CoolFace
Modelpublic

melcheikh/gemma-4-31B-it-qat-NVFP4-mse-Blackwell

sourceHugging Facegemmaupdated 3mo agoView on Hugging Face
3likes211downloads
Model Card

Gemma 4 31B IT QAT — NVFP4 (MSE scaling) for Blackwell

This is a second, better pass at quantizing Gemma 4 31B IT to NVFP4. My earlier upload (gemma-4-31B-it-qat-NVFP4-Blackwell) used ModelOpt's default absmax scaling. This one uses MSE-based weight scaling, and I'm publishing it as a separate repo so nobody's deployment changes under their feet.

What changed, and why

With absmax, each 16-weight block gets its scale from the single largest value in the block. One outlier and the other fifteen weights get crushed toward zero. That's the accuracy gap NVIDIA called out in their Nemotron 3 NVFP4 write-up, and their fix is what I adopted here: per-block scales chosen to minimize reconstruction error (MSE) plus an FP8 activation scale sweep. In ModelOpt terms, --qformat nvfp4_mse instead of --qformat nvfp4. Same ModelOpt release (0.44.0), same calibration set (cnndailymail, 512 samples), same excluded modules (`lmhead`, vision tower). Calibration takes about 7 minutes on an RTX 6000 Blackwell instead of 3 — the sweep does more work per batch.

One thing that will save you a bad afternoon

Do not bake calibrated FP8 KV-cache scales into a Gemma 4 NVFP4 export. I tried --kv_cache_qformat fp8 and the served model produced degenerate output — with two different weight quantizations producing byte-identical garbage, which pins the failure on the KV scales, not the weights. Gemma 4's mixed head dims (256 local / 512 global) push vLLM onto its Triton attention backend, and that path misapplies the per-layer checkpoint scales. This checkpoint ships with no KV scales (kv_cache_quant_algo: null); serve it with --kv-cache-dtype fp8 and let vLLM do dynamic scaling. Same memory savings, correct output.

Serving

bash
vllm serve melcheikh/gemma-4-31B-it-qat-NVFP4-mse-Blackwell \
    --quantization modelopt_fp4 \
    --kv-cache-dtype fp8 \
    --enforce-eager

Needs a Blackwell GPU for native FP4. Tested on an RTX 6000 Blackwell (96 GB) and a 24 GB RTX 5090 laptop — on 24 GB you'll want --max-model-len around 16000–20000 or the KV cache won't fit.

Speculative decoding

On vLLM ≥ 0.21 the Gemma 4 MTP draft path exists upstream (method: gemma4_mtp). Measured on the RTX 6000 Blackwell with this model as target, 300-token generations at temp 0.2: 23.0 tok/s without a draft, 52.1 tok/s with (2.26×, ~54% acceptance). One catch: that proposer loads the draft without applying its quant config, so the draft must be the unquantized BF16 assistant (google/gemma-4-31B-it-qat-q4_0-unquantized-assistant, 927 MB) — an NVFP4-quantized draft fails to load there. On older vLLM (0.20.x, legacy --speculative-model flags) the NVFP4 assistant works as the draft.

Why Hugging Face shows ~17B params

The FP4 weights are packed two per byte, so the tensor shapes are half-size and the parameter auto-count comes out near 17B. The architecture is the full 31B dense model.

License

Gemma Terms of Use, same as the base model.