path: deterministic fidxs for vit dataloader

torchtitan/experiments/path/config_registry.py

@@ -595,6 +595,7 @@ def _vit_dataloader_config(*, split: str) -> PathDataLoader.Config:
         _dataloader_config(split=split, fps=SUPERCOMBO_FPS, plan_only=True),
         dataset=os.path.join(XX_BASEDIR, dataset),
         pipeline_dir=BASE_DIR_GT_10M,
+        deterministic_fidxs=True,
     )
 
 

torchtitan/experiments/path/dataset.py

@@ -1,3 +1,9 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
 from __future__ import annotations
 
 import os
@@ -26,6 +32,7 @@ class Config(BaseDataLoader.Config):
         n_frames: int
         rgb: bool
         unvision: bool
+        deterministic_fidxs: bool = False
 
     def __init__(
         self,
@@ -41,11 +48,12 @@ def __init__(
         **kwargs: Any,
     ) -> None:
         del tokenizer, seq_len, snapshot_every_n_steps, kwargs
-        from gigashuffle import DataloaderConfig
         from xx.training.lib.dataloader import DataLoader
         from xx.training.path.config import DatasetConfig as XXPathDatasetConfig
         from xx.training.path.dataloader import get_dataset
 
+        from gigashuffle import DataloaderConfig
+
         self.config = config
         self.local_batch_size = local_batch_size
         self.dp_world_size = dp_world_size
@@ -72,6 +80,7 @@ def __init__(
             n_frames=config.n_frames,
             rgb=config.rgb,
             unvision=config.unvision,
+            deterministic_fidxs=config.deterministic_fidxs,
         )
         dataset = get_dataset(config.dataset, xx_config, val, self.local_rank)
         self.dataset = dataset
@@ -91,7 +100,9 @@ def __init__(
         self.loader = DataLoader(dataset, loader_config)
         self._iterator: Any | None = None
 
-    def __iter__(self) -> Iterator[tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]]:
+    def __iter__(
+        self,
+    ) -> Iterator[tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]]:
         iterator = iter(self.loader)
         self._iterator = iterator
         try: