[Dev] add support for deepep/hybridep dispatcher under thd format training

megatron/core/transformer/moe/fused_a2a.py

@@ -52,9 +52,12 @@ def get_buffer(group: torch.distributed.ProcessGroup, hidden_bytes: int):
         num_nvl_bytes = max(
             config.get_nvl_buffer_size_hint(hidden_bytes, group.size()), num_nvl_bytes
         )
-        num_rdma_bytes = max(
-            config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), num_rdma_bytes
-        )
+        # Local-only EP groups do not need an RDMA buffer, and DeepEP builds
+        # without internode support may not expose RDMA size hints.
+        if group.size() > torch.cuda.device_count():
+            num_rdma_bytes = max(
+                config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), num_rdma_bytes
+            )
 
     # Allocate buffer if not existed or not enough buffer
     # NOTES: the adaptive routing configuration of the network **must be off**
@@ -276,6 +279,9 @@ def set_deepep_num_sms(num_sms):
 
 _hybrid_ep_buffer = None
 
+# HybridEP dispatch/combine kernels use 64-token chunks for their public APIs.
+HYBRIDEP_TOKEN_ALIGNMENT = 64
+
 
 def init_hybrid_ep_buffer(
     group: torch.distributed.ProcessGroup,

megatron/core/transformer/moe/token_dispatcher.py

@@ -18,6 +18,7 @@
 )
 from megatron.core.transformer.enums import CudaGraphModule
 from megatron.core.transformer.moe.fused_a2a import (
+    HYBRIDEP_TOKEN_ALIGNMENT,
     fused_combine,
     fused_dispatch,
     hybrid_ep_combine,
@@ -1049,24 +1050,54 @@ def __init__(
 
         self.moe_expert_rank_capacity_factor = self.config.moe_expert_rank_capacity_factor
         self.over_budget = torch.zeros(1, dtype=torch.bool, device='cuda')
+        # THD sequence packing can produce different token counts per rank.
+        # HybridEP dispatch expects equal per-rank input sizes, so metadata and
+        # hidden states are padded to the group-wide max and trimmed in combine.
+        self._original_num_tokens: Optional[int] = None
+        self._padded_num_tokens: Optional[int] = None
 
     def setup_metadata(self, routing_map: torch.Tensor, probs: torch.Tensor):
         num_tokens = routing_map.shape[0]
-        self.routing_map = routing_map.reshape(num_tokens, self.num_experts)
-        self.token_probs = probs.reshape(num_tokens, self.num_experts)
+        self._original_num_tokens = num_tokens
+
+        padded_num_tokens = num_tokens
+        if self.config.sequence_packing_scheduler is not None:
+            # 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(
+                [num_tokens], device=routing_map.device, dtype=torch.long
+            )
+            torch.distributed.all_reduce(
+                max_num_tokens_across_ep, op=torch.distributed.ReduceOp.MAX, group=self.group
+            )
+            padded_num_tokens = int(max_num_tokens_across_ep.item())
+            padded_num_tokens += -padded_num_tokens % HYBRIDEP_TOKEN_ALIGNMENT
+        self._padded_num_tokens = padded_num_tokens
+
+        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:
+            pad_rows = padded_num_tokens - num_tokens
+            routing_map = torch.cat(
+                [routing_map, routing_map.new_zeros((pad_rows, self.num_experts))], dim=0
+            )
+            probs = torch.cat([probs, probs.new_zeros((pad_rows, self.num_experts))], dim=0)
+
+        self.routing_map = routing_map
+        self.token_probs = probs
 
         if self.moe_expert_rank_capacity_factor is not None:
             pad_multiple = get_align_size_for_quantization(self.config)
             budget = int(
-                routing_map.shape[0]
+                padded_num_tokens
                 * self.config.moe_router_topk
                 * self.moe_expert_rank_capacity_factor
             )
             budget += -budget % pad_multiple
             self.num_permuted_tokens = budget
         # Compute the capacity for each expert at the drop_and_pad mode
         if self.drop_and_pad:
-            num_out_tokens = num_tokens * self.config.moe_router_topk
+            num_out_tokens = padded_num_tokens * self.config.moe_router_topk
             # Drop and pad the input to capacity.
             self.capacity = get_capacity(
                 num_tokens=num_out_tokens,
@@ -1095,6 +1126,11 @@ def dispatch(
             self.token_probs = self.token_probs.float()  # downcast or upcast
         if self.config.fp8 or self.config.fp4:
             self.pad_multiple = get_align_size_for_quantization(self.config)
+        if self._padded_num_tokens is not None and hidden_states.shape[0] < self._padded_num_tokens:
+            pad_rows = self._padded_num_tokens - hidden_states.shape[0]
+            hidden_states = torch.cat(
+                [hidden_states, hidden_states.new_zeros((pad_rows, hidden_states.shape[-1]))], dim=0
+            )
         dispatched_hidden, self.dispatched_probs, _, tokens_per_expert, self.handle = (
             hybrid_ep_dispatch(
                 x=hidden_states,
@@ -1137,12 +1173,20 @@ def combine(
             pad_multiple=self.pad_multiple,
             fused=self.config.moe_permute_fusion_into_hybridep,
         )
+        if (
+            self._padded_num_tokens is not None
+            and self._original_num_tokens is not None
+            and hidden_states.shape[0] > self._original_num_tokens
+        ):
+            hidden_states = hidden_states[: self._original_num_tokens]
         # Release the used handle/num_permuted_tokens which could change in each iteration.
         # For drop_and_pad mode, we don't need to reset the num_permuted_tokens and
         # num_dispatched_tokens, because their values never change.
         self.handle = None
         if not self.drop_and_pad:
             self.num_permuted_tokens = None
+        self._original_num_tokens = None
+        self._padded_num_tokens = None
         return hidden_states
 
     def get_permuted_hidden_states_by_experts(self, hidden_states: torch.Tensor) -> torch.Tensor:

megatron/core/transformer/transformer_config.py

@@ -2851,10 +2851,9 @@ def _scope_to_str(s):
             # Needed for passing variable sequences between pp stages.
             self.variable_seq_lengths = True
 
-            # TODO(tailaim): add support for other dispatcher types
-            assert self.moe_token_dispatcher_type == "alltoall", (
-                f"sequence_packing only supports moe_token_dispatcher_type='alltoall', "
-                f"got '{self.moe_token_dispatcher_type}'"
+            assert self.moe_token_dispatcher_type in ("alltoall", "flex"), (
+                f"sequence_packing only supports moe_token_dispatcher_type in "
+                f"('alltoall', 'flex'), got '{self.moe_token_dispatcher_type}'"
             )
 
             supported_schedulers = ['dp_balanced', 'default_dynamic_cp']

tests/unit_tests/transformer/moe/test_token_dispatcher.py

@@ -2,16 +2,24 @@
 
 import copy
 import dataclasses
+import math
 from types import SimpleNamespace
 
 import pytest
 import torch
 
 from megatron.core import config, parallel_state
-from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_submodules
+from megatron.core.extensions.transformer_engine import get_thd_partitioned_indices
+from megatron.core.models.gpt.gpt_layer_specs import (
+    get_gpt_layer_local_submodules,
+    get_gpt_layer_with_transformer_engine_spec,
+)
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.transformer.moe.fused_a2a import reset_hybrid_ep_buffer
 from megatron.core.transformer.moe.moe_layer import MoELayer
 from megatron.core.transformer.moe.moe_utils import get_capacity
 from megatron.core.transformer.moe.token_dispatcher import MoETokenDispatcher
+from megatron.core.transformer.transformer_block import TransformerBlock
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.typed_torch import apply_module
 from megatron.core.utils import is_te_min_version
@@ -461,6 +469,171 @@ def is_hybrid_ep_available():
     return HAVE_HYBRIDEP
 
 
+def _round_up(value, divisor):
+    return value if divisor <= 1 else (value + divisor - 1) // divisor * divisor
+
+
+def _get_thd_padded_seqlens(seqlens, cp_size, tp_size):
+    # This follows the runtime packed-sequence path used by the Moonlight script:
+    # per-sequence lengths must be CP partitionable, and the packed token count
+    # must be even for TP/SP slicing.
+    cp_divisor = 2 * cp_size if cp_size > 1 else 1
+    padded_seqlens = [_round_up(seqlen, cp_divisor) for seqlen in seqlens]
+    total_seqlen = sum(padded_seqlens)
+    total_alignment = math.lcm(cp_divisor, tp_size)
+    padded_seqlens[-1] += _round_up(total_seqlen, total_alignment) - total_seqlen
+    return padded_seqlens
+
+
+def _to_cu_seqlens(seqlens):
+    cu_seqlens = torch.empty(len(seqlens) + 1, dtype=torch.int32, device="cuda")
+    cu_seqlens[0] = 0
+    cu_seqlens[1:] = torch.cumsum(torch.tensor(seqlens, dtype=torch.int32, device="cuda"), dim=0)
+    return cu_seqlens
+
+
+def _make_thd_packed_seq_params(seqlens, cp_size, tp_size):
+    padded_seqlens = _get_thd_padded_seqlens(seqlens, cp_size, tp_size)
+    cu_seqlens_padded = _to_cu_seqlens(padded_seqlens)
+    max_seqlen = max(padded_seqlens)
+    # Match get_batch_on_this_rank_for_sequence_packing(): TE consumes padded
+    # cumulative lengths as both cu_seqlens and cu_seqlens_padded for THD.
+    return PackedSeqParams(
+        qkv_format="thd",
+        cu_seqlens_q=cu_seqlens_padded,
+        cu_seqlens_kv=cu_seqlens_padded,
+        cu_seqlens_q_padded=cu_seqlens_padded,
+        cu_seqlens_kv_padded=cu_seqlens_padded,
+        max_seqlen_q=max_seqlen,
+        max_seqlen_kv=max_seqlen,
+    )
+
+
+def _make_sharded_thd_hidden_states(seqlens, hidden_size, cp_size, tp_size, dtype):
+    padded_seqlens = _get_thd_padded_seqlens(seqlens, cp_size, tp_size)
+    padded_sequences = []
+    for seqlen, padded_seqlen in zip(seqlens, padded_seqlens):
+        sequence = torch.randn(seqlen, hidden_size, device="cuda", dtype=dtype)
+        if padded_seqlen > seqlen:
+            sequence = torch.cat(
+                [
+                    sequence,
+                    torch.zeros(padded_seqlen - seqlen, hidden_size, device="cuda", dtype=dtype),
+                ],
+                dim=0,
+            )
+        padded_sequences.append(sequence)
+
+    hidden_states = torch.cat(padded_sequences, dim=0)
+    if cp_size > 1:
+        cu_seqlens_padded = _to_cu_seqlens(padded_seqlens)
+        cp_rank = parallel_state.get_context_parallel_rank()
+        index = get_thd_partitioned_indices(
+            cu_seqlens_padded, hidden_states.shape[0], cp_size, cp_rank
+        )
+        hidden_states = hidden_states.index_select(0, index)
+
+    tp_rank = parallel_state.get_tensor_model_parallel_rank()
+    sequence_parallel_length = hidden_states.shape[0] // tp_size
+    hidden_states = hidden_states[
+        tp_rank * sequence_parallel_length : (tp_rank + 1) * sequence_parallel_length
+    ]
+    return hidden_states.unsqueeze(1).contiguous().requires_grad_(True)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+@pytest.mark.skipif(
+    Utils.world_size % 8 != 0, reason="requires world size divisible by 8 for pp2/cp2/tp2/ep2/etp2"
+)
+@pytest.mark.internal
+@pytest.mark.parametrize("dispatcher", ["alltoall", "deepep", "hybridep"])
+def test_sequence_packing_thd_e2e_proxy_model(dispatcher):
+    """Run packed THD attention + MoE forward/backward with major parallelisms enabled."""
+    if not is_te_min_version("2.9.0"):
+        pytest.skip("SFT sequence packing requires Transformer Engine >= 2.9.0")
+    if dispatcher == "deepep" and not is_deep_ep_available():
+        pytest.skip("Deep EP is not available")
+    if dispatcher == "hybridep" and not is_hybrid_ep_available():
+        pytest.skip("Hybrid EP is not available")
+
+    tp_size, pp_size, cp_size, ep_size, etp_size = 2, 2, 2, 2, 2
+    Utils.initialize_model_parallel(
+        tensor_model_parallel_size=tp_size,
+        pipeline_model_parallel_size=pp_size,
+        context_parallel_size=cp_size,
+        expert_model_parallel_size=ep_size,
+        expert_tensor_parallel_size=etp_size,
+    )
+    _set_random_seed(seed_=123, data_parallel_random_init=False)
+
+    try:
+        spec = get_gpt_layer_with_transformer_engine_spec(num_experts=4, moe_grouped_gemm=False)
+        transformer_config = TransformerConfig(
+            num_layers=4,
+            hidden_size=1024,
+            ffn_hidden_size=2048,
+            moe_ffn_hidden_size=2048,
+            num_attention_heads=8,
+            tensor_model_parallel_size=tp_size,
+            pipeline_model_parallel_size=pp_size,
+            context_parallel_size=cp_size,
+            expert_model_parallel_size=ep_size,
+            expert_tensor_parallel_size=etp_size,
+            sequence_parallel=True,
+            sequence_packing_scheduler="dp_balanced",
+            max_seqlen_per_dp_cp_rank=1024,
+            cp_comm_type="p2p",
+            num_moe_experts=4,
+            moe_router_topk=2,
+            moe_router_load_balancing_type="aux_loss",
+            moe_token_dispatcher_type=(
+                "flex" if dispatcher in ("deepep", "hybridep") else dispatcher
+            ),
+            moe_flex_dispatcher_backend=(
+                dispatcher if dispatcher in ("deepep", "hybridep") else "deepep"
+            ),
+            moe_grouped_gemm=False,
+            moe_router_dtype="fp32",
+            params_dtype=torch.bfloat16,
+            pipeline_dtype=torch.bfloat16,
+            autocast_dtype=torch.bfloat16,
+            bf16=True,
+            add_bias_linear=False,
+            attention_dropout=0.0,
+            hidden_dropout=0.0,
+            use_cpu_initialization=True,
+        )
+        transformer_block = TransformerBlock(transformer_config, spec).cuda().to(torch.bfloat16)
+
+        torch.manual_seed(1000 + torch.distributed.get_rank())
+        seqlens = [257, 509, 1021]
+        hidden_states = _make_sharded_thd_hidden_states(
+            seqlens, transformer_config.hidden_size, cp_size, tp_size, torch.bfloat16
+        )
+        packed_seq_params = _make_thd_packed_seq_params(seqlens, cp_size, tp_size)
+
+        output = transformer_block(
+            hidden_states=hidden_states, attention_mask=None, packed_seq_params=packed_seq_params
+        )
+        assert output.shape == hidden_states.shape
+        assert torch.isfinite(output).all()
+
+        loss = output.float().square().mean()
+        loss.backward()
+
+        assert hidden_states.grad is not None
+        assert hidden_states.grad.shape == hidden_states.shape
+        assert torch.isfinite(hidden_states.grad).all()
+        assert any(
+            param.grad is not None and torch.isfinite(param.grad).all()
+            for param in transformer_block.parameters()
+            if param.requires_grad
+        )
+    finally:
+        reset_hybrid_ep_buffer()
+        Utils.destroy_model_parallel()
+
+
 @pytest.mark.skipif(
     not is_deep_ep_available() and not is_hybrid_ep_available(),
     reason="Deep EP and Hybrid EP are not available",
@@ -470,6 +643,7 @@ def setup_method(self, method):
         pass
 
     def teardown_method(self, method):
+        reset_hybrid_ep_buffer()
         Utils.destroy_model_parallel()
 
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")