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

* fix(test): repair legacy partial-rotary spec in vanilla-HF parity references

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>

* test(ci): use the shared relaxed resume envelope for MiniMax M2.7 LoRA

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>

* fix(moe): score sigmoid routing in fp32 by default

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>

* fix(test): pin fp32 reference tensors across vanilla-HF layout differences

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

* fix(minimax): restore the checkpoint's fp32 router end to end

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>

* fix(moe): allocate fp32-contract router gates in fp32 at construction

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>

* test(ci): gate MiniMax M2.7 cross-framework parity at the measured envelope

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>

* fix(minimax): extend the fp32 router contract and robustness gates to 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>

* fix(test): resolve in-tree reference configs past AutoModel registrations

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>

* test(ci): fit the 427B MiniMax-M3 vanilla reference on rank 0's node

The first M3-VL robustness attempt OOMed loading the 854 GiB bf16
reference under device_map=auto (GPU filled to 72 GiB, 9 GiB fused-expert
concat transient). Cap per-GPU memory at 55 GiB and spill the remainder to
host RAM with a longer reload timeout and job window, following the
Step-3.7 precedent at roughly double its scale.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

* fix(test): apply device-map memory caps to the source-load reference

hf_device_map_max_memory_gib / hf_device_map_cpu_max_memory_gib only
reached the phase-3 HF reload; the phase-0 source-load reference always
ran uncapped device_map=auto, so the 427B MiniMax-M3 reference filled
every GPU and OOMed on the fused-expert concat transient (pipeline
64586532; Step-3.7 never hit this because it skips phase-0 logit
parity). Thread the same caps through the source-load kwargs chain.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

* test(ci): skip the intractable 427B M3-VL vanilla-HF phases

The MiniMax-M3 vanilla reference (854 GiB bf16) cannot run economically
on rank 0's node: uncapped device_map OOMs on the fused-expert concat
transient, and the capped CPU-offload variant idles all 64 job GPUs for
hours per release run. Skip both vanilla-HF phases; MiniMax-M2.7 remains
the family's blocking cross-framework parity sentinel (identical measured
router lattice, AMINT-286), while this recipe keeps the AutoModel-side
gates blocking: train/save, bitwise reload, and relaxed-envelope resume
were all green on the first measuring run (pipeline 64577484).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

* fix(test): restrict fp32 leaf aliases to distinctive names

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>

* test(ci): right-size the M3-VL robustness job time to the measured runtime

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>

* refactor(moe): resolve the gate score dtype once at construction

Review feedback on PR #3674: score-path use sites repeated
'self.gate_precision or torch.float32' six times. Resolve it once in
Gate.__init__ into an always-concrete score_dtype; use sites now read
scores.to(dtype=self.score_dtype) with no None handling. gate_precision
itself stays tri-state because None means 'project in the input's
runtime dtype', which several reference routers require (AM-821:
Qwen3-MoE, GPT-OSS, DSV4, Laguna project bf16 while scoring fp32).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

---------

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>