Audit and align MoE router precision with reference implementations

Summary

A public MoE router precision audit compared AutoModel's effective router behavior with the pinned Hugging Face or checkpoint-owned reference implementations across 188 training, finetune, and benchmark YAMLs.

The audit checks four independent stages:

1. Param: stored router projection parameter dtype after checkpoint loading.
2. Proj: router linear/matmul arithmetic dtype.
3. Score: sigmoid/softmax, correction bias, normalization, and top-k arithmetic dtype.
4. Out: selected routing-weight dtype passed to expert computation.

At the audited head, 9 router families matched all four stages, 23 were mismatched or recipe-dependent, and Inkling could not be independently verified against public reference code. These counts describe the original audit baseline, not the current completion status. These are correctness-sensitive differences: projection or score precision can change top-k expert identity, while parameter/output dtype mismatches can change checkpoint fidelity or expert accumulation.

This issue tracks model-owned fixes and validation. It should not be solved by applying one global FP32 policy: several reference routers intentionally use BF16 for projection, scoring, or returned weights.

Reference and load-convention policy

Checkpoint tensor dtype, reference model code, and observed runtime behavior answer different questions and must be recorded separately. A Hugging Face model loaded with an explicit model-wide BF16 dtype can legitimately downcast an F32 checkpoint outlier; that loaded model is evidence for that named harness convention, not proof that the checkpoint intended the parameter to be BF16.

Production serving implementations such as vLLM, SGLang, and TensorRT-LLM are useful independent references when they have a model-specific implementation, but they are not universal precision authorities either. Record the exact runtime/version, load options, and observed Param / Proj / Score / Out chain instead of inferring per-parameter behavior from a global dtype=auto setting.

For MiniMax M2 specifically, vLLM's native implementation at 6a9c69f explicitly allocates the gate parameter in FP32, requests FP32 gate output/router logits, and keeps the correction bias in FP32. This independently corroborates the released checkpoint's F32 gate storage and the checkpoint-faithful FP32 router target. This is a model-specific override: vLLM's general dtype=auto policy otherwise selects a model-wide weight/activation dtype.

Progress — Mistral4 and GPT-OSS updated 2026-09-18

The September 18 refresh covers Mistral4 merge/validation status and GPT-OSS reference evidence. Other family entries retain their previous status pending a separate refresh.

Checked entries below record merged implementation scope. They do not imply full-depth HF parity has been validated for every variant; evidence limits are recorded separately.

Merged implementation

• Mistral 4 — routing semantics, precision, and checkpoint reference — NVIDIA-NeMo/Automodel#3865 merged September 14 as 4e00f6be078fbaf171272a775115023790284c78. It preserves selected routing weights and rotary state in FP32, corrects attention scaling, and disables default learning of DeepSeek-style correction biases that HF Mistral4 does not consume. The opt-in hf_reference_compute_fp32 covers HF RMSNorm, RoPE, router scoring, and expert accumulation. The final PR records a published-checkpoint 36-layer replay with identical weights and expert selections on shared inputs, focused output/gradient tests, and all six scoped checkpoint phases passing with exact AutoModel reload. Cross-framework tolerance exceptions remain; this is not bitwise vanilla-HF parity. Previously trained checkpoints with nonzero correction biases are not repaired by disabling future updates. Upstream HF improvements and validation for other families remain follow-up work.
• GLM-4.5/4.7 MoE, GLM-4.7 Flash/Lite, and GLM-5/5.1/5.2 DSA — fix(glm): align and diagnose cross-framework router parity merged on August 28. Model-owned defaults run projection/scoring in FP32 and preserve selected weights in FP32 while keeping BF16 checkpoint gate storage. All four audited stages match for normal ForCausalLM/recipe construction.
• MiniMax M2/M2.7 — fix(minimax): repair HF parity reference and gate at the measured envelope merged on August 27. The change allocates and retains the gate weight in F32, runs projection/scoring in FP32, preserves selected weights in FP32, and pins the HF reference to the checkpoint-faithful contract. The PR records full checkpoint-robustness validation for M2.7.
• MiniMax M3 — checkpoint router storage — fix(minimax): repair HF parity reference and gate at the measured envelope merged the F32 gate allocation/retention fix and AutoModel checkpoint-robustness coverage. This entry covers storage; selected-weight output precision remains open below.
• DeepSeek V3/V3.2 and Kimi K2 — router policy — fix: preserve selected routing weights in FP32 in deepseek merged on September 3. V3 and V3.2 preserve selected weights in FP32; V3.2 installs the FP32 gate default on its own construction path. Explicit overrides remain supported, caller-owned backend configs are not mutated, and focused tests cover Kimi K2's inherited policy. test(deepseek): update router parity for transformers 5.15, merged September 5, updates the parity test for Transformers 5.15.1.

