Fix sequence padding for DiT. Add support for DiT Context Parallel with THD.

Signed-off-by: sajadn <snorouzi@nvidia.com>

Author: sajadn

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

@@ -100,6 +100,7 @@ def cat(attr):
             __subflavors__=samples[0].__subflavors__,
             video=cat("video"),
             context_embeddings=cat("context_embeddings"),
+            context_mask=cat("context_mask"),
             loss_mask=cat("loss_mask"),
             seq_len_q=cat("seq_len_q"),
             seq_len_q_padded=cat("seq_len_q_padded"),

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

@@ -130,13 +130,14 @@ def encode_sample(self, sample: dict) -> DiffusionSample:
             "T H W d -> (T H W) d",
         )
 
-        if self.packing_buffer_size is None:
-            pos_ids = F.pad(pos_ids, (0, 0, 0, self.seq_length - seq_len))
-            loss_mask = torch.zeros(self.seq_length, dtype=torch.bfloat16)
-            loss_mask[:seq_len] = 1
-            video_latent = F.pad(video_latent, (0, 0, 0, self.seq_length - seq_len))
-        else:
-            loss_mask = torch.ones(seq_len, dtype=torch.bfloat16)
+        loss_mask = torch.ones(seq_len, dtype=torch.bfloat16)
+        sharding_factor = 64
+        seq_len_q_padded = ((seq_len + sharding_factor - 1) // sharding_factor) * sharding_factor
+
+        if seq_len < seq_len_q_padded:
+            video_latent = F.pad(video_latent, (0, 0, 0, seq_len_q_padded - seq_len))
+            loss_mask = F.pad(loss_mask, (0, seq_len_q_padded - seq_len))
+            pos_ids = F.pad(pos_ids, (0, 0, 0, seq_len_q_padded - seq_len))
 
         return DiffusionSample(
             __key__=sample["__key__"],
@@ -148,6 +149,7 @@ def encode_sample(self, sample: dict) -> DiffusionSample:
             context_mask=t5_text_mask,
             loss_mask=loss_mask,
             seq_len_q=torch.tensor([seq_len], dtype=torch.int32),
+            seq_len_q_padded=torch.tensor([seq_len_q_padded], dtype=torch.int32),
             seq_len_kv=torch.tensor([self.text_embedding_padding_size], dtype=torch.int32),
             pos_ids=pos_ids,
             latent_shape=torch.tensor([C, T, H, W], dtype=torch.int32),
@@ -168,6 +170,7 @@ def batch(self, samples: List[DiffusionSample]) -> dict:
             context_mask=sample.context_mask.unsqueeze_(0) if sample.context_mask is not None else None,
             loss_mask=sample.loss_mask.unsqueeze_(0) if sample.loss_mask is not None else None,
             seq_len_q=sample.seq_len_q,
+            seq_len_q_padded=sample.seq_len_q_padded,
             seq_len_kv=sample.seq_len_kv,
             pos_ids=sample.pos_ids.unsqueeze_(0) if sample.pos_ids is not None else None,
             latent_shape=sample.latent_shape,

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

@@ -13,16 +13,18 @@
 # limitations under the License.
 
 import torch
+from megatron.core import parallel_state as ps
 from megatron.core.packed_seq_params import PackedSeqParams
 
 
 def dit_data_step(qkv_format, dataloader_iter):
     # import pdb;pdb.set_trace()
     batch = next(iter(dataloader_iter.iterable))
-    batch = get_batch_on_this_cp_rank(batch)
-    batch = {k: v.to(device="cuda", non_blocking=True) if torch.is_tensor(v) else v for k, v in batch.items()}
     batch["is_preprocessed"] = True  # assume data is preprocessed
-    return encode_seq_length(batch, format=qkv_format)
+    batch = {k: v.to(device="cuda", non_blocking=True) if torch.is_tensor(v) else v for k, v in batch.items()}
+    batch = encode_seq_length(batch, format=qkv_format)
+    batch = get_batch_on_this_cp_rank(batch)
+    return batch
 
 
 def encode_seq_length(batch, format):
@@ -35,19 +37,20 @@ def encode_seq_length(batch, format):
         cu_seqlens_kv = batch["seq_len_kv"].cumsum(dim=0).to(torch.int32)
         cu_seqlens_kv = torch.cat((zero, cu_seqlens_kv))
 
+        cu_seqlens_q_padded = batch["seq_len_q_padded"].cumsum(dim=0).to(torch.int32)
+        cu_seqlens_q_padded = torch.cat((zero, cu_seqlens_q_padded))
+
         batch["packed_seq_params"] = {
             "self_attention": PackedSeqParams(
                 cu_seqlens_q=cu_seqlens_q,
                 cu_seqlens_kv=cu_seqlens_q,
-                cu_seqlens_q_padded=None,
-                cu_seqlens_kv_padded=None,
+                cu_seqlens_q_padded=cu_seqlens_q_padded,
                 qkv_format=format,
             ),
             "cross_attention": PackedSeqParams(
                 cu_seqlens_q=cu_seqlens_q,
                 cu_seqlens_kv=cu_seqlens_kv,
-                cu_seqlens_q_padded=None,
-                cu_seqlens_kv_padded=None,
+                cu_seqlens_q_padded=cu_seqlens_q_padded,
                 qkv_format=format,
             ),
         }
@@ -57,34 +60,26 @@ def encode_seq_length(batch, format):
 
 def get_batch_on_this_cp_rank(data):
     """Split the data for context parallelism."""
-    from megatron.core import mpu
-
-    cp_size = mpu.get_context_parallel_world_size()
-    cp_rank = mpu.get_context_parallel_rank()
-
-    t = 16
+    cp_size = ps.get_context_parallel_world_size()
     if cp_size > 1:
-        # cp split on seq_length, for video_latent, noise_latent and pos_ids
-        assert t % cp_size == 0, "t must divisibly by cp_size"
-        num_valid_tokens_in_ub = None
-        if "loss_mask" in data and data["loss_mask"] is not None:
-            num_valid_tokens_in_ub = data["loss_mask"].sum()
+        import transformer_engine_torch as tex
+
+        cp_rank = ps.get_context_parallel_rank()
+        for key in ["video", "loss_mask", "pos_ids"]:
+            if data[key] is not None:
+                index = tex.thd_get_partitioned_indices(
+                    data["packed_seq_params"]["self_attention"].cu_seqlens_q_padded,
+                    data[key].size(1),
+                    cp_size,
+                    cp_rank,
+                ).to(device=data[key].device, dtype=torch.long)
+                data[key] = data[key].index_select(1, index).contiguous()
 
-        for key, value in data.items():
-            if (value is not None) and (key in ["video", "video_latent", "noise_latent", "pos_ids"]):
-                if len(value.shape) > 5:
-                    value = value.squeeze(0)
-                B, C, T, H, W = value.shape
-                if T % cp_size == 0:
-                    # FIXME packed sequencing
-                    data[key] = value.view(B, C, cp_size, T // cp_size, H, W)[:, :, cp_rank, ...].contiguous()
-                else:
-                    # FIXME packed sequencing
-                    data[key] = value.view(B, C, T, cp_size, H // cp_size, W)[:, :, :, cp_rank, ...].contiguous()
-        loss_mask = data["loss_mask"]
-        data["loss_mask"] = loss_mask.view(loss_mask.shape[0], cp_size, loss_mask.shape[1] // cp_size)[
-            :, cp_rank, ...
-        ].contiguous()
-        data["num_valid_tokens_in_ub"] = num_valid_tokens_in_ub
+        for key in ["context_embeddings", "context_mask"]:
+            if data[key] is not None:
+                index = tex.thd_get_partitioned_indices(
+                    data["packed_seq_params"]["cross_attention"].cu_seqlens_kv, data[key].size(1), cp_size, cp_rank
+                ).to(device=data[key].device, dtype=torch.long)
+                data[key] = data[key].index_select(1, index).contiguous()
 
     return data

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

@@ -144,17 +144,11 @@ def _replace_no_cp_submodules(submodules):
         # Override Cross Attention to disable CP.
         # Disable TP Comm overlap as well. Not disabling will attempt re-use of buffer size same as Q and lead to
         # incorrect tensor shapes.
-        if submodules.cross_attention != IdentityOp:
-            cp_override_config = copy.deepcopy(config)
-            cp_override_config.context_parallel_size = 1
-            cp_override_config.tp_comm_overlap = False
-            self.cross_attention = build_module(
-                submodules.cross_attention,
-                config=cp_override_config,
-                layer_number=layer_number,
-            )
-        else:
-            self.cross_attention = None
+        self.cross_attention = build_module(
+            submodules.cross_attention,
+            config=self.config,
+            layer_number=layer_number,
+        )
 
         self.full_self_attention = build_module(
             submodules.full_self_attention,

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

@@ -18,6 +18,7 @@
 from typing import Iterable
 
 import torch
+import wandb
 from einops import rearrange
 from megatron.bridge.training.losses import masked_next_token_loss
 from megatron.bridge.training.state import GlobalState
@@ -41,7 +42,7 @@ def __init__(self):
         self.train = True
         self.validation_step = 0
 
-    def on_validation_start(self, batch, model, step):
+    def on_validation_start(self, state, batch, model, step):
         C, T, H, W = batch["latent_shape"][0]
         latent = self.diffusion_pipeline.generate_samples_from_batch(
             model,
@@ -81,6 +82,28 @@ def on_validation_start(self, batch, model, step):
             video_save_path=f"{image_folder}/validation_step={step}_rank={rank}.mp4",
         )
 
+        wandb_rank = parallel_state.get_data_parallel_world_size() - 1
+        if torch.distributed.get_rank() == wandb_rank:
+            gather_list = [None for _ in range(parallel_state.get_data_parallel_world_size())]
+        else:
+            gather_list = None
+
+        torch.distributed.gather_object(
+            obj=decoded_video[0],
+            object_gather_list=gather_list,
+            dst=wandb_rank,
+            group=parallel_state.get_data_parallel_group(),
+        )
+        if torch.distributed.get_rank() == wandb_rank:
+            if gather_list is not None:
+                videos = []
+                for video_data in gather_list:
+                    video_data_transposed = video_data.transpose(0, 3, 1, 2)
+                    videos.append(wandb.Video(video_data_transposed, fps=24, format="mp4"))
+
+                if state.wandb_logger is not None:
+                    state.wandb_logger.log({"prediction": videos})
+
     def __call__(
         self, state: GlobalState, data_iterator: Iterable, model: GPTModel, return_schedule_plan: bool = False
     ) -> tuple[torch.Tensor, partial]:
@@ -103,7 +126,7 @@ def __call__(
             self.train = False
             self.valid = True
             self.validation_step += 1
-            self.on_validation_start(batch, model, step=self.validation_step)
+            self.on_validation_start(state, batch, model, step=self.validation_step)
         return self.forward_step(state, batch, model, return_schedule_plan)
 
     def data_process(