[BUG] MoE router can potentially hack the aux loss when using sigmoid gating function
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**Describe the bug**
It's possible for the MoE router to "hack" the auxiliary load balancing loss by pushing sigmoid activations below 1e-20. While I haven't seen this happen in Megatron-LM yet, I have seen it in https://github.com/allenai/OLMo-core/pull/255 when we used the same normalizer epsilon of 1e-20, so I thought I'd raise this here just in case this happens to someone else.
https://github.com/NVIDIA/Megatron-LM/blob/55c968dabb6d32cf8c21b1e3fadeea47e2ceb78d/megatron/core/transformer/moe/router.py#L195-L199
**To Reproduce**
Train an MoE with auxiliary load balancing loss and the sigmoid router gating function. It's rather hard to reproduce as it doesn't happen all the time. The model needs to get somewhat lucky to "realize" that it can hack the aux loss.
**Expected behavior**
Aux loss stays >= 1.
**Environment (please complete the following information):**
- 55c968dabb6d32cf8c21b1e3fadeea47e2ceb78d
**Proposed fix**
Decrease the normalizing epsilon (e.g. use `torch.finfo(torch.float32).tiny` ~= `1.175e-38`) or remove it. I'm not sure it's necessary for stability.
Contributor guide
Assessment
This issue has not been assessed yet.