gemma-4-31B-it-FP8

Per-tensor FP8 (E4M3) post-training quantization of google/gemma-4-31B-it, produced with NVIDIA TensorRT Model Optimizer. Both the linear weights and the KV cache are FP8.

Summary

Base model	google/gemma-4-31B-it @ commit 419b2efe421994fdfd3394e621983d4cc511cd4f
Base license	Apache-2.0 (inherited)
Quantization tool	NVIDIA TensorRT Model Optimizer (nvidia-modelopt) 0.45.0
Weights	per-tensor FP8 E4M3, W8A8 static (weight_scale + input_scale)
KV cache	calibrated FP8 (k_scale / v_scale)
Excluded (kept BF16)	lm_head, model.embed_vision*, model.vision_tower* (Gemma-4 is a VLM; only the language layers are quantized). q_norm / k_norm also remain BF16.
Calibration set	nvidia/Nemotron-Post-Training-Dataset-v2 (in-domain math / code / stem / chat)
Calibration params	calib_size=1024, calib_seq=512, batch_size=auto — all ModelOpt defaults

How this checkpoint was generated

Quantized with ModelOpt's hf_ptq.py (from the examples/llm_ptq directory of the Model Optimizer repo, nvidia-modelopt==0.45.0):

# The Nemotron calibration dataset is gated; export a read token first.
export HF_TOKEN=<your_hf_token>

python3 hf_ptq.py \
  --pyt_ckpt_path google/gemma-4-31B-it \
  --qformat fp8 \
  --kv_cache_qformat fp8 \
  --export_fmt hf \
  --dataset nemotron-post-training-dataset-v2 \
  --trust_remote_code \
  --export_path <output_dir>

Notes on the flags:

• --qformat fp8 -> per-tensor FP8 E4M3 weights (static W8A8: a weight_scale and an input_scale per linear).
• --kv_cache_qformat fp8 -> calibrated FP8 KV cache (k_scale/v_scale). This is deliberate: ModelOpt's default KV mode is fp8_cast (constant amax), and none disables KV quantization entirely.
• --dataset nemotron-post-training-dataset-v2 -> gated nvidia/Nemotron-Post-Training-Dataset-v2; requires HF_TOKEN.
• calib_size / calib_seq / batch_size are left at ModelOpt defaults.

Reference: PTQ ran in ~90 s (calibrate + export) on 2 GPUs.

Evaluation

All modes: 4×H200, --max-model-len 262144, enable_thinking=true chat template (gemma_think.jinja), sampling per the NVFP4 card (temp=1.0, top_p=0.95, max_new_tokens=131072). Eval = simple_evals (eval-factory 26.03) called directly at the model; GSM8k via lm-eval chat.

• Eval — nvcr.io/nvidia/eval-factory/simple-evals:26.03: simple_evals for AIME / GPQA / MMLU-Pro; lm-eval (installed into the same image) for GSM8k.

Throughput note: these results were produced with TP=1 + DP=4 (4 model replicas) for faster inference. The vllm serve commands below show TP=1 and omit --data-parallel-size 4 — DP is throughput-only and does not affect accuracy. To reproduce our speed, add --data-parallel-size 4 and set --max-num-seqs to ceil(concurrent_requests / 4).

0. Two prerequisites the serve commands reference

0a. gemma_think.jinja — the thinking template (all 3 modes)

Gemma-4 ships its own chat_template.jinja, but reasoning is gated behind the template variable enable_thinking, which defaults to false. Off → the model under-reasons (AIME 2025 drops ~87%→67%). gemma_think.jinja is simply the model's own template with that flag forced on — generated once from the BF16 checkpoint's shipped template:

SRC=google/gemma-4-31B-it/chat_template.jinja   # the model's own template
DST=./gemma_think.jinja
printf '%s\n' '{%- set enable_thinking = true -%}' > "$DST"    # prepend: force thinking on
cat "$SRC" >> "$DST"

