Summary
When global_aux_loss is enabled, the router's auxiliary loss and its gradients depend on how many valid tokens each microbatch happens to contain. The running expert counts accumulate over every microbatch seen so far, but their denominator assumes that every earlier microbatch had the current microbatch's token count.
The denominator can be wrong when valid-token counts differ, including when microbatches are padded to the same physical size.
System Info
Environment: Megatron-Core core_v0.18.2, commit 571370c8, PyTorch 2.13.0+cu130, FP32 router computation.
The results below are from the pinned version. Source inspection on 2026-09-10 found the same calculation in d4550898; the full experiment has not been rerun on that commit.
The PR description records a separate CPU check of the extracted loss methods, run with Python 3.11.15 and PyTorch 2.13.0+cpu. That check does not execute the full CUDA router.
Steps/Code to reproduce bug
The test calls TopKRouter.forward(input, padding_mask) with four experts and top-1 routing. It feeds the same 12 ordered valid tokens with the same initial router weights and auxiliary coefficient, then applies a single SGD step. Only the per-microbatch valid-token counts differ:
- Balanced:
[4, 4, 4]
- Uneven:
[6, 2, 4]
Both cases use three microbatches padded to six physical positions each. Because the fixture fixes the routing distribution, the reference running-prefix calculation gives identical results for the two partitions. The native implementation first diverges at the second microbatch.
| Router gradient, balanced vs. uneven partition |
Relative L2 difference |
| Current implementation |
0.2500001 |
| Cumulative-token-count repair |
2.25e-7 |
| Independent FP64 reference |
1.33e-15 |
An equal-token-count control also matches the reference. Both the discrepancy and the repair were reproduced on a second host.
The regression code is in PR #7214. test_global_aux_loss_uneven_token_counts covers padded and unpadded inputs and a second window after reset. It requires the repository's CUDA/distributed test setup and has not been run on this CPU host; the pinned GPU results above come from the earlier fixture.
Expected behavior
For the fixed-routing-frequency fixture above, normalizing the same running expert counts should use the cumulative valid-token count. The balanced and uneven partitions should agree to numerical tolerance, including when both are padded to the same physical size. This expectation does not assume that the running-prefix approximation is invariant under arbitrary changes in routing distribution.
Root cause
_apply_global_aux_loss divides the accumulated expert counts by ga_steps and hands the result to a loss helper that normalizes by the current token count. The effective denominator is therefore
ga_steps * current_num_tokens
After microbatches with 6 and 2 valid tokens, the numerator covers 8 tokens while the denominator is 2 * 2 = 4. The estimated expert frequencies are inflated even though the routing decisions themselves are unchanged.
Proposed fix
Track the cumulative valid-token count alongside global_tokens_per_expert, reduced over the same scope, and normalize the running counts by it. Reset the new counter in reset_global_aux_loss_tracker.
The tested repair leaves the existing loss helper untouched by passing cumulative_expert_counts * current_token_count / cumulative_token_count as its count input. This removes the discrepancy above while keeping the current running-prefix approximation.
The scaling changes in #5007 concern DP compensation when attaching the loss. This report concerns the token-count denominator accumulated across microbatches.
Related pull request
#7214.
Summary
When
global_aux_lossis enabled, the router's auxiliary loss and its gradients depend on how many valid tokens each microbatch happens to contain. The running expert counts accumulate over every microbatch seen so far, but their denominator assumes that every earlier microbatch had the current microbatch's token count.The denominator can be wrong when valid-token counts differ, including when microbatches are padded to the same physical size.
System Info
Environment: Megatron-Core
core_v0.18.2, commit 571370c8, PyTorch 2.13.0+cu130, FP32 router computation.The results below are from the pinned version. Source inspection on 2026-09-10 found the same calculation in d4550898; the full experiment has not been rerun on that commit.
The PR description records a separate CPU check of the extracted loss methods, run with Python 3.11.15 and PyTorch 2.13.0+cpu. That check does not execute the full CUDA router.
Steps/Code to reproduce bug
The test calls
TopKRouter.forward(input, padding_mask)with four experts and top-1 routing. It feeds the same 12 ordered valid tokens with the same initial router weights and auxiliary coefficient, then applies a single SGD step. Only the per-microbatch valid-token counts differ:[4, 4, 4][6, 2, 4]Both cases use three microbatches padded to six physical positions each. Because the fixture fixes the routing distribution, the reference running-prefix calculation gives identical results for the two partitions. The native implementation first diverges at the second microbatch.
0.25000012.25e-71.33e-15An equal-token-count control also matches the reference. Both the discrepancy and the repair were reproduced on a second host.
The regression code is in PR #7214.
test_global_aux_loss_uneven_token_countscovers padded and unpadded inputs and a second window after reset. It requires the repository's CUDA/distributed test setup and has not been run on this CPU host; the pinned GPU results above come from the earlier fixture.Expected behavior
For the fixed-routing-frequency fixture above, normalizing the same running expert counts should use the cumulative valid-token count. The balanced and uneven partitions should agree to numerical tolerance, including when both are padded to the same physical size. This expectation does not assume that the running-prefix approximation is invariant under arbitrary changes in routing distribution.
Root cause
_apply_global_aux_lossdivides the accumulated expert counts byga_stepsand hands the result to a loss helper that normalizes by the current token count. The effective denominator is thereforeAfter microbatches with 6 and 2 valid tokens, the numerator covers 8 tokens while the denominator is
2 * 2 = 4. The estimated expert frequencies are inflated even though the routing decisions themselves are unchanged.Proposed fix
Track the cumulative valid-token count alongside
global_tokens_per_expert, reduced over the same scope, and normalize the running counts by it. Reset the new counter inreset_global_aux_loss_tracker.The tested repair leaves the existing loss helper untouched by passing
cumulative_expert_counts * current_token_count / cumulative_token_countas its count input. This removes the discrepancy above while keeping the current running-prefix approximation.The scaling changes in #5007 concern DP compensation when attaching the loss. This report concerns the token-count denominator accumulated across microbatches.
Related pull request
#7214.