diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index 52823ed88dd..74b09bfa9c1 100644 --- a/megatron/core/transformer/moe/router.py +++ b/megatron/core/transformer/moe/router.py @@ -872,8 +872,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, diff --git a/tests/unit_tests/transformer/moe/test_routers.py b/tests/unit_tests/transformer/moe/test_routers.py index 23a0cdc3b6a..f77adc304c2 100644 --- a/tests/unit_tests/transformer/moe/test_routers.py +++ b/tests/unit_tests/transformer/moe/test_routers.py @@ -737,6 +737,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):