fix(moe): preserve DeepEP autograd for empty expert ranks

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

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

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())

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)