Serve with --chat-template <DST>; verify a plain request opens with a thought block. (The FP8 checkpoint uses the same template — no need to regenerate.) Per-request alternative if your client supports it: chat_template_kwargs={"enable_thinking": true} (simple_evals/lm-eval can't send it, hence the baked template).

0b. patch_vllm_modelopt.py — FP8 tie-weights fix (FP8 modes on vLLM 0.24.0 only)

Gemma ties lm_head to embed_tokens, and the ModelOpt-FP8 checkpoint excludes lm_head from quantization. On vLLM 0.24.0, get_quant_method() in .../quantization/modelopt.py returns UnquantizedLinearMethod() for that excluded tied head — which has no tie_weights() → vllm serve crashes at load (NotImplementedError). The patch makes it return UnquantizedEmbeddingMethod() (which has tie_weights) for a ParallelLMHead, keeping UnquantizedLinearMethod() for regular linears:

# in vllm/model_executor/layers/quantization/modelopt.py, get_quant_method(), excluded-layer branch:
-  if isinstance(layer, (LinearBase, ParallelLMHead)):
-      return UnquantizedLinearMethod()
-  return None
+  if isinstance(layer, ParallelLMHead):
+      return UnquantizedEmbeddingMethod()   # has tie_weights
+  if isinstance(layer, LinearBase):
+      return UnquantizedLinearMethod()
+  return None
# (+ import UnquantizedEmbeddingMethod from vllm...vocab_parallel_embedding)

Apply once, as root, inside the vLLM container before serving: docker exec -u root <container> python3 patch_vllm_modelopt.py. Only needed for the FP8 modes on vLLM 0.24.0 — BF16 doesn't need it, and newer vLLM builds (incl. the FA4 build) auto-handle it.

1. vLLM serve — the 3 modes

(A) BF16 TRITON — image vllm/vllm-openai:v0.24.0

vllm serve google/gemma-4-31B-it \
  --served-model-name gemma --host 0.0.0.0 --port 8000 \
  --tensor-parallel-size 1 \
  --max-model-len 262144 --max-num-batched-tokens 8192 --enable-chunked-prefill --max-num-seqs 16 \
  --trust-remote-code \
  --attention-backend TRITON_ATTN \
  --chat-template ./gemma_think.jinja
# no --quantization (BF16); no tie-weights patch

(B) FP8 TRITON — image vllm/vllm-openai:v0.24.0 (+ tie-weights patch as root)

vllm serve jtchen0528/gemma-4-31B-it-FP8 \
  --served-model-name gemma --host 0.0.0.0 --port 8000 \
  --tensor-parallel-size 1 \
  --max-model-len 262144 --max-num-batched-tokens 8192 --enable-chunked-prefill --max-num-seqs 16 \
  --trust-remote-code \
  --quantization modelopt --kv-cache-dtype fp8 \
  --attention-backend TRITON_ATTN \
  --chat-template ./gemma_think.jinja
# requires patch_vllm_modelopt.py (FP8 Gemma tie-weights); FP8 KV → ~7x per-replica concurrency

(C) FP8 FA4 — image <vllm-fa4-build> (a vLLM build with FlashAttention-4)

vllm serve jtchen0528/gemma-4-31B-it-FP8 \
  --served-model-name gemma --host 0.0.0.0 --port 8000 \
  --tensor-parallel-size 1 \
  --max-model-len 262144 --max-num-batched-tokens 8192 --enable-chunked-prefill --max-num-seqs 16 \
  --trust-remote-code \
  --kv-cache-dtype fp8 \
  --attention-backend FLASH_ATTN --attention-config.flash_attn_version=3 --block-size 64 \
  --hf-overrides '{"text_config":{"use_bidirectional_attention":null}}' \
  --chat-template ./gemma_think.jinja
# FP8 auto-detected from ckpt (no --quantization flag, no tie patch on this newer build);
# requests FA v3, the flash-attention-dev build uses FA4 for full-attention layers.
# --hf-overrides disables Gemma's bidirectional attention (FLASH_ATTN requirement).

2. Accuracy

Benchmark	BF16 TRITON	FP8 TRITON	FP8 FA4
AIME 2025 (avg@64)	87.29	87.14	86.88
GPQA Diamond (avg@16)	85.42	85.39	86.55
MMLU-Pro (single)	86.44	86.25	86.64
GSM8k (flexible-extract)	94.62	94.62	94.47

License

Derivative of google/gemma-4-31B-it, which is released under Apache-2.0; this FP8 checkpoint is distributed under the same license.