Tutorial 04: RMSNorm — From Math to Code

Normalization layers stabilize transformer training and inference. The modern choice is RMSNorm (Root Mean Square Normalization), which is simpler and faster than LayerNorm. This tutorial derives the math, builds the pure-Python implementation, and shows how a Metal kernel replaces it.

---

1. Why Normalize?

Neural network activations can drift in scale as they pass through layers. If one layer outputs values in the range [0, 1] and the next expects [-10, 10], the model cannot learn effectively. Normalization rescales activations to a stable range at every layer.

The original transformer used Layer Normalization:

LayerNorm(x) = gamma * (x - mean(x)) / sqrt(var(x) + eps) + beta

LayerNorm subtracts the mean (centers the data) and divides by the standard deviation (scales the data). gamma and beta are learnable parameters.

RMSNorm simplifies this by dropping the mean centering:

RMSNorm(x) = gamma * x / sqrt(mean(x^2) + eps)

It only divides by the root mean square. Empirically, this works just as well as LayerNorm for transformers, but is ~10-15% faster because it avoids the mean subtraction.

---

2. The Math, Step by Step

Given an input vector x of dimension D:

Step 1: Square every element

x_squared = x_1^2, x_2^2, ..., x_D^2

Step 2: Compute the mean of the squares

mean_sq = (1/D) * sum(x_i^2)

This is the squared "energy" of the vector.

Step 3: Add epsilon for numerical stability

mean_sq_stable = mean_sq + eps

Epsilon prevents division by zero. If all inputs are zero, the RMS would be zero without epsilon, causing NaN. Typical values: eps = 1e-6 or 1e-5.

Step 4: Take the square root

rms = sqrt(mean_sq_stable)

Step 5: Divide and scale

normalized = (x / rms) * gamma

The division rescales the vector so its RMS is approximately 1. The learnable scale gamma allows the model to adjust the scale if needed.

Concrete Example

x = [3.0, 4.0]    # D = 2

x_squared = [9.0, 16.0]
mean_sq = (9 + 16) / 2 = 12.5
rms = sqrt(12.5) ≈ 3.536

normalized = [3.0/3.536, 4.0/3.536] ≈ [0.849, 1.131]

# Check: RMS of [0.849, 1.131] ≈ 1.0 ✓

---

3. The Pure-Python Implementation

import mlx.core as mx


class RMSNorm:
    def __init__(self, dim: int, weight: mx.array, eps: float = 1e-5):
        """
        Args:
            dim: The normalization dimension (hidden_size of the model).
            weight: The learnable scale parameter (gamma). Shape: (dim,)
            eps: Small constant for numerical stability.
        """
        self.dim = dim
        self.eps = eps
        self.weight = weight

    def __call__(self, x: mx.array) -> mx.array:
        """
        Apply RMSNorm to the input tensor.

        Args:
            x: Input tensor. Shape: (..., dim)
               The "..." means any number of leading dimensions (batch, seq, etc.)
               Only the last dimension (dim) is normalized over.

        Returns:
            Normalized tensor. Same shape as input.
        """
        # Remember the original dtype — we compute in float32 for stability,
        # then cast back to match the model's precision (float16 or bfloat16)
        orig_dtype = x.dtype

        # Cast to float32 to avoid underflow/overflow in intermediate computations
        x = x.astype(mx.float32)

        # Compute: x / rms(x)
        # rms(x) = sqrt(mean(x^2) + eps)
        #
        # Breaking this down:
        #   1. mx.square(x)        — element-wise x^2
        #   2. mx.mean(..., axis=-1, keepdims=True) — mean over the last dim
        #   3. + self.eps          — add epsilon
        #   4. mx.rsqrt(...)       — 1/sqrt (faster than 1/sqrt(x))
        #   5. x * result          — rescale
        x = x * mx.rsqrt(mx.mean(mx.square(x), axis=-1, keepdims=True) + self.eps)

        # Cast back to original precision and apply learnable scale
        x = x.astype(orig_dtype)
        return x * self.weight.astype(orig_dtype)