New contributions should start from current main; there is no longer a need to wait for NVIDIA-NeMo/Automodel#3674 or NVIDIA-NeMo/Automodel#3746 to merge.

Validation boundaries and follow-up

• DeepSeek-family full-depth source parity: NVIDIA-NeMo/Automodel#3746 records a full-model Moonlight distributed training smoke and an exact-head first-four-layer parity diagnostic. Full 27-layer Moonlight final-logits parity was explicitly still pending in that report. These results are not full-depth validation of the larger V3/V3.2/Kimi K2 checkpoints.
• MiniMax M3: NVIDIA-NeMo/Automodel#3674 records passing AutoModel train/save, reload, and resume coverage. Its vanilla-HF phases were skipped for documented reference-resource constraints; M2.7 serves as the family's cross-framework representative. Do not infer full M3 HF parity from the storage fix or AutoModel-only validation.

Any further shared score-arithmetic change must be checked against families with different reference precision contracts, including the DeepSeek V4 learned router. Prefer a model-owned default or an explicit typed control when router families intentionally differ.

Remaining precision work

Model-owned policy changes using existing router controls

• Ling / Bailing MoE V2 — @edjson is taking this scope; see the accepted scope and validation split. Add the model-owned FP32 projection default and preserve selected routing weights in FP32. Shared sigmoid scoring already defaults to FP32 on current main; verify Param / Proj / Score / Out independently rather than describing all three compute stages as currently BF16.
• Nemotron H / Nemotron 3 — preserve selected routing weights in FP32. Related NVIDIA-NeMo/Automodel#3845 is open and protects exact FP32 correction-bias values during force_hf checkpoint loading; that loading fix does not complete this Out-stage item.
• MiniMax M3 — selected routing weights — preserve FP32 selected weights at the Out stage; separate from the merged storage work in NVIDIA-NeMo/Automodel#3674.
• Step 3.5 and Step 3.7 — preserve selected routing weights in FP32 when need_fp32_gate=true.
• Kimi K2.5 VLM and Kimi VL A3B — verify wrapper construction inherits the DeepSeek-style FP32 router policy and preserves selected weights. Kimi K2 coverage in NVIDIA-NeMo/Automodel#3746 does not by itself establish these VLM wrapper paths.

F32 checkpoint router storage

The MiniMax storage fixes are merged above. Remaining storage work needs checkpoint-load and model-cast tests proving that both the required dtype and original values survive.

• MiMo V2 Flash — NVIDIA-NeMo/Automodel#3818 proposed FP32 router parameter/output preservation and focused numerical tests, but was closed without merging. This item remains unresolved.

Reference validation before BF16/de-promotion or routing-semantic changes

The entries below are audit discrepancies, not confirmed defects or approved precision reductions. “Over-promoted” means higher precision than the audited HF implementation/load convention; it does not establish that AutoModel is wrong. HF may not represent the model author's intended mixed-precision behavior, so matching HF alone is insufficient justification to decrease AutoModel's precision.

