RMSNorm

In the last chapter a block was a box: it took an x, ran it through attention and a feed-forward network, and handed back a better x of the same shape. You also saw that x was normalized before it entered each sublayer — self.norm1, self.norm2 — never after. This chapter opens that normalization. It is the first of the four boxes a block reaches for, and the simplest, which is why it comes first.

Start with the problem it solves. A block does not replace x; it adds to it. Attention proposes an adjustment and the block writes x + attention(x); the feed-forward network proposes another and the block writes x + ff(x). Do that twice per block, twenty-eight blocks deep, and the numbers in x have no reason to stay in any particular range — each add can push them larger. Attention and the feed-forward network are built assuming their input arrives at a sane, familiar scale. Hand them a vector that has drifted ten times too large and their internal sums saturate, their softmaxes spike, and the arithmetic that was tuned during training comes apart. Something has to reset the scale of x before each sublayer reads it. That something is RMSNorm.

3.1 What it computes

Here is the whole of it, from forge’s model.py.

class RMSNorm(nn.Module):
    # the fp32 upcast is load-bearing: computing this in bf16 fails parity
    def __init__(self, emb_dim: int, eps: float = 1e-6) -> None:
        super().__init__()
        self.eps = eps
        self.scale = nn.Parameter(torch.ones(emb_dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        dtype = x.dtype
        x32 = x.float()
        var = x32.pow(2).mean(dim=-1, keepdim=True)
        return (x32 * torch.rsqrt(var + self.eps) * self.scale.float()).to(dtype)

Listing 3.1 · forge’s RMSNorm. One learned vector, one scale computed from the input, one multiply.

Read forward line by line. It takes a vector x — one token’s emb_dim numbers — and returns a rescaled version of the same shape. The first move is to upcast to float32; forge marks that line as load-bearing, and it is: the next step squares numbers and averages them, and in half precision — the format the real 0.6B model’s weights arrive in — the small errors in that sum grow enough to break parity. The float32 driver you build below already runs in single precision, so the upcast is inert there; it earns its keep on the real model. Then x32.pow(2).mean(dim=-1, keepdim=True) squares every entry and averages the squares across the feature dimension. Take the reciprocal square root of that average, torch.rsqrt, and you have one over the root-mean-square of the vector — the number that, multiplied through, brings the vector to a root-mean-square of one. That is the whole idea in the name: root, mean, square. The + self.eps is a floor so the reciprocal square root never divides by zero on an all-zero vector.

The last multiply does two things at once. x32 * torch.rsqrt(...) is the normalization — it fixes the vector’s scale. * self.scale.float() then applies a learned per-feature gain: scale is a vector the model trained, one number per feature, so after normalizing to a common size the model can still turn individual features up or down. Normalize to a known scale, then let training decide how loud each feature should be from there.

3.2 Why root-mean-square, and not the mean

The older normalization this replaced, LayerNorm, does one more thing: before dividing by the scale, it subtracts the mean of the vector, recentering it on zero. RMSNorm deliberately skips that step. It divides by the root-mean-square and never subtracts the mean — it fixes the vector’s size without moving its center. That is the whole difference, and it is a deliberate trade: dropping the mean-subtraction removes a reduction and a subtraction from a computation that runs twice per block across every token, and the released weights were trained without it, so forge must match that choice exactly to reproduce them.

Skipping the mean is safe. Squaring the values is not optional. The square is what guarantees the quantity under the root is never negative — a mean of squares cannot be less than zero, so its reciprocal square root is always a real number. Take the mean of the raw values instead, and for a vector that sits near zero on average, that mean can land slightly negative; the reciprocal square root of a negative number is not a number at all. The next section makes that failure concrete.

3.3 Correct, to the bit

The same problem from the last chapter applies here in miniature: RMSNorm can run, return a correctly shaped vector, and still compute the wrong thing. So it gets the same treatment — an independent second implementation of the same formula, the same weights loaded into both, the same input fed to each, and a required agreement between the two outputs. The companion driver builds forge’s RMSNorm and the reference’s, hands both one random (2, 9, 32) batch, and prints the largest disagreement:

The gap is exactly zero — not near zero, zero. That is worth pausing on. In the last chapter forge and the reference disagreed by a rounding-sized 5.96e-07, because they did the same math in a different order and floating-point addition is not perfectly associative. Here the two implementations do the same operations in the same order: square, mean, reciprocal square root, multiply by scale. Identical arithmetic on identical bits produces an identical result, down to the last digit. RMSNorm is simple enough to leave no rounding noise to hide in, which makes the check as sharp as a check can be — any nonzero number here is a real divergence, not float dust.

RMSNorm resets the scale of x, but it treats every position identically — it has no idea which token in the sequence it is looking at, or in what order they arrived. Attention is where positions finally read each other, and for that to mean anything, each vector first has to be told where it sits. That is the next box: rotary position embeddings.