fix(minimax): repair HF parity reference and gate at the measured envelope by yuhezhang-ai · Pull Request #3674 · NVIDIA-NeMo/Automodel

yuhezhang-ai · 2026-08-25T21:02:45Z

What

Closes the MiniMax M2.7 checkpoint-robustness failures (AMINT-286): the catastrophic cross-framework KL (13.1) was an invalid reference, the resume drift matches the shared routed-MoE envelope, and the residual ~0.09 KL is knife-edge-router amplification of bf16 kernel noise — proven not to be an implementation defect — now gated at its measured envelope.

Changelog

Reference repair (root cause of the 13.1): Transformers 5.12's in-tree MiniMaxM2Config drops the legacy rotary_dim field, so the vanilla reference silently ran full-head RoPE instead of partial (64/128). The harness now derives rope_parameters["partial_rotary_factor"] = rotary_dim / head_dim for legacy configs (13.1 → 0.086). Upstreamed as MiniMaxM2 silently applies full-head RoPE: config drops the checkpoints' legacy rotary_dim field huggingface/transformers#48241 / #48242.
Checkpoint-faithful fp32 router: the released checkpoint stores gate.weight and e_score_correction_bias in fp32; AutoModel now scores sigmoid routing in fp32, allocates fp32-contract gate weights in fp32 at construction (keeps FSDP dtype groups uniform), and the harness pins the HF reference's fp32 tensors across vanilla-layout differences. Gate metrics are unchanged by this (the residual was never router precision), but the model now matches the checkpoint contract.
Recipe: shared relaxed resume envelope (restored state and first forward are exact; 0.012–0.020 loss drift at continuation steps 6–7 across three runs), and measured parity_threshold_overrides for the two cross-framework gates on the full 2048-token document (mean_kl 0.15 / p95_kl 0.5 / cos 0.95 vs measured 0.086–0.094 / 0.346–0.375 / 0.964–0.965 across scoped CI 64182587, 64344786, 64351027). Overriding all three metrics numerically shadows the profile, so no relaxed-profile selection is used.
MiniMax-M3-VL: the released M3 checkpoint stores its router tensors exactly like M2.7 (fp32 gate.weight; fp32 1e-3-quantized bias lattices at magnitudes 4.7–7.9, measured via ranged Hub reads). The same fp32 contract is ported to minimax_m3_vl (fp32 gate allocation + weight pin), and the previously uncovered M3-VL LoRA recipe is enrolled in the checkpoint-robustness gates for the AutoModel side: train/save, bitwise reload, and relaxed-envelope resume. Green end to end in the shipped configuration: nemo-ci pipeline 64755685 (job 413579005, 31 min: train 401s / reload 243s bitwise-exact / resume 359s, harness status PASS on head 0728023). The two vanilla-HF phases are skipped: the 427B (854 GiB bf16) single-node reference OOMs uncapped and, capped, CPU-offloads ~400 GiB while idling all 64 job GPUs for hours per run — M2.7 remains the family'''s blocking cross-framework sentinel with the identical measured router lattice.
Harness fixes surfaced by the new coverage: (1) AutoModel'''s CONFIG_MAPPING registrations shadow in-tree config classes, so built-in references paired a nemo-owned config with the in-tree model and crashed (temporal_patch_size); _replace_nemo_owned_reference_config now falls back to Transformers''' own CONFIG_MAPPING_NAMES. GLM-4.7-Flash is verified unaffected (its checkpoint is glm4_moe_lite, never hijacked). DeepSeek-V4-Flash is the other affected type, and a local meta-device A/B proves the change strictly recovers a broken reference there too: the in-tree DSV4 model cannot even construct with the nemo-owned config ('DeepseekV4Config' object has no attribute 'layer_types'), while the in-tree config constructs cleanly — no working gate changes numbers. (2) hf_device_map_max_memory_gib/_cpu_max_memory_gib only reached the phase-3 reload; they now also apply to the phase-0 source-load reference (unit-tested).

Why the envelope is a closure, not a workaround

Full-weight local investigation on the real 230 GB checkpoint (62-layer streaming sweep, single GPU):

The in-tree HF reference against itself (eager vs sdpa attention) reaches mean KL 0.237 / cos 0.973 at full depth through the same expert-flip cascade (~90 % of tokens re-routed by layer 59) that drives AutoModel's 0.320 / 0.962 in the same protocol — the reference's own kernel-choice variance is the same order as the AutoModel gap.
Per-layer matched-input deltas are at bf16 ulp scale in every arm; TE and SDPA attention are bit-identical at layer 0; an end-to-end fp32 router leaves the numbers unchanged; HF-rounding-order RMSNorms make the first layernorm bitwise equal and cut 1-layer KL 6×, yet the effect is fully erased by depth 62.
Router evidence (pipeline 64182587): 68 % of tokens flip an expert within five layers, all flips inside the measured score-noise band, replacements symmetric, no token unflipped across 62 layers.

Real conversion or model-math regressions stay loud: the rope bug measured mean KL 13.1, two orders above the envelope.

Validation

MiniMax M2.7 job green end to end at commit 695a9bc (nemo-ci pipeline 64546388); the commits added since (M3-VL contract/coverage, reference-config resolution, phase-0 memory caps, fp32-alias fix) do not touch the M2.7 path — its minimax_m2 model_type is not registered, so the config-resolution change never fires for it. Full gate readings: (source_load 0.0856 / 0.3565 / 0.9661 and hf_export_reload 0.0890 / 0.3396 / 0.9651 against the 0.15 / 0.5 / 0.95 envelope; all self-repeats and the AutoModel reload bitwise-exact; resume within the shared relaxed envelope).
Scoped CI green end to end (nemo-ci pipeline 64481928, exact job minimax_m2.7_hellaswag_lora): all six phases pass — source_load 0.0860 / 0.3448 / 0.9644 and hf_export_reload 0.0869 / 0.3524 / 0.9646 against the 0.15 / 0.5 / 0.95 envelope; every self-repeat and the AutoModel reload bitwise-exact; resume within the shared relaxed envelope.
Unit tests: 218 passed, 3 skipped (harness kwargs, MiniMax router precision, MoE layers).

The router/shape diagnostics used during the investigation are deliberately NOT enabled in the recipe: the diagnosis is complete and its evidence is archived (AMINT-286, the 62-layer sweep data, and the router-flip reports from pipeline 64182587), so report-only capture would add per-run cost without decision value. The MiniMax router-capture extension remains recoverable from the investigation branch if ever needed again.

Note: this PR and #3659 add functions in adjacent regions of the same harness file; whichever merges second needs a trivial context-only rebase.

🤖 Generated with Claude Code

copy-pr-bot · 2026-08-25T21:02:49Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

yuhezhang-ai · 2026-08-25T21:06:06Z

/ok to test 695a9bc

…erences Transformers 5.x in-tree configs keep MiniMax-M2's legacy rotary_dim field as a plain attribute while the in-tree model reads only rope_parameters["partial_rotary_factor"], so the checkpoint-robustness vanilla-HF reference silently rotates the full 128-dim head for a checkpoint trained with 64-dim partial RoPE. The reference is deterministic but invalid, which produced the catastrophic AMINT-286 source/export cross-framework failures (mean KL 13.1/6.3, cosine 0.02) while AutoModel matched the checkpoint's own modeling semantics exactly (tiny CPU fp32 parity: KL 3e-17; injecting the factor makes unmodified in-tree Transformers match to 7e-7). Derive the missing factor as rotary_dim / head_dim when a loaded reference config carries the legacy field without a partial factor, for both the source-load and consolidated-export HF reference loads. The repair is a strict no-op for configs without the legacy spec, and the repaired config is only passed explicitly to from_pretrained when a repair actually happened, so every other model's load call is unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

Native-resume restores exact model/buffer/optimizer state and an exact first forward, but three independent scoped-CI runs measured 0.012-0.020 routed-MoE loss drift at continuation steps 6-7 against the standard 0.0117 allowance (pipelines 63311388, 63340676, 64093591), always within the shared relaxed envelope (0.0432). Select the relaxed resume profile for this routed-MoE LoRA config, matching the Nemotron hybrid-MoE chat precedent; all logit gates stay on the standard profile (AMINT-286). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

The softmax scoring path already defaults to fp32 (softmax(dtype=gate_precision or float32)), but both sigmoid paths applied sigmoid to the raw bf16 gate output. HF sigmoid-router references (MiniMax-M2, GLM, DeepSeek families) compute sigmoid(logits.float()), and bf16 sigmoid quantizes scores at ~2e-3 — the same order as the 1e-3 e_score_correction_bias lattices those checkpoints ship — flipping 7-13% of top-k selections per layer against the fp32 reference on MiniMax-M2.7's real bias tensors. This is one of the two contributors to the stable ~0.086 cross-framework mean-KL floor in AMINT-286 (softmax-router models in the same gate set sit at 0.001-0.003). Score sigmoid and sigmoid_with_bias in gate_precision-or-fp32, exactly mirroring the softmax branch; selection, gathering, and top-k normalization now run in fp32 and the final weights keep the existing cast back to the input dtype, which matches the HF references' top_k_weights.to(hidden dtype). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

…ences _hf_fp32_module_names mirrors AutoModel's _keep_in_fp32_modules_strict onto the vanilla-HF reference load, but the entries use AutoModel module paths: in-tree Transformers MiniMax-M2 keeps e_score_correction_bias directly on mlp (not mlp.gate), so the injected name never matched and the reference's router bias was silently cast to bf16. MiniMax-M2.7's biases are fp32 lattices with ~1e-3 gaps at magnitudes up to 8.6, where a bf16 ulp is 0.031: 215-242 of 255 adjacent expert pairs collide exactly after the cast, scrambling 30-73% of top-8 selections per layer against the checkpoint's true ordering — the dominant contributor to the weight-independent ~0.086 cross-framework KL floor in AMINT-286 (phase 0 = 0.0860 vs phase 3 = 0.0885). Register the distinctive leaf component of each dotted strict name as well, so any vanilla layout keeps the tensor in fp32; generic weight/bias leaves are excluded. Verified at tiny scale: the AutoModel-path entry alone leaves the in-tree bias bf16, the leaf entry pins it fp32 while the gate weight stays bf16. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

Released MiniMax-M2 checkpoints store the router gate weight in fp32 (verified from the M2.5/M2.7 safetensors headers) and the HF reference projects with hidden_states.to(weight.dtype), so the checkpoint-faithful router is fp32 at every stage: fp32 parameter, fp32 projection, fp32 scoring, fp32 selected weights. AutoModel allocated the gate in model dtype and downcast the fp32 checkpoint tensor to bf16 at load. Follow the ERNIE pattern: default backend.gate_precision to fp32 (explicit overrides preserved), keep mlp.gate.weight in the strict fp32 contract next to the existing e_score_correction_bias entry (which also pins the vanilla-HF reference's gate weight through the harness's mirrored contract), and return fp32 selected weights (router_weights_fp32, as GLM4-MoE-Lite already runs in CI). Measured motivation (AMINT-286, pipeline 64308157): with bf16 routers on both sides, cross-framework KL at the 128-token gate is 0.106 — 13x the HF reference's own shape-noise floor (0.0083) — because knife-edge bias-lattice routing flips saturate within the first five layers; an fp32 router on both sides removes projection- and weight-rounding from that flip-noise budget. Adds a router-precision contract test (fp32 gate parameter and bias survive the model-wide bf16 cast; fp32 selected weights; override preserved). The router-diagnostics roundtrip test now seeds and initializes its AutoModel properly and forces eager compilation for the newly-specialized fp32-weight expert kernel on CPU. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

The fp32-router pilot failed EP sharding on all 32 ranks (pipeline 64344786): 'FSDP could not isolate parameters with a distinct dtype from siblings in the same module: mlp.gate.weight'. Root cause, reproduced exactly on CPU: FSDP shards the freshly constructed module, where the gate weight still carries its allocation dtype (MoEConfig.dtype, bf16) while its compute dtype is pinned fp32 and its sibling e_score_correction_bias buffer is fp32 storage — a mixed (storage, compute) group inside one module that the dtype-aware sharder cannot isolate. The fp32 storage contract only materialized later (initialize_weights / checkpoint cast), which is why local single-process construction paths passed. Add MoEConfig.gate_dtype (default None inherits dtype, backward compatible) so models whose checkpoints store the gate in fp32 allocate it fp32 from birth on every construction path, and set it for MiniMax-M2. A regression test shards a freshly constructed block through fully_shard_by_dtype and pins the construction-time dtypes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

…velope Full-weight local investigation (AMINT-286) localized the residual cross-framework KL to knife-edge-router amplification of bf16 kernel noise rather than an implementation defect. A 62-layer streaming sweep on the real checkpoint shows the in-tree HF reference against itself (eager vs sdpa attention) reaches mean KL 0.237 / cos 0.973 through the same expert-flip cascade (~90% of tokens re-routed by layer 59) that drives AutoModel's 0.320 / 0.962; per-layer matched-input deltas stay at bf16 ulp scale in every arm, TE and SDPA attention are bit-identical at layer 0, and an end-to-end fp32 router leaves the gate metrics unchanged. Gate the full 2048-token document at the measured envelope (scoped CI 64182587/64344786/64351027: mean KL 0.086-0.094, p95 0.346-0.375, cos 0.964-0.965) via numeric threshold overrides; overriding all three metrics shadows the profile, so no relaxed-profile selection is needed. Real conversion or model-math regressions remain loud: the repaired rope reference bug measured mean KL 13.1, two orders above this envelope. Validated green end to end on nemo-ci pipeline 64481928 (all six phases). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

… M3-VL The released MiniMax-M3 checkpoint stores its router tensors exactly like MiniMax-M2.7: gate.weight in fp32 and e_score_correction_bias as fp32 1e-3-quantized lattices at magnitudes 4.7-7.9 (measured via ranged reads of the Hub shards; ~70-76% of adjacent expert biases collide in bf16). Port the M2.7 checkpoint-faithful contract to minimax_m3_vl: allocate the gate weight fp32 (gate_dtype) so FSDP dtype groups stay uniform, and pin mlp.gate.weight in _keep_in_fp32_modules_strict alongside the existing bias pin. Enroll the M3-VL LoRA recipe in the checkpoint-robustness gates (previously uncovered; the VL wrapper, processor-driven text parity path, and mixed dense/sparse decoder are distinct coverage from the M2.7 LLM). Cross-framework thresholds start from M2.7's measured envelope given the identical router lattice, to be calibrated to this recipe's own scoped-run measurements. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

…ions AutoModel's CONFIG_MAPPING registrations also shadow in-tree config classes: for built-in references (trust_remote_code=false) the harness paired an AutoModel-owned config with the in-tree model, which crashes on attribute contracts the local class does not carry — the first M3-VL robustness run failed constructing the in-tree vision tower ('MiniMaxM3VLVisionConfig' object has no attribute 'temporal_patch_size', pipeline 64577484). Extend _replace_nemo_owned_reference_config to resolve the in-tree class from Transformers' CONFIG_MAPPING_NAMES (which registration cannot shadow) when the remote-code auto_map path does not apply. Besides minimax_m3_vl this changes the vanilla-reference config resolution for the robustness-gated glm_moe_dsa and deepseek_v4 recipes (the other registered-and-in-tree types; mistral4 already resolves in-tree, hy_v3 is not gated). In-tree model + in-tree config is the self-consistent vanilla pair a transformers user gets without AutoModel imports, so this is the faithful reference for all of them; their gate readings should be watched on the next full run. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

Deriving a leaf alias from every _keep_in_fp32_modules_strict entry leaked generic names into the vanilla reference's dtype plan: Gemma4's router.proj / router.scale entries produced 'proj' and 'scale', and Transformers matches these as unanchored substrings, so q_proj, down_proj, and friends would load fp32 inside a bf16 reference (dtype errors, doubled memory, invalid parity). Register leaf aliases only from an explicit allowlist — e_score_correction_bias, the one leaf whose vanilla parent path actually differs (in-tree MiniMax-M2 hangs it on mlp, not mlp.gate). Full AutoModel paths still pass through unchanged. Addresses the P1 review finding on PR #3674. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

yuhezhang-ai · 2026-08-26T17:22:45Z

/ok to test 6fb4b3d

…ntime The robustness matrix measured 1028s of phase time on the first scoped run (job 412098422; ~32 min GitLab-side including setup outside the Slurm window). 45 min covers it with headroom, versus the 04:00:00 window sized for the now-skipped CPU-offloaded HF phases. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

yuhezhang-ai · 2026-08-26T17:40:47Z

/ok to test 0728023

akoumpa · 2026-08-26T22:21:28Z

+        # and fp32 selected weights. Keep that default while preserving an
+        # explicit backend override (see AMINT-286; ERNIE follows the same
+        # pattern for its fp32 router).
+        if backend.gate_precision is None:


🤔 would it make sense to have this if -> set default in the __post_init__ of BackendConfig? I'm wondering if that would break anything, otherwise looks like it would us some headache.

I think None needs to remain the BackendConfig default. It means the caller did not specify a gate-compute dtype, allowing each model to install its reference-compatible policy while preserving an explicit override.

Making FP32 the global default in BackendConfig.__post_init__ would change projection arithmetic for router families whose references intentionally project in BF16, including Qwen3 MoE, DeepSeek V4, GPT-OSS, Mistral 4, and Laguna. #3650 specifically found that a global FP32 router policy would be incorrect.

For MiniMax, the F32 checkpoint gate and FP32 routing chain are model-owned, so I think this conditional belongs here. We could consider a follow-up helper to reduce the repeated immutable replace(...) pattern, but it should still receive the model-specific default explicitly and preserve None globally.

yuhezhang-ai requested a review from a team as a code owner August 25, 2026 21:02

yuhezhang-ai mentioned this pull request Aug 25, 2026

fix(minimax): repair HF parity reference and gate at the measured envelope #3643

Closed

yuhezhang-ai mentioned this pull request Aug 25, 2026

fix(test): unblock the Kimi-Linear vanilla-HF parity reference #3659

Merged

copy-pr-bot Bot temporarily deployed to nemo-ci August 25, 2026 21:06 Inactive

copy-pr-bot Bot temporarily deployed to test August 25, 2026 21:06 Inactive

copy-pr-bot Bot deployed to public August 25, 2026 21:06 Active

copy-pr-bot Bot deployed to public August 25, 2026 21:09 Active

copy-pr-bot Bot temporarily deployed to nemo-ci August 25, 2026 21:11 Inactive

yuhezhang-ai and others added 8 commits August 25, 2026 17:18

yuhezhang-ai force-pushed the yuhez/fix/amint286-minimax-rope-reference branch from 695a9bc to a83a083 Compare August 26, 2026 00:47

copy-pr-bot Bot temporarily deployed to nemo-ci August 26, 2026 14:55 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci August 26, 2026 17:23 Error

copy-pr-bot Bot temporarily deployed to nemo-ci August 26, 2026 17:23 Inactive

copy-pr-bot Bot had a problem deploying to test August 26, 2026 17:23 Error

copy-pr-bot Bot deployed to public August 26, 2026 17:23 Active

copy-pr-bot Bot deployed to public August 26, 2026 17:25 Active

copy-pr-bot Bot deployed to public August 26, 2026 17:26 Active

copy-pr-bot Bot temporarily deployed to nemo-ci August 26, 2026 17:26 Inactive

yuhezhang-ai force-pushed the yuhez/fix/amint286-minimax-rope-reference branch from 957ecfc to 0728023 Compare August 26, 2026 17:39

copy-pr-bot Bot deployed to public August 26, 2026 17:41 Active

copy-pr-bot Bot deployed to test August 26, 2026 17:41 Active

copy-pr-bot Bot temporarily deployed to nemo-ci August 26, 2026 17:42 Inactive

akoumpa reviewed Aug 26, 2026

View reviewed changes

yuhezhang-ai mentioned this pull request Aug 26, 2026

Audit and align MoE router precision with reference implementations #3650

Open

19 tasks

Navigation Menu

Uh oh!

fix(minimax): repair HF parity reference and gate at the measured envelope - #3674

fix(minimax): repair HF parity reference and gate at the measured envelope#3674
yuhezhang-ai wants to merge 14 commits into
mainfrom
yuhez/fix/amint286-minimax-rope-reference

yuhezhang-ai commented Aug 25, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Aug 25, 2026

Uh oh!

yuhezhang-ai commented Aug 25, 2026

Uh oh!

yuhezhang-ai commented Aug 26, 2026

Uh oh!

yuhezhang-ai commented Aug 26, 2026

Uh oh!

akoumpa Aug 26, 2026

Uh oh!

yuhezhang-ai Aug 26, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

yuhezhang-ai commented Aug 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What

Changelog

Why the envelope is a closure, not a workaround

Validation

Uh oh!

copy-pr-bot Bot commented Aug 25, 2026

Uh oh!

yuhezhang-ai commented Aug 25, 2026

Uh oh!

yuhezhang-ai commented Aug 26, 2026

Uh oh!

yuhezhang-ai commented Aug 26, 2026

Uh oh!

akoumpa Aug 26, 2026

Choose a reason for hiding this comment

Uh oh!

yuhezhang-ai Aug 26, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

yuhezhang-ai commented Aug 25, 2026 •

edited

Loading