clean up, remove all CP for sbhd, CP now is only for thd

Author: Huy Vu2

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

@@ -27,7 +27,6 @@
 from dfm.src.megatron.model.wan.utils.utils import grid_sizes_calculation, patchify
 from megatron.core import parallel_state
 from torch.nn import functional as F
-from dfm.src.megatron.model.wan.utils.utils import split_inputs_cp, cat_outputs_cp
 
 import math
 from typing import Tuple, Union
@@ -99,10 +98,8 @@ def __init__(
         wan_checkpoint_dir = self._select_checkpoint_dir(checkpoint_dir, checkpoint_step)
         self.model = self.setup_model_from_checkpoint(wan_checkpoint_dir)
 
-        # DEBUGGING thd
-        # set qkv_format to to "thd" for context parallelism
-        # self.model.config.qkv_format = "sbhd"
-        self.model.config.qkv_format = "thd"
+        # if we use context parallelism, we need to set qkv_format to "thd" for context parallelism
+        self.model.config.qkv_format = "thd" # "sbhd"
 
         # set self.sp_size=1 for later use, just to respect the original Wan inference code
         self.sp_size = 1
@@ -487,31 +484,13 @@ def noop_no_sync():
                 timestep = [t] * batch_size
                 timestep = torch.stack(timestep)
 
-                # DEBUGGING thd
-                # # run context parallelism slitting
-                # if parallel_state.get_context_parallel_world_size() > 1:
-                #     latent_model_input = split_inputs_cp(latent_model_input, 0)
-                #     arg_c['context'] = split_inputs_cp(arg_c['context'], 0)
-                #     arg_null['context'] = split_inputs_cp(arg_null['context'], 0)
-
                 self.model.to(self.device)
                 noise_pred_cond = self.forward_pp_step(
                     latent_model_input, grid_sizes=grid_sizes, max_video_seq_len=max_video_seq_len, timestep=timestep, arg_c=arg_c)
 
                 noise_pred_uncond = self.forward_pp_step(
                     latent_model_input, grid_sizes=grid_sizes, max_video_seq_len=max_video_seq_len, timestep=timestep, arg_c=arg_null)
 
-                # DEBUGGING thd
-                # # run context parallelism gathering
-                # if parallel_state.get_context_parallel_world_size() > 1:
-                #     arg_c['context'] = cat_outputs_cp(arg_c['context'], 0) # we need to cat the context back together for the next timestep
-                #     arg_null['context'] = cat_outputs_cp(arg_null['context'], 0) # we need to cat the context back together for the next timestep
-                #     # TODO: does this step slow down speed???
-                #     noise_pred_cond = noise_pred_cond.contiguous()
-                #     noise_pred_uncond = noise_pred_uncond.contiguous()
-                #     noise_pred_cond = cat_outputs_cp(noise_pred_cond, 0)
-                #     noise_pred_uncond = cat_outputs_cp(noise_pred_uncond, 0)
-
                 # run unpatchify
                 unpatchified_noise_pred_cond = noise_pred_cond
                 unpatchified_noise_pred_cond = unpatchified_noise_pred_cond.transpose(0, 1) # bring sbhd -> bshd

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

@@ -20,7 +20,7 @@
 from torch import Tensor
 from diffusers import WanPipeline
 from dfm.src.megatron.model.wan.flow_matching.time_shift_utils import compute_density_for_timestep_sampling
-from dfm.src.megatron.model.wan.utils.utils import patchify, split_inputs_cp, thd_split_inputs_cp
+from dfm.src.megatron.model.wan.utils.utils import patchify, thd_split_inputs_cp
 
 class FlowPipeline:
 
@@ -150,12 +150,6 @@ def training_step(
         # ========================================================================
 
         if parallel_state.get_context_parallel_world_size() > 1:
-            # DEBUGGING thd
-            # video_latents = split_inputs_cp(video_latents, 0)
-            # noisy_latents = split_inputs_cp(noisy_latents, 0)
-            # noise = split_inputs_cp(noise, 0)
-            # context_embeddings = split_inputs_cp(context_embeddings, 0)
-            # split_loss_mask = split_inputs_cp(loss_mask, 0)
             video_latents = thd_split_inputs_cp(video_latents, packed_seq_params['self_attention'].cu_seqlens_q, parallel_state.get_context_parallel_group())
             noisy_latents = thd_split_inputs_cp(noisy_latents, packed_seq_params['self_attention'].cu_seqlens_q, parallel_state.get_context_parallel_group())
             noise = thd_split_inputs_cp(noise, packed_seq_params['self_attention'].cu_seqlens_q, parallel_state.get_context_parallel_group())

dfm/src/megatron/model/wan/rope_utils.py

@@ -1,6 +1,5 @@
 import torch
 from torch.cuda import amp
-from dfm.src.megatron.model.wan.utils.utils import split_inputs_cp
 from megatron.core import parallel_state
 
 class Wan3DRopeEmbeddings(torch.nn.Module):
@@ -57,10 +56,9 @@ def forward(self, n_head, dim_head, max_seq_len, grid_sizes, device):
         # We concatenate them along dim=1 to get (max_seq_len, batch_size, 1, dim_head)
         freqs_real = torch.cat(freqs_real, dim=1)
 
-        # DEBUGGING thd
-        # # TODO: if run context/sequence related parallel, then we need to scatter 
-        # # the freqs_real to the context parallel region, using specific cp_rank split method
-        # if parallel_state.get_context_parallel_world_size() > 1:
-        #     freqs_real = split_inputs_cp(freqs_real, 0)
+        # Note:
+        # when running context_parallel, which must use "thd" for qkv_format,
+        # we don't need to scatter the freqs to the context parallel region,
+        # because mcore rope_utils will automatically retrieve the correct freqs for each context parallel region
 
         return freqs_real

dfm/src/megatron/model/wan/utils/utils.py

@@ -129,7 +129,7 @@ def cat_outputs_cp(x: torch.Tensor, seq_dim: int) -> torch.Tensor:
     else:
         return x
 
-# DEBUGGING thd
+
 def thd_split_inputs_cp(x: torch.Tensor,
                            cu_seqlens_q_padded: torch.Tensor,
                            cp_group: dist.ProcessGroup) -> torch.Tensor: