Add Nemotron-H Puzzle model support

[sxuff]

Addresses #1506.

Summary

This adds native loading and inference support for checkpoints whose model_type is nemotron_h_puzzle by reusing the existing Nemotron-H implementation and applying Puzzle's blockwise configuration.

The change:

• remaps nemotron_h_puzzle to the Nemotron-H model implementation
• derives the layer count from block_configs when the official config omits num_hidden_layers
• supports per-layer moe_intermediate_size and num_experts_per_tok
• remaps official model.* checkpoint tensors to MLX's backbone.* hierarchy
• keeps Puzzle's output projection outside ordinary low-bit quantization
• matches NVIDIA's FP32 RMSNorm and router behavior
• preserves existing nemotron_h bias, norm, router, timestep, and SSM behavior by scoping Puzzle-specific numerical semantics to Puzzle configs
• preserves any explicitly configured time_step_limit, including finite upper bounds
• preserves the projected activation dtype for Mamba softplus(dt + dt_bias) while retaining FP32 state-transition and recurrence operands
• uses the official time_step_min lower bound
• uses the official projection bias setting

Why timestep precision is configurable

The existing shared SSM utility promotes dt to FP32 before softplus. NVIDIA's Puzzle implementation preserves the projected activation dtype at that operation, then promotes the recurrence-sensitive state-space operands.

With identical fused-CUDA first-layer intermediates, promoting dt before softplus produced an SSM cosine similarity of 0.8832. Preserving BF16 at that boundary produced 0.999998. This establishes fidelity to that reference boundary; an independent same-process A/B on the released 6-bit checkpoint found only small end-to-end differences between the two modes, with per-layer cosine similarity of at least 0.997 and coherent generation in both modes (details). The new promote_dt argument defaults to True, so other callers retain the previous behavior.

Tests

Added Puzzle regressions for:

• model registration
• official layer-count derivation
• heterogeneous MoE dimensions and top-k
• official prefix remapping
• routed expert stacking
• reduced forward execution
• strict quantized loading
• save/load behavior
• output-head quantization policy
• Mamba prefill versus cached decoding
• timestep activation precision and minimum clamp
• explicit timestep-limit passthrough for base Nemotron-H and Puzzle
• FP32 state-transition behavior with native-dtype timesteps

Targeted M2 Max result:

88 passed, 1 skipped, 51 subtests passed

Full repository result on the M2 Max:

218 passed, 1 skipped, 1 deselected, 79 subtests passed

The deselected test was the pre-existing Hugging Face billsum integration test. The environment had datasets 5.0.0 and huggingface_hub 1.23.0, whose current repository-ID validation rejects the test's legacy bare billsum identifier before loading. Running the complete suite without the deselection produced the same 218 passes plus that single unrelated failure.

Scope

MTP/speculative decoding is intentionally not included. This PR implements ordinary autoregressive base-model loading and inference.

[pierre427]

Author of the competing #1536 here — since the two PRs landed the same day, I tried to settle the main technical difference (promote_dt) with measurements rather than have them race.

I ported this PR's exact dt semantics (bf16-native softplus via compute_dt_native, plus A following A_log's dtype) behind a runtime toggle on top of #1536, and A/B'd both modes in the same process, on the same inputs, on the released georgeis55/Nemotron-Labs-3-Puzzle-75B-A9B-MLX-6bit (mlx 0.32.0.dev, M5):

metric	promoted-fp32 (upstream ssm.py)	native-bf16 (this PR)
teacher-forced ppl (1536-token prose pack)	7.0828	7.0749

• Per-layer output cosine between the two modes: worst layer 0.9971, layer 0 (first Mamba block) 0.9999970, all 88 layers ≥ 0.997.
• Final logits cosine 0.99994, greedy top-1 agreement 96.6%.
• Greedy generations both fully coherent with the same structure/content (wording diverges after ~60 tokens, as expected from ~4% token-level disagreement).
• Same picture on plain nemotron_h (Nemotron-3-Nano-30B-4bit): logits cosine 0.9993, identical 60-token greedy output.

