Quantization-Aware Training: When the Model Learns Its Own Precision

Kimi K2 Thinking ships its experts in INT4, gpt-oss in MXFP4, and Kimi K3 in MXFP4 with MXFP8 activations; K2 Thinking and gpt-oss state that their benchmark scores were measured in those formats, not in BF16. Their makers did not quantize a finished model and hope for the best, the way most 4-bit checkpoints on Hugging Face are made. They trained the model knowing it would be rounded: quantization-aware training (QAT). This piece explains the trick, reads what each of the three actually did, and measures it on a small model: the same fine-tuning with and without quantization in the loop, at 4 and 3 bits.

After the fact, or during training

Post-training quantization (PTQ) takes trained weights and rounds them to a grid of at most 16 levels (4 bits) or 8 (3 bits), with one scale per group of weights. Each weight lands on the nearest level, and the error is whatever the rounding left. GPTQ, AWQ and rotations are ways of making that rounding hurt less, but the weights themselves were never chosen with the grid in mind.

QAT keeps training, with the rounding inside the forward pass. The model computes its outputs with the quantized weights, measures its loss on those outputs, and updates its weights to lower it. Over training, the weights move to values whose rounded versions work well. At the end, rounding them one last time changes nothing the model has not already seen.

The trick: fake quantization and a straight-through gradient

Rounding has a gradient of zero almost everywhere, so a naive backward pass through it would never update anything. QAT uses two ideas, both a few lines of PyTorch.

Fake quantization. The weights stay in full precision (the "latent" or master weights). In the forward pass, each one is quantized and immediately dequantized: rounded to its group's grid, then multiplied back by the scale. The layer computes in ordinary floating point, but with exactly the values the quantized model will have.

The straight-through estimator (STE). In the backward pass, the rounding is treated as if it were the identity: the gradient with respect to the quantized weight is applied to the latent weight unchanged (Bengio et al., 2013). In PyTorch that is one line:

def fake_quant(w, bits, group=32):
    """Symmetric, one scale per group of 32 input weights, round to nearest."""
    qmax = 2 ** (bits - 1) - 1
    g = w.reshape(w.shape[0], -1, group)
    scale = g.abs().amax(-1, keepdim=True).clamp(min=1e-8) / qmax
    return (torch.clamp(torch.round(g / scale), -qmax - 1, qmax) * scale).reshape(w.shape)

class FakeQuantSTE(torch.nn.Module):
    def __init__(self, bits):
        super().__init__()
        self.bits = bits

    def forward(self, w):
        # forward: the quantized weight; backward: d(out)/d(w) = 1
        return w + (fake_quant(w, self.bits) - w).detach()

# every linear layer of the transformer blocks computes with its quantized weight
for m in model.model.layers.modules():
    if isinstance(m, torch.nn.Linear):
        torch.nn.utils.parametrize.register_parametrization(m, "weight", FakeQuantSTE(4))

Train as usual. When done, remove the parametrization and quantize the latent weights with the same fake_quant: the checkpoint is exactly what the forward pass used. This is the scheme of PyTorch's torchao QAT (its prepare inserts the fake quantization, its convert applies the real one) and of NVIDIA's Model Optimizer. The labs below say they used QAT but don't publish how their gradients pass the rounding.

What three frontier models did

Kimi K2 Thinking (Moonshot AI, 1T parameters, 32B active). The model card says QAT was applied "in post-training stage", with "INT4 weight-only quantization to the MoE components", and that "all benchmark results are reported under INT4 precision". The checkpoint's config.json gives the scheme: 4-bit integers, symmetric, one scale per group of 32 weights, on the routed experts only; attention, shared experts, the dense MLP layers and the output layer stay in BF16. The stated reason: thinking models "use excessive decoding lengths", and thus quantization "often results in substantial performance drops". The reward is speed: Moonshot reports a roughly 2× faster generation in low-latency serving, where decode is limited by reading the experts' weights.

gpt-oss (OpenAI, 117B and 21B parameters). The model card (§2.1): "We post-trained the models with quantization of the MoE weights to MXFP4 format", 4.25 bits per parameter (4-bit values plus one shared 8-bit scale per 32), on weights that are "90+% of the total parameter count". That is what lets the 120B model fit on one 80 GB GPU and the 20B one in 16 GB.

Kimi K3 (Moonshot AI, 2.8T parameters, 104B active). The technical report (§4.1.4) goes further: QAT "throughout the entire post-training stage, covering both SFT and RL", with MoE expert weights in MXFP4 and activations in MXFP8, while attention projections, shared experts and routers stay in higher precision. It adds a reason specific to reinforcement learning: "rollout and training share the same quantization scheme — eliminating the train–inference mismatch". The model that generates the RL rollouts is the quantized one that will be served, so the policy is optimized for the model users get.

