[Dev] fix(moe): align MoE router loss scale semantics

[Victarry]

Summary

This PR aligns the effective scale of MoE router auxiliary losses across aux-loss variants, --calculate-per-token-loss, and TP/CP/DP parallel strategies.

The conservative semantic baseline is kept unchanged: load_balancing_type=aux_loss with --calculate-per-token-loss=false. Other paths are aligned to that baseline.

Fixes #3672.

Problem Manifestation

Before this fix, equivalent MoE training configurations could apply different effective router regularization strengths even when using the same moe_aux_loss_coeff.

The most visible symptoms were:

• global_aux_loss did not have the same effective scale as the standard aux_loss. In practice, the active load-balancing loss and router gradient strength could drift when switching only the aux-loss type from aux_loss to global_aux_loss.
• With aux_loss and --calculate-per-token-loss enabled, changing TP/CP sizes changed the effective aux-loss scale. The same model and data could therefore show different router-loss/router-gradient behavior under different TP/CP parallel strategies.
• The same per-token normalization issue also affected seq_aux_loss/z_loss style router-loss attachment paths, so aux type, per-token mode, and TP/CP/DP strategy were not cleanly comparable.

The expected behavior is that these configurations should preserve the same coefficient semantics: changing aux-loss type or parallel decomposition should not systematically rescale router gradients, except where the loss definition itself intentionally differs.

Root Cause

There were two scale mismatches in the router-loss attachment path:

• global_aux_loss builds tokens_per_expert from a DP-aggregated tp_dp_cp_group, while router probabilities remain local. Since the final gradient is averaged across DP ranks, the attached gradient was diluted by dp_size.
• Per-token aux/z-loss paths attached loss scaled by the local valid-token count, but finalize_model_grads normalizes by the global token count. With TP/CP enabled, this missed the tp_cp_group.size() factor.

As a result, otherwise equivalent configurations could show systematic router-loss and router-gradient scale drift when switching aux-loss type, per-token-loss mode, or TP/CP/DP strategy.

Fix

• Scale global_aux_loss by dp_size = tp_dp_cp_group.size() // tp_cp_group.size() so its backward strength matches aux_loss.
• Keep the logged global aux loss normalized through the same scale factor, so reported metrics stay comparable.
• Scale per-token aux-loss and z-loss attachment by local_valid_tokens * tp_cp_group.size() to match global-token-count gradient normalization.
• Preserve the non-per-token aux_loss baseline behavior.

Validation

Public W&B report: https://api.wandb.ai/links/megatron-core-moe-dev/xafmkree

Original W&B report: https://wandb.ai/megatron-core-moe-dev/moe-aux-loss-scale-alignment/reports/MoE-Aux-Loss-Scale-Alignment-E2E-Ablation---2026-05-25--VmlldzoxNzAwODkxMw==

E2E validation uses the same model/config family and compares loss ratios against the baseline instead of forcing load balancing. Horizontal aux-type comparisons use a deterministic mock_repeat dataset so each global batch is identical across aux types. Vertical parallel-strategy comparisons use the real SlimPajama data path and load all compared TP/CP/DP jobs from the same one-iteration checkpoint.

Key validation results:

Comparison	Before fix active LB ratio	After fix active LB ratio	After fix LM ratio
aux_loss vertical	0.9994 - 1.0455	0.9952 - 1.0063	0.9985 - 1.0053
seq_aux_loss vertical	0.9929 - 1.0419	0.9948 - 1.0053	0.9997 - 1.0057
global_aux_loss vertical	0.9483 - 1.0080	0.9981 - 1.0086	0.9986 - 1.0035
z_loss vertical	0.9844 - 2.6181	0.9254 - 1.0281	0.9968 - 1.0067

Additional checks:

• Horizontal deterministic repeat-data jobs: 12/12 completed.
• aux_loss + --calculate-per-token-loss=false is bitwise identical before/after fix, confirming the baseline is unchanged.
• aux_loss and seq_aux_loss fixed TP1/CP1 horizontal runs are bitwise identical before/after fix where the scale correction should be neutral.
• global_aux_loss changes after the first step as expected because the fix intentionally changes its backward scale.
• Added unit coverage for global_aux_loss gradient scaling.

Local sanity checks:

git diff --check
python3 -m py_compile megatron/core/transformer/moe/router.py megatron/core/transformer/moe/moe_utils.py tests/unit_tests/transformer/moe/test_aux_loss.py

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[Victarry]

/ok to test 90fcb53

[Victarry]

/claude review

[claude]

LGTM

[zyeric]

why use dp_size instead of self.tp_dp_cp_group.size(). Since global_tokens_per_expert is reduced at tp_dp_cp_group.

[Victarry]

why use dp_size instead of self.tp_dp_cp_group.size(). Since global_tokens_per_expert is reduced at tp_dp_cp_group.

It's to keep consistent with the standard aux loss, otherwise the effective scale of global aux loss will be increased.

[zyeric]

why use dp_size instead of self.tp_dp_cp_group.size(). Since global_tokens_per_expert is reduced at tp_dp_cp_group.

It's to keep consistent with the standard aux loss, otherwise the effective scale of global aux loss will be increased.

I think all aux loss should align with --calculate-per-token-loss's behavior: activation = MoEAuxLossAutoScaler.apply(activation, aux_loss * activation.shape[0]) (mean loss -> sum loss). If dp_size is used when tp * cp > 1, I think the semantic is confusing.

[Victarry]

why use dp_size instead of self.tp_dp_cp_group.size(). Since global_tokens_per_expert is reduced at tp_dp_cp_group.

It's to keep consistent with the standard aux loss, otherwise the effective scale of global aux loss will be increased.

I think all aux loss should align with --calculate-per-token-loss's behavior: activation = MoEAuxLossAutoScaler.apply(activation, aux_loss * activation.shape[0]) (mean loss -> sum loss). If dp_size is used when tp * cp > 1, I think the semantic is confusing.

@zyeric Thanks for your comments.
I agreed aux loss with/without --calculate-per-token-loss should have the same bahavior.
I think the aux loss without calculate per token loss is the most commonly used loss, so we don't want change its semantics for now, which will influence the optimal hyper parameters. So I prefer to keep the current formula. There are some historical issues for the current "confusing" implementation, like https://github.com/NVIDIA/Megatron-LM/pull/2217/changes#r2521887449

For the --calculate-per-token-loss, there are still some issues which may cause loss diff with different parallel sizes, see #3105 and #1652 for details. I'll revisit the implementations to make sure all of different combinations(parallel size, per-token-loss, different aux loss types) follow the same semantics.

Marked this PR as draft until resolving all aux-loss related issues.

[Victarry]

/ok to test 7810a78

[Victarry]

/ok to test 25642dd

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/26496906565

[Victarry]

Closed this in favor of #5047