llama-nemotron-rerank-1b-v2-mlx

A hand-written MLX / Metal inference path for NVIDIA's nvidia/llama-nemotron-rerank-1b-v2 cross-encoder reranker, built to run natively on Apple Silicon.

No weights are redistributed in this repo. The model weights are NVIDIA's and are loaded from the base_model (nvidia/llama-nemotron-rerank-1b-v2) out of your local Hugging Face cache at runtime. This repo ships only inference code (an MLX forward pass) and documentation — nothing here downloads or re-hosts the weights for you.

What this is

nvidia/llama-nemotron-rerank-1b-v2 is not a standard causal LLM, and mlx_lm cannot load it as-is. Its architecture is LlamaBidirectionalForSequenceClassification (config model_type: llama_bidirec): a Llama-3.2-1B backbone with bidirectional (non-causal) attention (is_causal=False on every layer — encoder-style, every token attends to every other token), followed by masked average pooling over the non-pad tokens and a linear score head (score.weight, shape [1, 2048], temperature 1.0) that projects the pooled vector to a single relevance logit.

Because of the bidirectional mask + pool + score-head, there is no off-the-shelf loader for it in the MLX ecosystem. This repo provides the forward pass written by hand: it reuses mlx_lm's Llama backbone modules for the transformer layers, swaps in a bidirectional (padding-only) attention mask, does the masked-mean pool, and applies the popped score.weight head — reproducing NVIDIA's reranker on Metal.

• Input prompt format: question:{q} \n \n passage:{p} (one (query, passage) pair per sequence).
• Output: one relevance logit per pair; sort descending to rerank.

Derivation — how we got here

• Weights: NVIDIA's nvidia/llama-nemotron-rerank-1b-v2, used as-is. They are read from the HF-cached parent snapshot and cast to f16 for the Metal forward. No fine-tuning, no surgery — same weights, faithfully run.
• The MLX forward is hand-written. There is no automatic conversion: the bidirectional attention, masked-mean pooling, and linear score head are implemented explicitly. The weights load via safetensors safe_open(..., "pt") (the bf16 tensors decode through the torch path, then cast to f16 MLX arrays — mlx/numpy loaders cannot decode bf16 directly).
• Two implementations:
  • Fast Rust engine (production): mlx-rerank — a torch-free mlx-rs 0.25 + tokenizers 0.20 runner (single rerank binary) that serves a /rerank HTTP endpoint. This is the engine you'd actually deploy.
  • Python reference: reference_mlx.py in this repo — a compact, readable MLX forward that loads weights from the HF-cached parent, does length-sorted sub-batching, masked-mean pooling, and the score-head projection. Use it to understand or validate the forward.

Files

File	Purpose
README.md	This model card.
reference_mlx.py	Python MLX reference forward + a 100-doc micro-benchmark; loads weights from the HF-cached base_model.
requirements.txt	Python deps for the reference (mlx, mlx-lm, torch, safetensors, transformers, numpy).

The production Rust engine lives in the sibling repo mlx-rerank, not here.

Benchmarks

Speed is a 100-doc real payload measured on the machine below.

Stack	ms / 100 docs
Nemotron-1B MLX (this code, local)	~1.43 s (sub_batch=10 optimum)
Nemotron-1B PyTorch-MPS (same model)	~5.5 s

• ~5x faster than the PyTorch path end-to-end on the same hardware.
• Runs fully local on Apple Silicon (no API, no per-query cost).

Notes on speed: length-sorting sequences before batching is the one tuning lever that survived clean measurement. Quantization (int8/int4) is a loss on Metal here (40–63% slower) — this forward is compute-bound, not memory-bound, so shrinking weights can't help. The earlier B=25 batching measured ~1925 ms; the current optimum is sub_batch=10 at ~1.43 s / 100 docs. The bundled reference_mlx.py ships with range(0, len(docs), 25); lower it to 10 to reproduce the optimum.

Numerical validation

The MLX port is faithful to NVIDIA's original weights — the only difference from PyTorch is float accumulation order:

• Single-doc P(yes): 0.14934 (MLX) vs 0.15001 (PyTorch) → diff 0.0007.
• Batched 100-doc ranking: top-3 identical = [2, 26, 0] in both MLX and PyTorch.

This proves the forward is correct — it is the same model, faithfully run on Metal.

Machine

All speed and validation figures were measured on:

Model:   MacBook Pro (MacBookPro18,4)
Chip:    Apple M1 Max  (10 CPU cores: 8P + 2E)
GPU:     Apple Silicon unified-memory GPU, driven via Metal (MLX runs natively)
RAM:     64 GB unified
macOS:   26.5.1 (build 25F80)

Toolchain: mlx.core 0.31.2 (Python reference); the Rust engine uses mlx-rs 0.25 + tokenizers 0.20 on rustc 1.93.0.

Usage

Python reference (this repo)

pip install -r requirements.txt
# Ensure the base model is in your HF cache first, e.g.:
#   huggingface-cli download nvidia/llama-nemotron-rerank-1b-v2
python reference_mlx.py

reference_mlx.py locates the base_model snapshot in ~/.cache/huggingface/hub/..., loads the weights as f16, and runs the hand-written forward. The included run(...) function takes a query and a list of documents and returns one relevance score per document (sort descending to rerank). The __main__ block benchmarks a 100-doc payload.

Fast Rust engine (production)

For real workloads use the companion repo mlx-rerank: a torch-free mlx-rs runner that serves a /rerank endpoint.

git clone https://github.com/jak-pan/mlx-rerank
cd mlx-rerank
cargo run --release --bin rerank
# POST {"query": "...", "documents": ["...", ...]} to /rerank

License

The weights are governed by NVIDIA's license for nvidia/llama-nemotron-rerank-1b-v2 (published under the NVIDIA Open Model License, an NVIDIA-specific community license — not Apache/MIT). The local snapshot does not pin a license string, so verify the exact license name and terms on the NVIDIA model page before relying on it. This repo redistributes no weights — only inference code and documentation. The backbone lineage is Llama-3.2-1B (_name_or_path: nvidia/llama-3.2-nv-rerankqa-1b-v2).