PigeonColumba compressa

SmoothQuant (weights + activations)

Activations with a few huge channels can't be quantized to 8 bits; move that difficulty into the weights with an exact per-channel rescaling.

Runs hereRun it in your browser

On this page (9 sections)

The idea

To run a matrix multiply on 8-bit integer hardware, both the weights and the activations (the inputs flowing into the layer) must be 8-bit. LLM activations are hard: a few fixed channels are up to 100 times larger than the rest, so a scale per tensor, or even per token, spends almost all its levels on them. Weights are easy. SmoothQuant moves the difficulty from activations to weights with an exact per-channel rescaling, y = (x / s) · (s · W), so both become quantizable.

How it works

SmoothQuant writes Y = X W with X (tokens x in) and W (in x out), the same orientation this repo stores.

  1. On calibration data, record each input channel's maximum activation max|X_j| and maximum weight max|W_j| (row j of W).
  2. Set s_j = max|X_j|^α / max|W_j|^(1-α). The migration strength α sets how much difficulty moves. At α = 0.5, both sides end up with the same per-channel maximum sqrt(max|X_j| · max|W_j|).
  3. Divide input channel j by s_j and multiply weight row j by s_j. The product is unchanged.
  4. Fold 1/s into whatever produced X, offline, so no extra work happens at run time.
  5. Quantize weights per output channel and activations per tensor (or per token) to 8 bits.

Why per-channel activation scales are not an option. An integer matmul sums over input channels, so a per-input-channel scale sits inside the sum and can't be pulled out afterwards. Per-token and per-output-channel scales can.

In a transformer, the outlier channels appear in the LayerNorm outputs that feed the attention QKV projection and the first MLP layer. 1/s is folded into that LayerNorm's gain and bias. The paper uses α = 0.5 for most models, more for harsher outliers.

The code, line by line

inject_outliers(net, layer, k, factor) builds the problem on purpose. It multiplies k neurons of hidden layer layer (weights and bias) by factor and divides the next layer's matching weight rows by factor. Because relu(c·z) = c·relu(z) for c > 0, the function is unchanged, but in the page's run three channels become outlier size times louder.

smoothquant(net, X, alpha):

  • _, hid, _ = net.forward(X) records the calibration activations once. hid[i - 1] is layer i's input. One pass is enough: rescaling layer i's inputs never changes its outputs, which are the next layer's inputs.
  • if net.res[i - 1] or net.res[i]: continue skips residual layers. Rescaling a residual stream would also rescale the skip path.
  • a = np.abs(hid[i - 1]).max(0) and w = np.abs(net.W[i]).max(1) + 1e-8 are max|X_j| and max|W_j|. .max(1) on the (in, out) array runs over a whole input row.
  • s = a ** alpha / w ** (1 - alpha) is step 2.
  • silent = a < 1e-3 * a.max() finds ReLU neurons that never fire on calibration data. The formula would give them s ≈ 0, multiplying their incoming weights by hundreds, so any test input that switches one on would produce a giant activation. np.where(silent, 1.0, np.clip(s, 1e-3, 1e3)) leaves them unscaled. On the 512 calibration inputs, 1 of layer 2's 256 inputs counts as silent (2 once injected outliers raise the maximum), and 9 of the output layer's.
  • net.W[i] = net.W[i] * s[:, None] scales the weight rows. net.W[i - 1] / s[None, :] and net.b[i - 1] / s fold 1/s into the previous layer, which is legal because ReLU is positively homogeneous.

quant_act(h, bits) is dynamic per-token quantization: scale = np.abs(h).max(1, keepdims=True) / qmax, one scale per row (per example here, per token in an LLM), computed on the fly. One huge channel still sets its row's scale, which is exactly the problem. forward_wa(net, X, abits) quantizes the input of every layer after the first, then runs the layer.

With weight bits below 16, run_smoothquant passes both models through apply_layerwise(..., rtn, skip_last=True): per-output-row RTN after smoothing, so the difficulty moved into the weights is charged. The first layer is quantized; the output layer stays float, like lm_head.

What the simulations show

3x256 teacher with 3 outlier channels injected into layer 2's input, at outlier size 30. Weights are 8-bit (W8), and the float model scores 0.9150.

activation bitsplain, W8SmoothQuant α 0.5, W8
160.91460.9146
80.91370.9154
60.90490.9155
50.85170.9152
40.56610.9079

At outlier size 100, plain falls to 0.6850 at A6 and 0.2157 at A4; SmoothQuant gives 0.9155 and 0.9079, identical to size 30, because s absorbs the injected factor exactly. Without injected outliers, plain A4 already gives 0.9096 and SmoothQuant 0.9089: there was nothing to fix.

The weight side costs little here. At W4A6, plain gives 0.8963 and SmoothQuant 0.9112. Even α = 1, which moves all of the activation range into the weights, gives 0.9101 at W4A6. Across α = 0, 0.5, 0.8 and 1 at W8A4, SmoothQuant gives 0.9035, 0.9079, 0.9090 and 0.9099. So this toy does not show the paper's trade-off, where too large an α makes the weights hard to quantize, probably because its weights have no outliers of their own for the migrated range to collide with.

When it breaks

  • Channels that are silent during calibration have no measured range. Scaling them by the formula blows up their weights for unseen inputs; the code leaves them at s = 1.
  • Large α pushes outliers into the weights, which in real LLMs eventually costs weight accuracy (barely, here).
  • s is static. A channel that is an outlier on only some tokens is not fully tamed by it.
  • The fold needs a positively homogeneous op (ReLU, or LayerNorm's affine gain) right before the layer. The paper smooths only layers fed by a LayerNorm, where the fold is free.

Try this

  • inject outlier channels on, weight bits 8, activation bits 4: compare the plain row with SmoothQuant.
  • Raise outlier size to 100 at activation bits 6: look at how far plain falls while SmoothQuant doesn't move.
  • weight bits 4, then migration strength 0, 0.5 and 1: look for whether moving range into the weights costs anything at W4.
  • Turn inject outlier channels off: look for SmoothQuant having nothing left to fix.

Going further

SmoothQuant enabled W8A8 inference of models up to 530B parameters with negligible accuracy loss. Rotation methods (QuaRot, SpinQuant) push the same idea further: an orthogonal rotation spreads every outlier across all channels, which makes 4-bit activations practical.

Source code

The exact code that runs, read from the repository when this page was built.

smoothquant lab/layerwise.py:247
def smoothquant(net: MLP, X: np.ndarray, alpha: float = 0.5) -> MLP:
    """SmoothQuant: per input channel j of layer i, s_j = max|x_j|^a / max|w_j|^(1-a).
    Divide the activation by s (fold into the previous layer: legal because
    relu is positively homogeneous) and multiply the weights by s. Same
    function; activation outliers shrink, weights absorb the difficulty."""
    net = net.copy()
    _, hid, _ = net.forward(X)   # scaling layer i's inputs never changes layer i's own inputs
    for i in range(1, net.n_layers):
        if net.res[i - 1] or net.res[i]:
            continue
        a = np.abs(hid[i - 1]).max(0)
        w = np.abs(net.W[i]).max(1) + 1e-8
        s = a ** alpha / w ** (1 - alpha)
        # a channel that never fired on calibration data has no range to fix;
        # scaling it by a tiny s would blow up its weights for unseen inputs
        silent = a < 1e-3 * a.max()
        s = np.where(silent, 1.0, np.clip(s, 1e-3, 1e3)).astype(F)
        net.W[i] = net.W[i] * s[:, None]
        net.W[i - 1] = net.W[i - 1] / s[None, :]
        net.b[i - 1] = net.b[i - 1] / s
        hid[i - 1] = hid[i - 1] / s
    return net
inject_outliers lab/layerwise.py:230
def inject_outliers(net: MLP, layer: int = 0, k: int = 3, factor: float = 30.0,
                    seed: int = 0) -> MLP:
    """Create LLM-style activation outlier channels *without changing the
    function*: multiply k neurons' outputs by ``factor`` and divide the next
    layer's matching weights by it. Because relu(c z) = c relu(z) for c > 0,
    the network computes exactly the same thing, but its activations now have
    a few huge channels, which is what breaks per-tensor 8-bit activations
    in real LLMs."""
    net = net.copy()
    rng = np.random.default_rng(seed)
    idx = rng.choice(net.sizes[layer + 1], k, replace=False)
    net.W[layer][:, idx] *= factor
    net.b[layer][idx] *= factor
    net.W[layer + 1][idx, :] /= factor
    return net
forward_wa lab/layerwise.py:218
def forward_wa(net: MLP, X: np.ndarray, abits: int) -> np.ndarray:
    """Forward pass that also quantizes the input of every layer after the
    first (W8A8-style). Weights are whatever ``net.eff`` gives; quantize them
    first (e.g. with ``apply_layerwise`` + ``rtn``) for the full W+A setting."""
    h = X
    for i in range(net.n_layers):
        if i > 0:
            h = quant_act(h, abits)
        h = layer_forward(net, i, h)
    return h
quant_act lab/layerwise.py:207
def quant_act(h: np.ndarray, bits: int) -> np.ndarray:
    """Dynamic per-token activation quantization: one scale per row (per
    example here, per token in an LLM), computed on the fly. A single huge
    channel still sets every row's scale, which is the problem SmoothQuant fixes."""
    if bits >= 16:
        return h
    qmax = qmax_for(bits)
    scale = np.maximum(np.abs(h).max(1, keepdims=True), 1e-12) / qmax
    return np.clip(np.round(h / scale), -qmax, qmax) * scale

Simulation figures

No simulation in the lab covers this method yet.

Papers

  • Xiao et al. SmoothQuant (2022)
  • Dettmers et al. LLM.int8() (2022)

Run it in your browser

16 = leave weights in float; 8 = the W8A8 setting

Not run yet.

A run loads the lab's trained model, compresses it and scores the result on the lab's toy problem, here in your browser. Nothing is trained. The first run downloads Python (about 20 MB); most runs then take 5 to 60 seconds. Every model a run makes smaller ends in a bird.

Pigeon's warning

Channels that are silent during calibration have no measured range. Scaling them by the formula blows up their weights for unseen inputs; the code leaves them at s = 1.

When it breaks: all 4 points

Build it yourself

No Flight School lesson builds this method yet.

All Flight School lessons