The pattern is the same in all three: quantize the expert weights, which are almost all of the bytes, keep the small and sensitive parts (attention, routers, shared experts) in higher precision, and do it during post-training, not pre-training. None of the three publishes a QAT-versus-PTQ comparison, so the size of the gain is their claim, not a measurement we can check. We measured one ourselves, on a small scale.

Our experiment: PTQ, fine-tuning, and QAT on a small model

The setup, on the GX10 (an NVIDIA GB10), in about an hour: Qwen3-0.6B-Base, every linear layer of its transformer blocks quantized to INT4 or INT3, symmetric, one scale per group of 32 weights (the format and group size of K2 Thinking's checkpoint), embeddings and output layer kept in BF16. Five versions of the model, all scored on the same 81,880 tokens of WikiText-2 (perplexity) and on held-out web text, by the KL divergence of their next-token distribution from the full-precision model they stand in for:

• the original model, and the original quantized after the fact (plain PTQ);
• the model fine-tuned for 400 steps on 3.3 million tokens of FineWeb-Edu, in full precision: the reference;
• that fine-tuned model quantized after the fact;
• the same fine-tuning, same data, same steps, with fake quantization in every forward pass: QAT.

The control matters. Fine-tuning by itself changes the model, so the fair comparison is not QAT against the original model, but QAT against the same fine-tuning quantized afterwards.

(Full precision: 12.35 after fine-tuning, 11.82 before; the original model quantized after the fact: 14.22 at 4 bits, 56.15 at 3 bits.)

At 4 bits, QAT closed about half of the gap PTQ left, and its next-token distribution moved 61% less from the full-precision model's. At 3 bits, where rounding to seven levels (−3 to 3 times the scale) breaks the model after the fact (perplexity four times higher), QAT brought it back to within 45% of full precision. The PyTorch team saw a similar pattern on Llama 3 8B in different setups: QAT recovered 68% of the WikiText perplexity degradation with INT4 weights and INT8 activations, and 99% with 2-bit weights once the most sensitive layers were left unquantized (torchao). Where PTQ hurts most, QAT has the most to recover.

Two limits. This is one small model, one seed, and a short fine-tuning; a frontier model's post-training runs on far more data, and how much QAT recovers there is the labs' claim, not something this experiment shows. And the after-the-fact quantization here is plain round-to-nearest; GPTQ would narrow the 4-bit PTQ gap from the other side.

What QAT costs

Fake quantization is cheap next to a training step: our QAT runs took 998 and 1,000 seconds against 953 for the same fine-tuning without it, 5% more. The PyTorch team measured 34% slower fine-tuning on Llama 3 8B with INT4 weights and fake-quantized INT8 activations; on our 0.6B model, the output layer over a 151,936-token vocabulary dominates each step, which hides the overhead. Memory grows a little, because each layer keeps a quantized copy of its weights for the forward pass.

The real cost is having a training run at all: the data, the pipeline, and the GPUs. That is why QAT shows up in labs' post-training, not in the community checkpoints quantized from a finished model.

When QAT is worth it

• You are training anyway. QAT is a change to a fine-tuning or post-training run, not a separate step. Labs that ship in 4 bits fold it into post-training; if you fine-tune a model for your product and will serve it in 4 bits, fine-tune it with fake quantization.
• The format is aggressive. At 8 bits, PTQ is already close to lossless; at 4 bits it costs a little, and at 3 bits a lot. QAT pays most where PTQ hurts most.
• The model generates long outputs. Moonshot's reason for K2 Thinking: with very long decoding, quantization often costs substantial performance.
• Not when you only have the checkpoint. Without the training data and a training setup, GPTQ, AWQ or rotations are what you can do. NVIDIA's variant, quantization-aware distillation (QAD), trains the quantized model to match the full-precision model's output distribution (a KL loss), and its authors report it recovers accuracy without the full training data or a rerun of the SFT and RL pipeline (arXiv 2601.20088).

Gotchas

Train and serve with the same quantizer. The gain exists only if the inference engine rounds exactly as training did: same format, same group size, same symmetric or asymmetric choice, same scale computation. K2 Thinking's checkpoint ships in compressed-tensors format with its scheme in config.json, so vLLM and SGLang load it as trained.

Compare against fine-tuning without quantization, not against the original model. Fine-tuning alone changes perplexity; a QAT gain measured against the original model mixes the two.

Weight-only QAT does nothing for activations. Training with fake-quantized activations helps (torchao's INT8 activations, K3's MXFP8), but 4-bit activations still lean on per-block scales (MXFP4/NVFP4) or rotations.