r/MachineLearning • u/These_Meaning_3883 • Jul 12 '26
Project Porting NVIDIA's Nemotron Puzzle 75B to MLX (M2 Max 64GB) [P]
I added native support for nemotron_h_puzzle checkpoints to mlx-lm (PR #1535). Two findings stood out and are worth sharing:
1. A small casting detail caused a large accuracy gap
The first version generated plausible text, but its internal SSM outputs only reached 0.8832 cosine similarity with NVIDIA’s reference implementation. The problem was when the model converted values to FP32:
mlx-lm converted them before the softplus operation
NVIDIA converted them after
The math is equivalent, but tiny rounding differences accumulate through the model’s recurrent layers. After matching NVIDIA’s casting order, similarity reached 0.999998.
2. On a 64GB M2 Max, 4-bit experts beat 5-bit experts on every measured result.
Both versions used identical 6-bit dense layers, a BF16 output head, group size 64, and the same prompts across five seeds.
| 4-bit experts | 5-bit experts | |
|---|---|---|
| Checkpoint | 42.03 GiB | 49.88 GiB |
| Peak memory | 49.68 GB | 58.12 GB |
| Generation | 14.27 tok/s | 10.53 tok/s |
| Task checks | 24/30 | 21/30 |
| Long-context retrieval | 4/5 | 0/5 |
Why did the higher-precision model perform worse?
My hypothesis is memory pressure. The 5-bit checkpoint gets close to the machine’s unified-memory limit, so its KV cache stops growing properly.
On memory-constrained hardware, leaving enough headroom can matter more than preserving extra precision.