[FeSens]

What does this PR do?

What's wrong

The DeepSeek-V2/V3-style routers compute their logits as F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32)), and 23 other models copy or inherit these routers. Under torch.autocast(dtype=torch.bfloat16), autocast runs F.linear in bf16 and undoes the casts. The logits come back in bf16, so the sigmoid/softmax, bias addition and top-k all run on bf16 scores. Trainer and Accelerate turn on autocast for bf16=True (outside DeepSpeed). As a result, a bf16 fine-tune of these models picks different experts than inference does with the same weights.

Reproduction

import torch
from transformers import DeepseekV3Config
from transformers.models.deepseek_v3.modeling_deepseek_v3 import DeepseekV3TopkRouter

torch.manual_seed(0)
config = DeepseekV3Config(hidden_size=256, n_routed_experts=64, num_experts_per_tok=8, n_group=8, topk_group=4)
router = DeepseekV3TopkRouter(config).eval()
torch.nn.init.normal_(router.weight, std=0.02)
torch.nn.init.normal_(router.e_score_correction_bias, std=0.01)
hidden_states = torch.randn(4096, config.hidden_size)

with torch.no_grad():
    logits, weights, indices = router(hidden_states)
    with torch.autocast("cpu", dtype=torch.bfloat16):
        logits_ac, weights_ac, indices_ac = router(hidden_states)

same_set = (indices.sort(-1).values == indices_ac.sort(-1).values).all(-1)
print("logits dtype:", logits.dtype, "| under autocast:", logits_ac.dtype)
print(f"tokens with a different expert set under autocast: {(~same_set).float().mean():.1%}")

On main (50e35cc):

logits dtype: torch.float32 | under autocast: torch.bfloat16
tokens with a different expert set under autocast: 10.5%

With this change: torch.float32 and 0.0%.

Impact

The bug affects any bf16-autocast training or evaluation of axk1, axk2, bailing_hybrid, beam, deepseek_ocr2, deepseek_v2, deepseek_v3, deepseek_v32, dots1, exaone_moe, glm4_moe, glm4_moe_lite, glm4v_moe, glm5_next, glm_moe_dsa, hy_v3, hy_v4, kimi_linear, kolibri1, longcat_flash, mimo_v2_flash, muse_spark, nemotron_h, solar_open and unlimited_ocr. Ernie 4.5 MoE (text and VL) and Flint already disable autocast around their router matmul, and they behave correctly.

I measured this on zai-org/GLM-4.7-Flash (glm4_moe_lite, 64 experts, top-4, sigmoid router) with these settings:

• the first 6 decoder layers, which include 5 MoE layers
• bf16 weights, on an RTX 5070 Ti
• 8,192 tokens of the wikitext-2-raw-v1 test split

Each cell is the share of tokens whose top-4 expert set changes under torch.autocast("cuda", dtype=torch.bfloat16):

MoE layer	same router input	full forward
1	1.87%	1.87%
2	2.81%	3.17%
3	1.44%	2.67%
4	2.01%	3.76%
5	1.76%	3.83%

• When the expert set does not change, routing weights differ by 2e-3 to 4.5e-3 on average (absolute).
• Repeating the forward pass without autocast changes 0.00% of tokens, so GPU nondeterminism is not the cause.
• With this change, under autocast, 0.00% of tokens change in every layer.

Fix

The router matmul now runs under maybe_autocast(device_type=device_type, enabled=False), the same way Ernie4_5_MoeTopKRouter and FlintTopKRouter already do it. The explicit casts stay, because bf16 weights still need them when autocast is off.

• I edited the 8 modular roots that define these forwards: deepseek_v2, deepseek_v3, axk2, beam, hy_v3, kolibri1, longcat_flash and muse_spark. The other generated modeling_*.py files were regenerated from them.
• Only the matmul needs wrapping. Autocast does not cast sigmoid, topk or gather, and it already runs softmax in float32.
• maybe_autocast returns a nullcontext when autocast is off. Without autocast, outputs are unchanged and torch.compile gets no extra nodes. A fullgraph=True compile of the router works both with and without autocast.
• For Kolibri 1, only the F.linear branch is wrapped. Its torch.mm(..., out_dtype=torch.float32) fast path already returns float32 under CUDA autocast (checked).

Not changed here: routers that compute in the activations' dtype (Qwen3-MoE, Mixtral) don't ask for float32. MiniMaxM2TopKRouter, MiniMaxM3VLTopKRouter and Step3p7TopKRouter (step3p7 is generated from minimax_m3_vl) compute F.linear(hidden_states.to(self.weight.dtype), self.weight), so they run in the gate's dtype.

• MiniMax checkpoints store the router weight in float32, and Keep the MoE correction bias (and MiniMax-M2's router) in float32 when loading in bfloat16 #50002 keeps it in float32 on load. With an fp32 gate under CPU bf16 autocast, these routers return bf16 logits, and 1.8–2.4% of tokens change expert set (random weights).
• The same wrap works for them unchanged. It runs the matmul in the gate's dtype and leaves bf16 gates as they are (checked). That PR adds it together with the float32 gate. muse_spark subclasses MiniMaxM2TopKRouter, but it has its own float32 forward, which is fixed here.

Tests

• I added test_router_ignores_autocast to tests/models/deepseek_v3 and tests/models/deepseek_v2. The test checks that the router output (logits, weights, indices) under torch.autocast(torch_device, dtype=torch.bfloat16) equals the output without autocast, and that the logits are float32.
  • On main, both tests fail on CPU with AssertionError: torch.bfloat16 != torch.float32.
  • With this change, both pass on CPU. The same test body also passes on CUDA, and fails there on main.
• I ran the existing tests of all 25 affected models (-k "not slow", CPU, macOS): 3889 passed, 2830 skipped, 52 failed.
  • All 52 failures are the TP/EP tests of glm4_moe, glm4_moe_lite, glm4v_moe and glm_moe_dsa. They raise "Tensor parallelism is not supported on MPS devices", and main fails the exact same 52 tests on this machine.
• utils/check_modular_conversion.py reports the 25 models as consistent. utils/check_copies.py and ruff check / ruff format --check (ruff 0.14.10) pass on the changed files.

Note for run-slow: the unlimited_ocr integration tests wrap generate in bf16 autocast to match the reference implementation. The reference MoEGate has the same F.linear(... float32 ...) pattern, so under that autocast it actually routes in bf16. Those expected outputs may need an update. Suggested: run-slow: deepseek_v2, deepseek_v3, glm4_moe, glm4_moe_lite, kolibri1, unlimited_ocr.

Code Agent Policy

The Transformers repo is currently being overwhelmed by a large number of PRs and issue comments written by code agents. These often are low-quality, or fix extremely minor issues that occur rarely or never in practice. As a result, we're instituting a rule that first-time contributors should not use code agents to submit PRs or issues. We'd also ask autonomous "OpenClaw"-like agents not to open any PRs or issues.

Issues/PRs from first-time contributors that violate this rule will probably just be closed without review, and we might block you, especially if you open more than one or appear to be deliberately ignoring this. We especially do not want new contributors to jump in on random issues to contribute an agent-written fix. This creates lots of noise for reviewers and other users and will almost certainly get you blocked.

For more information, please read CONTRIBUTING.md.

• (First-time contributors only): I confirm that this PR description and code is not written by an LLM or code agent

The code, the tests and this description were drafted with an AI coding agent (Claude Code), so the box above is not checked.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline and the Pull Request checks?
• Was this discussed/approved via a Github issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes according to the guidelines?
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag members/contributors who may be interested in your PR.

[github-actions]

Thank you for your contribution 🤗!

CI Security Gate — waiting for maintainer approval

CI is approved automatically only for authors who already have a contribution merged in this repository. A maintainer will review this PR and approve CI to run on it.

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: axk1, axk2, bailing_hybrid, beam, deepseek_ocr2, deepseek_v2, deepseek_v3, deepseek_v32, dots1, exaone_moe, glm4_moe, glm4_moe_lite, glm4v_moe, glm5_next, glm_moe_dsa, hy_v3