Skip to content

Fix MoE aux_loss / z_loss gradient scaling with TP > 1 - #5047

Merged
ko3n1g merged 3 commits into
NVIDIA:mainfrom
deepakn94:dnarayanan/fix_moe_loss_tp_scaling
Jun 3, 2026
Merged

Fix MoE aux_loss / z_loss gradient scaling with TP > 1#5047
ko3n1g merged 3 commits into
NVIDIA:mainfrom
deepakn94:dnarayanan/fix_moe_loss_tp_scaling

test: TP/CP-invariance of aux_loss under --calculate-per-token-loss

d2633d8
Select commit
Loading
Failed to load commit list.
Sign in for the full log view
check-approval
succeeded May 31, 2026 in 21s