[FeSens]

What does this PR do?

What's wrong

DeepseekV4TopKRouter and DeepseekV4HashRouter compute logits = F.linear(flat, self.weight) in the dtype of their inputs. With dtype=torch.bfloat16, almost everything is bf16:

• The router input is bf16, because post_attention_layernorm casts its output back to the activation dtype.
• gate.weight is bf16. The checkpoint stores it in bf16, and it is not in _keep_in_fp32_modules_strict.
• Only e_score_correction_bias is float32.

So the logits, the sqrtsoftplus scores and the normalised routing weights are all bf16.

DeepSeek's reference Gate.forward (inference/model.py in deepseek-ai/DeepSeek-V4-Flash) computes scores = linear(x.float(), self.weight.float()) and keeps the scores and weights in float32. This library's DeepseekV3TopkRouter does the same, with F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32)).

Near 1, consecutive bf16 values are 2⁻⁷ ≈ 0.0078 apart. The scores of competing experts can be closer than that. The float32 bias is then added to rounded scores, so the top-k depends on the rounding, not only on the router.

Reproduction

import torch
import torch.nn.functional as F
from transformers import DeepseekV4Config
from transformers.models.deepseek_v4.modeling_deepseek_v4 import DeepseekV4TopKRouter

# Two experts with logits 1 and 1.0078125 (both exact in bf16), correction bias 0.001 and 0, top-1.
config = DeepseekV4Config(hidden_size=64, n_routed_experts=2, num_experts_per_tok=1)
router = DeepseekV4TopKRouter(config)
with torch.no_grad():
    router.weight.zero_()
    router.weight[:, 0] = torch.tensor([1.0, 1.0078125])
    router.e_score_correction_bias.copy_(torch.tensor([0.001, 0.0]))
router.weight.data = router.weight.data.bfloat16()  # dtype=torch.bfloat16 load; the bias stays float32
hidden_states = torch.zeros(1, 64, dtype=torch.bfloat16)
hidden_states[0, 0] = 1.0

# DeepSeek's reference Gate: sqrt(softplus(linear(x.float(), weight.float()))), then topk(scores + bias).
scores = F.softplus(F.linear(hidden_states.float(), router.weight.float())).sqrt()
print("reference scores:", scores.tolist(), "-> expert", (scores + router.e_score_correction_bias).argmax(-1).tolist())

with torch.no_grad():
    logits, weights, indices = router(hidden_states)
print("router logits:", logits.dtype, "scores:", F.softplus(logits).sqrt().tolist(), "-> expert", indices.tolist())

On main (95bafdc):

reference scores: [[1.1459763050079346, 1.1484681367874146]] -> expert [1]
router logits: torch.bfloat16 scores: [[1.1484375, 1.1484375]] -> expert [[0]]

With this PR:

reference scores: [[1.1459763050079346, 1.1484681367874146]] -> expert [1]
router logits: torch.float32 scores: [[1.1459763050079346, 1.1484681367874146]] -> expert [1]

Impact

I measured this on DeepSeek-V4-Flash's layer-5 router:

• The tensors are the real checkpoint tensors: layers.5.ffn.gate.weight (bf16), layers.5.ffn.gate.bias (float32) and layers.5.ffn_norm.weight.
• They go through this library's DeepseekV4TopKRouter, with the dtypes that a dtype=torch.bfloat16 load gives.
• The inputs are synthetic: 32,768 tokens of RMSNorm(N(0, 1)) times the real ffn_norm weight, cast to bf16. I did not run the full model on text, so real hidden states may give a different rate.
• Each row is compared with the reference Gate formula in float32.

router	tokens with a different top-6 set	routing-weight relative error on the other tokens (mean / max)
main	8.23% (2,697 / 32,768)	2.86e-3 / 1.37e-2
main, under torch.autocast("cpu", dtype=torch.bfloat16)	8.23% (2,697 / 32,768)	2.86e-3 / 1.37e-2
this PR	0 (0 / 32,768)	1.29e-8 / 3.42e-7
this PR, under torch.autocast("cpu", dtype=torch.bfloat16)	0 (0 / 32,768)	1.29e-8 / 3.42e-7

The remaining 3.42e-7 is float32 rounding. The router takes topk(..., sorted=False), so it adds up the six weights in a different order from the reference.

The hash layers (the first 3 in V4-Flash) pick their experts from tid2eid, so only their weights are affected.

Fix

• Both routers now compute F.linear(flat.type(torch.float32), self.weight.type(torch.float32)), as DeepseekV3TopkRouter does. The scores, the bias addition, the top-k and the normalisation then run in float32, and the routers return float32 logits and weights.
• The matmul runs under maybe_autocast(device_type=device_type, enabled=False), as in Ernie4_5_MoeTopKRouter.
  • Without it, torch.autocast(dtype=torch.bfloat16) runs the matmul in bf16 again. Trainer, through Accelerate, turns that autocast on with bf16=True, outside DeepSpeed.
  • I checked this: with the cast alone (maybe_autocast patched to a no-op), the new test still fails on its autocast case, with torch.bfloat16 != torch.float32.
  • Without autocast, maybe_autocast is a nullcontext.
• The edit is in modular_deepseek_v4.py, and modeling_deepseek_v4.py is regenerated.
  • No other model inherits these two routers.
  • beam, glm5_next, hy_v4, minimax_m3_vl and step3p7 import other V4 classes. They regenerate unchanged.

The routers now hand float32 top_k_weights to the experts. DeepSeek-V3's router already does, so all experts backends take them:

