fix(data): scope batch-sampler shuffling to training

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

Author: filaretov

src/megatron/bridge/data/base.py

@@ -43,8 +43,12 @@ 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."""
+    train_shuffle: bool = True
+    """Whether the training global-batch sampler reshuffles samples each epoch.
+
+    Only used when ``dataloader_type="batch"``. Validation and test samples are
+    never shuffled.
+    """
 
     persistent_workers: bool = True
     """Whether dataloader workers persist between iterations."""

src/megatron/bridge/data/loaders.py

@@ -334,7 +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,
+        shuffle=cfg.dataset.train_shuffle,
     )
     eval_gbs = (
         cfg.validation.eval_global_batch_size
@@ -363,7 +363,7 @@ def worker_init_fn(_):
             global_batch_size=eval_gbs,
             drop_last=not (isinstance(cfg.dataset, GPTSFTDatasetConfig) and cfg.dataset.dataloader_type == "batch"),
             seed=sampler_seed,
-            shuffle=cfg.dataset.shuffle,
+            shuffle=False,
         )
     elif cfg.validation.eval_iters > 0:
         val_dataloader_type = "cyclic" if isinstance(cfg.dataset, GPTDatasetConfig) else cfg.dataset.dataloader_type
@@ -383,7 +383,7 @@ def worker_init_fn(_):
             global_batch_size=eval_gbs,
             drop_last=not (isinstance(cfg.dataset, GPTSFTDatasetConfig) and val_dataloader_type == "batch"),
             seed=sampler_seed,
-            shuffle=cfg.dataset.shuffle,
+            shuffle=False,
         )
 
     if cfg.validation.eval_iters > 0:
@@ -403,7 +403,7 @@ def worker_init_fn(_):
             global_batch_size=eval_gbs,
             drop_last=not (isinstance(cfg.dataset, GPTSFTDatasetConfig) and cfg.dataset.dataloader_type == "batch"),
             seed=sampler_seed,
-            shuffle=cfg.dataset.shuffle,
+            shuffle=False,
         )
 
     # Flags to know if we need to do training/validation/testing.

tests/functional_tests/test_groups/data/test_loaders.py

@@ -320,7 +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,
+                train_shuffle=False,
                 num_workers=0,
                 persistent_workers=False,
             )
@@ -356,14 +356,14 @@ 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
+        assert train_call.kwargs["shuffle"] is cfg.dataset.train_shuffle
+        assert valid_call.kwargs["shuffle"] is False
+        assert test_call.kwargs["shuffle"] is False
 
     @mock.patch("torch.distributed.broadcast")
     @mock.patch("torch.distributed.get_world_size", return_value=1)
     @mock.patch("torch.distributed.get_rank", return_value=0)
-    def test_gpt_sft_batch_eval_preserves_split_smaller_than_global_batch(
+    def test_gpt_sft_batch_eval_is_unshuffled_and_preserves_split_smaller_than_global_batch(
         self, _mock_rank, _mock_world_size, _mock_broadcast
     ):
         class PaddingAwareDataset:
@@ -404,8 +404,8 @@ def __getitem__(self, index):
         test_batch = next(iter(test_dataloader))
         assert valid_batch.shape == (cfg.validation.eval_global_batch_size,)
         assert test_batch.shape == (cfg.validation.eval_global_batch_size,)
-        assert -1 in valid_batch
-        assert -1 in test_batch
+        assert valid_batch.tolist() == [0, 1, 2, -1]
+        assert test_batch.tolist() == [0, 1, 2, -1]
 
     @mock.patch("torch.distributed.broadcast")
     @mock.patch("torch.distributed.get_world_size", return_value=1)

tests/unit_tests/data/builders/test_gpt_sft_config.py

@@ -59,7 +59,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,
+        train_shuffle=False,
         enable_offline_packing=True,
         offline_packing_specs=specs,
         do_test=False,
@@ -74,7 +74,7 @@ def test_config_round_trip_is_declarative_and_serializable(tmp_path):
     assert restored.offline_packing_specs.packed_sequence_size == 128
     assert restored.offline_packing_specs.max_single_sequence_length == 120
     assert isinstance(restored.preprocessing, PromptCompletionSFTPreprocessingConfig)
-    assert restored.shuffle is False
+    assert restored.train_shuffle is False
     assert "tokenizer" not in serialized