Phase 2: math-accuracy regression on tq3a-tq2e g32 (Nemotron-3 Super) under non-zero rep_penalty

[manjunathshiva]

Summary

The Phase-1 hybrid quantization for 48 GB Apple Silicon — tq3a-tq2e at group_size=32 on nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 — passes 5 of 6 stress tests cleanly but fails step-by-step arithmetic when run with the recommended sampler-B config (temp=0.7 rep_penalty=1.04 rep_ctx=256). The 2-bit experts cause small slips in numeric reasoning that the repetition penalty doesn't compensate for.

This issue tracks the Phase-2 plan to fix it permanently.

Reproduction

python -m turboquant_mlx.generate \
    --model ~/models/nemotron-3-super-120b-tq3a-tq2e-g32 \
    --prompt \"A train leaves Boston at 9:00 AM going 60 mph. A second train leaves New York at 10:00 AM going 80 mph toward Boston, 215 miles away. At what time do they meet?\" \
    --max-tokens 2048 --min-tokens 50 \
    --temp 0.7 --rep-penalty 1.04 --rep-ctx 256

Reasoning chain looks coherent, final number lands wrong. Removing --rep-penalty recovers arithmetic on most prompts but reintroduces long-gen tail loops.

Phase-1 mitigation (shipped, documented)

• Model card Phase-1 known limitation section recommends omitting --rep-penalty for numeric prompts.
• Stress harness flags math as PARTIAL rather than FAIL — all other paths pass.
• Users wanting reliable math are pointed at the standard tq3 variant.

This is a workaround, not a fix. Phase 2 needs to recover math accuracy under sampler-B so the hybrid is unconditionally usable.

Phase-2 candidates (ranked, cheapest first)

1. Calibration-data Lloyd-Max codebook (preferred)

• Idea: replace the data-free Lloyd-Max codebook for 2-bit experts with one fit to a small calibration set (a few hundred prompts of math + reasoning text). Keep the codebook fixed and embedded in config.json as today, so the inference path doesn't change.
• Cost: a few hours of calibration on the 120B model, no kernel work.
• Risk: drifts away from the data-free promise of TurboQuant; the codebook becomes domain-flavored. Mitigation: include enough prose + code + math in calibration so it stays general.
• Why first: smallest blast radius. If it works, ship.

2. First/last-layer bit protection

• Idea: keep layer 0 and layer N at 3-bit (or BF16) regardless of the per-path policy. Architectural prior — first/last layers are known to be the most sensitivity-critical.
• Cost: ~1 GB extra on disk, ~1 GB extra peak. Still under the 48 GB cap.
• Risk: low. Already a standard trick in PTQ literature.
• Why second: cheap to try, but doesn't address the root cause if the regression isn't concentrated at the layer extremes.

3. Fused QJL Metal kernel for PolarQuantizedSwitchLinear

• Idea: enable --use-qjl for experts via a fused decode kernel, recovering tq3-quality at tq2-storage cost. QJL is currently unfused on SwitchLinear and unusably slow (28 min on the first stress test vs hybrid's 2 min — see project memory qjl_decode_unusably_slow).
• Cost: significant kernel work; nanobind plumbing; correctness verification.
• Risk: high engineering cost, but potentially recovers full math accuracy at the same 36 GB / 40.8 GB peak / 27.2 tok/s budget.
• Why third: only justified if (1) and (2) don't recover math.

Acceptance criteria for Phase 2 closeout

The chosen fix must satisfy all of:

• 6/6 stress tests pass under sampler-B (temp=0.7 rep_penalty=1.04 rep_ctx=256), including the train-meeting arithmetic prompt.
• Peak unified memory stays at or below 48 GB (default iogpu.wired_limit_mb).
• Decode speed stays at or above 25 tok/s on M-series MacBook (current Phase-1 baseline is 27.2 tok/s).
• On-disk size stays at or below 40 GB (current Phase-1 baseline is 36 GB).
• No regression on the 4000-token needle-in-a-haystack long-context test.

References

• Model card: https://huggingface.co/manjunathshiva/Nemotron-3-Super-120B-A12B-tq3a-tq2e-g32
• Standard tq3 variant (math-correct, needs sysctl bump on 64 GB): https://huggingface.co/manjunathshiva/Nemotron-3-Super-120B-A12B-tq3
• Phase-1 long-context kernel fix: Long-prompt prefill crashes in PolarQuantizedSwitchLinear when mlx-lm triggers do_sort #1 (closed)