Fix MoE routers returning probabilities instead of logits

[yacinemebarki]

What does this PR do?

Fixes issue #45120: Several MoE routers returned softmaxed probabilities as router_logits, which caused load_balancing_loss_func to compute softmax(softmax(logits)), flattening routing distributions and weakening gradient signals during fine-tuning. This PR fixes it by keeping router_logits as raw logits and computing router_probs separately for top-k routing.

---

Code Agent Policy

The Transformers repo is currently being overwhelmed by a large number of PRs and issue comments written by code agents. We are currently bottlenecked by our ability to review and respond to them. As a result, we ask that new users do not submit pure code agent PRs at this time.

• I confirm that this is not a pure code agent PR.

For more information, please read CONTRIBUTING.md.

---

Before submitting

• This PR fixes the issue and does not break inference
• Did you read the contributor guideline?
• Documentation does not need changes, as behavior/API is unchanged
• No new tests are required; logic change only affects auxiliary loss

How to test / verify

1. Train a model using output_router_logits=True.
2. Check that router_logits returned by routers are raw logits, not probabilities.
3. Verify that load_balancing_loss_func receives proper logits and computes meaningful gradients.
4. Confirm that inference behavior is unchanged.

[vasqu]

Careful approval but this is correct in my eyes, we definitely shouldn't have a double softmax imo.

2 things

1. You need to apply this to actual modeling files and their inheritances via make fix-repo
2. We need to adjust our test to catch this, i.e. let's add another check in the test_load_balancing_loss tests (we should be easily check via sum != 1 which should be hard to achieve without a softmax)

Also cc @ArthurZucker since you have the most exp here with the moe loss

[vasqu]

@yacinemebarki you applied the template a bit wrong and basically pinged every member of the transformers repo. Please be careful about this the next time; for now I commented it out again 🤗

[vasqu]

Looks like a bad merge/rebase happened 😓

[vasqu]

Will likely be superceded by #45346

[github-actions]

[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

[HuggingFaceDocBuilderDev]

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.

[ArthurZucker]

let's go