[FeSens]

What does this PR do?

What's wrong

These models route like DeepSeek-V3: they add e_score_correction_bias to the sigmoid scores, then take the top-k:

• MiniMax-M2 / M2.1 / M2.5
• MiniMax-M3 (through minimax_m3_vl)
• Step-3.7-Flash
• LongCat-Flash

More than 20 models with that router already declare _keep_in_fp32_modules_strict = ["e_score_correction_bias"], including DeepSeek-V3, GLM-4.5/4.6/4.7, dots1 and MiMo-V2. These four don't. With dtype=torch.bfloat16, the bias is rounded to bf16.

For MiniMax and Step that breaks routing. Their checkpoints store the bias in float32, and its values are large but differ little between experts:

checkpoint	bias mean (per layer)	spread between experts (std)	distinct values per layer, after bf16
MiniMax-M2 / M2.1 / M2.5	1.5 – 8.6	0.02 – 0.18	9 – 41 of 256
MiniMax-M3	4.9 – 7.9		8 – 12 of 128
Step-3.7-Flash (router_bias)	5.8 – 8.7		6 – 12 of 288

Near 8, bf16 values are 0.0625 apart, so the rounding error (up to 0.031) is larger than the gap between the sigmoid scores of competing experts. Expert selection then follows the rounding of the bias rather than the router.

MiniMax also stores the router weight (gate.weight) in float32. Its quantization config keeps it out of FP8 (modules_to_not_convert: ["gate", "e_score_correction_bias", "lm_head"]). The router already casts its input to the weight's dtype (F.linear(hidden_states.to(self.weight.dtype), self.weight)). Under torch.autocast, which Trainer turns on for bf16=True, that matmul still runs in bf16.

Impact

Measured on MiniMax-M2.5's first MoE layer, using the real checkpoint tensors on 16,384 tokens (4 × 4096, prose and code):

• The embeddings and layer-0 attention come from the checkpoint (fp8, dequantized).
• Hidden states are computed once in bf16 with this library's MiniMax-M2 modules.
• Only the router's precision changes between rows.
• Each row is compared with routing that uses the checkpoint's float32 gate and bias.

router	tokens with a different top-8 set	experts swapped per token (of 8)	busiest expert load / mean	idle experts
checkpoint (fp32 gate, fp32 bias) = this PR	0	0	5.15	1
main, dtype=torch.bfloat16	99.7%	3.81	17.4	16
bf16 gate, fp32 bias	14.1%	0.14	5.16	1
fp32 gate, bf16 bias	99.7%	3.81	17.3	16

MiniMax-M2 gives the same picture:

• On main: 99.7% of tokens get a different set, 3.92 experts are swapped per token, the busiest expert carries 16.5× the mean load (6.03× from the checkpoint), and 17 experts get no tokens (1 from the checkpoint).
• With a bf16 gate and an fp32 bias: 12.5% of tokens get a different set.

So on main, almost every token has about half of its 8 expert slots filled by experts the float32 router would not pick. The bias causes almost all of this. The bf16 gate causes the remaining 13–14%, and the effect differs by checkpoint:

• M2.5: the gate weights are exactly representable in bf16. The loss comes from the router computing, and returning, its logits in bf16. A bf16 GEMM with float32 output gives 0%.
• MiniMax-M3: 99.998% of the gate weights change when rounded to bf16.

The router computes in the weight's dtype, so keeping the weight in float32 covers both cases.

LongCat-Flash's biases are centred near zero, so bf16 costs little there. It is included so the checkpoint's float32 values are kept as for the other models.

Fix

• MiniMaxM2PreTrainedModel: _keep_in_fp32_modules_strict = ["e_score_correction_bias", "mlp.gate.weight"]. MiniMaxM3VL inherits it through the modular file.
• Step3p7PreTrainedModel and LongcatFlashPreTrainedModel: _keep_in_fp32_modules_strict = ["e_score_correction_bias"].
• MiniMaxM2TopKRouter and MiniMaxM3VLTopKRouter: the matmul runs under maybe_autocast(device_type, enabled=False), so it computes in the gate's dtype even under autocast. Ernie4_5_MoeTopKRouter already does this. Step-3.7's router is generated from M3-VL and gets the same wrapper; its gate is bf16 there, so it computes in bf16 as before. Without autocast, maybe_autocast is a nullcontext, and the outputs are unchanged.

The modeling files are regenerated from the modular files, and utils/check_modular_conversion.py passes. Step's router_bias is renamed to e_score_correction_bias before the keep-in-fp32 match, which the test covers through save_pretrained / from_pretrained.

Tests

One new test per model (test_router_and_correction_bias_kept_in_float32 for both MiniMax models, test_e_score_correction_bias_kept_in_float32 for Step-3.7 and LongCat):

1. Fill the biases with 8.5 + 0.02 * arange(n).
2. save_pretrained, then from_pretrained(..., dtype=torch.bfloat16).
3. Check that the biases are float32 and bitwise equal to what was saved, and for MiniMax that the gate weight is float32 and that the router returns float32 logits under torch.autocast("cpu", dtype=torch.bfloat16).

Results:

• On main, all four fail (torch.bfloat16 != torch.float32); with the fix, all four pass.
• With the keep-in-fp32 lists but without the router wrapper, the two MiniMax tests still fail on the autocast check.
• The full test files of the four models: 703 passed, 677 skipped.
• utils/check_modular_conversion.py passes, and ruff 0.14.10 check and format are clean.

Related

• Ensure e_score_correction_bias dtype of DeepSeek-V3/R1 is FP32 #42580 (merged 2025-12-10) did the same for DeepSeek-V3, dots1, GLM4-MoE and GLM4V-MoE. Models added since then didn't get it.
• Rename only MiniMax-M3-VL's MoE block, so a save keeps the dense and vision mlp keys #49786 (open) changes MiniMax-M3-VL's key renaming and also edits test_modeling_minimax_m3_vl.py. Both PRs only add code there, so they shouldn't conflict. The keep-in-fp32 match runs on the renamed keys (mlp.gate.weight), which that PR keeps.
• I found no existing issue or PR for MiniMax-M2, Step-3.7 or LongCat.
• Keep float32 MoE router logits in float32 under autocast #50003 adds the same autocast wrapper to the routers that cast to float32 explicitly (DeepSeek-V3, GLM-4.x and 23 others).

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: longcat_flash, minimax_m2, minimax_m3_vl, step3p7