#84 · FP8 Precision

running

◷train_gpt.py▸mm_t_op — FP8 matmul◎learner＄no GPU

Blocks

Quick summary

A custom Torch operator wraps _scaled_mm with E4M3 (4 exponent, 3 mantissa) quantization for the forward pass. Weights…

Full explanation below the code →

fba-lab — train_gpt.py · mm_t_op — FP8 matmulexecuting

// block: mm_t_op — FP8 matmul · lines 60–95$ study train_gpt.py --block fp8-mm-opA custom Torch operator wraps _scaled_mm with E4M3 (4 exponent, 3 mantissa) quantization for the for… ✓

Explanation

A custom Torch operator wraps _scaled_mm with E4M3 (4 exponent, 3 mantissa) quantization for the forward pass. Weights are divided by a learned scale before quantizing — this prevents exponent overflow. The transposed layout (.T.contiguous().T) tricks _scaled_mm into using the column-major path it needs for fast H100 FP8 Tensor Cores.

Think about

Why does quantizing to E4M3 (float8) speed up training rather than just reducing memory?

// architecture

Live diagram

Record #84 · PR #306 · FP8 MLP up-projection · −0.8 s

Hardware precision

H100 has dedicated FP8 tensor cores — half the bits, double the throughput

H100 GPUs contain separate hardware units for different floating-point precisions. BF16 tensor cores process 16-bit numbers. FP8 tensor cores process 8-bit numbers — and run at roughly 2× the throughput. The catch: FP8 has only 3 mantissa bits (vs 7 in BF16) and a tiny dynamic range. Without careful scaling, model weights quantized to FP8 would lose too much precision and diverge.

PR #306 targets the MLP up-projection — the widest matmul in the transformer block. Weights are quantized to E4M3 FP8 format with a learned per-layer scale factor that keeps values in the representable range. The BF16 master weights live in the optimizer; only the compute copy is FP8. After each optimizer step, the BF16 weights are re-quantized to FP8 and cached for the next forward pass.

# lines 60–95 — E4M3 quantized matmul for MLP up-projection
w_fp8 = w_bf16.to(torch.float8_e4m3fn) / scale
y = torch._scaled_mm(x, w_fp8, scale_a, scale_b)

# lines 158–172 — cache FP8 copy; re-quantize after each opt step
model._mlp_up_proj_f8 = quantize_weights_fp8(model)

Why not train entirely in FP8?

Optimizer state (momentum buffers) and weight updates need high precision to accumulate correctly across thousands of steps. FP8 has only 8 representable exponent values — rounding errors compound and the model diverges. The standard solution: BF16 master weights for the optimizer, FP8 compute copies for the forward pass matmuls only.

Why only the MLP up-projection?

The MLP up-projection is the largest matmul in the transformer (d_model → 4×d_model). It dominates matmul time. Attention QKV projections and the down-projection are smaller and more sensitive to quantization error — particularly attention where small numeric differences can shift which tokens receive focus. FP8 is applied surgically where the gain is largest and the risk lowest.

What changed vs #62

-         
-         
-             w_f8.T, 

H100 has dedicated FP8 tensor cores — half the bits, double the throughput

Why not train entirely in FP8?

Why only the MLP up-projection?

Lab

#84 · FP8 Precision

Live diagram

H100 has dedicated FP8 tensor cores — half the bits, double the throughput

What changed vs #62

mm_t_op FP8 matmul

quantize_weights_fp8 cache

#84 · FP8 Precision

Live diagram

H100 has dedicated FP8 tensor cores — half the bits, double the throughput

What changed vs #62

mm_t_op FP8 matmul

quantize_weights_fp8 cache