[Bug] global_aux_loss uses the current microbatch token count to normalize accumulated expert counts
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
### 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](https://github.com/NVIDIA/Megatron-LM/blob/571370c829ca768fe37244f4e2e7f28d8accc4ab/megatron/core/transformer/moe/router.py), 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](https://github.com/NVIDIA/Megatron-LM/blob/d4550898cc1a0fd32c6defe1f79de6d5183cf6d0/megatron/core/transformer/moe/router.py); 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](https://github.com/NVIDIA/Megatron-LM/blob/ae9eaf14d5d3a963518e56257dbfa3c5222e7115/tests/unit_tests/transformer/moe/test_aux_loss.py). `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
```text
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.
Contributor guide
Research direction
Start in megatron/core/transformer/moe/router.py at _apply_global_aux_loss and reset_global_aux_loss_tracker; compare the regression coverage in tests/unit_tests/transformer/moe/test_aux_loss.py. Run test_global_aux_loss_uneven_token_counts with the repository's CUDA/distributed setup. Done means balanced and uneven padded or unpadded cases agree within tolerance and the reset/window case remains covered.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning, testing-qa
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 25/100