diff --git a/nemo_automodel/components/moe/experts.py b/nemo_automodel/components/moe/experts.py index 224bb7edb8..8cef4b5796 100644 --- a/nemo_automodel/components/moe/experts.py +++ b/nemo_automodel/components/moe/experts.py @@ -1098,9 +1098,22 @@ def forward( None if self.config.apply_router_weight_after_down else permuted_probs, ) else: - output1 = torch.matmul(x[0] * 0, gate_and_up_projs[0]) - output1_ = self.expert_activation(output1, activation_probs) - output2 = torch.matmul(output1_, down_projs[0]) + # Keep the zero-token path connected to the tensors produced by + # DeepEP dispatch. Basing the dummy computation on ``x`` bypasses + # FusedDispatch autograd, so this rank never enters the reverse + # dispatch collective while peers with local tokens do. Preserve + # the dispatched empty shape and attach zero-valued dependencies + # for every local expert parameter as grouped GEMM would. + zero_dependency = permuted_local_hidden_states.sum() * 0 + zero_dependency = zero_dependency + permuted_probs.sum() * 0 + zero_dependency = zero_dependency + gate_and_up_projs.reshape(-1)[0] * 0 + zero_dependency = zero_dependency + down_projs.reshape(-1)[0] * 0 + if self.expert_bias: + gate_up_proj_bias = self.gate_up_proj_bias.to_local().to(compute_dtype) + down_proj_bias = self.down_proj_bias.to_local().to(compute_dtype) + zero_dependency = zero_dependency + gate_up_proj_bias.reshape(-1)[0] * 0 + zero_dependency = zero_dependency + down_proj_bias.reshape(-1)[0] * 0 + output2 = (permuted_local_hidden_states * 0 + zero_dependency).to(compute_dtype) if self.config.apply_router_weight_after_down: # HybridEP/DeepEP combine expects the expert activation dtype. Keep diff --git a/tests/functional_tests/moe/L2_MoE_DeepEP_EmptyRank_Backward.sh b/tests/functional_tests/moe/L2_MoE_DeepEP_EmptyRank_Backward.sh new file mode 100755 index 0000000000..6df0730f85 --- /dev/null +++ b/tests/functional_tests/moe/L2_MoE_DeepEP_EmptyRank_Backward.sh @@ -0,0 +1,23 @@ +#!/bin/bash +# 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. + +set -xeuo pipefail + +export PYTHONPATH=${PYTHONPATH:-}:$(pwd) +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1}" +export TORCH_NCCL_USE_COMM_NONBLOCKING=0 + +timeout 180s torchrun --nproc_per_node=2 --nnodes=1 \ + tests/functional_tests/moe/run_deepep_empty_rank_backward.py diff --git a/tests/functional_tests/moe/run_deepep_empty_rank_backward.py b/tests/functional_tests/moe/run_deepep_empty_rank_backward.py new file mode 100755 index 0000000000..340ed1f70c --- /dev/null +++ b/tests/functional_tests/moe/run_deepep_empty_rank_backward.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python +# 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. + +"""Two-rank DeepEP regression for an EP rank that receives no expert tokens. + +All tokens are routed to expert 0, so rank 1 takes the zero-token expert path. +Backward must still traverse DeepEP's reverse-dispatch collective on both ranks +and materialize explicit zero gradients for rank 1's local expert parameters. +""" + +from __future__ import annotations + +import os +import sys +from datetime import timedelta + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor import Shard, distribute_tensor + +from nemo_automodel.components.moe.config import MoEConfig +from nemo_automodel.components.moe.experts import GroupedExpertsDeepEP +from nemo_automodel.components.moe.megatron.fused_a2a import free_buffer + + +def _config() -> MoEConfig: + return MoEConfig( + n_routed_experts=2, + n_shared_experts=0, + n_activated_experts=1, + n_expert_groups=1, + n_limited_groups=1, + train_gate=True, + gate_bias_update_factor=0.0, + aux_loss_coeff=0.0, + score_func="softmax", + route_scale=1.0, + dim=16, + inter_dim=32, + moe_inter_dim=32, + norm_topk_prob=False, + expert_bias=True, + expert_activation="swiglu", + dtype=torch.bfloat16, + ) + + +def main() -> int: + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group("nccl", timeout=timedelta(seconds=90)) + + try: + if dist.get_world_size() != 2: + if dist.get_rank() == 0: + print("ERROR: this regression requires exactly two ranks", file=sys.stderr) + return 1 + + ep_mesh = init_device_mesh("cuda", (2,), mesh_dim_names=("ep",)) + experts = GroupedExpertsDeepEP(_config()).to(device=device, dtype=torch.bfloat16) + with torch.no_grad(): + experts.init_weights(device) + for name, parameter in list(experts.named_parameters(recurse=False)): + sharded = nn.Parameter(distribute_tensor(parameter.detach(), ep_mesh, [Shard(0)])) + sharded.requires_grad = parameter.requires_grad + experts.register_parameter(name, sharded) + experts.init_token_dispatcher(ep_mesh) + + generator = torch.Generator(device=device).manual_seed(20260825 + dist.get_rank()) + hidden_states = torch.randn(4, 16, generator=generator, device=device, dtype=torch.bfloat16) + hidden_states.requires_grad_(True) + routing_probs = torch.ones(4, 1, device=device, dtype=torch.float32, requires_grad=True) + expert_indices = torch.zeros(4, 1, device=device, dtype=torch.long) + token_mask = torch.ones(4, device=device, dtype=torch.bool) + + output = experts(hidden_states, token_mask, routing_probs, expert_indices) + assert output.shape == hidden_states.shape + assert torch.isfinite(output).all() + output.float().square().sum().backward() + + assert hidden_states.grad is not None and torch.isfinite(hidden_states.grad).all() + assert routing_probs.grad is not None and torch.isfinite(routing_probs.grad).all() + has_nonzero_expert_grad = False + for parameter in experts.parameters(): + assert parameter.grad is not None + local_grad = parameter.grad.to_local() + assert torch.isfinite(local_grad).all() + if dist.get_rank() == 1: + assert torch.count_nonzero(local_grad) == 0 + else: + has_nonzero_expert_grad |= bool(torch.count_nonzero(local_grad)) + if dist.get_rank() == 0: + assert has_nonzero_expert_grad + + dist.barrier() + if dist.get_rank() == 0: + print("PASS: DeepEP backward completed with an empty expert rank") + return 0 + finally: + free_buffer() + if dist.is_initialized(): + dist.destroy_process_group() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/unit_tests/moe/test_experts.py b/tests/unit_tests/moe/test_experts.py index 74b0f8824f..91cdcf3751 100644 --- a/tests/unit_tests/moe/test_experts.py +++ b/tests/unit_tests/moe/test_experts.py @@ -844,6 +844,68 @@ def test_empty_routing_probs_match_activation_dtype(self): assert stabilized.dtype == torch.bfloat16 assert _stabilize_empty_routing_probs_dtype(nonempty_probs, torch.bfloat16) is nonempty_probs + @pytest.mark.parametrize("expert_bias", [False, True], ids=["no-bias", "bias"]) + def test_zero_local_tokens_preserve_dispatch_and_parameter_gradients(self, expert_bias): + """The zero-token fallback preserves its empty shape and explicit parameter gradients.""" + config = MoEConfig( + n_routed_experts=2, + n_shared_experts=0, + n_activated_experts=1, + n_expert_groups=1, + n_limited_groups=1, + train_gate=True, + gate_bias_update_factor=0.0, + aux_loss_coeff=0.0, + score_func="softmax", + route_scale=1.0, + dim=4, + inter_dim=8, + moe_inter_dim=8, + norm_topk_prob=False, + expert_bias=expert_bias, + expert_activation="swiglu", + dtype=torch.float32, + ) + experts = GroupedExpertsDeepEP(config) + with torch.no_grad(): + experts.init_weights(torch.device("cpu")) + experts.ep_size = 1 + experts.n_routed_experts = config.n_routed_experts + + # Production weights are DTensors after parallelization. A plain CPU + # unit test can exercise the same forward path with identity to_local. + for parameter in experts.parameters(): + parameter.to_local = lambda parameter=parameter: parameter + + class EmptyRankDispatcher: + def token_permutation2(self, hidden_states, num_local_tokens, token_probs, token_indices): + del num_local_tokens, token_indices + return ( + hidden_states[:0], + torch.zeros(config.n_routed_experts, dtype=torch.long), + token_probs.reshape(-1)[:0], + ) + + def token_unpermutation(self, hidden_states): + return hidden_states + + experts.token_dispatcher = EmptyRankDispatcher() + x = torch.randn(3, config.dim, requires_grad=True) + weights = torch.rand(3, config.n_activated_experts, requires_grad=True) + indices = torch.zeros(3, config.n_activated_experts, dtype=torch.long) + token_mask = torch.ones(3, dtype=torch.bool) + + output = experts(x, token_mask, weights, indices) + assert output.shape == (0, config.dim) + assert torch.isfinite(output).all() + output.sum().backward() + + assert x.grad is not None + assert weights.grad is not None + for parameter in experts.parameters(): + assert parameter.grad is not None + assert torch.count_nonzero(parameter.grad) == 0 + def test_grouped_experts_deepep_token_dispatcher_init(self, moe_config): """Test token dispatcher initialization.""" experts = GroupedExpertsDeepEP(moe_config)