Sync with torchtitan

.ci/docker/requirements.txt

@@ -1,6 +1,5 @@
-torch >= 2.3.0
 torchdata >= 0.8.0
-datasets >= 2.19.0
+datasets >= 2.21.0
 tomli >= 1.1.0 ; python_version < "3.11"
 tensorboard
 sentencepiece

torchtitan/datasets/hf_datasets.py

@@ -11,18 +11,13 @@
 from torch.distributed.checkpoint.stateful import Stateful
 from torch.utils.data import IterableDataset
 
-try:
-    from torchdata.stateful_dataloader import StatefulDataLoader
-except ImportError as e:
-    raise ImportError(
-        "Please install the latest torchdata nightly to use StatefulDataloader via:"
-        "pip3 install --pre torchdata --index-url https://download.pytorch.org/whl/nightly"
-    ) from e
+from torchdata.stateful_dataloader import StatefulDataLoader
+
 
 from torchtitan.datasets.tokenizer import Tokenizer
 from torchtitan.logging import logger
 
-from datasets import load_dataset
+from datasets import Dataset, load_dataset
 from datasets.distributed import split_dataset_by_node
 
 # map from dataset name to a local directory, or
@@ -102,7 +97,7 @@ def __init__(
         else:
             ds = load_dataset(dataset_path, split="train")
 
-        # TODO: support shuffling and checkpointing
+        # TODO: support shuffling
         self.dataset_name = dataset_name
         self._data = split_dataset_by_node(ds, rank, world_size)
         self._tokenizer = tokenizer
@@ -143,17 +138,11 @@ def _get_data_iter(self):
         if self._sample_idx == 0:
             return iter(self._data)
 
-        # Skip samples
-        if isinstance(self._data, IterableDataset):
-            it = iter(self._data)
-            # Naively iterate through the samples as skip may not be supported
-            for _ in range(self._sample_idx):
-                next(it)
-            return it
-
         # As skipping to the end throws an error in case of map-style dataset, return an empty iterator
-        if self._sample_idx == len(self._data):
+        if isinstance(self._data, Dataset) and self._sample_idx == len(self._data):
             return iter([])
+
+
         return iter(self._data.skip(self._sample_idx))
 
     def load_state_dict(self, state_dict):
@@ -179,7 +168,7 @@ def state_dict(self) -> Dict[str, Any]:
         return {self._rank_id: pickle.dumps(super().state_dict())}
 
     def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
-        # State being empty is valid, don't log a warning
+        # State being empty is valid
         if not state_dict:
             return
 

train.py

@@ -204,6 +204,11 @@ def loss_fn(pred, labels):
     checkpoint_loaded = checkpoint.load()
 
     if parallel_dims.pp_enabled and not checkpoint_loaded:
+        if job_config.experimental.pipeline_parallel_split_mode == "tracer":
+            raise RuntimeError(
+                "Pipeline parallelism with tracer mode is not supported without a seed checkpoint."
+            )
+
         # TODO: fix this by allowing each rank to set their own seed
         logger.warning(
             "Pipeline Parallelism is being used without a seed checkpoint. "