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

Commits on May 31, 2026

Claude's suggestion to fix seq_loss discrepancy with different TP
deepakn94
committed
c939862
Clarify TP/CP gradient scaling comments for MoE aux/z loss

deepakn94
and
claude
committed
96664cc
test: TP/CP-invariance of aux_loss under --calculate-per-token-loss

deepakn94
and
claude
committed
d2633d8