make log clear when dataset is run out and re-looped

Author: rakkit

torchtitan/hf_datasets/text_datasets.py

@@ -154,6 +154,7 @@ def __init__(
         ds = dataset_loader(path)
 
         self.dataset_name = dataset_name
+        self.dataset_path = dataset_path
         self._data = split_dataset_by_node(ds, dp_rank, dp_world_size)
         self._tokenizer = tokenizer
         self.infinite = infinite
@@ -201,12 +202,16 @@ def __iter__(self):
                 yield sample_tokens
 
             if not self.infinite:
-                logger.warning(f"Dataset {self.dataset_name} has run out of data")
+                logger.warning(
+                    f"HuggingFaceDataset {self.dataset_name} from {self.dataset_path} has run out of data"
+                )
                 break
             else:
                 # Reset offset for the next iteration
                 self._sample_idx = 0
-                logger.warning(f"Dataset {self.dataset_name} is being re-looped")
+                logger.warning(
+                    f"HuggingFaceDataset {self.dataset_name} from {self.dataset_path} is being re-looped"
+                )
                 # Ensures re-looping a dataset loaded from a checkpoint works correctly
                 if not isinstance(self._data, Dataset):
                     if hasattr(self._data, "set_epoch") and hasattr(
@@ -344,6 +349,10 @@ def __init__(
     def dataset_name(self):
         return self._data.dataset_name
 
+    @property
+    def dataset_path(self):
+        return self._data.dataset_path
+
     def _get_data_iter(self):
         # We don't use the sample index because we defer skipping to the
         # sub-dataset.
@@ -367,13 +376,15 @@ def __iter__(self):
 
             if not self.infinite:
                 logger.warning(
-                    f"Packed dataset {self.dataset_name} has run out of data"
+                    f"GreedyPackedDataset {self.dataset_name} from {self.dataset_path} has run out of data"
                 )
                 break
             else:
                 # Reset offset for the next iteration
                 self._sample_idx = 0
-                logger.warning(f"Packed dataset {self.dataset_name} is being re-looped")
+                logger.warning(
+                    f"GreedyPackedDataset {self.dataset_name} from {self.dataset_path} is being re-looped"
+                )
                 # Ensures re-looping a dataset loaded from a checkpoint works correctly
                 if not isinstance(self._data, Dataset):
                     if hasattr(self._data, "set_epoch") and hasattr(