Skip to content

Audit and align MoE router precision with reference implementations #3650

Description

@yuhezhang-ai

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 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.

Active work

Any shared sigmoid-scoring default changed by the MiniMax work must be rechecked against BF16-score references such as the DeepSeek V4 learned router. Prefer a model-owned default or an explicit typed control when router families intentionally differ.

Remaining families

Likely small model-owned policy changes using existing router controls

  • DeepSeek V3 and V3.2 — preserve selected routing weights in FP32. V3 already owns the FP32 projection/scoring default; V3.2 bypasses the V3 ForCausalLM constructor, so its normal construction path must also install the model-owned FP32 projection/scoring default while preserving an explicit caller override.
  • Ling / Bailing MoE V2 — default projection/scoring to FP32 and preserve selected weights in FP32.
  • Nemotron H / Nemotron 3 — preserve selected routing weights in FP32.
  • Step 3.5 and Step 3.7 — preserve selected routing weights in FP32 when need_fp32_gate=true.
  • Kimi K2 — follows the DeepSeek V3 router policy; verify it with the shared fix.
  • Kimi K2.5 VLM and Kimi VL A3B — ensure wrapper construction inherits the DeepSeek-style FP32 router policy and preserves selected weights.

F32 checkpoint router storage

The existing strict-FP32 parameter/buffer retention machinery should make these focused changes, but each still needs a checkpoint-load and model-cast test proving the gate remains F32.

Model-specific BF16/de-promotion or routing-semantic work

These are not simple “turn on FP32” fixes. AutoModel currently over-promotes at least one stage or has a router variant whose arithmetic differs from the reference. They may require a model-owned score/projection policy or a narrowly typed shared control.

  • DeepSeek V4 learned and hash routers.
  • Gemma 4 MoE.
  • DiffusionGemma MoE.
  • GPT-OSS.
  • Mistral 4 text router.
  • Laguna S 2.1.
  • HYV3 / Hy3 Preview and Hy-MT2.

Reference required

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

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, projection/scoring arithmetic dtype, selected-weight dtype, and default/explicit-override behavior.
  • Run a tiny numerical router/component comparison against the reference, followed by the relevant end-to-end checkpoint/parity test.
  • 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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions