fix(moe): equalize and align per-rank token counts for HybridEP dispatch

nemo_automodel/components/moe/megatron/token_dispatcher.py

@@ -352,6 +352,11 @@ def forward(self, token_indices: torch.Tensor, token_probs: torch.Tensor) -> tup
         return routing_map, multihot_probs
 
 
+# DeepEP's hybrid-ep metadata allgather asserts bytes_per_rank % 16 == 0 on a
+# 4-byte-per-token array, so per-rank token counts must be multiples of 4.
+_HYBRIDEP_TOKEN_ALIGNMENT = 4
+
+
 class _HybridEPManager(_DispatchManager):
     """
     A manager class to handle fused all-to-all communication processes for MoE models using
@@ -390,6 +395,9 @@ def __init__(
         # Handle used for combine operation
         self.handle = None
         self.pad_multiple = None
+        # Set by dispatch() when this rank padded its tokens up to the EP-group
+        # maximum; combine() slices the padding back off.
+        self.num_unpadded_tokens: int | None = None
 
         if hybrid_ep_dispatch is None:
             raise ImportError(
@@ -443,6 +451,27 @@ def dispatch(
         self.num_permuted_tokens = None
         if self.token_probs.dtype != torch.float32:
             self.token_probs = self.token_probs.float()
+
+        # HybridEP's fused all-to-all exchanges fixed-extent buffers, so every rank
+        # in the EP group must dispatch the same number of tokens, and the kernel's
+        # metadata allgather additionally requires that count to be 16-byte aligned
+        # (4 tokens). Unequal or unaligned counts (variable-length or packed dynamic
+        # batches) abort or deadlock the collective. Pad this rank's tokens up to
+        # the aligned group-wide maximum: padded rows route to no expert (all-False
+        # routing map), and combine() slices them back off.
+        self.num_unpadded_tokens = None
+        if torch.distributed.is_initialized() and torch.distributed.get_world_size(self.group) > 1:
+            num_tokens = hidden_states.shape[0]
+            group_max = torch.tensor(num_tokens, device=hidden_states.device)
+            torch.distributed.all_reduce(group_max, op=torch.distributed.ReduceOp.MAX, group=self.group)
+            target_tokens = -(-int(group_max) // _HYBRIDEP_TOKEN_ALIGNMENT) * _HYBRIDEP_TOKEN_ALIGNMENT
+            pad_tokens = target_tokens - num_tokens
+            if pad_tokens > 0:
+                self.num_unpadded_tokens = num_tokens
+                hidden_states = nn.functional.pad(hidden_states, (0, 0, 0, pad_tokens))
+                self.routing_map = nn.functional.pad(self.routing_map, (0, 0, 0, pad_tokens))
+                self.token_probs = nn.functional.pad(self.token_probs, (0, 0, 0, pad_tokens))
+
         dispatched_hidden, self.dispatched_probs, _, tokens_per_expert, self.handle = hybrid_ep_dispatch(
             x=hidden_states,
             routing_map=self.routing_map,
@@ -474,6 +503,9 @@ def combine(
         )
         self.handle = None
         self.num_permuted_tokens = None
+        if self.num_unpadded_tokens is not None:
+            hidden_states = hidden_states[: self.num_unpadded_tokens]
+            self.num_unpadded_tokens = None
         return hidden_states
 
     def get_dispatched_metadata(self) -> torch.Tensor:

tests/functional_tests/moe/run_hybridep_unequal_tokens.py

@@ -0,0 +1,102 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""HybridEP dispatch/combine parity with unequal per-rank token counts.
+
+Every rank in a HybridEP group must dispatch the same token extent; the
+dispatcher now pads shorter ranks up to the group maximum. With identity
+"experts", each token's combined output depends only on its own routing, so
+running the same data once with equal counts (no padding path) and once with
+rank 1 truncated (padding path) must produce identical outputs for the common
+tokens.
+
+Run:
+    torchrun --standalone --nproc_per_node=2 \
+        tests/functional_tests/moe/run_hybridep_unequal_tokens.py
+"""
+
+import os
+
+import torch
+import torch.distributed as dist
+
+from nemo_automodel.components.moe.megatron.token_dispatcher import (
+    MoEFlexTokenDispatcher,
+    TokenDispatcherConfig,
+)
+
+HIDDEN = 256
+NUM_EXPERTS = 4
+TOPK = 2
+FULL_TOKENS = int(os.environ.get("FULL_TOKENS", "5"))
+SHORT_TOKENS = int(os.environ.get("SHORT_TOKENS", "3"))  # rank 1 in the unequal run
+
+
+def run_dispatch_combine(dispatcher: MoEFlexTokenDispatcher, hidden, indices, probs):
+    hidden = hidden.detach().requires_grad_(True)
+    out, _tokens_per_expert, _permuted_probs = dispatcher.token_permutation2(
+        hidden_states=hidden,
+        num_local_tokens=hidden.shape[0],
+        token_probs=probs,
+        token_indices=indices,
+    )
+    # Identity experts: combine returns each token's prob-weighted sum of its
+    # own dispatched copies, independent of every other token.
+    combined = dispatcher.token_unpermutation(out)
+    combined.float().square().sum().backward()
+    return combined.detach(), hidden.grad
+
+
+def main():
+    rank = int(os.environ["RANK"])
+    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+    dist.init_process_group("nccl")
+    torch.manual_seed(1234 + rank)
+
+    ep_group = dist.new_group(ranks=list(range(dist.get_world_size())))
+    config = TokenDispatcherConfig(
+        moe_flex_dispatcher_backend="hybridep",
+        num_moe_experts=NUM_EXPERTS,
+        moe_router_topk=TOPK,
+        moe_share_token_dispatcher=False,
+    )
+    num_local = NUM_EXPERTS // dist.get_world_size()
+    dispatcher = MoEFlexTokenDispatcher(
+        num_local_experts=num_local,
+        local_expert_indices=list(range(rank * num_local, (rank + 1) * num_local)),
+        config=config,
+        ep_group=ep_group,
+    )
+
+    hidden = torch.randn(FULL_TOKENS, HIDDEN, dtype=torch.bfloat16, device="cuda")
+    indices = torch.stack([torch.randperm(NUM_EXPERTS, device="cuda")[:TOPK] for _ in range(FULL_TOKENS)])
+    probs = torch.rand(FULL_TOKENS, TOPK, dtype=torch.float32, device="cuda")
+    probs = probs / probs.sum(dim=-1, keepdim=True)
+
+    # Reference: every rank dispatches FULL_TOKENS (equal counts).
+    reference, reference_grad = run_dispatch_combine(dispatcher, hidden.clone(), indices, probs)
+
+    # Unequal: rank 1 truncates to SHORT_TOKENS, forcing the padding path.
+    keep = FULL_TOKENS if rank == 0 else SHORT_TOKENS
+    unequal, unequal_grad = run_dispatch_combine(dispatcher, hidden[:keep].clone(), indices[:keep], probs[:keep])
+
+    assert unequal.shape == (keep, HIDDEN), f"rank {rank}: got {tuple(unequal.shape)}"
+    torch.testing.assert_close(unequal, reference[:keep], rtol=0, atol=0)
+    torch.testing.assert_close(unequal_grad, reference_grad[:keep], rtol=0, atol=0)
+    print(f"[rank {rank}] OK: unequal-count forward AND backward bitwise-match the equal-count run")
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

