refactor inference_configs + mockdatamodule

Author: root

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

@@ -73,7 +73,7 @@ def __init__(
         self.patch_temporal = patch_temporal
         self.seq_length = seq_length
 
-    ## actual encode_sample() for production
+
     def encode_sample(self, sample: dict) -> dict:
         video_latent = sample["pth"]
         context_embeddings = sample["pickle"]
@@ -103,24 +103,6 @@ def encode_sample(self, sample: dict) -> dict:
             video_metadata=video_metadata,
         )
 
-    ## mock encode_sample() for debugging
-    # def encode_sample(self, sample: dict) -> dict:
-
-    #     # mock encode sample
-    #     F_latents = 24
-    #     H_latents = 104
-    #     W_latents = 60
-    #     video_latent = torch.tensor(torch.randn(16, F_latents, H_latents, W_latents), dtype=torch.float32)
-    #     grid_size = torch.tensor([video_latent.shape[1] // self.patch_temporal, video_latent.shape[2] // self.patch_spatial, video_latent.shape[3] // self.patch_spatial], dtype=torch.int32)
-    #     context_embeddings = torch.tensor(torch.randn(512, 4096), dtype=torch.float32)
-    #     video_metadata = {}
-
-    #     return dict(
-    #         video_latent=video_latent,
-    #         grid_size=grid_size,
-    #         context_embeddings=context_embeddings,
-    #         video_metadata=video_metadata,
-    #     )
 
     def batch(self, samples: list[dict]) -> dict:
         # process video latents

dfm/src/megatron/model/wan/flow_matching/flow_inference_pipeline.py

@@ -71,7 +71,7 @@ def _encode_text(
 class FlowInferencePipeline:
     def __init__(
         self,
-        config,
+        inference_cfg,
         model_id="Wan-AI/Wan2.1-T2V-14B-Diffusers",
         checkpoint_dir=None,
         checkpoint_step=None,
@@ -90,8 +90,8 @@ def __init__(
         Initializes the FlowInferencePipeline with the given parameters.
 
         Args:
-            config (EasyDict):
-                Object containing model parameters initialized from config.py
+            inference_cfg (dict):
+                Object containing inference configuration.
             checkpoint_dir (`str`):
                 Path to directory containing model checkpoints
             t5_checkpoint_dir (`str`, *optional*, defaults to None):
@@ -106,7 +106,7 @@ def __init__(
                 Whether to place T5 model on CPU. Only works without t5_fsdp.
         """
         self.device = torch.device(f"cuda:{device_id}")
-        self.config = config
+        self.inference_cfg = inference_cfg
         self.model_id = model_id
         self.rank = rank
         self.t5_cpu = t5_cpu
@@ -115,25 +115,26 @@ def __init__(
         self.pipeline_parallel_size = pipeline_parallel_size
         self.sequence_parallel = sequence_parallel
         self.pipeline_dtype = pipeline_dtype
-        self.num_train_timesteps = config.num_train_timesteps
-        self.param_dtype = config.param_dtype
+        self.num_train_timesteps = inference_cfg.num_train_timesteps
+        self.param_dtype = inference_cfg.param_dtype
+        self.text_len = inference_cfg.text_len
 
         self.text_encoder = UMT5EncoderModel.from_pretrained(
             model_id,
             subfolder="text_encoder",
-            torch_dtype=config.t5_dtype,
+            torch_dtype=inference_cfg.t5_dtype,
         )
         self.tokenizer = AutoTokenizer.from_pretrained(
             model_id,
             subfolder="tokenizer",
         )
 
-        self.vae_stride = config.vae_stride
-        self.patch_size = config.patch_size
+        self.vae_stride = inference_cfg.vae_stride
+        self.patch_size = inference_cfg.patch_size
         self.vae = AutoencoderKLWan.from_pretrained(
             model_id,
             subfolder="vae",
-            torch_dtype=config.param_dtype,
+            torch_dtype=inference_cfg.param_dtype,
         )
         self.vae.to(self.device)
 
@@ -150,7 +151,7 @@ def __init__(
             dist.barrier()
         self.model.to(self.device)
 
-        self.sample_neg_prompt = config.sample_neg_prompt
+        self.sample_neg_prompt = inference_cfg.sample_neg_prompt
 
     def setup_model_from_checkpoint(self, checkpoint_dir):
         provider = WanModelProvider()
@@ -362,7 +363,7 @@ def generate(
         # we implement similar to Wan's diffuser setup
         # (https://github.com/huggingface/diffusers/blob/0f252be0ed42006c125ef4429156cb13ae6c1d60/src/diffusers/pipelines/wan/pipeline_wan.py#L157)
         # in which we pad the text to 512, pass through text encoder, and truncate to the actual tokens, then pad with 0s to 512.
-        context_max_len = 512
+        context_max_len = self.text_len
         context_lens = []
         contexts = []
         contexts_null = []

dfm/src/megatron/model/wan/inference/configs/__init__.py

@@ -28,7 +28,7 @@
 from megatron.core.transformer.transformer_block import TransformerBlock
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import make_sharded_tensor_for_checkpoint
-from nemo.collections.diffusion.models.dit.dit_embeddings import ParallelTimestepEmbedding
+from dfm.src.megatron.model.common.dit_embeddings import ParallelTimestepEmbedding
 from torch import Tensor
 
 from dfm.src.megatron.model.wan.wan_layer_spec import (

dfm/src/megatron/model/wan/inference/configs/shared_config.py

@@ -31,6 +31,7 @@
 from megatron.core.distributed import DistributedDataParallelConfig
 
 from dfm.src.megatron.data.wan.wan_energon_datamodule import WanDataModuleConfig
+from dfm.src.megatron.data.wan.wan_mock_energon_datamodule import WanMockDataModuleConfig
 from dfm.src.megatron.model.wan.wan_provider import WanModelProvider
 
 
@@ -158,6 +159,28 @@ def pretrain_config(
 
     precision_config.grad_reduce_in_fp32 = False
 
+    if mock:
+        dataset = WanMockDataModuleConfig(
+            path=None,
+            seq_length=1024, # we don't need to use this value, just add because Bridge training requires for LLMs
+            F_latents=3,
+            H_latents=104,
+            W_latents=60,
+            context_seq_len=512,
+            context_embeddings_dim=4096,
+            micro_batch_size=micro_batch_size,
+            global_batch_size=global_batch_size,
+            num_workers=10,
+        )
+    else:
+        dataset = WanDataModuleConfig(
+            path=None,
+            seq_length=1024, # we don't need to use this value, just add because Bridge training requires for LLMs
+            micro_batch_size=micro_batch_size,
+            global_batch_size=global_batch_size,
+            num_workers=10,
+        )
+
     # Config Container
     cfg = ConfigContainer(
         model=model_cfg,
@@ -182,13 +205,7 @@ def pretrain_config(
             use_distributed_optimizer=True,
             use_megatron_fsdp=use_megatron_fsdp,  # need use_distributed_optimizer=True
         ),
-        dataset=WanDataModuleConfig(
-            path=None,
-            seq_length=1024,  # we don't need to use this value, just add because Bridge training requires for LLMs
-            micro_batch_size=micro_batch_size,
-            global_batch_size=global_batch_size,
-            num_workers=10,
-        ),
+        dataset=dataset,
         logger=LoggerConfig(
             log_interval=10,
             tensorboard_dir=tensorboard_dir,