From 8f1fb1fd7c71d3de81249467c4c74eb1a42f655b Mon Sep 17 00:00:00 2001 From: tdophung Date: Wed, 22 Jul 2026 11:21:52 -0700 Subject: [PATCH] [JAX] Keep MoE aux-loss cotangent scalar --- transformer_engine/jax/router.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/transformer_engine/jax/router.py b/transformer_engine/jax/router.py index 80ec42a95f..170c74fe5f 100644 --- a/transformer_engine/jax/router.py +++ b/transformer_engine/jax/router.py @@ -342,12 +342,11 @@ def _fused_moe_aux_loss_fwd(probs, tokens_per_expert, topk, coeff): def _fused_moe_aux_loss_bwd(topk, coeff, residuals, g): del topk, coeff const_buf, tokens_per_expert, num_tokens = residuals - grad_aux_loss = g.reshape(1) grad_probs = fused_moe_aux_loss_bwd( const_buf, tokens_per_expert, - grad_aux_loss, + g, num_tokens, ) return grad_probs, None