Before de-promotion, establish the model-specific contract from author guidance/reference code and checkpoint evidence, then cross-check pinned vLLM/SGLang or other relevant inference implementations. Record Param / Proj / Score / Out separately, including load options and backend. If references disagree, retain the existing higher precision while resolving the discrepancy; a BF16 change needs model-specific numerical/gradient and relevant end-to-end evidence.

The NVIDIA-NeMo/Automodel#3865 motivates this recheck: HF 5.15.1 uses BF16 router scoring under a BF16 load, while the audited vLLM CUDA router scores in FP32 and SGLang's normal CUDA path produces FP32 router logits. This is evidence to question HF as the sole reference for every candidate below, not proof that every listed HF implementation is incorrect or that FP32 is always the intended contract.

For checkpoint CI, prefer an explicitly documented reference comparison with aligned precision when practical. Otherwise, a measured model-specific HF tolerance exception may be appropriate, with a clear rationale and follow-up; do not lower production precision solely to pass HF parity. Preserve independent checkpoint round-trip/resume checks and focused FP32 numerical coverage. Mistral4's implementation and scoped validation are now merged as recorded above. Upstream HF reference improvements and further reduction of its documented tolerance exceptions remain follow-up work.

• DeepSeek V4 learned and hash routers.
• Gemma 4 MoE.
• DiffusionGemma MoE.
• GPT-OSS — scoring contract confirmed different; full-model cause unresolved. Under the audited BF16 load/default backend, HF 5.15.1 uses BF16 / BF16 / BF16 / BF16 for Param / Proj / Score / Out; AutoModel at e4f0cc9cd01f0f8ba51184dc3ec09578309deda6 uses BF16 / BF16 / FP32 / BF16. This is specifically selected-softmax precision, not FP32 router projection or FP32 weights passed to experts. HF source, AutoModel score default and routing. These are tensor/operation dtype contracts; BF16 softmax may still use FP32 internal kernel accumulation. OpenAI's own PyTorch reference at 7b583341 also uses a BF16 gate and selected softmax without FP32 promotion, so this is not established as an HF precision defect. Both select top-k logits before softmax; changing only that softmax dtype cannot change that router's expert indices for identical logits. Both return BF16 weights, so differing score precision need not produce different forward weights, while backward arithmetic can differ. Validate the pinned GPU forward/backward paths and full-model effects before reducing production precision, changing the reference, or attributing HF-reload KL to this discrepancy.
• Mistral 4 text router — model-owned policy and precision-aligned checkpoint-reference work merged in #3865; see merged implementation and remaining upstream-reference/tolerance follow-up above.
• Laguna S 2.1.
• HYV3 / Hy3 Preview and Hy-MT2.

Follow-up: higher-precision Hugging Face references

Mistral4 merged example: NVIDIA-NeMo/Automodel#3865 merged September 14. The explicit hf_reference_compute_fp32 mode promotes HF RMSNorm, RoPE, router scoring, and expert accumulation while retaining native parameter/projection dtypes. Pinned vLLM CUDA/FlashInfer source paths corroborate the selected precision policy. The merged change also fixes attention scaling, rotary preservation, selected routing-weight dtype, and training-time correction-bias semantics; it is not evidence that the original discrepancy was purely benign BF16 drift. Focused component/gradient tests, full-depth published-checkpoint diagnostics, all six scoped checkpoint phases, and final PR CI are recorded as passing. This mode currently supports Mistral4 only; other families require their own contract and numerical validation.

The harness already has _hf_fp32_module_names / _keep_hf_modules_in_fp32 in tests/functional_tests/checkpoint_robustness/test_checkpoint_robustness_llm.py, used by both source-load and HF-reload paths to preserve declared FP32 tensors during loading. AutoModel's cast_model_to_dtype separately preserves declared FP32 values across model casts. Reuse this infrastructure; the missing control is the reference's compute policy, not an entirely new loader.

