diff --git a/torchtitan/experiments/path/config_registry.py b/torchtitan/experiments/path/config_registry.py index 05281ae4fe..42cf248c21 100644 --- a/torchtitan/experiments/path/config_registry.py +++ b/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, ) diff --git a/torchtitan/experiments/path/dataset.py b/torchtitan/experiments/path/dataset.py index 592a3a1809..54330c0f95 100644 --- a/torchtitan/experiments/path/dataset.py +++ b/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: