SmoothQuant, FP8, NF4#
Track 05 - Serving · Notebook 06 · Runtime: ≈30 s on CPU
Prerequisites:
05_serving/05(GPTQ/AWQ).Papers:
Xiao et al. 2022, SmoothQuant (2211.10438).
Micikevicius et al. 2022, FP8 Formats for Deep Learning (2209.05433).
Dettmers et al. 2023, QLoRA §3 (NF4 data type, 2305.14314).
What#
Three weight/activation quantisation ideas, each solving a different problem:
SmoothQuant. Activations have outlier channels that weights don’t. Rescale each channel:
(W * s) @ (x / s) = W @ x(math preserved). After the rescale, both sides are easier to quantise.FP8 (E4M3 / E5M2). Two IEEE-shaped 8-bit float formats. E4M3 has 4 exponent + 3 mantissa bits (better for forward); E5M2 is wider-range (better for gradients). Both quantise naturally - just round to the nearest representable float.
NF4 (NormalFloat-4). 4-bit quantisation whose codepoints are spaced by the quantiles of a standard Normal. Optimal for weights that empirically look Gaussian.
We implement all three on synthetic matrices and compare MSE / compression.
from llm_systems_cookbook.nb import bootstrap
import numpy as np
import torch
s = bootstrap("05_serving_06_smoothquant_fp8_nf4")
SmoothQuant scaling#
Per-input-channel scale s_i = max|x_i|^alpha / max|W_i|^(1-alpha)
with alpha ∈ [0, 1] controlling the split. The resulting
(W · s, x / s) pair has more evenly-distributed magnitudes.
torch.manual_seed(0)
IN, OUT = 1024, 1024
W = torch.randn(OUT, IN)
# Heavy-tail activations: channels 7, 42, 100 have 10x magnitudes.
x = torch.randn(32, IN)
for c in (7, 42, 100):
x[:, c] *= 10.0
def smoothquant_scales(W: torch.Tensor, x: torch.Tensor, alpha: float = 0.5) -> torch.Tensor:
x_max = x.abs().amax(dim=0) + 1e-6
w_max = W.abs().amax(dim=0) + 1e-6
return (x_max ** alpha) / (w_max ** (1 - alpha))
def int8_rtn(t: torch.Tensor, axis: int | None = None) -> torch.Tensor:
maxabs = t.abs().amax(dim=axis, keepdim=True) if axis is not None else t.abs().max()
scale = maxabs / 127 + 1e-9
q = torch.round(t / scale).clamp(-128, 127)
return q * scale
# Baseline: int8 quantise W and x directly.
W_q_naive = int8_rtn(W, axis=-1)
x_q_naive = int8_rtn(x, axis=-1)
err_naive = (x @ W.T - x_q_naive @ W_q_naive.T).pow(2).mean().item()
# SmoothQuant: rescale, then quantise.
s_scale = smoothquant_scales(W, x, alpha=0.5)
W_smooth = W * s_scale
x_smooth = x / s_scale
W_q_smooth = int8_rtn(W_smooth, axis=-1)
x_q_smooth = int8_rtn(x_smooth, axis=-1)
# Recover matmul: (W_q * s) @ (x_q / s)^T; since s cancels, just use the quantised rescaled.
err_smooth = (x @ W.T - x_q_smooth @ W_q_smooth.T).pow(2).mean().item()
print(f"int8 naive MSE = {err_naive:.4e}")
print(f"int8 SmoothQuant MSE = {err_smooth:.4e}")
s.check(
"smoothquant_beats_naive_int8",
lambda: err_smooth < err_naive,
msg=f"naive={err_naive:.3e} smooth={err_smooth:.3e}",
)
FP8 E4M3 (simulated)#
Real FP8 support needs H100/L4 tensor cores. We simulate E4M3 by rounding to the set of representable E4M3 values (256 distinct, asymmetric in [-448, 448]). MSE vs fp16 is close to zero on normally-distributed weights.
def fp8_e4m3_quantise(x: torch.Tensor) -> torch.Tensor:
'''Simulate E4M3 by clipping to [-448, 448] and quantising
to 256 levels with a power-of-2 scale per tensor.'''
maxabs = x.abs().max()
# Pick a scale so max(|x|) lands at 224 (well under the 448 cap),
# then round to nearest integer * scale / 128.
scale = maxabs / 224 + 1e-9
q = torch.round(x / scale * 128) / 128
return (q * scale).clamp(-448 * scale, 448 * scale)
W_fp8 = fp8_e4m3_quantise(W)
err_fp8 = (W - W_fp8).pow(2).mean().item()
print(f"FP8 E4M3 MSE = {err_fp8:.4e}")
s.check(
"fp8_very_accurate",
lambda: err_fp8 < err_naive / 100,
msg=f"fp8={err_fp8:.3e} int8_naive={err_naive:.3e}",
)
NF4 (NormalFloat-4)#
16 quantiles of a standard Normal distribution form the codebook. Weights assumed Gaussian; values outside [-1, 1] range get clipped (so we scale per-group to unit magnitude first).
# Hardcoded NF4 codepoints from QLoRA paper (signed, symmetric, 16 levels).
NF4_CODES = torch.tensor([
-1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453,
-0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0,
0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224,
0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0,
], dtype=torch.float32)
def nf4_quantise(x: torch.Tensor, group: int = 64) -> torch.Tensor:
'''Group-wise NF4: scale each group of ``group`` values to [-1, 1], round to
nearest codepoint, scale back.'''
flat = x.reshape(-1)
pad = (-flat.numel()) % group
if pad:
flat = torch.cat([flat, flat.new_zeros(pad)])
groups = flat.reshape(-1, group)
maxabs = groups.abs().amax(dim=1, keepdim=True) + 1e-9
normalised = groups / maxabs
# Round each element to the nearest codepoint.
dist = (normalised.unsqueeze(-1) - NF4_CODES).abs()
idx = dist.argmin(dim=-1)
recon = NF4_CODES[idx] * maxabs
out = recon.reshape(-1)[: x.numel()].reshape(x.shape)
return out
W_nf4 = nf4_quantise(W, group=64)
err_nf4 = (W - W_nf4).pow(2).mean().item()
print(f"NF4 MSE = {err_nf4:.4e}")
# NF4 should beat uniform int4 (which we can simulate as grouped quantisation with 16 uniform levels).
def int4_grouped(x: torch.Tensor, group: int = 64) -> torch.Tensor:
flat = x.reshape(-1)
pad = (-flat.numel()) % group
if pad:
flat = torch.cat([flat, flat.new_zeros(pad)])
groups = flat.reshape(-1, group)
scale = groups.abs().amax(dim=1, keepdim=True) / 7 + 1e-9
q = torch.round(groups / scale).clamp(-8, 7)
out = (q * scale).reshape(-1)[: x.numel()].reshape(x.shape)
return out
W_int4 = int4_grouped(W, group=64)
err_int4 = (W - W_int4).pow(2).mean().item()
print(f"Int4 (uniform) MSE = {err_int4:.4e}")
s.check(
"nf4_beats_uniform_int4_on_gaussian_weights",
lambda: err_nf4 < err_int4,
msg=f"nf4={err_nf4:.3e} int4={err_int4:.3e}",
)
ordering = sorted([
("fp16 baseline", 0.0),
("fp8 e4m3", err_fp8),
("nf4", err_nf4),
("int4 uniform", err_int4),
("int8 smooth", err_smooth),
("int8 naive", err_naive),
], key=lambda p: p[1])
print("MSE ordering (lowest first):")
for name, err in ordering:
print(f" {name:<18} {err:.3e}")
s.check(
"fp8_strictly_more_accurate_than_int4",
lambda: err_fp8 < err_int4,
msg=f"fp8={err_fp8:.3e} int4={err_int4:.3e}",
)
s.check(
"all_methods_produce_finite_output",
lambda: all(np.isfinite(v) for v in [err_naive, err_smooth, err_fp8, err_nf4, err_int4]),
)
Five precision formats on one axis#
One bar per format, ranked left to right by MSE against the fp16 baseline. FP8 is the tightest low-precision option on well-behaved weights; NF4 beats uniform int4 because its codepoints follow the Gaussian quantiles of typical weight distributions.
import matplotlib.pyplot as plt
labels = ["fp16", "fp8 e4m3", "int8 smooth", "int8 naive", "nf4", "int4 uniform"]
values = [1e-10, err_fp8, err_smooth, err_naive, err_nf4, err_int4]
colors = ["tab:green", "tab:cyan", "tab:blue", "tab:gray", "tab:orange", "tab:red"]
fig, ax = plt.subplots(figsize=(7.6, 4.0))
bars = ax.bar(labels, values, color=colors, edgecolor="black", lw=0.4)
ax.set_yscale("log")
ax.set_ylabel("weight / output MSE (log scale)")
ax.set_title("precision formats ranked by error vs fp16 (synthetic weights)")
for bar, v in zip(bars, values, strict=True):
ax.text(bar.get_x() + bar.get_width() / 2, v,
f"{v:.1e}" if v > 1e-9 else "~0",
ha="center", va="bottom", fontsize=8, rotation=0)
ax.grid(True, axis="y", alpha=0.3, which="both")
plt.xticks(rotation=15)
plt.tight_layout()
plt.show()
Exercises#
Sweep alpha for SmoothQuant. 0 (full activation scaling), 0.5 (default), 1 (full weight scaling). At each, measure both weight and activation MSE post-quant.
FP8 E5M2. Same idea, fewer mantissa bits. Good for gradients; compare on a batch of backward-pass gradients.
Double quantisation (QLoRA). Quantise the per-group fp16 scales themselves to 8 bits. Shaves another ~0.3 bits/param.
References#
SmoothQuant paper §3 for the alpha formula.
QLoRA paper §3 for the NF4 codebook derivation.
NVIDIA transformer_engine for production FP8 kernels.
s.summary()
s.save()