Tensor storage and router arithmetic must remain separate. FP32 weights alone do not force FP32 projection with BF16 activations. FP32 scoring can instead start at softmax/sigmoid while projection remains BF16; for example, an explicit softmax(..., dtype=torch.float32) also keeps subsequent score normalization in FP32 unless cast back. Full FP32 projection additionally needs compatible input/weight dtypes and an autocast policy that does not lower them.

• Prefer upstream HF fixes: once author/reference evidence establishes the intended per-stage precision, submit focused model-specific Transformers PRs with regression tests for routing selections/weights and gradients. Do not propose a global FP32 policy or assume full FP32 projection is required when only scoring differs.
• Provide a scoped harness reference mode where needed: reuse load-time FP32 preservation, then apply explicit, model-specific HF compute handling in the test harness while the rest of the reference stays BF16. Cover both source and exported-checkpoint HF inference, including actual device-map and dequantization paths. Label this as a precision-aligned reference, retain the vanilla-HF comparison separately, and avoid silent/global arithmetic patches or replacing HF math with AutoModel math.
• Validate the complete contract: prove exact precision-critical checkpoint values survive loading and verify actual Param / Proj / Score / Out dtypes, expert choices, outputs, and input/gate gradients. Exercise BF16 activations, relevant autocast behavior, and scoped end-to-end checkpoint CI; component-only FP32 passes are insufficient.
• Revisit temporary tolerances: after the higher-precision reference is validated, remeasure and tighten/remove the model-specific HF exceptions where justified. Keep production precision and independent checkpoint round-trip/resume checks intact.

Reference required

• Inkling — obtain an authoritative router implementation or numerical oracle before declaring its current BF16/FP32 behavior correct or changing it.

Ling validation plan

• Contributor coverage: tiny Mini/Flash/1T construction tests through the normal model path, default and explicit-override behavior, and deterministic grouped-router comparisons of selected experts, routing weights, and input/gate gradients against a pinned reference. Preserve the checkpoint's parameter-storage contract separately from the compute-precision policy.
• Real-checkpoint representative: Ling Mini is sufficient for the contributor's initial checkpoint coverage. _real_forward_smoke.py is supplemental: it passes an explicit MoEConfig and checks loading/finite output, so it does not by itself prove the new model defaults.
• Maintainer GPU coverage: @yuhezhang-ai can help validate the final PR commit with full-depth Ling Mini HF-vs-AutoModel source parity using Phase 0 of the checkpoint-robustness harness, followed by a short distributed training check. Verify reference load conventions and exact correction-bias values, including the HF expert_bias to AutoModel e_score_correction_bias naming difference. The remaining train/save/reload/resume phases provide broader regression coverage when checkpoint or casting behavior changes.
• Larger variants: use targeted router/checkpoint-tensor or reduced-layer checks as needed; the contributor need not run full Flash or 1T models. A reduced-layer 1T check must include a MoE layer: its first four layers are dense. Label reduced-layer evidence separately from full-depth parity.

Acceptance criteria per family

• Identify an immutable reference implementation and, when relevant, checkpoint tensor dtype evidence.
• State every reference load convention explicitly; do not treat a framework loaded with a model-wide dtype override as a checkpoint-storage oracle.
• Where a model-specific production implementation exists, cross-check a pinned serving runtime and record its actual loaded parameter and routing arithmetic dtypes.
• Match Param / Proj / Score / Out independently; do not infer full correctness from FP32 logits alone.
• Keep model-specific defaults in the owning model package. Preserve explicit caller overrides where supported.
• Add focused tests for stored gate dtype and exact precision-critical values, projection/scoring arithmetic dtype, selected-weight dtype, and default/explicit-override behavior.
• Run a tiny numerical router/component comparison against the reference, including input/gate gradients, followed by the relevant end-to-end checkpoint/parity test. Finite training loss or first-few-layer agreement alone does not establish full-depth logits parity.
• Update this checklist and link the fixing PR plus exact-head validation evidence.

