[dev] [fix] Skip auxiliary loss for hash MoE layers

megatron/core/transformer/moe/router.py

@@ -904,8 +904,15 @@ def routing(
                 pad_to_capacity=self.config.moe_pad_expert_input_to_capacity,
             )
 
-        # Apply each aux loss type and attach aux loss autograd function to probs
-        if self.training and torch.is_grad_enabled() and self.is_aux_loss_enabled():
+        # Hash routing assignments come from tid2eid rather than the learned logits, so a
+        # learned top-k aux loss would optimize routing decisions that were never dispatched.
+        # Keep aux loss enabled for the non-hash MoE layers that share this configuration.
+        if (
+            not self.is_hash_layer
+            and self.training
+            and torch.is_grad_enabled()
+            and self.is_aux_loss_enabled()
+        ):
             # Calculate scores and routing_map for aux loss
             routing_map_for_aux_loss, scores_for_aux_loss = compute_routing_scores_for_aux_loss(
                 logits,

tests/unit_tests/transformer/moe/test_routers.py

@@ -759,6 +759,40 @@ def test_is_hash_layer_logic(self):
         assert hash_r.enable_expert_bias is False
         assert normal_r.enable_expert_bias is True
 
+    @pytest.mark.internal
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    @pytest.mark.parametrize("aux_loss_type", ["aux_loss", "seq_aux_loss", "global_aux_loss"])
+    def test_hash_routing_skips_aux_loss(self, aux_loss_type, monkeypatch):
+        """Hash layers must not attach learned-routing aux-loss gradients."""
+        config = _hash_routing_config(
+            moe_router_load_balancing_type=aux_loss_type, moe_aux_loss_coeff=1.0
+        )
+        router = TopKRouter(
+            config=config, pg_collection=get_default_pg_collection(), layer_number=1
+        )
+
+        def fail_if_called(*args, **kwargs):
+            pytest.fail("hash routing must not recompute a learned top-k map for aux loss")
+
+        monkeypatch.setattr(
+            "megatron.core.transformer.moe.router.compute_routing_scores_for_aux_loss",
+            fail_if_called,
+        )
+
+        logits = (
+            torch.tensor([3.0, 2.0, 1.0, 0.0], device="cuda")
+            .expand(4, 2, -1)
+            .clone()
+            .requires_grad_()
+        )
+        input_ids = torch.arange(8, device="cuda").reshape(2, 4)
+
+        probs, _ = router.routing(logits, input_ids=input_ids)
+        probs.sum().mul_(0).backward()
+
+        assert logits.grad is not None
+        torch.testing.assert_close(logits.grad, torch.zeros_like(logits.grad))
+
     @pytest.mark.internal
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     def test_moe_layer_hash_routing_integration(self):