[data] fix: repeat FLUX mock data across training (#5534)

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

Author: yaoyu-33

src/megatron/bridge/diffusion/data/flux/flux_mock_datamodule.py

@@ -15,6 +15,7 @@
 """Mock data module for FLUX model training."""
 
 from dataclasses import dataclass
+from itertools import chain, repeat
 
 import torch
 from torch.utils.data import DataLoader, Dataset
@@ -219,7 +220,7 @@ class FluxMockDataModuleConfig(DatasetProvider):
 
     def __post_init__(self):
         """Initialize the mock dataset and dataloader."""
-        mock_ds = _MockT2IDataset(
+        self._mock_ds = _MockT2IDataset(
             image_H=self.image_H,
             image_W=self.image_W,
             length=self.num_train_samples,
@@ -232,13 +233,19 @@ def __post_init__(self):
             vae_channels=self.vae_channels,
         )
 
+        self._train_dl = self._build_dataloader()
+        self._valid_dl = self._build_dataloader()
+        self._test_dl = self._build_dataloader()
+        self.sequence_length = self.seq_length
+
+    def _build_dataloader(self) -> DataLoader:
         kwargs = {}
         if self.num_workers > 0:
             kwargs["prefetch_factor"] = 8
             kwargs["persistent_workers"] = True
 
-        self._train_dl = DataLoader(
-            mock_ds,
+        return DataLoader(
+            self._mock_ds,
             batch_size=self.micro_batch_size,
             num_workers=self.num_workers,
             collate_fn=_collate_fn,
@@ -247,10 +254,10 @@ def __post_init__(self):
             pin_memory=True,
             **kwargs,
         )
-        self._train_dl_iter = iter(self._train_dl)
-        self.sequence_length = self.seq_length
 
     def build_datasets(self, _context: DatasetBuildContext):
         """Build and return train/val/test dataloaders."""
-        # Return iterator for external dataloader type
-        return self._train_dl_iter, self._train_dl_iter, self._train_dl_iter
+        # Reiterate each loader without caching generated batches in host memory.
+        return tuple(
+            chain.from_iterable(repeat(dataloader)) for dataloader in (self._train_dl, self._valid_dl, self._test_dl)
+        )

tests/unit_tests/diffusion/data/flux/test_flux_mock_datamodule.py

@@ -38,7 +38,8 @@ def test_flux_mock_datamodule_build_and_batch_shapes():
         num_train_samples=100,
     )
     train_dl, val_dl, test_dl = cfg.build_datasets(_context=None)
-    assert train_dl is val_dl and val_dl is test_dl
+    assert train_dl is not val_dl
+    assert val_dl is not test_dl
 
     batch = next(iter(train_dl))
     expected_keys = {
@@ -77,6 +78,31 @@ def test_flux_mock_datamodule_build_and_batch_shapes():
     assert batch["loss_mask"].dtype == torch.bfloat16
 
 
+def test_flux_mock_datamodule_repeats_with_independent_splits():
+    cfg = FluxMockDataModuleConfig(
+        micro_batch_size=1,
+        global_batch_size=2,
+        num_workers=0,
+        image_H=8,
+        image_W=8,
+        vae_channels=1,
+        vae_scale_factor=8,
+        prompt_seq_len=2,
+        context_dim=2,
+        pooled_prompt_dim=2,
+        num_train_samples=2,
+    )
+
+    train_dl, val_dl, test_dl = cfg.build_datasets(_context=None)
+
+    for _ in range(cfg.num_train_samples + 1):
+        next(train_dl)
+    next(val_dl)
+    next(test_dl)
+    assert train_dl is not val_dl
+    assert val_dl is not test_dl
+
+
 def test_flux_mock_datamodule_without_precaching():
     """Test the mock datamodule with non-precached data."""
     cfg = FluxMockDataModuleConfig(