NVIDIA / NVIDIA/Megatron-LM

[BUG] MoE router can potentially hack the aux loss when using sigmoid gating function

Open
#1,545 8 comments 3 reactions 0 assignees View on GitHub
bug community-request module: moe
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.