[data] fix: Keep cyclic sampler ranks epoch-aligned (#5944)

Signed-off-by: Yu Yao <yaoyu.094@gmail.com>

Author: yaoyu-33

src/megatron/bridge/data/samplers.py

@@ -506,12 +506,11 @@ def __iter__(self) -> Iterator[list[int]]:
             random_idx = torch.randperm(bucket_size, generator=g).tolist()
             idx_range = [start_idx + x for x in random_idx[bucket_offset:]]
         else:
-            full_bucket_size = (self.total_samples // self.micro_batch_size) * self.micro_batch_size
             full_bucket_offset = current_epoch_samples
             g = torch.Generator()
             g.manual_seed(self.epoch)
-            idx_range_total = torch.randperm(full_bucket_size, generator=g).tolist()
-            idx_range_active = idx_range_total[full_bucket_offset:]
+            idx_range_total = torch.randperm(self.total_samples, generator=g).tolist()
+            idx_range_active = idx_range_total[full_bucket_offset:active_total_samples]
             idx_range = idx_range_active[self.data_parallel_rank :: self.data_parallel_size]
 
         batch = []

tests/unit_tests/data/test_samplers.py

@@ -5,6 +5,7 @@
 from torch.utils.data import Dataset
 
 from megatron.bridge.data.samplers import (
+    MegatronPretrainingRandomSampler,
     MegatronPretrainingSampler,
     RandomSeedDataset,
     build_pretraining_data_loader,
@@ -108,3 +109,30 @@ def test_cyclic_sampler_resume_seeds_worker_dataset_for_current_epoch(
     expected_value = torch.randint(0, 1_000_000, (1,), generator=generator).item()
 
     assert actual_value.item() == expected_value
+
+
+def test_cyclic_sampler_non_sharded_tail_keeps_data_parallel_epochs_aligned() -> None:
+    """Every DP rank must cross a truncated cyclic epoch on the same step."""
+
+    def make_sampler(rank: int, consumed_samples: int = 0) -> MegatronPretrainingRandomSampler:
+        return MegatronPretrainingRandomSampler(
+            list(range(10)),
+            total_samples=10,
+            consumed_samples=consumed_samples,
+            micro_batch_size=2,
+            data_parallel_rank=rank,
+            data_parallel_size=3,
+            data_sharding=False,
+        )
+
+    def cyclic_batches(sampler: MegatronPretrainingRandomSampler):
+        while True:
+            yield from sampler
+
+    uninterrupted = [cyclic_batches(make_sampler(rank)) for rank in range(3)]
+    _ = [next(rank_batches) for rank_batches in uninterrupted]
+    second_step = [next(rank_batches) for rank_batches in uninterrupted]
+    resumed_step = [next(iter(make_sampler(rank, consumed_samples=6))) for rank in range(3)]
+
+    assert len({index for batch in second_step for index in batch}) == 6
+    assert second_step == resumed_step