Numerics & Hand-Written Backward · AI infra interview lesson

Why It Matters

exp(1000) overflows, but softmax([1000,1001]) equals softmax([-1,0]). Explain why subtracting the maximum preserves the result before testing outputs and gradients.

One concrete number: the Qwen vocabularies have about 152k entries. If the softmax backward built the Jacobian explicitly, one token would need 152k * 152k ≈ 2.3e10 elements, about 92 GB in FP32; the VJP formula dx = y * (g - sum(g*y)) needs only O(V) extra memory. Another: BF16 spacing just above 1.0 is 2^-7 = 0.0078, so an update of 1e-3 (learning rate times gradient) is rounded away entirely. That is why mixed-precision training must keep FP32 master weights.

Walk the Shapes First

Stable softmax / LSE on one row of logits x = [1000, 1001, 1002] (a direct exp(1000) is inf in FP32):

Shapes in the RMSNorm backward, x: (B=2, T=3, D=4), w: (4,):

Rule read off the table: every numerical trick is "subtract a per-row statistic first, then compute"; every backward formula contains one per-row reduction (sum(g*y), mean(gh*xhat)), and weight gradients are sums over the batch and sequence dimensions.

Running the examples: execute this page in order in one Python environment. Later snippets may reuse earlier definitions. CPU examples validate calculations and state changes, not GPU timing. Download the notebook for the complete exercise workspace.

Core Implementations

Step 1: stable softmax, logsumexp, log_softmax

What's new: subtract the max to avoid overflow, and a convention for all -inf rows.

lse(x) = m + log(sum(exp(x - m))), log_softmax = (x - m) - log(sum(exp(x - m))) (mathematically x - lse), softmax = exp(log_softmax). For a row that is entirely -inf (for example fully masked by causal + padding) we use lse = -inf and softmax = 0. Outputting zeros matches how FlashAttention treats fully masked rows; what LSE value such a row gets differs between implementations, so check the source when you depend on it.

import torch
torch.manual_seed(0)

def logsumexp_stable(x, dim=-1):
    m = x.amax(dim=dim, keepdim=True)
    m = torch.where(torch.isinf(m) & (m < 0), torch.zeros_like(m), m)  # replace -inf max by 0
    return (m + torch.log(torch.exp(x - m).sum(dim=dim, keepdim=True))).squeeze(dim)

def log_softmax_stable(x, dim=-1):
    m = x.amax(dim=dim, keepdim=True)
    m = torch.where(torch.isinf(m) & (m < 0), torch.zeros_like(m), m)
    z = x - m                                   # shift first, then subtract log(sum)
    return z - torch.log(torch.exp(z).sum(dim=dim, keepdim=True))

def softmax_stable(x, dim=-1):
    out = torch.exp(log_softmax_stable(x, dim))
    return torch.nan_to_num(out, nan=0.0)   # all -inf row: -inf - (-inf) = NaN -> 0

x = torch.tensor([[1000., 1001., 1002.],
                  [-1000., 0., 1000.],
                  [0., float("-inf"), 1.]])
assert torch.allclose(logsumexp_stable(x), torch.logsumexp(x, -1))
assert torch.allclose(log_softmax_stable(x), torch.log_softmax(x, -1))
assert torch.allclose(softmax_stable(x), torch.softmax(x, -1))
print(softmax_stable(x)[0])        # tensor([0.0900, 0.2447, 0.6652])

# the naive version overflows at 1000
naive = torch.exp(x[0]) / torch.exp(x[0]).sum()
assert torch.isnan(naive).all()

# all -inf rows
z = torch.full((2, 5), float("-inf"))
assert torch.isnan(torch.softmax(z, -1)).all()          # torch gives NaN
assert (softmax_stable(z) == 0).all()                   # our convention gives 0
assert torch.isneginf(logsumexp_stable(z)).all()

# other dims
y = torch.randn(3, 4, 5) * 50
assert torch.allclose(softmax_stable(y, dim=1), torch.softmax(y, dim=1), atol=1e-6)

torch.logsumexp itself returns -inf for an all -inf row, which is correct; the problem is the hand-written x - x.max() step.