Key Details

keepdims=True: The mean must have shape (..., 1) so it broadcasts correctly against x's shape (..., D). Without keepdims, the mean would collapse the last dimension and the shapes would not align.

mx.rsqrt instead of 1/mx.sqrt: The reciprocal square root is a single hardware instruction on most GPUs. Computing 1/sqrt(x) directly is faster than 1 / sqrt(x).

Float32 intermediate: Even if the model runs in float16 or bfloat16, the normalization statistics are computed in float32. This prevents precision loss when computing the mean of squared values (which can be very small).

---

4. Verification

# Verify RMSNorm produces outputs with RMS ≈ 1
x = mx.random.normal((2, 4, 128))  # batch=2, seq=4, hidden=128
gamma = mx.ones((128,))             # unit scale
norm = RMSNorm(dim=128, weight=gamma, eps=1e-6)

y = norm(x)

# Compute RMS of each vector — should be approximately 1.0
rms_values = mx.sqrt(mx.mean(mx.square(y), axis=-1))
print(rms_values)
# Should print values close to 1.0 for every (batch, seq) position

---

5. Where RMSNorm Appears in the Transformer

A transformer block applies RMSNorm twice:

input
  │
  ├──→ RMSNorm → Attention → (+ input) ──→ h
  │                                            │
  │                                            ├──→ RMSNorm → FFN → (+ h) ──→ output

In code:

class TransformerBlock:
    def __call__(self, x, offset, cache, mask=None):
        # First norm: normalize before attention
        r = self.self_attn(self.input_layernorm(x), offset, cache, mask)
        h = x + r   # residual connection

        # Second norm: normalize before feed-forward
        r = self.mlp(self.post_attention_layernorm(h))
        return h + r  # residual connection

The pattern is always: norm → compute → residual add. This is called Pre-LN (pre-normalization) and is the standard for modern transformers.

---

6. RMSNorm vs LayerNorm: When It Matters

Property	LayerNorm	RMSNorm
Centers data (subtracts mean)	Yes	No
Parameters	gamma + beta (2×D)	gamma only (D)
Compute cost	Higher	~10-15% faster
Quality	Baseline	Equivalent for transformers

For a model with hidden_size=4096, the extra beta parameter would add 4096 additional floats per layer (4B × 32 layers = 128MB). RMSNorm avoids this overhead with no quality loss.

---

7. From Pure Python to Metal

The pure-Python implementation works correctly but calls three separate MLX operations: square, mean, and rsqrt. A custom Metal kernel can fuse these into a single GPU pass:

# core/kernels.py — Metal-backed RMSNorm
class FastRMSNorm:
    def __init__(self, dim: int, weight: mx.array, eps: float = 1e-5):
        self.dim = dim
        self.weight = weight
        self.eps = eps

    def __call__(self, x: mx.array) -> mx.array:
        # One Metal kernel call instead of three MLX operations
        return tiny_llm_ext_ref.rms_norm(
            mx.contiguous(x),
            mx.contiguous(self.weight.astype(x.dtype)),
            self.eps,
        )

The Metal kernel does the same math — square, mean, rsqrt, multiply — but in a single GPU dispatch. This eliminates two intermediate tensor allocations and two memory round-trips.

The engine chooses between them at initialization time:

# In model.py — selecting between pure Python and Metal paths
norm_cls = FastRMSNorm if features.kernels else RMSNorm

self.input_layernorm = norm_cls(
    config.hidden_size,
    weights.tensor(f"{prefix}.input_layernorm.weight"),
    eps=config.rms_norm_eps,
)

---

Summary

• RMSNorm normalizes by the root mean square, avoiding mean centering.
• The formula: x * rsqrt(mean(x^2) + eps) * gamma
• Pure Python: three MLX ops (square → mean → rsqrt → multiply).
• Metal kernel: single fused GPU dispatch for the same math.
• Applied twice per transformer block: before attention and before feed-forward.

Next: Tutorial 05 — Rotary Position Embeddings (RoPE)