moe_hybridep_pad_variable_tokens

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

Author: zhongbozhu

megatron/core/transformer/moe/token_dispatcher.py

@@ -1061,7 +1061,11 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor):
         self._original_num_tokens = num_tokens
 
         padded_num_tokens = num_tokens
-        if self.config.sequence_packing_scheduler is not None:
+        equalize_thd_token_counts = (
+            self.config.sequence_packing_scheduler is not None
+            or self.config.moe_hybridep_pad_variable_tokens
+        )
+        if equalize_thd_token_counts:
             # Use the actual tp_ep max so all ranks in the MoE communication
             # group pass the same token count to HybridEP.
             max_num_tokens_across_ep = torch.tensor(
@@ -1076,7 +1080,7 @@ def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor):
 
         routing_map = routing_map.reshape(num_tokens, self.num_experts)
         probs = probs.reshape(num_tokens, self.num_experts)
-        if self.config.sequence_packing_scheduler is not None and padded_num_tokens > num_tokens:
+        if equalize_thd_token_counts and padded_num_tokens > num_tokens:
             pad_rows = padded_num_tokens - num_tokens
             routing_map = torch.cat(
                 [routing_map, routing_map.new_zeros((pad_rows, self.num_experts))], dim=0

megatron/core/transformer/transformer_config.py

@@ -866,6 +866,13 @@ class TransformerConfig(ModelParallelConfig):
     moe_permute_fusion_into_hybridep: bool = False
     """Fuse token rearrangement ops during token dispatching for HybridEP."""
 
+    moe_hybridep_pad_variable_tokens: bool = False
+    """Pad uneven local token counts to the HybridEP group maximum before dispatch.
+
+    This is needed when the frontend supplies locally packed THD inputs whose token counts
+    can differ across ranks, without using Megatron Core's sequence_packing_scheduler.
+    """
+
     moe_per_layer_logging: bool = False
     """Enable per-layer logging for MoE, currently supports auxiliary loss and z loss."""