gemma-4-31B-it-qat-q4_0-mlx

A 4-bit MLX conversion of Google's Gemma 4 31B QAT checkpoint that keeps Google's own quantization grid instead of deriving a new one.

Loads with stock mlx-lm — no patched version required.

Why this exists

Google trains these checkpoints so they survive being squashed to 4 bits, and ships the intended 4-bit result as a GGUF. Standard MLX conversion ignores that and re-derives its own grid from the weights, discarding the alignment the model was trained for.

This conversion reproduces Google's grid exactly. Same 4 bits, same file size, same speed — the only difference is which 4-bit values get written.

How much it helps

Measured against this model's own bfloat16 reference over 200 chat-formatted prompts, teacher-forced.

Prediction drift — how far the 4-bit model's next-word probabilities wander from the original. Lower is better; 0 would be identical.

Agreement — how often the 4-bit model picks the exact same next word as the original. Higher is better.

	standard MLX 4-bit	this model
prediction drift	0.0321	0.0205
agreement	96.5%	96.9%

Across the dense Gemma 4 family the improvement runs 21–48% less drift, at identical file size:

model	drift, standard	drift, this method	better by
E2B	0.0338 (group 64) / 0.0220 (group 32)	0.0175	21%
E4B	0.0225 / 0.0160	0.0112	30%
12B	0.1048 / 0.1219	0.0631	48%
31B	0.0314 / 0.0321	0.0205	36%

How the grid is recovered

Per 32-weight group, the scale comes from the signed element of largest magnitude:

extremum = w[argmax(|w|)]        # signed, not abs
d        = extremum / -8
code     = clip(trunc(w/d + 8.5), 0, 15)
scales   = d
biases   = -8 * d

Stored as ordinary affine 4-bit with group size 32, which is why stock MLX loads it.

This is verified, not asserted. Running that derivation on Google's -qat-q4_0-unquantized weights reproduces the codes and scales in Google's shipped Q4_0 GGUF exactly — 100% of 721,551,360 weights matched on 12B, with zero scale error. Google generated their GGUF deterministically from these weights, and this conversion lands on the same grid.

Use with mlx

pip install mlx-lm

mlx_lm.generate --model goodolclint/gemma-4-31B-it-qat-q4_0-mlx --prompt "Explain the Doppler effect briefly."

from mlx_lm import load, generate

model, tokenizer = load("goodolclint/gemma-4-31B-it-qat-q4_0-mlx")
prompt = tokenizer.apply_chat_template(
    [{"role": "user", "content": "Explain the Doppler effect briefly."}],
    add_generation_prompt=True,
    tokenize=False,
)
print(generate(model, tokenizer, prompt=prompt, max_tokens=512))

Gemma 4 is a thinking-mode model — use the chat template. Bare completions produce degenerate output regardless of quantization.

Reproducing it

mlx_lm.convert \
  --hf-path google/gemma-4-31B-it-qat-q4_0-unquantized \
  --mlx-path gemma-4-31B-it-qat-q4_0-mlx \
  -q --q-calibration q4_0

--q-calibration is proposed upstream in ml-explore/mlx-lm and ml-explore/mlx-swift-lm. Until it merges, use the branch linked from those pull requests. The published weights need none of that — they are ordinary affine 4-bit and load with released mlx-lm.

Honest limits

• Only linear layers are calibrated. Embeddings and norms use standard derivation at the same 4-bit/group-32 geometry, because MLX has no way to hand a precomputed grid to a quantized embedding. In Google's GGUF the embedding ships at Q6_K, so it is outside the grid claim either way.

• Do not expect this on mixture-of-experts. The same method on Gemma 4 26B-A4B produced no significant improvement, so no MoE conversion is published here.

• Measurements are single-run on one prompt corpus with a fixed seed. They measure agreement with the bfloat16 reference — not whether the model is good. That is a benchmark question, not answered here.

• This costs speed and memory. Group size 32 is needed to hit Google's grid, and it is more expensive than the group-64 conversions most MLX models use. Measured on an M4-class machine against mlx-community 4-bit builds, 200-token decode:

  model	this model	typical group-64 build	cost
  E2B	183.7 tok/s, 3.0 GB	192.4 tok/s, 2.7 GB	4.5% slower, 11% more memory
  E4B	116.5 tok/s, 4.8 GB	125.8 tok/s, 4.3 GB	7.4% slower, 12% more memory
  12B	56.8 tok/s, 7.6 GB	62.2 tok/s, 6.8 GB	8.7% slower, 12% more memory
  31B	24.4 tok/s, 19.5 GB	27.0 tok/s, 17.6 GB	9.6% slower, 11% more memory

  That cost is the group size, not the calibration — at equal group size the calibration is free. If throughput matters more to you than fidelity, a group-64 build is the right choice.

Licence and modification notice

Licensed under the Apache License, Version 2.0, inherited from the base model — see Google's Gemma 4 licence and the full licence text. Google's Gemma Terms of Use and Prohibited Use Policy also apply to how you use these weights.

Modification notice (Apache-2.0 §4b): the weight files in this repository are modified copies of google/gemma-4-31B-it-qat-q4_0-unquantized. Every .safetensors file has been changed: bfloat16 weights were quantized to 4-bit using the grid recovery described above, and config.json records the resulting quantization block. No other change was made — no fine-tuning, no merging, no distillation, no vocabulary or architecture change.

All credit for the model and for the quantization-aware training belongs to Google DeepMind. This repository changes only how the weights are converted for MLX.