45. Int8 Fake-Quantization of a Linear Layer

Problem

Serving large models cheaply often means quantizing weights to int8: storing an 8-bit integer per weight plus a small per-tensor float "scale" that maps integers back to the original range, instead of a full 32-bit float per weight. Before shipping real int8 kernels, it's standard practice to prototype with "fake quantization" -- round weights through int8 and immediately dequantize back to float -- so you can measure the accuracy impact of quantization while still running ordinary float matmuls.

Implement symmetric per-tensor int8 fake-quantization of a Linear layer's weight, then run the forward pass with the fake-quantized weight (bias stays in full precision).

Source: src/45_int8_fake_quant_linear.py

def int8_fake_quant_linear(
    x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor,
) -> torch.Tensor: ...

>>> x = torch.randn(2, 4)
>>> w = torch.randn(3, 4)
>>> b = torch.randn(3)
>>> int8_fake_quant_linear(x, w, b).shape
torch.Size([2, 3])

Step-by-Step Approach

1. Compute the per-tensor scale: scale = weight.abs().max() / 127 -- symmetric quantization maps the single largest-magnitude weight to ±127, leaving -128 unused so the integer grid is symmetric around zero.
2. Divide the weight by the scale to bring it onto (roughly) the integer grid.
3. Round to the nearest integer with torch.round.
4. Clamp to [-127, 127] -- rounding alone can still leave values slightly outside range from floating-point error at the extremes.
5. Dequantize by multiplying back by scale, returning to the original float units.
6. Run the ordinary linear forward pass, x @ w_dequant.T + bias, using the dequantized weight; the bias is never quantized.

The key insight is that "fake" quantization keeps everything as float tensors the whole way through -- the int8 constraint is enforced only by forcing values onto the rounded, clamped grid, which is exactly what lets you measure quantization error using ordinary float matmuls before committing to real int8 storage and kernels.

Reference solution

def int8_fake_quant_linear(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor,
) -> torch.Tensor:
    # Symmetric per-tensor scale: the single largest-magnitude weight maps
    # to +/-127 (int8's range, keeping -128 unused so the grid is
    # symmetric around zero -- avoids a small positive/negative bias).
    scale = weight.abs().max() / 127.0

    # round() then clamp(): clamp alone is not enough, since values must
    # also land on integer grid points before being treated as "int8".
    # This stays a float tensor throughout (fake quant), so gradients
    # could still flow through it in a straight-through-estimator setup,
    # unlike a real torch.qint8 tensor.
    w_int8 = torch.clamp(torch.round(weight / scale), -127, 127)
    w_dequant = w_int8 * scale  # back to float, same units as the original weight

    # bias is never quantized in practice -- it's a tiny fraction of total
    # parameters and quantizing it buys no meaningful memory/compute win.
    return x @ w_dequant.T + bias

Key Functions & Tricks

• weight.abs().max() — the per-tensor scale anchor for symmetric quantization.
• torch.round — snaps a value to the nearest integer; note PyTorch uses round-half-to-even, so exact .5 boundaries round to the nearest even integer, not always "up."
• torch.clamp(t, min, max) — bounds values into the valid int8 range after rounding.
• x @ w_dequant.T — matches nn.Linear's convention of storing weight as (out_features, in_features), transposed for the matmul.

How to Recognize This Pattern