log_softmax is not written as x - lse: FP32 spacing near 1002 is about 6e-5, so lse = 1002.4076 is rounded first and the subtraction loses significant digits (measured 2e-5 away from torch.log_softmax, beyond allclose's default tolerance); exp(x - lse) has the same issue. Shifting to z = x - m keeps the numbers near 0, and subtracting log(sum(exp(z))) from there avoids it.

Step 2: CE gradient, softmax VJP, linear backward

What's new: the three most common hand-written backward passes, all checked against autograd in float64.

• Cross entropy (mean reduction over N valid tokens): dlogits = (softmax(logits) - onehot(target)) / N.
• Softmax VJP: y = softmax(x), upstream g, dx = y * (g - sum(g * y, -1, keepdim=True)). No need for the V×V Jacobian diag(y) - y yᵀ.
• y = x @ W.T + b, W: (out, in): dx = dy @ W, dW = dy.T @ x, db = dy.sum(0) (flatten leading dims for multi-dimensional inputs).

import torch
import torch.nn.functional as F
torch.manual_seed(0)

def ce_grad(logits, target, ignore_index=-100):
    valid = target != ignore_index
    n = valid.sum()
    p = torch.softmax(logits, -1)
    onehot = F.one_hot(target.clamp_min(0), logits.shape[-1]).to(logits.dtype)
    return (p - onehot) * valid.unsqueeze(-1) / n

def softmax_vjp(y, g):
    return y * (g - (g * y).sum(-1, keepdim=True))

Propagate gradients through a linear layer

After an upstream gradient arrives, the linear layer needs separate input and weight gradients. Write the shapes of inputs, weights and upstream gradients before checking each matrix product.

def linear_backward(x, W, dy):
    x2, dy2 = x.reshape(-1, x.shape[-1]), dy.reshape(-1, dy.shape[-1])
    return (dy2 @ W).reshape(x.shape), dy2.T @ x2, dy2.sum(0)

# CE with ignore_index
logits = torch.randn(6, 7, dtype=torch.float64, requires_grad=True)
target = torch.tensor([0, 3, -100, 6, 2, -100])
F.cross_entropy(logits, target, ignore_index=-100).backward()
assert torch.allclose(ce_grad(logits.detach(), target), logits.grad)

# softmax VJP
x = torch.randn(4, 9, dtype=torch.float64, requires_grad=True)
g = torch.randn(4, 9, dtype=torch.float64)
y = torch.softmax(x, -1)
y.backward(g)
assert torch.allclose(softmax_vjp(y.detach(), g), x.grad)

# linear backward with a 3-D input
x = torch.randn(2, 3, 5, dtype=torch.float64, requires_grad=True)
W = torch.randn(4, 5, dtype=torch.float64, requires_grad=True)
b = torch.randn(4, dtype=torch.float64, requires_grad=True)
dy = torch.randn(2, 3, 4, dtype=torch.float64)
(x @ W.T + b).backward(dy)
dx, dW, db = linear_backward(x.detach(), W.detach(), dy)
assert torch.allclose(dx, x.grad) and torch.allclose(dW, W.grad) and torch.allclose(db, b.grad)

The CE derivation is two lines: loss = lse(z) - z[t], d lse / dz = softmax(z), d z[t] / dz = onehot(t). That is also why a fused CE kernel can write the gradient out during the forward pass: it only needs the softmax and the target.

Step 3: LayerNorm / RMSNorm forward + backward, fused add + RMSNorm

What's new: the backward of normalization layers, and the "residual add + norm" fusion common in inference. (Pro lesson)

Compare LayerNorm and RMSNorm

LayerNorm subtracts a mean; RMSNorm does not. The backward pass must include that mean dependency, not just change the denominator. (Pro lesson)

Step 4: Welford variance and merging

What's new: single-pass, numerically stable mean and variance, plus merging two partial statistics (needed for parallel reductions). (Pro lesson)

Step 5: BF16 swallowing small updates, master weights, global gradient clipping

What's new: what goes wrong when accumulating in low precision and how to fix it, plus clipping gradients by their global norm. (Pro lesson)

In Real Systems

nano-vllm nanovllm/layers/layernorm.py , RMSNorm : rms_forward does x.float() before rsqrt(var + eps) and casts back to the input dtype befo… (Pro lesson)

Common Pitfalls

Written: exp(x) / exp(x).sum() . Problem: FP32 exp overflows once logits exceed about 88.7, and FP16 much earlier ( exp(11.1) is already abo… (Pro lesson)

Say It in the Interview

"Stable softmax subtracts the row max before exponentiating, and a fully masked row needs a special case because −inf minus −inf is NaN; (Pro lesson)

Hands-on Notebook

Each exercise in the notebook has a problem statement and a function skeleton; write your version, then run the test cell below it. (Pro lesson)