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

yuhezhang-ai · 2026-08-24T16:40:17Z

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.

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

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

Follow-up (stacked on #3635 once it lands): re-enable capture_router_diagnostics / shape_diagnostic for this recipe plus the MiniMax router-capture extension — that machinery lives in #3635.

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-24T16:40:20Z

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.

* test(checkpoint): expand parity metrics and phase coverage Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(checkpoint): calibrate all parity phases on long input Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(checkpoint): use unique long-context parity input Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(checkpoint): stabilize long-context parity input Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(checkpoint): define representative parity cohort Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(checkpoint): calibrate parity profiles Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(checkpoint): relax Step resume drift Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(checkpoint): calibrate parity profiles from scoped CI Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(ci): capture checkpoint repeatability metrics Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * docs(ci): explain repeatability metric diagnostics Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(ci): add high-variance checkpoint parity profile Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(checkpoint): avoid redundant resume checkpoint writes Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * refactor(checkpoint): use targeted Step parity overrides Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * fix(test): repair checkpoint parity CI regressions Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(checkpoint): cover Nemotron Flash HF reload Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * fix(test): stabilize remote checkpoint parity Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * fix(ci): use cached Nemotron family tokenizer Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(ci): enable routed MoE resume coverage Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * fix(checkpoint): map Nemotron PEFT export namespace Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * fix(ci): reach configured resume boundary Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(ci): calibrate Nemotron chat resume drift Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * docs(ci): clarify checkpoint robustness compatibility Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(ci): enable Nemotron resume coverage Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(ci): preserve Nemotron tokenizer issue gate Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(ci): retire stale Nemotron KL overrides Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(ci): forbid legacy max KL recipe thresholds Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(ci): remove unmeasured Mistral FP8 overrides Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(ci): remove legacy checkpoint parity fields Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(ci): unify checkpoint parity overrides Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(ci): support per-comparison parity profiles Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * fix(ci): keep parity threshold numeric across YAML loaders Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> --------- Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> Co-authored-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com>

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

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

Superseded: GitHub closed this PR when its merged base branch was deleted and the stacking metadata prevents retargeting. Replacement (same head branch, rebased onto main): see the new PR linked below.

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

Replacement: #3674

yuhezhang-ai mentioned this pull request Aug 24, 2026

fix(test): unblock Kimi-Linear HF reference and gate MiniMax routed parity #3644

Closed

yuhezhang-ai force-pushed the yuhez/fix/amint286-minimax-rope-reference branch from 4843722 to 2359f9d Compare August 25, 2026 14:50

yuhezhang-ai mentioned this pull request Aug 25, 2026

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

Merged

yuhezhang-ai changed the title ~~fix(test): repair the MiniMax M2.7 vanilla-HF parity reference~~ fix(minimax): repair HF parity reference and gate at the measured envelope Aug 25, 2026

yuhezhang-ai marked this pull request as ready for review August 25, 2026 15:49

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

yuhezhang-ai force-pushed the yuhez/fix/amint286-minimax-rope-reference branch from 2359f9d to 79bb6b2 Compare August 25, 2026 18:43

An error occurred while trying to automatically change base from yuhez/test/checkpoint-parity-metrics to main August 25, 2026 20:34

yuhezhang-ai and others added 7 commits August 25, 2026 13:56

yuhezhang-ai force-pushed the yuhez/fix/amint286-minimax-rope-reference branch from 79bb6b2 to 695a9bc Compare August 25, 2026 20:59

yuhezhang-ai requested review from a team, HuiyingLi, akoumpa, athitten and snowmanwwg as code owners August 25, 2026 20:59

yuhezhang-ai deleted the branch yuhez/test/checkpoint-parity-metrics August 25, 2026 21:01

yuhezhang-ai closed this Aug 25, 2026

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 temporarily deployed to public August 25, 2026 21:06 Inactive

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

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

Navigation Menu

Uh oh!

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

fix(minimax): repair HF parity reference and gate at the measured envelope#3643
yuhezhang-ai wants to merge 8 commits into
yuhez/test/checkpoint-parity-metricsfrom
yuhez/fix/amint286-minimax-rope-reference

yuhezhang-ai commented Aug 24, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Aug 24, 2026

Uh oh!

yuhezhang-ai commented Aug 25, 2026

Uh oh!

yuhezhang-ai commented Aug 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

yuhezhang-ai commented Aug 24, 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 24, 2026

Uh oh!

yuhezhang-ai commented Aug 25, 2026

Uh oh!

yuhezhang-ai commented Aug 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

yuhezhang-ai commented Aug 24, 2026 •

edited

Loading