Enhance DiT and Wan mock data modules by initializing training data l… (#65)

* Enhance DiT and Wan mock data modules by initializing training data loaders as iterators. Update test configuration for DiT pretraining to include checkpoint save and load paths.

* Remove unnecessary DataLoader assertion from Wan mock datamodule tests to streamline test execution.

Author: abhinavg4

dfm/src/megatron/data/dit/dit_mock_datamodule.py

@@ -158,6 +158,7 @@ def __post_init__(self):
             shuffle=False,
             drop_last=False,
         )
+        self._train_dl = iter(self._train_dl)
         self.sequence_length = self.seq_length
 
     def build_datasets(self, _context: DatasetBuildContext):

dfm/src/megatron/data/wan/wan_mock_datamodule.py

@@ -136,6 +136,7 @@ def __post_init__(self):
             shuffle=False,
             drop_last=False,
         )
+        self._train_dl = iter(self._train_dl)
         self.sequence_length = self.seq_length
 
     def build_datasets(self, _context: DatasetBuildContext):

tests/functional_tests/mcore/recipes/test_dit_pretrain.py

@@ -55,6 +55,8 @@ def test_DiT_pretrain_mock(self, tmp_path):
             "model.context_parallel_size=1",
             "model.qkv_format=thd",
             "model.num_attention_heads=16",
+            f"checkpoint.save={checkpoint_dir}",
+            f"checkpoint.load={checkpoint_dir}",
             "dataset.task_encoder_seq_length=4608",
             "dataset.seq_length=4608",
             "train.global_batch_size=2",

tests/unit_tests/megatron/data/wan/test_wan_mock_datamodule.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import torch
-from torch.utils.data import DataLoader
 
 from dfm.src.megatron.data.wan.wan_mock_datamodule import WanMockDataModuleConfig
 
@@ -37,7 +36,6 @@ def test_wan_mock_datamodule_build_and_batch_shapes():
         context_embeddings_dim=64,
     )
     train_dl, val_dl, test_dl = cfg.build_datasets(_context=None)
-    assert isinstance(train_dl, DataLoader)
     assert train_dl is val_dl and val_dl is test_dl
 
     batch = next(iter(train_dl))