tests/unit_tests/moe/test_token_dispatcher.py

@@ -132,3 +132,64 @@ def test_token_unpermutation_applies_async_setting_to_deepep_combine(enabled):
 
     torch.testing.assert_close(actual, hidden_states)
     manager.combine.assert_called_once_with(hidden_states, enabled, enabled)
+
+
+class TestHybridEPTokenCountEqualization:
+    """dispatch() must pad unequal per-rank token counts up to the EP-group max."""
+
+    def _run(self, hybrid_ep_manager, monkeypatch, num_tokens, group_max):
+        import nemo_automodel.components.moe.megatron.token_dispatcher as td
+
+        monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+        monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: 2)
+
+        def fake_all_reduce(tensor, op=None, group=None):
+            tensor.fill_(group_max)
+
+        monkeypatch.setattr(torch.distributed, "all_reduce", fake_all_reduce)
+
+        dispatched = {}
+
+        def fake_dispatch(x, routing_map, probs, **kwargs):
+            dispatched.update(x=x, routing_map=routing_map, probs=probs)
+            return x, probs, None, routing_map.sum(dim=0), "handle"
+
+        monkeypatch.setattr(td, "hybrid_ep_dispatch", fake_dispatch)
+        monkeypatch.setattr(td, "hybrid_ep_combine", lambda x, **kwargs: dispatched["x"])
+
+        hidden = torch.randn(num_tokens, 4)
+        hybrid_ep_manager.routing_map = torch.ones(num_tokens, 8, dtype=torch.bool)
+        hybrid_ep_manager.token_probs = torch.full((num_tokens, 8), 0.125)
+        out = hybrid_ep_manager.dispatch(hidden)
+        combined = hybrid_ep_manager.combine(out)
+        return hidden, dispatched, combined
+
+    def test_shorter_rank_pads_to_aligned_group_max_and_slices_back(self, hybrid_ep_manager, monkeypatch):
+        # group max 5 rounds up to the 4-token kernel alignment -> 8.
+        hidden, dispatched, combined = self._run(hybrid_ep_manager, monkeypatch, num_tokens=3, group_max=5)
+
+        assert dispatched["x"].shape[0] == 8
+        assert dispatched["routing_map"].shape[0] == 8
+        assert dispatched["probs"].shape[0] == 8
+        # Padded rows carry zero hidden state and route to no expert.
+        assert torch.equal(dispatched["x"][3:], torch.zeros(5, 4))
+        assert not dispatched["routing_map"][3:].any()
+        assert not dispatched["probs"][3:].any()
+        # combine() returns only this rank's real tokens.
+        assert combined.shape[0] == 3
+        assert torch.equal(combined, hidden)
+        assert hybrid_ep_manager.num_unpadded_tokens is None
+
+    def test_equal_unaligned_counts_pad_to_alignment(self, hybrid_ep_manager, monkeypatch):
+        hidden, dispatched, combined = self._run(hybrid_ep_manager, monkeypatch, num_tokens=6, group_max=6)
+
+        assert dispatched["x"].shape[0] == 8
+        assert combined.shape[0] == 6
+        assert torch.equal(combined, hidden)
+
+    def test_equal_aligned_counts_do_not_pad(self, hybrid_ep_manager, monkeypatch):
+        hidden, dispatched, combined = self._run(hybrid_ep_manager, monkeypatch, num_tokens=4, group_max=4)
+
+        assert dispatched["x"].shape[0] == 4
+        assert combined.shape[0] == 4
+        assert torch.equal(combined, hidden)