What happens to each norm layer before 8-bit inference.
On this page
Before a model runs in 8-bit integers, each norm layer either disappears into the layer before it, keeps running with its gain and shift moved, or forces a fix for the huge activations around it. Figure 1 maps the three cases.
Each change on this page is done once, offline, before export; the only part that still runs on every token is LayerNorm's centering and division, and RMSNorm's division [2, 5, 6].
| Change | When it runs | What it uses and stores |
|---|---|---|
| BN fold | once, offline, after training and before export | the trained γ, β, μ, σ²; stores nothing new and removes a layer [2] |
| LN gain fold | once, offline | γ moves into the next linear weights and β into its bias; centering and division still run every forward pass [6] |
| SmoothQuant | once, offline, from calibration samples | per-channel factors fold into the previous LayerNorm and the next linear weights; no extra runtime operation [5] |
At inference, BatchNorm with fixed statistics only multiplies by a number and adds a number per channel [1]. A conv is already a multiply and add, so the two collapse into one conv with new weights, and nothing is left to quantize between them. Take one output channel of a 1 × 1 conv followed by BatchNorm2d, the Conv, BN pair that static quantization (Quantization, section 6) fuses first.
With ε = 0.00001 both paths give −5.0999. Put the conv output into BN's inference line and regroup.
Term by term, γ/σ is the per-channel scale, and it grows when σ is small, which can widen the range of weights to quantize: in MobileNet v1 the folded weights show larger outliers than the unfolded ones [3]. μ and β only move the bias. Folding removes the BN compute, its data movement and the quantization of its output [2]. In PyTorch, by hand and with the library call:
import torch
from torch import nn
from torch.nn.utils import fusion
conv = nn.Conv2d(3, 8, 3, bias=False).eval()
bn = nn.BatchNorm2d(8).eval()
with torch.no_grad():
sd = torch.sqrt(bn.running_var + bn.eps)
s = bn.weight / sd
w = conv.weight * s.view(-1, 1, 1, 1)
b = bn.bias - bn.running_mean * s
fused = fusion.fuse_conv_bn_eval(conv, bn)s is γ/σ, and b uses b = 0 because bias=False.During quantization-aware training the batch statistics would make these folded weights change every step, so BN switches to its running statistics after a delay; the schedule and the PyTorch recipe live in Freezing observers and BatchNorm statistics in QAT (Quantization, section 13).
LayerNorm and RMSNorm take their statistics from the token in front of them, so the mean and spread are new for every token and cannot be baked into weights. What can move is the gain: γ multiplies each feature right before the next linear layer, so it folds into that layer's weights as diag(γ) Win, where diag(γ) is the matrix with γ on its diagonal and Win the next layer's weights [6]. β folds into that layer's bias as β Win. That works only when the norm's output feeds linear layers alone, the Pre-LN case.
LLM.int8() measured a few huge activation features, the outliers, in language models [4]; earlier work it cites traces them to LayerNorm.
| LLM.int8() measurement [4] | Result |
|---|---|
| Models | GPT-2, Fairseq and GPT-J, 125M to 13B parameters; 8-bit inference on OPT (Meta's open GPT-style LLMs) up to 175B |
| Outlier size | up to 20 times the other dimensions |
| Where | about 25% of layers at first; all layers from about 6.7B parameters |
| At 6.7B | 150,000 outliers per sequence, in only 6 feature dimensions |
| Zeroing them | perplexity (how surprised the model is by real text, lower is better) rises 600 to 1000% |
One outlier channel stretches the range for the whole activation tensor, so every other channel gets only a few of the 256 levels. Why not give each activation channel its own scale, like the per-channel weight scales in Quantization? Per-channel activation scales do not map well to hardware-accelerated matrix-multiply (GEMM) kernels [5].
Weights are easy to quantize and activations are not, so shift part of the size to the weights and fold the division into the LayerNorm gain from Section 2. Take 2 illustrative tokens shaped like an OPT layer's input, entering a linear layer, with input channel 1 standing for an outlier channel (real outliers are about 100 times larger than most activations [5]). With one int8 range for the whole activation tensor, channel 2 uses only 8 of the 256 levels.
Divide each activation channel by a factor and multiply its weight row by the same factor.
Term by term, diag(s)−1 shrinks the outlier channel, and channel 2 rises from 8 to about 45 levels; diag(s) grows the matching weights, whose tensor max rises from 0.5 to 5.66, so weight row 2 also gets about 45 levels instead of all 256. The difficulty is split, not removed. Because X comes out of a LayerNorm, the division by s folds offline into that LayerNorm's γ and β, so no extra operation runs [5], the same gain-folding move as Section 2. Slide α to move the difficulty:
α = 0.5: activation channel 2 gets 45 of 256 levels, weight row 2 gets 45
The dashed ticks ② mark the levels before smoothing. At α = 0, s = 1/max|W| = 2 halves both activation maxima, so the levels stay at ②. In PyTorch, smoothing one LayerNorm and linear pair from recorded activations:
import torch
from torch import nn
ln, lin = nn.LayerNorm(2), nn.Linear(2, 2, bias=False)
acts = torch.tensor([[64., 1.], [-60., 2.]])
x_max = acts.abs().amax(0)
w_max = lin.weight.abs().amax(0)
alpha = 0.5
s = x_max**alpha / w_max**(1 - alpha)
with torch.no_grad():
ln.weight.div_(s)
ln.bias.div_(s)
lin.weight.mul_(s)acts, the LayerNorm outputs recorded on 2 calibration tokens (Figure 4's X).nn.Linear stores its weight as out × in, so amax(0) gives one max per input channel [7]. The weights here are random, so s differs from (11.31, 2.0).mul_(s) broadcasts s along the input dimension.| Symptom | Cause | Fix | What to log |
|---|---|---|---|
fuse_conv_bn_eval raises "Fusion only for eval!" | the modules are still in training mode [8] | call eval() on both first | nothing; it fails loudly |
| an LLM of about 6.7B parameters or more loses accuracy under plain 8-bit activation quantization | outlier features set the range [4] | mixed-precision decomposition (LLM.int8()) or SmoothQuant | per-channel activation max on calibration data |
| outputs change after folding a LayerNorm gain | the norm's output also fed the residual stream (Post-LN) | fold γ only into layers that are the norm's only reader (Figure 3) | outputs before and after the fold |
BN folding has no knob. SmoothQuant has one:
| Knob | Default | What it does |
|---|---|---|
| α, migration strength | 0.5 | larger α moves more of the outlier's size into the weights (Figure 5); SmoothQuant uses 0.5 for all OPT and BLOOM (BigScience's open multilingual LLM) models and a larger α such as 0.75 when activation outliers are more significant [5] |
When accuracy still drops after folding and smoothing, Debugging a Quantized Model shows how to measure where the error comes from.
torch/nn/utils/fusion.py (fuse_conv_bn_eval, fuse_conv_bn_weights). GitHub.