Fix softmaxing router logits#45315
Conversation
|
(Doing the fix in a separate PR because the existing PRs were stale or had rebase issues) |
|
cc @vasqu - this has popped up in multiple issues since January so it'd be good to finally put it to rest! |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
33dd54b to
b1c8150
Compare
|
[For maintainers] Suggested jobs to run (before merge) run-slow: flex_olmo, minimax, mixtral, olmoe, qwen2_moe, qwen3_5_moe, qwen3_moe, qwen3_next, qwen3_omni_moe, qwen3_vl_moe |
|
Answered on the other PR but replying here just for viz, that really has become messy 😬 sorry |
|
Yep, closing in favour of the other PR to keep things neater! |
Reusing a variable name meant that we returned a softmaxed value instead of the original logits in some MoE routers. This generally did not affect inference, but could affect the auxiliary loss on MoE logits in training when the coefficient for that loss was > 0. First reported in #43542. Fixes #45120