One observation on the 0.8832-vs-0.999998 cosine result: if both variants are compared against a common reference and one sits at 0.8832 while the other sits at 0.999998, the two variants must differ from each other by roughly cosine 0.88. On real activations we measure ≥ 0.997 at every layer — so the 0.88 figure appears to be a property of the injected fused-CUDA intermediate harness (inputs off the natural activation distribution, or a dt regime real prompts don't reach) rather than something that occurs end-to-end on the released checkpoint.

That said, I think bf16-at-softplus is still the right call where it matches NVIDIA's reference semantics — it's zero-cost and strictly more faithful; it just doesn't appear to be load-bearing for output quality, so it needn't decide between the two ports.

On consolidation: happy to converge rather than compete. This PR's MTP-tensor exclusion and output-projection quant policy are things #1536 should have regardless; #1536 is the smaller diff for the core remap. Whichever direction the maintainers prefer, I'll adapt mine or review/support this one.

[sxuff]

@pierre427 Thanks for running the same-process A/B, this is very helpful (i also like the coincidence of both our PRs haha).

I agree it reframes promote_dt: preserving BF16 at softplus is still preferable because it matches NVIDIA’s reference semantics, I’ll soften that wording and reference your results.

One small correction after checking the current base: mtp.* filtering is already present on main, so I shouldn’t list that as new in #1535. The additional pieces here include the official model.* → backbone.* remap, preserving lm_head outside low-bit quantization, and the reference-aligned numerical behavior and tests.

I’m happy to consolidate either way. I’ll leave #1535 open and wait for maintainer preference. If #1536 is preferred, I’m happy to help port the remaining pieces and tests; if #1535 is preferred, I’d appreciate your review/support.

[pierre427]

Appreciate the thoughtful reply — and yeah, funny we both landed the same port the same day.

I'm happy to consolidate onto #1535 as the base. It's the more complete port — the official model.*→backbone.* remap, keeping lm_head out of low-bit quant, the fp32 reference-norm/router, and the per-layer block_configs handling are all the right calls. With the promote_dt wording softened I don't think there's a meaningful numerics gap between us: my same-process A/B has both modes at ≥0.997 per-layer cosine and a perplexity wash (7.075 vs 7.083) on the released 6-bit checkpoint, so bf16-at-softplus is a free fidelity nicety rather than a correctness fix. I'll close #1536 in favor of yours once a maintainer signals preference, and I'm glad to review and help land it.

Two things I hit while testing that might be worth folding in:

1. Consider gating the Puzzle-specific paths on block_configs is not None rather than model_type == "nemotron_h_puzzle". The public conversion I tested — georgeis55/Nemotron-Labs-3-Puzzle-75B-A9B-MLX-6bit — ships config.model_type = "nemotron_h" but does carry block_configs. On it, is_puzzle is False, so the fp32 norm/router, native-softplus, and (notably) the lm_head quant-exclusion all no-op. The per-layer MoE path still works since it keys off block_configs, but a 4-bit conversion of a checkpoint labeled nemotron_h would miss the lm_head protection and could hit exactly the repetitive-generation failure your comment describes. Keying the niceties off block_configs would make them robust to how a converter set model_type.
2. A small MoEGate fallback. Puzzle isn't group-routed; if a config omits n_group/topk_group/routed_scaling_factor, group_expert_select's n_group > 1 throws on None. Defaulting them to 1/1/1.0 when absent avoids that. (The georgeis55 config happens to supply them, so it's just hardening for other conversions.)

Separately, happy to contribute the numerics probe as a reference/regression check for the softplus-dtype boundary if useful — single model load, same inputs through both modes, reporting per-layer cosine, final-logits cosine/top-1, and teacher-forced ppl.

Whatever the maintainers prefer, count on my review and support to get it merged.

[sxuff]

Thanks Pierre, agreed on both points.

block_configs is the better capability signal. I confirmed the current model_type checks would miss conversions like the georgeis55 checkpoint even though they contain Puzzle’s blockwise configuration. I’ll gate the Puzzle-specific norm/router, Mamba timestep behavior, projection bias, and lm_head quantization exclusion on block_configs is not None, with regression coverage for model_type="nemotron_h" plus block_configs. Ordinary Nemotron-H configs without block_configs will remain unchanged.

I’ll also add the MoEGate fallbacks (n_group=1, topk_group=1, and routed_scaling_factor=1.0) and test an ungrouped configuration.
The numerics probe sounds useful, though I’d keep it as a follow-up unless the maintainers want it included here, so this PR stays focused.
Thanks again for consolidating on #1535 and offering to review

[zcbenz]

Would mind rewriting the PR description following our new template?
https://github.com/ml-explore/mlx-lm/blob/main/.github/PULL_REQUEST_TEMPLATE/new_model.md?plain=1

It would help me verify the implementation much quicker.