Dit unit tests

dfm/src/common/utils/save_video.py

@@ -44,5 +44,4 @@ def save_video(
         "output_params": ["-f", "mp4"],
     }
 
-    print("video_save_path", video_save_path)
     imageio.mimsave(video_save_path, grid, "mp4", **kwargs)

dfm/src/megatron/data/common/diffusion_energon_datamodule.py

@@ -55,7 +55,11 @@ def __post_init__(self):
         self.sequence_length = self.dataset.seq_length
 
     def build_datasets(self, context: DatasetBuildContext):
-        return self.dataset.train_dataloader(), self.dataset.val_dataloader(), self.dataset.test_dataloader()
+        return (
+            iter(self.dataset.train_dataloader()),
+            iter(self.dataset.val_dataloader()),
+            iter(self.dataset.val_dataloader()),
+        )
 
 
 class DiffusionDataModule(EnergonMultiModalDataModule):

dfm/src/megatron/data/common/diffusion_task_encoder_with_sp.py

@@ -56,7 +56,7 @@ def __init__(
         self,
         *args,
         max_frames: int = None,
-        text_embedding_padding_size: int = 512,
+        text_embedding_max_length: int = 512,
         seq_length: int = None,
         patch_spatial: int = 2,
         patch_temporal: int = 1,
@@ -65,7 +65,7 @@ def __init__(
     ):
         super().__init__(*args, **kwargs)
         self.max_frames = max_frames
-        self.text_embedding_padding_size = text_embedding_padding_size
+        self.text_embedding_max_length = text_embedding_max_length
         self.seq_length = seq_length
         self.patch_spatial = patch_spatial
         self.patch_temporal = patch_temporal

dfm/src/megatron/data/common/sequence_packing_utils.py

@@ -71,35 +71,3 @@ def first_fit_decreasing(seqlens: List[int], pack_size: int) -> List[List[int]]:
     """
     sorted_seqlens = sorted(seqlens, reverse=True)
     return first_fit(sorted_seqlens, pack_size)
-
-
-def concat_pad(tensor_list, max_seq_length):
-    """
-    Efficiently concatenates a list of tensors along the first dimension and pads with zeros
-    to reach max_seq_length.
-
-    Args:
-        tensor_list (list of torch.Tensor): List of tensors to concatenate and pad.
-        max_seq_length (int): The desired size of the first dimension of the output tensor.
-
-    Returns:
-        torch.Tensor: A tensor of shape [max_seq_length, ...], where ... represents the remaining dimensions.
-    """
-    import torch
-
-    # Get common properties from the first tensor
-    other_shape = tensor_list[0].shape[1:]
-    dtype = tensor_list[0].dtype
-    device = tensor_list[0].device
-
-    # Initialize the result tensor with zeros
-    result = torch.zeros((max_seq_length, *other_shape), dtype=dtype, device=device)
-
-    current_index = 0
-    for tensor in tensor_list:
-        length = tensor.shape[0]
-        # Directly assign the tensor to the result tensor without checks
-        result[current_index : current_index + length] = tensor
-        current_index += length
-
-    return result

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

@@ -31,9 +31,9 @@ class DiTTaskEncoder(DiffusionTaskEncoderWithSequencePacking):
     Attributes:
         cookers (list): A list of Cooker objects used for processing.
         max_frames (int, optional): The maximum number of frames to consider from the video. Defaults to None.
-        text_embedding_padding_size (int): The padding size for text embeddings. Defaults to 512.
+        text_embedding_max_length (int): The maximum length for text embeddings. Defaults to 512.
     Methods:
-        __init__(*args, max_frames=None, text_embedding_padding_size=512, **kwargs):
+        __init__(*args, max_frames=None, text_embedding_max_size=512, **kwargs):
             Initializes the BasicDiffusionTaskEncoder with optional maximum frames and text embedding padding size.
         encode_sample(sample: dict) -> dict:
             Encodes a given sample dictionary containing video and text data.
@@ -71,7 +71,6 @@ def encode_sample(self, sample: dict) -> DiffusionSample:
             // self.patch_spatial**2
             // self.patch_temporal
         )
-        is_image = T == 1
 
         if seq_len > self.seq_length:
             print(f"Skipping sample {sample['__key__']} because seq_len {seq_len} > self.seq_length {self.seq_length}")
@@ -100,8 +99,8 @@ def encode_sample(self, sample: dict) -> DiffusionSample:
         t5_text_embeddings = torch.from_numpy(sample["pickle"]).to(torch.bfloat16)
         t5_text_embeddings_seq_length = t5_text_embeddings.shape[0]
 
-        if t5_text_embeddings_seq_length > self.text_embedding_padding_size:
-            t5_text_embeddings = t5_text_embeddings[: self.text_embedding_padding_size]
+        if t5_text_embeddings_seq_length > self.text_embedding_max_length:
+            t5_text_embeddings = t5_text_embeddings[: self.text_embedding_max_length]
         t5_text_mask = torch.ones(t5_text_embeddings_seq_length, dtype=torch.bfloat16)
 
         pos_ids = rearrange(

dfm/src/megatron/model/dit/dit_layer_spec.py

@@ -156,7 +156,9 @@ def _replace_no_cp_submodules(submodules):
             layer_number=layer_number,
         )
 
-        self.adaLN = AdaLN(config=self.config, n_adaln_chunks=9 if self.cross_attention else 6)
+        self.adaLN = AdaLN(
+            config=self.config, n_adaln_chunks=9 if not isinstance(self.cross_attention, IdentityOp) else 6
+        )
 
     def forward(
         self,
@@ -176,7 +178,7 @@ def forward(
     ):
         timestep_emb = attention_mask
 
-        if self.cross_attention:
+        if not isinstance(self.cross_attention, IdentityOp):
             shift_full, scale_full, gate_full, shift_ca, scale_ca, gate_ca, shift_mlp, scale_mlp, gate_mlp = (
                 self.adaLN(timestep_emb)
             )
@@ -192,7 +194,7 @@ def forward(
             packed_seq_params=None if packed_seq_params is None else packed_seq_params["self_attention"],
         )
 
-        if self.cross_attention:
+        if not isinstance(self.cross_attention, IdentityOp):
             hidden_states, pre_cross_attn_layernorm_output_ada = self.adaLN.scaled_modulated_layernorm(
                 residual=hidden_states,
                 x=attention_output,
@@ -210,7 +212,7 @@ def forward(
         hidden_states, pre_mlp_layernorm_output_ada = self.adaLN.scaled_modulated_layernorm(
             residual=hidden_states,
             x=attention_output,
-            gate=gate_ca if self.cross_attention else gate_full,
+            gate=gate_ca if not isinstance(self.cross_attention, IdentityOp) else gate_full,
             shift=shift_mlp,
             scale=scale_mlp,
         )

dfm/src/megatron/model/dit/dit_model_provider.py

@@ -14,7 +14,6 @@
 
 import logging
 from dataclasses import dataclass
-from typing import Callable
 
 import torch
 from megatron.bridge.models.model_provider import ModelProviderMixin
@@ -39,14 +38,14 @@ class DiTModelProvider(TransformerConfig, ModelProviderMixin[VisionModule]):
     add_bias_linear: bool = False
     gated_linear_unit: bool = False
 
-    num_layers: int = 28
-    hidden_size: int = 1152
+    num_layers: int = 12
+    hidden_size: int = 384
     max_img_h: int = 80
     max_img_w: int = 80
     max_frames: int = 34
     patch_spatial: int = 2
     patch_temporal: int = 1
-    num_attention_heads: int = 16
+    num_attention_heads: int = 6
     layernorm_epsilon = 1e-6
     normalization = "RMSNorm"
     add_bias_linear: bool = False
@@ -110,52 +109,27 @@ def configure_vae(self):
 
 
 @dataclass
-class DiT7BModelProvider(DiTModelProvider):
-    hidden_size: int = 4096
-    max_img_h: int = 240
-    max_img_w: int = 240
-    max_frames: int = 128
-    num_attention_heads: int = 32
+class DiTBModelProvider(DiTModelProvider):
+    """DiT-B"""
 
-    apply_rope_fusion: bool = True  # TODO: do we support this?
-    additional_timestamp_channels = None  # TODO: do we support this?
-    vae_module: str = None
-    vae_path: str = None
+    num_layers: int = 12
+    hidden_size: int = 768
+    num_attention_heads: int = 12
 
 
 @dataclass
-class DiT14BModelProvider(DiTModelProvider):
-    num_layers: int = 36
-    hidden_size: int = 5120
-    max_img_h: int = 240
-    max_img_w: int = 240
-    max_frames: int = 128
-    num_attention_heads: int = 40
-    apply_rope_fusion: bool = True
-    layernorm_zero_centered_gamma: bool = False
-    additional_timestamp_channels = None
-    vae_module: str = None
-    vae_path: str = None
-    loss_add_logvar: bool = True
+class DiTLModelProvider(DiTModelProvider):
+    """DiT-L"""
+
+    num_layers: int = 24
+    hidden_size: int = 1024
+    num_attention_heads: int = 16
 
 
 @dataclass
-class DiTLlama30BConfig(DiTModelProvider):
-    num_layers: int = 48
-    hidden_size: int = 6144
-    ffn_hidden_size: int = 16384
-    num_attention_heads: int = 48
-    num_query_groups: int = 8
-    gated_linear_unit: int = True
-    bias_activation_fusion: int = True
-    activation_func: Callable = torch.nn.functional.silu
-    layernorm_epsilon: float = 1e-5
-    max_frames: int = 128
-    max_img_h: int = 240
-    max_img_w: int = 240
-    init_method_std: float = 0.01
-    add_bias_linear: bool = False
-    seq_length: int = 256
-    masked_softmax_fusion: bool = True
-    persist_layer_norm: bool = True
-    bias_dropout_fusion: bool = True
+class DiTXLModelProvider(DiTModelProvider):
+    """DiT-XL"""
+
+    num_layers: int = 28
+    hidden_size: int = 1152
+    num_attention_heads: int = 16

dfm/src/megatron/model/dit/dit_step.py

@@ -50,7 +50,7 @@ def on_validation_start(self, state, batch, model):
             num_steps=model.config.val_generation_num_steps,
             is_negative_prompt=True if "neg_context_embeddings" in batch else False,
         )
-        caption = batch["video_metadata"][0]["caption"]
+        caption = batch["video_metadata"][0]["caption"] if "caption" in batch["video_metadata"][0] else "no caption"
         latent = latent[0, None, : batch["seq_len_q"][0]]
         latent = rearrange(
             latent,
@@ -157,7 +157,6 @@ def forward_step(self, state, batch, model, return_schedule_plan: bool = False):
 
         check_for_nan_in_loss = state.cfg.rerun_state_machine.check_for_nan_in_loss
         check_for_spiky_loss = state.cfg.rerun_state_machine.check_for_spiky_loss
-        # import pdb;pdb.set_trace()
         straggler_timer = state.straggler_timer
         with straggler_timer:
             if parallel_state.is_pipeline_last_stage():

dfm/src/megatron/model/dit/edm/edm_pipeline.py

@@ -179,21 +179,17 @@ def training_step(
         # import pdb; pdb.set_trace()
         # Get the input data to noise and denoise~(image, video) and the corresponding conditioner.
         self.net = model
-        x0_from_data_batch, x0, condition = self.get_data_and_condition(data_batch)
+        x0, condition = self.get_data_and_condition(data_batch)
 
         # Sample pertubation noise levels and N(0, 1) noises
         sigma, epsilon = self.draw_training_sigma_and_epsilon(x0.size(), condition)
 
         if parallel_state.is_pipeline_last_stage():
-            output_batch, pred_mse, edm_loss = self.compute_loss_with_epsilon_and_sigma(
-                data_batch, x0_from_data_batch, x0, condition, epsilon, sigma
-            )
+            output_batch, pred_mse, edm_loss = self.compute_loss_with_epsilon_and_sigma(x0, condition, epsilon, sigma)
 
             return output_batch, edm_loss
         else:
-            net_output = self.compute_loss_with_epsilon_and_sigma(
-                data_batch, x0_from_data_batch, x0, condition, epsilon, sigma
-            )
+            net_output = self.compute_loss_with_epsilon_and_sigma(x0, condition, epsilon, sigma)
             return net_output
 
     def denoise(self, xt: torch.Tensor, sigma: torch.Tensor, condition: dict[str, torch.Tensor]):
@@ -232,8 +228,6 @@ def denoise(self, xt: torch.Tensor, sigma: torch.Tensor, condition: dict[str, to
 
     def compute_loss_with_epsilon_and_sigma(
         self,
-        data_batch: dict[str, torch.Tensor],
-        x0_from_data_batch: torch.Tensor,
         x0: torch.Tensor,
         condition: dict[str, torch.Tensor],
         epsilon: torch.Tensor,
@@ -294,14 +288,14 @@ def get_per_sigma_loss_weights(self, sigma: torch.Tensor):
 
     def get_condition_uncondition(self, data_batch: Dict):
         """Returns conditioning and unconditioning for classifier-free guidance."""
-        _, _, condition = self.get_data_and_condition(data_batch, dropout_rate=0.0)
+        _, condition = self.get_data_and_condition(data_batch, dropout_rate=0.0)
 
         if "neg_context_embeddings" in data_batch:
             data_batch["context_embeddings"] = data_batch["neg_context_embeddings"]
             data_batch["context_mask"] = data_batch["context_mask"]
-            _, _, uncondition = self.get_data_and_condition(data_batch, dropout_rate=1.0)
+            _, uncondition = self.get_data_and_condition(data_batch, dropout_rate=1.0)
         else:
-            _, _, uncondition = self.get_data_and_condition(data_batch, dropout_rate=1.0)
+            _, uncondition = self.get_data_and_condition(data_batch, dropout_rate=1.0)
 
         return condition, uncondition
 
@@ -419,13 +413,14 @@ def get_data_and_condition(self, data_batch: dict[str, Tensor], dropout_rate=0.2
             Raw data, latent data, and conditioning information.
         """
         # Latent state
-        raw_state = data_batch["video"] * self.sigma_data
-        # assume data is already encoded
-        latent_state = raw_state
-
-        # Condition
-        data_batch["crossattn_emb"] = self.random_dropout_input(
+        latent_state = data_batch["video"] * self.sigma_data
+        condition = {}  # Create a new dictionary for condition
+        # Copy all keys from data_batch except 'video'
+        for key, value in data_batch.items():
+            if key not in ["video", "context_embeddings"]:
+                condition[key] = value
+
+        condition["crossattn_emb"] = self.random_dropout_input(
             data_batch["context_embeddings"], dropout_rate=dropout_rate
         )
-
-        return raw_state, latent_state, data_batch
+        return latent_state, condition

dfm/src/megatron/recipes/dit/dit.py

@@ -32,7 +32,7 @@
 
 from dfm.src.megatron.data.common.diffusion_energon_datamodule import DiffusionDataModuleConfig
 from dfm.src.megatron.data.dit.dit_mock_datamodule import DiTMockDataModuleConfig
-from dfm.src.megatron.model.dit.dit_model_provider import DiTModelProvider
+from dfm.src.megatron.model.dit.dit_model_provider import DiTModelProvider, DiTXLModelProvider
 
 
 def model_config(
@@ -57,7 +57,7 @@ def model_config(
     Returns:
         DiTModelProvider: Configuration for the DiT-S model.
     """
-    return DiTModelProvider(
+    return DiTXLModelProvider(
         tensor_model_parallel_size=tensor_parallelism,
         pipeline_model_parallel_size=pipeline_parallelism,
         pipeline_dtype=pipeline_parallelism_dtype,