nemotron-omni-mlx

The problem

NVIDIA's Nemotron-3-Nano-Omni-30B-A3B takes text, images, video and audio and answers in text. A 4-bit MLX quantization already existed on Hugging Face (mlx-community, quantized by yayr), but only its text backbone loaded with the standard MLX tools. The vision tower (C-RADIOv4-H) and audio tower (Parakeet Conformer) needed "a multimodal runtime" that didn't exist for MLX (README.md). Matt wrote that runtime.

Constraints

• Apple Silicon, 32 GB or more. Peak use is about 22 GB, measured on one machine, an M5 Max with 128 GB (README.md).
• macOS gives the GPU roughly two thirds of RAM by default. On a 32 GB Mac that ceiling sits under the 22.1 GB image peak, so the README documents raising iogpu.wired_limit_mb.
• Metal fp32 matmul and FFT lose precision (about 7.7e-4 relative error against float64, README.md), which matters for an audio front end.
• Small dependency list: mlx, mlx-lm, numpy, pillow, huggingface_hub. No librosa or soundfile, so OmniProcessor._load_audio_stdlib reads WAV with the standard library.

What he built

Everything under src/ is his; the model, weights, reference PyTorch code and processor are NVIDIA's, and the text backbone is mlx-lm's nemotron_h.

• src/vision.py: the RADIO ViT-H tower, including position embedding interpolation, pixel shuffle and the projector, plus a video path that stacks frames into channels the way the reference does.
• src/audio.py: a log-mel featurizer that uses the checkpoint's own filterbank and runs on the CPU stream (mx.stream(mx.cpu)) to avoid the Metal precision loss, and the Conformer encoder with relative position attention.
• src/processing.py: a re-implementation of torch's antialiased bicubic resize and NVIDIA's dynamic resolution tiling.
• src/omni.py: splice_embeddings() places tower outputs at the context token positions. _make_omni_lm() is the key decision. mlx-lm's nemotron_h has no way to accept precomputed embeddings, so instead of forking mlx-lm the wrapper temporarily swaps the embedding layer for a stub that returns the spliced embeddings and restores it in a finally block. PLAN.md explains why that is safe: MLX builds the graph eagerly and only evaluation is lazy.
• scripts/setup_weights.py builds a filtered text-only/ checkpoint, because mlx-lm's load() rejects the full checkpoint's vision and audio keys.

PLAN.md states the rule he worked by: "Port the WEIGHT GRAPH, not the framework."

Hard bugs

The git history has one code commit (c4188bd, about 3,700 lines), so these come from the gotcha log in PLAN.md rather than fix commits.

• Silent garbage from vision. The reference normalizes pixels outside the tower. Feeding raw [0,1] pixels produces plausible-looking wrong output with no error.
• NaNs on padded audio batches. A -inf attention mask poisons padded rows with NaN, in the NVIDIA and Hugging Face reference too. He used -1e9; padded batch parity came out at minimum cosine 0.99994.
• Bicubic coefficient. torch's antialiased bicubic uses a = -0.5, not -0.75. The wrong value cost up to 27/255 per pixel. Computing weights in float32 instead of float64 mattered as well.

Results (from the repo)

• 23 parity tests against NVIDIA's PyTorch reference (7 vision, 2 audio, 14 processing). README.md reports 23/23 passing. They need the reference code and a 3.7 GB weight shard, so they don't run in CI.
• Parity: audio cosine 0.99999130 (worst frame), vision cosine 0.99996227 (worst token), processor token ids exact and pixels within 1.55e-06 (README.md).
• End-to-end on the M5 Max (PLAN.md): text-only 109 tok/s at 17.9 GB, with an image 67.7 tok/s at 22.1 GB, with audio 110.5 tok/s at 21.0 GB.
• The 152.5 tok/s text figure in PLAN.md is the text backbone alone through plain mlx-lm (first smoke test).
• Parity was checked fp32 against fp32; the shipped runtime runs the towers in bf16, and PLAN.md lists re-checking bf16 quality as open work.

What it shows

This is model porting the careful way: read the reference, port the graph, and prove each piece numerically before moving on. He used cosine similarity rather than absolute difference for vision because of outlier activations, moved one stage to CPU for precision, and extended mlx-lm without forking it.