AI Engineer Exhaustive
EasyPlusLLaMA

RMSNorm

LLaMA

Easy

Root Mean Square Layer Normalization

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

RMSNorm (Zhang and Sennrich, 2019) rescales a vector by its root-mean-square and a learned gain. It does not subtract a mean and has no bias. Llama 3 inherits this from Llama / Llama 2; the 2024 herd paper never writes the formula. You infer the operator from the official RMSNorm module and from the earlier Llama papers the authors say they did not deviate from.

In the released model, RMSNorm appears three times per depth: attention_norm and ffn_norm inside every block (pre-norm), plus a final Transformer.norm before the output projection. Each instance owns its own gain of length (d).

How it works

For a residual activation (x) with last-axis width (d) (dim), Zhang and Sennrich define

\mathrm{RMS}(x) = \sqrt{\frac{1}{d}\sum_{i=1}^{d} x_i^{2}+\varepsilon},\qquad \mathrm{RMSNorm}(x) = \frac{x}{\mathrm{RMS}(x)}\odot g.

(g\in\mathbb{R}^{d}) is weight, initialized to ones. There is no (\beta) and no centering term (\bar x). That is the whole difference from LayerNorm: the mean of (x) is allowed to pass through.

The official _norm is the same expression in multiply-rsqrt form,

\hat x = x \cdot \bigl(\mathrm{mean}(x^{2},\;\text{axis}=-1)+\varepsilon\bigr)^{-1/2},

then forward does weight * hat{x}. Reduction is only over the feature axis; batch and sequence axes stay independent. eps on the module defaults to (10^{-6}), but every constructed Llama 3 norm is passed args.norm_eps, and ModelArgs.norm_eps defaults to (10^{-5}). Checkpoints use the ModelArgs value.

Numerics: _norm runs in float32 (x.float()), then the result is type_as(x) before the gain multiply. The public generate path typically materializes the model in BF16 or FP16, so this upcast is what keeps mean(x^2) from underflowing on long-tailed hidden states.

Shapes: if (x) is (B\times S\times d), the RMS tensor is (B\times S\times 1) and (g) broadcasts over batch and time. The paper’s Table 3 widths (4096 / 8192 / 16384) are exactly these (d) values. FFN inner dimensions are not RMSNorm-ed in the official block — only the residual stream.

Specified vs inferred: the use of RMSNorm (pre-attention, pre-FFN, pre-unembed) is inferred from llama/model.py and the Llama 2 lineage. The 2024 paper discusses architecture at the level of GQA, SwiGLU, RoPE (\theta), and vocab size, and is silent on (\varepsilon), float32, and the absence of bias.

Official code

llama/model.py — class RMSNorm (_norm, forward). Call sites: TransformerBlock.attention_norm / ffn_norm and Transformer.norm. No other normalization appears in that file (no LayerNorm, no QK-norm).

Watch-outs

Sources