feat(data): expose shuffle in DataloaderConfig

Signed-off-by: Hristo Filaretov <hfilaretov@nvidia.com>

Author: filaretov

src/megatron/bridge/data/base.py

@@ -43,6 +43,9 @@ class DataloaderConfig:
     drop_last: bool = True
     """Whether dataloaders drop the last incomplete batch."""
 
+    shuffle: bool = True
+    """Whether the global-batch sampler used by ``dataloader_type="batch"`` reshuffles samples each epoch."""
+
     persistent_workers: bool = True
     """Whether dataloader workers persist between iterations."""
 

src/megatron/bridge/data/loaders.py

@@ -334,6 +334,7 @@ def worker_init_fn(_):
         global_batch_size=cfg.train.global_batch_size,
         drop_last=drop_last,
         seed=sampler_seed,
+        shuffle=cfg.dataset.shuffle,
     )
     eval_gbs = (
         cfg.validation.eval_global_batch_size
@@ -361,6 +362,7 @@ def worker_init_fn(_):
             data_parallel_size=eval_dp_size,
             global_batch_size=eval_gbs,
             seed=sampler_seed,
+            shuffle=cfg.dataset.shuffle,
         )
     elif cfg.validation.eval_iters > 0:
         val_dataloader_type = "cyclic" if isinstance(cfg.dataset, GPTDatasetConfig) else cfg.dataset.dataloader_type
@@ -379,6 +381,7 @@ def worker_init_fn(_):
             data_parallel_size=eval_dp_size,
             global_batch_size=eval_gbs,
             seed=sampler_seed,
+            shuffle=cfg.dataset.shuffle,
         )
 
     if cfg.validation.eval_iters > 0:
@@ -397,6 +400,7 @@ def worker_init_fn(_):
             data_parallel_size=eval_dp_size,
             global_batch_size=eval_gbs,
             seed=sampler_seed,
+            shuffle=cfg.dataset.shuffle,
         )
 
     # Flags to know if we need to do training/validation/testing.

src/megatron/bridge/data/samplers.py

@@ -26,6 +26,7 @@ def build_pretraining_data_loader(
     drop_last: Optional[bool] = True,
     global_batch_size: Optional[int] = None,
     seed: int | None = None,
+    shuffle: bool = True,
 ) -> Optional[DataLoader]:
     """Build a dataloader for pretraining.
 
@@ -53,6 +54,8 @@ def build_pretraining_data_loader(
         seed: Explicit shuffle seed for the global-batch sampler. Supplying the
               dataset seed keeps pipeline stages on the same sample order even
               though their model RNG seeds differ.
+        shuffle: Whether the global-batch sampler reshuffles samples each epoch.
+                 Only used when dataloader_type is 'batch'.
 
     Returns:
         A PyTorch DataLoader instance, or the dataset itself if dataloader_type is
@@ -101,6 +104,7 @@ def build_pretraining_data_loader(
             data_parallel_size=data_parallel_size,
             drop_last=drop_last,
             pad_samples_to_global_batch_size=not drop_last,
+            shuffle=shuffle,
             seed=seed,
         )
     elif dataloader_type == "external":

tests/functional_tests/test_groups/data/test_loaders.py

@@ -320,6 +320,7 @@ def test_build_train_valid_test_data_loaders_uses_eval_dp_group(
                 dataset_root="/tmp/dataset",
                 seq_length=512,
                 seed=4321,
+                shuffle=False,
                 num_workers=0,
                 persistent_workers=False,
             )
@@ -355,6 +356,9 @@ def test_build_train_valid_test_data_loaders_uses_eval_dp_group(
         assert train_call.kwargs["seed"] == expected_seed
         assert valid_call.kwargs["seed"] == expected_seed
         assert test_call.kwargs["seed"] == expected_seed
+        assert train_call.kwargs["shuffle"] is cfg.dataset.shuffle
+        assert valid_call.kwargs["shuffle"] is cfg.dataset.shuffle
+        assert test_call.kwargs["shuffle"] is cfg.dataset.shuffle
 
     @mock.patch("torch.distributed.broadcast")
     @mock.patch("torch.distributed.get_world_size", return_value=1)

tests/functional_tests/test_groups/data/test_samplers.py

@@ -202,6 +202,7 @@ def test_batch_sampler_initialization(self):
         assert sampler.micro_batch_size == 4
         assert sampler._global_batch_size == 16
         assert sampler.data_parallel_size == 2
+        assert sampler.shuffle is True
 
     def test_batch_sampler_length(self):
         """Test length calculation for batch sampler.
@@ -845,6 +846,23 @@ def mock_dataloader():
 class TestBatchDataloaderIntegration:
     """Integration tests for batch dataloader type."""
 
+    def test_build_batch_dataloader_propagates_shuffle(self):
+        """The loader builder must pass shuffle through to the batch sampler."""
+        dataloader = build_pretraining_data_loader(
+            dataset=list(range(16)),
+            consumed_samples=0,
+            dataloader_type="batch",
+            micro_batch_size=1,
+            num_workers=0,
+            data_sharding=False,
+            global_batch_size=4,
+            seed=1234,
+            shuffle=False,
+        )
+
+        assert dataloader.batch_sampler.shuffle is False
+        assert next(iter(dataloader.batch_sampler)) == [0, 1, 2, 3]
+
     def test_build_batch_dataloader_explicit_seed_is_independent_of_model_rng(self):
         """Pipeline-stage model seeds must not change the fine-tuning sample order."""
         import torch

tests/unit_tests/data/builders/test_gpt_sft_config.py

@@ -55,6 +55,7 @@ def test_config_round_trip_is_declarative_and_serializable(tmp_path):
         hf_output_root=str(tmp_path),
         hf_validation_proportion=0.1,
         seed=5678,
+        shuffle=False,
         enable_offline_packing=True,
         offline_packing_specs=specs,
         do_test=False,
@@ -68,6 +69,7 @@ def test_config_round_trip_is_declarative_and_serializable(tmp_path):
     assert restored.hf_dataset.dataset_name == "squad"
     assert restored.offline_packing_specs.packed_sequence_size == 128
     assert isinstance(restored.preprocessing, PromptCompletionSFTPreprocessingConfig)
+    assert restored.shuffle is False
     assert "tokenizer" not in serialized