Skip to content

feat(models): vendor Nemotron-H Puzzle checkpoint support - #1228

Open
pierre427 wants to merge 2 commits into
raullenchai:mainfrom
pierre427:codex/port-mlx-lm-1536
Open

feat(models): vendor Nemotron-H Puzzle checkpoint support#1228
pierre427 wants to merge 2 commits into
raullenchai:mainfrom
pierre427:codex/port-mlx-lm-1536

Conversation

@pierre427

@pierre427 pierre427 commented Jul 25, 2026

Copy link
Copy Markdown
Contributor

feat(models): vendor Nemotron-H Puzzle checkpoint support

Source: ml-explore/mlx-lm#1536

Review state: validated — real 75B server/tool-call gate passed; MTP unavailable in the tested artifact

Why

Nemotron-H Puzzle matters to Rapid users because it is a practical 75B / ~9B-active hybrid MoE that can run locally on Apple Silicon with a 262k context window and strong tool-calling behavior. The port lets Rapid instantiate the heterogeneous per-layer Puzzle variant while keeping hybrid SSM and speculative-decode safety gates conservative.

Summary

Add native-first nemotron_h compatibility for the Nemotron Labs Puzzle
checkpoint while preserving Rapid-MLX's conservative hybrid-model safety
gates.

AI assistance disclosure

OpenAI Codex Terra implemented the initial adaptation. The coordinator
independently inspected the diff, repaired process-global test cleanup,
verified public checkpoint metadata, corrected the tool parser profile to the
checkpoint's Nemotron XML template, and reran focused tests, Ruff, format, and
whitespace checks.

Implementation

  • Vendor the Nemotron-H architecture only when the installed mlx-lm has no
    native module.
  • Register the architecture in Rapid's direct loader fallback.
  • Add a text-only Nemotron Puzzle profile with the verified Nemotron tool-call
    parser and Qwen3 reasoning parser.
  • Preserve the checkpoint's MTP sidecar metadata without claiming safe
    speculative decoding: Rapid still lacks validated hybrid SSM
    rollback/replay for this architecture.

Test plan

  • pytest -q tests/test_nemotron_puzzle_vendored.py tests/test_aliases_contract.py — 2545 passed.
  • Ruff check on all changed Python files.
  • Ruff format check on all changed Python files.
  • git diff --check.
  • Public checkpoint metadata verified: nemotron_h, 88 block
    descriptors, one configured MTP layer, 13 weight shards, and zero MTP
    tensor keys in the tested community 6-bit artifact.
  • Parser profile corrected to match the checkpoint's
    <function> / <parameter> Nemotron template.
  • Real 75B GPU server validation on Apple Silicon: Rapid loaded the
    community 6-bit artifact, /v1/models advertised 262144 context with
    text/tools capabilities, ordinary chat returned Paris, and a forced
    Nemotron-format weather tool request produced
    get_weather({"city":"Toronto"}).

Remaining validation

The tested community 6-bit artifact has no MTP tensor keys, so MTP decode was
not a valid validation target for this PR. mlx-lm#1533 remains deferred until a
Puzzle checkpoint with MTP weights is available and Rapid has a safe hybrid SSM
rollback/replay contract.

Attribution

The vendor module preserves Apple copyright and Apache-2.0 attribution and is
adapted from mlx-lm #1536 for Rapid-MLX's native-first vendoring seam.


Draft opened from pierre427 fork after duplicate-checking fresh raullenchai/Rapid-MLX@b3ccb921cd2d435b7238d89b079134258f7d3485.

@pierre427
pierre427 marked this pull request as ready for review July 27, 2026 17:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant