diff --git a/nemo_automodel/components/moe/megatron/token_dispatcher.py b/nemo_automodel/components/moe/megatron/token_dispatcher.py index f81e3791d8..e509d9fd41 100644 --- a/nemo_automodel/components/moe/megatron/token_dispatcher.py +++ b/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: diff --git a/tests/functional_tests/moe/run_hybridep_unequal_tokens.py b/tests/functional_tests/moe/run_hybridep_unequal_tokens.py new file mode 100644 index 0000000000..8e943ef3a7 --- /dev/null +++ b/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() diff --git a/tests/unit_tests/moe/test_token_dispatcher.py b/tests/unit_tests/moe/test_token_dispatcher.py index 6409bbe8c1..3b3caf7f49 100644 --- a/tests/unit_tests/moe/test_token_dispatcher.py +++ b/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)