Already exact in the audit

No action is currently required for Qwen2 MoE fallback, ERNIE 4.5 MoE, Qwen3 MoE, Qwen3 Next, Qwen3.5-family MoE, Qwen3 Omni MoE, Qwen3 VL MoE, Kimi K3, or Kimi Linear. Re-audit them if shared router behavior changes.

Deferred follow-up: rotary-frequency precision

The Mistral4 checkpoint-parity investigation also identified a rotary-buffer precision issue. NVIDIA-NeMo/Automodel#3865 merged the model-specific fix together with attention-scaling and router-policy changes. The broader rotary audit below remains open; Mistral4's completed work does not validate other model families.

Why include this in the precision audit

The Mistral 4 multimodal initializer rebuilt rotary frequencies in FP32 and then rounded them to BF16 when casting the wrapper. Preserving their dtype only after rounding cannot recover the original values. The new regression exercises the real initializer with the published model's 64 rotary dimensions; the previous attention-only fixture used two rotary dimensions and bypassed wrapper initialization. The final merged Mistral4 PR records passing scoped coverage of all six checkpoint phases; the broader audit below remains open.

The absence of freqs_cis from a preservation list is not sufficient evidence of a bug. Other implementations protect the parent rotary_emb, preserve values in _apply, recompute frequencies after casting, or cache frequencies outside registered buffers. Final cos/sin outputs may legitimately use the activation dtype.

Starting evidence — 2026-09-11

Representative CPU probes at AutoModel 6effd77b32d2944dd7a9574ab4b67f102c3e341e, with Transformers 5.15.1 and PyTorch 2.10.0:

• DeepSeek V3/R1 and V3.2: real BF16 initializers reconstruct FP32 frequencies with exact values. This covers the tested initialization paths, not every subsequent cast/load path.
• Llama/Qwen2 rotary implementation: the buffer can become BF16, but cache construction recomputes from configuration; the Llama probe retained exact positional outputs at positions 0, 2047, and 8192.
• GPT-OSS rotary implementation: the cached FP32 frequency tensor survives a module-wide BF16 cast because it is not a registered buffer.
• Qwen3-Omni, MiMo-V2-Flash, DiffusionGemma: their three existing rotary-precision regression tests passed.
• Ministral 3 — follow-up candidate: a real tiny Ministral3ForCausalLM followed by cast_model_to_dtype(..., bfloat16) rounds inv_freq; positional cosine/sine differences reach approximately 1.98/1.38 across the tested positions. This is an isolated cast-path reproduction, not proof that normal checkpoint loading or CI is affected. Establish whether those paths reinitialize the buffer afterward before proposing a fix.

Audit checklist

• Inventory precision-critical rotary buffers and their actual construction, initialization, casting, and checkpoint-loading paths, including VLM wrappers and distributed materialization.
• Verify original FP32 values survive, or are reconstructed from configuration before use; checking final dtype or preservation declarations alone is insufficient.
• Compare positional outputs with a pinned reference after the real lifecycle operations, using nontrivial rotary widths and positions near and beyond the original context boundary.
• Trace the Ministral 3 candidate through normal checkpoint loading and determine whether the cast vulnerability is reachable in supported training/reload flows.
• For confirmed defects, add model-owned fixes and focused regression coverage, followed by relevant forward/backward and checkpoint-parity evidence. Record component-only versus full-model evidence separately.

Retain each model's reference precision contract; do not apply a global FP32 policy to all rotary outputs or router stages.

Start from current main and the accepted Ling/Bailing scope in the linked comment, then inspect the model-owned router construction paths and existing need_fp32_gate controls. Verify Param, Proj, Score, and Out independently, including checkpoint loading and model-cast behavior; done means the specified families have focused numerical or checkpoint tests and documented reference evidence.