Skip to content

FP4 Quantization-Aware Training (Exp10)

Paper Section 5.2.1 · Code: src/deepseek_v4_lab/quant/fake_quant.py

Evidence tags: PAPER stated by DeepSeek · DERIVED mathematically implied · IMPLEMENTED reproduced in this project · MEASURED observed in our experiments · INFERRED our interpretation · UNKNOWN not established. Full definitions: evidence ledger

Why this exists

Problem: serving cost is dominated by weight bytes and memory movement; FP16/FP8 leave cost unclaimed.

Response: MXFP4-style tiles with straight-through estimation so quantization lives inside the training objective instead of being applied after it.

New cost: numerical headroom shrinks to the E2M1 grid; post-training-only quantization collapses; STE must be implemented exactly right.

Our experiment: Exp10's matched-budget control shows QAT-FP4 ≈ fp32-control while PTQ-FP4 never reaches even the untouched baseline [MEASURED].

Exchange rate: LOW PRECISION buys serving cost and cache bytes; spends numerical headroom unless trained-in.

Implemented [IMPLEMENTED]

  • MXFP4-style emulation: E2M1 grid {0, ±0.5, ±1, ±1.5, ±2, ±3, ±4, ±6} over 1×32 tiles with per-tile amax scaling; FP8 as intermediate format.
  • Straight-through estimator: forward = quantized values; backward = identity to master weights [PAPER 5.2.1].
  • Tile isolation verified: an outlier in one tile cannot destroy precision of neighbors (the fine-grained-scaling argument).
  • FakeQuantizedLinear for QAT insertion; quantize_model_weights for PTQ.

Measured (Markov LM, 400 pretrain + 250 finetune steps)

variant eval loss
fp32 2.1430
PTQ-FP8 2.1410
PTQ-FP4 2.1453
fp32 + same extra finetune (control) 2.1167
QAT-FP4 2.1183

[MEASURED]

Reading

  • PTQ-FP4 never reaches even the original fp32 baseline.
  • QAT-FP4 matches the extra-steps fp32 control almost exactly: with the quantizer active during training, weights move to positions that work on the E2M1 grid; precision became part of the objective.
  • Magnitudes are toy-scale; the paper's version additionally exploits the lossless FP4→FP8 dequant trick (E4M3 exponent range absorbs sub-block scales) [PAPER Section 5.2.1], which we emulate implicitly via per-tile scales.

Lab finding

First STE implementation leaked gradient through the per-tile scale computation; correct STE requires full detach of the quantized branch.