• eager (DeepseekV4Experts.forward), grouped_mm and batched_mm multiply the bf16 expert output by the weights. That promotes to float32, and the result is cast to the hidden dtype (.to(final.dtype) / .to(hidden_states.dtype)).
• deepgemm, the finegrained-FP8 forwards and sonicmoe first cast the weights to the activation dtype (sorted_weights.to(out.dtype), sample_weights.to(proj_out.dtype), top_k_weights.to(hidden_states.dtype)). For them only the final multiply stays in bf16. Expert selection and normalisation are still float32.
• deepgemm_megamoe copy_s the weights into its buffer.

With output_router_logits=True, the recorded router logits are now float32 instead of bf16. load_balancing_loss_func already returned a float32 aux loss. I checked a tiny bf16 V4 with labels and output_router_logits=True, on main and with this PR:

• The loss and the aux loss are float32 on both.
• The backward runs, and the bf16 gate gradient is finite.

Tests

I added DeepseekV4ModelTest::test_router_computes_in_float32:

1. It builds both routers with the dtypes of a bf16 load: bf16 weight, float32 bias.
2. Their outputs (logits, weights, indices) must equal a float32 copy of the router, with and without torch.autocast(torch_device, dtype=torch.bfloat16).
3. It also runs the two-expert case above and checks that expert 1 is picked.

Results (CPU, macOS arm64, Python 3.13.12, torch 2.13.0):

• On main, the test fails with AssertionError: torch.bfloat16 != torch.float32. With this PR, it passes.
• The full test_modeling_deepseek_v4.py gives 122 passed, 195 skipped and 13 failed.
  • All 13 failures are TP/EP tests that raise Tensor parallelism is not supported on MPS devices.
  • Main fails the same 13 on this machine, plus the new test.
• utils/checkers.py passes with ruff_check, ruff_format, init_isort, sort_auto_mappings, noisy_comments, modular_conversion, copies. I did not run make typing (ty is not installed).
• I ran the numbers above on 95bafdc. This branch is on 8090b6a, where no deepseek_v4 file has changed. On 8090b6a, the full test file and the checks give the same results.

I also ran the test file on Linux (Python 3.13.3, torch 2.13.0+cu130), on 8090b6a with this PR's test file and on this PR's head:

• CPU: main gives 1 failed (the new test), 134 passed and 195 skipped. This PR gives 135 passed and 195 skipped. The 13 TP/EP tests that fail on MPS pass on both.
• CUDA, one L4: main gives 1 failed (the new test), 126 passed and 203 skipped. This PR gives 127 passed and 203 skipped.
• On the L4, the impact measurement above, with the same tensors and inputs:
  • Main: 8.24% of tokens (2,699 / 32,768) get a different top-6 set.
  • Main under torch.autocast("cuda", dtype=torch.bfloat16): 2.44% (801 / 32,768). The logits are bf16 and the weights float32.
  • This PR: 0 in both cases, with a max weight relative error of 3.57e-7.

Related

Duplicate search (read-only, 2026-10-27, huggingface/transformers):

query	result
gh search issues --include-prs "DeepseekV4TopKRouter"	#45643 (merged, "Add DeepSeek V4")
gh search issues --include-prs "DeepseekV4HashRouter"	none
gh search issues --include-prs "deepseek v4 router float32" / "… bf16" / "… fp32"	none
gh search issues --include-prs "deepseek v4 routing precision" / "deepseek v4 gate float" / "deepseek v4 expert selection"	none
gh search issues --include-prs "deepseek_v4 router"	#45782 (merged, WeightConverter regex for shared experts)
gh pr list --state open --search "deepseek_v4"	8 open PRs
gh pr list --state open --search "DeepSeek-V4"	25 open PRs
gh pr list --state open --search "deepseek v4 router"	5 open PRs

For each of the 27 distinct open PRs, I listed the changed files (gh pr view --json files). For the PRs that touch deepseek_v4 files, I read the diff (gh pr diff):

• [WIP] modular cleanup #49266 ("[WIP] modular cleanup", conflicting) renames the variables in both V4 router forwards (flat → hidden_states, logits → router_logits). It keeps F.linear(hidden_states, self.weight), so the bf16 router remains. It is a textual conflict only: whichever lands second needs a small rebase.
• DSv4 static cache #46384 (DSv4 static cache) and Finegrained FP8 v4: fused grouepd/batched #46957 (Finegrained FP8 v4) change modular_deepseek_v4.py, but not the routers.
• These change other V4 files, but not the routers:
  • Update deepseek v4 snippet: model id, quant config #46257: the docstring snippet.
  • FP8 on the fine-grained machinery (Fp8Config) #49777: the config.
  • 🚨 MXFP4 on the fine-grained machinery, and DeepSeek-V4's FP8 + MXFP4 plan #49778, fine-grained kernels: monolith PR #49658 and Finegrained kernels: unified quantization kernels api (separate quantizers) #49015: the docs, the config or test_modeling_deepseek_v4.py.
  • I did not check whether their test-file edits sit next to the new test.
• Keep float32 MoE router logits in float32 under autocast #50003 (mine) adds the same maybe_autocast(..., enabled=False) wrapper to the routers that already cast to float32. That is 25 models, including DeepSeek-V3. It changes no deepseek_v4 file: V4 didn't cast, so it was not in that PR.

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: deepseek_v4

[Rocketknight1]

cc @ArthurZucker since you added V4! I checked this out and I think they're correct that upstream uses float32, so this does make router logits noisy for us