[BUG] the function seq_aux_loss_load_balancing seems not align with DeepSeekV2 paper
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 3h
- Merged PRs (30d)
- 272
Description
**Describe the bug**

if the capacity is not None, computing `fi` should counts the score '1' base on the topk_mask instead of capacity_mask, otherwise the aux loss will be less than the true aux loss about ~30%,(1.1 vs 0.8).
(test under setting num_experts=80, topk_experts = 8, capcacity=2)
The aim of aux loss is to balancing expert selection, so this kind of capacity drop indeed weaken the regulation of balancing.
In the deepseekv3 paper, the declaration is more clear:
the score is 1 as long as `s'` in topk score of experts `i` with respect to token`t` , not `Token 𝑡 selects Expert`

**Environment (please complete the following information):**
- Megatron-LM commit ID
- core_r0.11.0
- [b1022a3](https://github.com/NVIDIA/Megatron-LM/commit/b1022a323087584fdfe51d9690da162f3767afc7)
- PyTorch version
- CUDA version
- NCCL version
**Proposed fix**
If you have a proposal for how to fix the issue state it here or link to a PR.
the last return value `tokens_per_expert` of fucntion `topk_softmax_with_capacity` is just as well `fi` before capactiy mask.
aux loss can be computed based on this tensor

Contributor guide
Research direction
Start by tracing seq_aux_loss_load_balancing and topk_softmax_with_capacity, focusing on how fi and tokens_per_expert are computed with a capacity mask. Compare that behavior with the DeepSeekV2/V3 paper and the proposed use of the pre-capacity value; done means the auxiliary loss reflects top-k expert selection rather than dropped tokens.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100