A training loop from scratch: AdamW, accumulation, mixed precision, BPE

Before this lesson: lesson 06 (the from-scratch routine, stable softmax and cross-entropy). Background: Adam and AdamW and learning-rate schedules and clipping in the Optimization track, backpropagation (what loss.backward() computes), and floating point (machine epsilon, ulp). Time: about 80 minutes, plus the notebook.

Why this matters #

ML engineer and research engineer interviews for LLM roles often include a round like "write the training step", "implement AdamW" or "implement a BPE tokenizer". Stanford CS336, assignment 1, asks for exactly these pieces, and many interviewers know it. The round tests two things: the exact order of the operations, and the numerics behind them. The bugs here are silent. A loss that is not divided by the number of micro-batches, clipping before unscaling, or weight decay inside Adam's gradient all train without an error, but to a different result. The same loop, scaled up with data parallelism, trains LLMs, vision models and scientific foundation models (lesson 12 covers the memory and the parallelism).

1. Predict first #

2. The picture #

2.1 The five-line loop #

One training step estimates the gradient of the mean loss over a batch, then moves the parameters. The simplest loop:

import torch

torch.manual_seed(0)
model = torch.nn.Linear(3, 1)
opt = torch.optim.SGD(model.parameters(), lr=0.1)
X, y = torch.randn(64, 3), torch.randn(64, 1)
for step in range(100):
    loss = ((model(X) - y) ** 2).mean()   # forward: build the graph, compute the loss
    opt.zero_grad(set_to_none=True)       # clear old gradients (backward ADDS into .grad)
    loss.backward()                       # gradients of the loss into p.grad
    opt.step()                            # update the parameters from p.grad
print(round(loss.item(), 3))              # 0.766

The key fact for this lesson: backward() adds into p.grad; it does not overwrite it. This is why zero_grad is needed, and it is also what makes gradient accumulation work. Each of the four pieces below fixes one practical problem on top of these lines.

2.2 AdamW: the step size comes from the learning rate, not from the gradient #

SGD moves by $\gamma g$: a gradient 1000 times larger gives a step 1000 times larger. Adam divides a running mean of the gradient by the square root of a running mean of the squared gradient. If the gradient is multiplied by $c$, both the top and the bottom are multiplied by $c$, so the step stays the same. The step size is about $\gamma$, the learning rate, in every coordinate.

Weight decay pulls every weight towards zero a little at each step. AdamW applies it as a separate step, $\theta \leftarrow (1 - \gamma\lambda)\theta$, outside Adam's division. Section 3 shows why this matters.

2.3 Gradient accumulation: many small backward passes, one step #

A batch of 512 sequences may not fit in GPU memory. Split it into $K$ micro-batches, run forward and backward on each one, and let p.grad add up the gradients. Then take one optimizer step. Gradients are linear in the loss, so if each micro-batch loss is divided by $K$, the sum equals the gradient of the full-batch mean.

For language models there is a trap. The loss is a mean over tokens, and micro-batches can hold different numbers of target tokens (padding, packing, different lengths). Then "the mean of the micro-batch means" gives different tokens different weights.

2.4 Mixed precision: 16-bit numbers are coarse #

A 16-bit number has either few mantissa bits (bf16: 8 significant bits) or a small range (fp16: largest value 65504, and tiny numbers become 0). So the matrix products run in 16 bits, but the weights, the optimizer state and long sums stay in fp32.

2.5 BPE: merge the most frequent pair, again and again #

A tokenizer turns text into integer ids. Byte-level BPE (byte-pair encoding) starts with the 256 byte values as the vocabulary. It counts every pair of neighbouring tokens in the training text, merges the most frequent pair into one new token, and repeats. Each merge makes the sequences shorter, and shorter sequences make the model cheaper to run.

3. Definitions and main results #

3.1 AdamW #

Notation (as in the PyTorch documentation): $\theta$ are the parameters, $g_t$ is the gradient at step $t$, $\gamma$ is the learning rate, $\lambda$ is the weight decay, $\beta_1, \beta_2 \in [0, 1)$ are the decay rates of the two running averages, and $\varepsilon$ is a small constant. The optimizer state of each parameter tensor is $m$ and $v$ (same shape as $\theta$, both start at 0) and the step count $t$.

AdamW, as in torch.optim.AdamW. For $t = 1, 2, \dots$:

1. $\theta \leftarrow \theta - \gamma\lambda\theta$ (decoupled decay: multiply by $1 - \gamma\lambda$);
2. $m \leftarrow \beta_1 m + (1 - \beta_1)\, g_t$ (the first moment, a running mean of the gradient);
3. $v \leftarrow \beta_2 v + (1 - \beta_2)\, g_t^2$ (the second moment, elementwise square);
4. $\hat m = m / (1 - \beta_1^t)$, $\quad \hat v = v / (1 - \beta_2^t)$ (bias correction);
5. $\theta \leftarrow \theta - \gamma\, \hat m / (\sqrt{\hat v} + \varepsilon)$.

PyTorch defaults: $\gamma = 10^{-3}$, $\beta = (0.9, 0.999)$, $\varepsilon = 10^{-8}$, $\lambda = 0.01$. LLM pretraining usually takes $\beta_2 = 0.95$ (GPT-3, LLaMA) and $\lambda = 0.1$, applied to weight matrices only, not to biases and normalization gains.

3.2 Bias correction #

Why it is needed, in one line. $m$ starts at 0, so after one step $m_1 = (1 - \beta_1) g_1 = 0.1\,g_1$: far too small. The correction divides by exactly the missing factor.

Theorem (bias correction). If every $g_s$ has the same mean $\mu$, then $\mathbb{E}[m_t] = (1 - \beta_1^t)\,\mu$, so $\hat m_t = m_t/(1 - \beta_1^t)$ is an unbiased estimate of $\mu$.

Proof.

1. $m_0 = 0$, so unrolling the update gives $m_t = (1 - \beta_1)\sum_{s=1}^{t} \beta_1^{t-s} g_s$.
2. Take expectations: $\mathbb{E}[m_t] = \mu (1 - \beta_1)\sum_{k=0}^{t-1}\beta_1^k$.
3. The geometric sum is $(1 - \beta_1^t)/(1 - \beta_1)$, so $\mathbb{E}[m_t] = (1 - \beta_1^t)\,\mu$.
4. If every $g_s^2$ also has the same mean $\mathbb{E}[g^2]$, the same argument gives $\mathbb{E}[v_t] = (1 - \beta_2^t)\,\mathbb{E}[g^2]$, so $\hat v_t$ is unbiased. $\blacksquare$

Two consequences.

• The first step has size $\gamma$ in every coordinate. At $t = 1$, $\hat m_1 = g_1$ and $\hat v_1 = g_1^2$, so $\hat m_1/\sqrt{\hat v_1} = g_1/|g_1| = \pm 1$. Without the correction the first step is $m_1/\sqrt{v_1} = (1 - \beta_1)/\sqrt{1 - \beta_2} \approx 3.16$ times $\operatorname{sign}(g)$ for the default betas.
• Adam is invariant to the scale of the loss. Replace $g$ by $cg$ with $c > 0$: $\hat m$ is multiplied by $c$, $\sqrt{\hat v}$ by $c$, and the ratio does not change (ignoring $\varepsilon$).

In code, $t$ is the step count after the increment, so $t = 1$ at the first step. With $t = 0$, $1 - \beta^0 = 0$ and you divide by zero.

3.3 Decoupled weight decay vs L2 #

For SGD, an L2 term $\tfrac{\lambda}{2}\|\theta\|^2$ in the loss and decoupled decay are the same: $\theta - \gamma(g + \lambda\theta) = (1 - \gamma\lambda)\theta - \gamma g$.

For Adam they differ. With L2, the term $\lambda\theta$ is added to $g$, goes into $m$, and is then divided by $\sqrt{\hat v}$. A weight with a large gradient history has a large $\sqrt{\hat v}$, so its decay is divided by a large number: it is barely regularized. A weight with tiny gradients has a tiny $\sqrt{\hat v}$ and is pulled to zero fast. AdamW applies the decay outside the adaptive scaling, so every weight shrinks by the same factor $1 - \gamma\lambda$ per step (Loshchilov and Hutter, 2019). Because PyTorch multiplies $\lambda$ by $\gamma$, the decay also follows the learning-rate schedule.

Memory. AdamW keeps two fp32 tensors per parameter, $m$ and $v$: 8 bytes. With mixed precision the usual count is 16 bytes per parameter: bf16 weights (2) + bf16 gradients (2) + fp32 master weights (4) + $m$ (4) + $v$ (4). A 7B-parameter model needs $16 \cdot 7 \cdot 10^9 = 112$ GB before any activations (lesson 12 splits this over GPUs).

3.4 How a PyTorch optimizer is built #

An interviewer may ask for AdamW as a subclass of torch.optim.Optimizer. The class below is not AdamW; it is sign descent, $\theta \leftarrow \theta - \gamma\operatorname{sign}(g)$. It shows the four parts of the API: parameter groups, lazy per-parameter state, no_grad, and in-place updates.

import torch

class SignSGD(torch.optim.Optimizer):
    """p <- p - lr * sign(grad). Only here to show the Optimizer API."""

    def __init__(self, params, lr=1e-3):
        if lr <= 0:
            raise ValueError(f"invalid lr: {lr}")
        super().__init__(params, defaults={"lr": lr})   # builds self.param_groups

    @torch.no_grad()                                     # the update is not part of the graph
    def step(self, closure=None):
        loss = None if closure is None else closure()
        for group in self.param_groups:                  # a list of dicts: params + hyperparameters
            for p in group["params"]:
                if p.grad is None:                       # frozen or unused parameter
                    continue
                state = self.state[p]                    # per-parameter dict, empty at first
                if not state:                            # lazy init on the first step
                    state["step"] = 0                    # AdamW also adds exp_avg, exp_avg_sq here
                state["step"] += 1
                p.add_(torch.sign(p.grad), alpha=-group["lr"])   # in place: keep the same tensor
        return loss

p = torch.nn.Parameter(torch.tensor([0.0]))
opt = SignSGD([p], lr=0.25)
for _ in range(20):
    opt.zero_grad()
    ((p - 3.0) ** 2).sum().backward()
    opt.step()
print(p.item(), opt.state[p]["step"])   # 3.0 20

For AdamW, add exp_avg and exp_avg_sq (torch.zeros_like(p)) to the state, and replace the update line by steps 1–5 of section 3.1, with $t$ = state["step"] after the increment.

3.5 Gradient accumulation #

Theorem (accumulation is exact). Split a batch into $K$ micro-batches of equal size, with mean losses $L_1, \dots, L_K$. The full-batch mean loss is $L = \frac{1}{K}\sum_k L_k$. If you call backward() on $L_k / K$ for every $k$ without zeroing in between, then p.grad equals $\nabla L$ (up to float rounding).

Proof.

1. The gradient is linear: $\nabla L = \frac{1}{K}\sum_k \nabla L_k = \sum_k \nabla (L_k/K)$.
2. Each backward() adds $\nabla (L_k/K)$ into p.grad, which starts at zero (or None).
3. After $K$ calls, p.grad $= \sum_k \nabla(L_k/K) = \nabla L$. $\blacksquare$

The one exception is BatchNorm: it computes statistics per micro-batch, so its result changes.

Token-level losses. For language models the right loss is the mean over all target tokens in the accumulation window. Let micro-batch $k$ hold $n_k$ target tokens. The mean of the micro-batch means gives each token of micro-batch $k$ the weight $1/(K n_k)$ instead of $1/\sum_j n_j$; the two agree only when all $n_k$ are equal. The fix: sum the token losses (reduction="sum") and divide by the total token count of the window, counted before the backward passes (and summed over all GPUs with data parallelism). Hugging Face Transformers fixed this bug in its Trainer in 2024.

With DDP (distributed data parallel), every backward() would start an all-reduce of the gradients over the GPUs. Wrap micro-batches $0, \dots, K-2$ in model.no_sync(), so the all-reduce runs once per optimizer step instead of $K$ times.

3.6 Floating-point formats and the AMP recipe #

Machine epsilon is the gap between 1 and the next number of the format. The ulp (unit in the last place) of $x$, $\operatorname{ulp}(x)$, is the gap between $x$ and the next larger number. Two failure modes follow:

• Underflow (fp16). Values below about $3 \cdot 10^{-8}$ (half the smallest subnormal $2^{-24} \approx 5.96 \cdot 10^{-8}$) round to 0, and values below $6.1 \cdot 10^{-5}$ lose precision as subnormals. Small gradients disappear.
• Stagnation (bf16, and fp16 for large sums). Adding $x$ to $S$ does nothing when $|x| < \operatorname{ulp}(S)/2$. In bf16, $\operatorname{ulp}(1) = 2^{-7}$, so an update of $10^{-3}$ to a weight of size 1 is lost.

The AMP recipe (automatic mixed precision):

1. Master weights and optimizer state stay in fp32; the optimizer updates fp32 numbers.
2. Under torch.autocast, matrix products and convolutions take 16-bit inputs; the tensor cores accumulate the products in fp32.
3. Under CUDA autocast, ops with exp, log or long reductions (softmax, log-softmax, layer norm, sums, losses) run in fp32.
4. fp16 only: dynamic loss scaling. Multiply the loss by a scale $S$ (start $2^{16}$), so small gradients do not underflow in the backward pass. Unscale the gradients before clipping and before the step. If any gradient is inf or NaN, skip the step and halve $S$; after 2000 good steps in a row, double $S$ (the GradScaler defaults).
5. bf16 has the fp32 range, so it needs no loss scaling. It has only 8 significant bits, so long sums and small updates must not happen in bf16.

3.7 The full training step #

The right order is: zero the gradients → forward under autocast → scaled backward (for each micro-batch) → unscale → clip → step → update the scaler → step the schedule. The code runs on a CPU with bf16 autocast. On an older GPU with fp16, set amp_dtype = torch.float16; the same code then uses the scaler.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR

torch.manual_seed(0)
device = "cuda" if torch.cuda.is_available() else "cpu"
amp_dtype = torch.bfloat16                 # torch.float16 only on GPUs without bf16
use_scaler = amp_dtype == torch.float16    # bf16 has the fp32 range: no loss scaling

W_true = torch.randn(16, 4, device=device)  # toy task: a 4-class linear teacher
def get_batch(n=32):
    x = torch.randn(n, 16, device=device)
    return x, (x @ W_true).argmax(-1)

model = nn.Sequential(nn.Linear(16, 64), nn.GELU(), nn.Linear(64, 4)).to(device)
decay = [p for p in model.parameters() if p.dim() >= 2]      # weight matrices
no_decay = [p for p in model.parameters() if p.dim() < 2]    # biases (and norm gains)
opt = torch.optim.AdamW([{"params": decay, "weight_decay": 0.1},
                         {"params": no_decay, "weight_decay": 0.0}],
                        lr=3e-3, betas=(0.9, 0.95))
n_steps, warmup, K, max_norm = 200, 20, 4, 1.0              # K micro-batches per step
sched = SequentialLR(opt, [LinearLR(opt, 0.05, 1.0, warmup),
                           CosineAnnealingLR(opt, n_steps - warmup, eta_min=3e-4)], [warmup])
scaler = torch.amp.GradScaler(device, enabled=use_scaler)

for step in range(n_steps):
    opt.zero_grad(set_to_none=True)          # once per optimizer step, not per micro-batch
    total = 0.0
    for micro in range(K):                   # with DDP: model.no_sync() for micro < K - 1
        x, y = get_batch()
        with torch.autocast(device_type=device, dtype=amp_dtype):
            logits = model(x)                # matmuls in bf16, weights stay fp32
        loss = F.cross_entropy(logits.float(), y)    # loss in fp32
        scaler.scale(loss / K).backward()    # outside autocast; grads add up in p.grad (fp32)
        total += loss.item() / K
    scaler.unscale_(opt)                     # true gradients, before clipping
    norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
    scaler.step(opt)                         # skips the step if a gradient is inf/NaN
    scaler.update()                          # adapts the scale (fp16 only)
    sched.step()                             # once per optimizer step
    if step in (0, n_steps - 1):
        print(step, round(total, 3), round(sched.get_last_lr()[0], 5))
# 0 1.42 0.00029
# 199 0.261 0.0003

The loss falls, and the learning rate first rises (warmup: linear, from 5% of the peak over 20 steps) and then decays along a cosine to its minimum. Gradient clipping rescales all gradients together when their global norm $\sqrt{\sum_p \|\nabla_p\|^2}$ is above max_norm.

3.8 Byte-level BPE #

Training (GPT-2 style, as in CS336 §2.4):

1. Vocabulary: the 256 single bytes, plus special tokens such as <|endoftext|>.
2. Pre-tokenize: split the text into chunks (GPT-2 uses a regular expression; this lesson splits on spaces) and count each distinct chunk. Merges never cross chunk boundaries or special tokens.
3. Count pairs: write each chunk as a tuple of tokens (bytes at first). Count adjacent pairs, each weighted by the chunk's count.
4. Pick a pair: the most frequent one. Ties: CS336 takes the lexicographically greater pair (compare the byte strings, first element first). Libraries differ, so state your rule.
5. Merge: replace every occurrence of the pair by one new token, append the pair to the list of merges, and go back to step 3. Stop when the vocabulary reaches the target size (256 + specials + number of merges).

Encoding: pre-tokenize the same way, then inside each chunk apply the merges in the order they were learned. Decoding: concatenate the bytes of the ids and decode UTF-8 with errors="replace", because an arbitrary id sequence may not be valid UTF-8.

No unknown token. Every string is a sequence of bytes 0–255, and all 256 bytes are in the vocabulary, so byte-level BPE can encode any text.

Cost. The simple version recounts all pairs for every merge: $O(MN)$ for $M$ merges and $N$ symbols in the table of distinct chunks. The fast version keeps the pair counts and an index from each pair to the chunks that contain it, and updates only the chunks that a merge touches.

Answers to section 1. (1) With a constant gradient, $\hat m = g$ and $\hat v = g^2$ exactly at every step, so the step is $\gamma g/(|g| + \varepsilon) \approx 0.1$ for both $g = 0.5$ and $g = 0.0005$. Adam's step size comes from $\gamma$, not from the gradient. (2) SGD: the gradient is 4 times larger, which acts like a 4 times larger learning rate. AdamW: $\hat m/\sqrt{\hat v}$ does not depend on the gradient scale, so the update hardly changes. But clipping now sees a norm 4 times larger: it starts when the true norm is $0.25$ instead of $1.0$. So the steps where clipping is active change, and with them the training run.

4. Worked example #

(a) AdamW against Adam with L2. One parameter, $\theta_0 = 1$, constant gradient $g = 0.5$, $\gamma = 0.1$, $\lambda = 0.1$, $\beta = (0.9, 0.999)$, $\varepsilon$ negligible.

AdamW, step 1.

1. Decay: $\theta = 1 \cdot (1 - 0.1 \cdot 0.1) = 0.99$.
2. $m_1 = 0.1 \cdot 0.5 = 0.05$, $\quad v_1 = 0.001 \cdot 0.25 = 0.00025$.
3. $\hat m = 0.05/0.1 = 0.5$, $\quad \hat v = 0.00025/0.001 = 0.25$.
4. Step $= 0.1 \cdot 0.5/\sqrt{0.25} = 0.1$. So $\theta_1 = 0.99 - 0.1 = 0.89$.

AdamW, step 2.

1. Decay: $0.89 \cdot 0.99 = 0.8811$.
2. $m_2 = 0.9 \cdot 0.05 + 0.05 = 0.095$, $\quad v_2 = 0.999 \cdot 0.00025 + 0.00025 = 0.00049975$.
3. $\hat m = 0.095/0.19 = 0.5$, $\quad \hat v = 0.00049975/0.001999 = 0.25$.
4. Step $0.1$. So $\theta_2 = 0.7811$.

Adam with L2 ($\lambda = 0.1$ added to the gradient). Step 1: $g' = 0.5 + 0.1 \cdot 1 = 0.6$; with a single gradient, $\hat m/\sqrt{\hat v} = 1$ again, so $\theta_1 = 0.9$. Step 2: $g' = 0.5 + 0.1 \cdot 0.9 = 0.59$, and $\theta_2 = 0.8000473$. Plain Adam without any decay gives $0.9$ and $0.8$. The L2 term almost vanished inside the adaptive scaling; AdamW shrank the weight by 1% per step, as intended.

(b) Accumulation with unequal token counts. Micro-batch 1 has 2 target tokens with losses 4 and 2; micro-batch 2 has 4 tokens with loss 1 each.

• Correct: total loss $4 + 2 + 1 + 1 + 1 + 1 = 10$ over 6 tokens: $10/6 = 1.667$.
• Mean of means: $(3 + 1)/2 = 2.0$. Tokens of micro-batch 1 get weight $1/4$, tokens of micro-batch 2 get $1/8$, instead of $1/6$ each.

(c) 16-bit numbers.

• A gradient of $10^{-8}$ in fp16 is 0: it is below half the smallest subnormal, $2^{-24}/2 \approx 2.98 \cdot 10^{-8}$. Scaled by $2^{16}$ it is $6.55 \cdot 10^{-4}$, a normal fp16 number. Unscaled in fp32 it is $9.997 \cdot 10^{-9}$ (fp16 keeps about 3 decimal digits).
• Add 1.0 a thousand times to a bf16 accumulator. The sum stops at 256: there $\operatorname{ulp} = 2$, and $256 + 1$ is a tie that rounds to the even neighbour, 256. An fp16 accumulator stops at 2048 for the same reason.
• A bf16 weight $w = 1$ plus an update of $10^{-3}$ stays 1, because $10^{-3} < \operatorname{ulp}(1)/2 = 0.0039$.

(d) BPE by hand. Corpus after pre-tokenization: {low: 3, lot: 2, slow: 3, slot: 1}.

In round 3, (lo, t) and (s, low) both have count 3. Compare the first elements: s > lo (the byte s is greater than the byte l), so (s, low) wins.

Merges: [(l,o), (lo,w), (s,low), (lo,t), (s,lot)]; vocabulary size $256 + 5$ plus specials. Encoding new words:

• slots: s l o t s → merge 1 → s lo t s → merges 2 and 3 do not apply → merge 4 → s lot s → merge 5 → slot s: 2 tokens.
• lowest → low e s t.
• owl → o w l: the pair (o, w) is not a merge; only (lo, w) is.

5. In code #

PyTorch checks the numbers above (the notebook does the same in NumPy, and has the exercises).

import torch
import torch.nn.functional as F
from collections import Counter

f64 = torch.float64

# (a) AdamW vs Adam + L2 vs Adam, constant gradient 0.5
def run(opt_cls, **kw):
    p = torch.nn.Parameter(torch.tensor([1.0], dtype=f64))
    opt, out = opt_cls([p], lr=0.1, **kw), []
    for _ in range(2):
        p.grad = torch.tensor([0.5], dtype=f64)
        opt.step()
        out.append(round(p.item(), 7))
    return out
print(run(torch.optim.AdamW, weight_decay=0.1))   # [0.89, 0.7811]
print(run(torch.optim.Adam, weight_decay=0.1))    # [0.9, 0.8000473]   (Adam's weight_decay is L2)
print(run(torch.optim.Adam))                      # [0.9, 0.8]

# Adam is (almost) invariant to the gradient scale; SGD is not (prediction 2)
def final(opt_cls, scale):
    torch.manual_seed(0)
    p = torch.nn.Parameter(torch.randn(5, dtype=f64))
    opt = opt_cls([p], lr=0.01)
    for _ in range(20):
        opt.zero_grad()
        (scale * (p ** 2).sum()).backward()
        opt.step()
    return p.detach()
for cls in (torch.optim.AdamW, torch.optim.SGD):
    print(cls.__name__, f"{(final(cls, 1.0) - final(cls, 4.0)).abs().max().item():.1e}")
# AdamW 2.6e-09
# SGD 1.0e+00

# (b) per-micro-batch means weight tokens unequally
mb1, mb2 = torch.tensor([4.0, 2.0]), torch.tensor([1.0, 1.0, 1.0, 1.0])
print((mb1.mean() + mb2.mean()).item() / 2, (mb1.sum() + mb2.sum()).item() / 6)  # 2.0 1.666...

# accumulation with equal micro-batches is exact
torch.manual_seed(0)
X, y = torch.randn(32, 3, dtype=f64), torch.randn(32, dtype=f64)
w = torch.zeros(3, dtype=f64, requires_grad=True)
((X @ w - y) ** 2).mean().backward()
g_full, w.grad = w.grad.clone(), None
for Xb, yb in zip(X.chunk(4), y.chunk(4)):
    (((Xb @ w - yb) ** 2).mean() / 4).backward()            # grads add up in w.grad
print(f"{(g_full - w.grad).abs().max().item():.1e}")         # 1.1e-16

# (c) 16-bit formats
g = torch.tensor(1e-8)
print(g.half().item(), (g * 2**16).half().item())          # 0.0 0.0006551742553710938
s = torch.tensor(0.0, dtype=torch.bfloat16)
for _ in range(1000):
    s += 1.0
print(s.item())                                             # 256.0
w = torch.tensor(1.0, dtype=torch.bfloat16)
print((w + 1e-3).item(), torch.finfo(torch.bfloat16).eps)   # 1.0 0.0078125

# which dtypes autocast picks (CPU, bf16)
lin = torch.nn.Linear(8, 4)
with torch.autocast("cpu", dtype=torch.bfloat16):
    out = lin(torch.randn(2, 8))
    loss = F.cross_entropy(out, torch.tensor([0, 1]))
loss.backward()
print(lin.weight.dtype, out.dtype, loss.dtype, lin.weight.grad.dtype)
# torch.float32 torch.bfloat16 torch.float32 torch.float32

# (d) BPE round-1 pair counts, weighted by word counts
corpus = {"low": 3, "lot": 2, "slow": 3, "slot": 1}
pairs = Counter()
for word, c in corpus.items():
    for a, b in zip(word, word[1:]):
        pairs[a + b] += c
print(pairs.most_common())   # [('lo', 9), ('ow', 6), ('sl', 4), ('ot', 3)]

Only the matrix-product output is bf16; the weights, the gradients and the loss stay fp32. On the CPU, autocast does not move softmax and layer norm to fp32; the CUDA op lists do.

6. Where it shows up #

• LLM pretraining. Karpathy's build-nanogpt uses exactly this loop: AdamW with $\beta = (0.9, 0.95)$, weight decay 0.1 on parameters with dim() >= 2, gradient accumulation up to 524,288 tokens per step, bf16 autocast, clipping at 1.0, linear warmup then cosine decay, and gradient synchronization only on the last micro-step (the manual form of no_sync).
• Large-scale training with FSDP. The mixed-precision policy sets the dtype of the parameters, of the gradient reduction and of the buffers separately; reducing gradients in fp32 is the safe default. On A100 and H100 GPUs, bf16 removes the need for loss scaling.
• Scientific ML. PDE residuals and spectral loss terms can span many orders of magnitude. Compute them in fp32, even when the network runs in bf16; an fp16 residual overflows above 65504.
• Tokenizers. BPE (GPT-2, LLaMA 3 and most LLMs) or a close relative decides the sequence length and therefore the cost. Science models often replace tokenizers by patches of a grid, but the trade-off is the same: fewer, larger tokens make attention cheaper and each token harder to model.

7. Common mistakes #

• Loss not divided by $K$, or per-micro-batch means with unequal token counts. Divide the summed token loss by the total token count of the window.
• zero_grad inside the micro-batch loop. Only the last micro-batch's gradient survives. Zero once per optimizer step.
• Clipping before unscale_. The threshold then acts on gradients multiplied by $S$. Unscale, then clip.
• scaler.update() or the scheduler stepped every micro-batch. Both belong once per optimizer step.
• Backward or the optimizer step inside autocast. Only the forward pass belongs there.
• Adam(weight_decay=...) when AdamW was meant, and decay on biases, norm gains or embeddings. Use AdamW with two parameter groups.
• Bias correction with $t$ starting at 0. $1 - \beta^0 = 0$: a division by zero. Increment first.
• Optimizer state bugs: updating p outside torch.no_grad(), assigning a new tensor (p = p - ...) instead of an in-place update, or keeping $m$ and $v$ in bf16 instead of fp32.
• fp16 without a scaler, or long sums and weight updates in bf16.
• BPE bugs: merges across chunks or special tokens, pair counts without the chunk weights, a tie rule that depends on dictionary order, and encoding by greedy longest match instead of merge order.

8. Check yourself #

1. Adam without bias correction, $\beta = (0.9, 0.999)$. How large is the first step $m_1/\sqrt{v_1}$, relative to $\operatorname{sign}(g)$?
   Answer

   $m_1 = 0.1g$ and $v_1 = 0.001g^2$, so $m_1/\sqrt{v_1} = 0.1/\sqrt{0.001} \approx 3.16$ times $\operatorname{sign}(g)$: about 3 times the intended step.
2. AdamW with $\gamma = 10^{-3}$ and $\lambda = 0.1$, and zero gradient. By which factor does a weight shrink per step, and after 10,000 steps?
   Answer

   $1 - \gamma\lambda = 1 - 10^{-4}$ per step; $(1 - 10^{-4})^{10000} \approx e^{-1} \approx 0.368$ after 10,000 steps.
3. DDP on 16 GPUs, micro-batches of 4 sequences of 1024 tokens, $K = 8$. How many tokens per optimizer step, and how many gradient all-reduces per step with no_sync used correctly?
   Answer

   $16 \cdot 8 \cdot 4 \cdot 1024 = 524{,}288 = 2^{19}$ tokens. One all-reduce of the gradients per step (sent in buckets), during the backward pass of the last micro-batch.
4. fp16 training with scale $S = 2^{16}$. One step produces an inf gradient. What does GradScaler do?
   Answer

   `scaler.step` skips `optimizer.step()` for this step, and `scaler.update()` halves the scale to $2^{15}$. After 2000 steps in a row without inf or NaN, it doubles the scale again.
5. Add the chunk lots with count 4 to the corpus of section 4(d). What are the first three merges?
   Answer

   (l, o) with count 13; then (lo, t) with 7, which now beats (lo, w) with 6; then (lo, w).
6. Why does BPE training count pairs inside pre-tokenized chunks, weighted by chunk counts, instead of scanning the raw text?
   Answer

   Speed: each distinct chunk is processed once, times its count, instead of once per occurrence. Quality: merges cannot cross word or punctuation boundaries, so `dog.` and `dog!` share the token for `dog`.

9. Interview angle #

Recognize the task from the wording.

What to say aloud before coding AdamW (about 30 seconds): "The state per parameter is the step count and two fp32 tensors, $m$ and $v$, initialized to zero. Each step: decay the weight by $1 - \gamma\lambda$, update $m$ and $v$, correct both by $1 - \beta^t$ with $t$ starting at 1, and step by $\gamma \hat m/(\sqrt{\hat v} + \varepsilon)$. All updates in place, under no_grad. I will test 20 steps on a random quadratic against torch.optim.AdamW in float64."

A plan for 45 minutes.

• 0–5 min: clarify (may you subclass torch.optim.Optimizer? use nn.Module? the regex package?), write the signatures, list the state.
• 5–20 min: the core (AdamW steps 1–5, or the BPE count–pick–merge loop).
• 20–30 min: a test. AdamW: the same random quadratic, 20 steps, against torch.optim.AdamW in float64. BPE: the merges of section 4(d) and the 12 merges of CS336's small example.
• 30–40 min: extensions (parameter groups with no decay on biases; special tokens; incremental pair counts; accumulation and AMP in the loop).
• 40–45 min: complexity, memory (16 bytes per parameter), failure modes.

Typical questions and short model answers.

1. "Why AdamW and not Adam with L2?" — "With L2, the decay gradient $\lambda\theta$ goes into $m$ and is divided by $\sqrt{\hat v}$. So weights with a large gradient history are barely regularized, and weights with tiny gradients are decayed too much. AdamW applies the decay directly to the weight, $\theta \leftarrow (1 - \gamma\lambda)\theta$, so every weight shrinks at the same rate. The paper shows that this also makes the best weight decay less dependent on the learning rate."
2. "Why bias correction?" — "$m$ and $v$ start at zero, so in the first steps they are too small: $\mathbb{E}[m_t] = (1 - \beta_1^t)\mu$. Dividing by $1 - \beta^t$ makes them unbiased. Without it the first step is $0.1/\sqrt{0.001} \approx 3.16$ times too large for the default betas, and the error lasts for hundreds of steps through $v$."
3. "Why $\beta_2 = 0.95$ for LLMs?" — "$v$ averages over about $1/(1 - \beta_2)$ steps: 20 instead of 1000. When the gradient scale grows suddenly, a stale, too small $v$ makes very large steps, which can start a loss spike. A shorter memory follows the change faster."
4. "bf16 or fp16?" — "bf16 when the hardware supports it: it has the fp32 range, so there is no loss scaling and no overflow. It has only 8 significant bits, so I keep master weights, optimizer state and reductions in fp32. fp16 has more precision but needs dynamic loss scaling."
5. "How do you make accumulation exact for language models?" — "Count the target tokens of the whole window first, summed over GPUs. Then use the summed token loss of each micro-batch divided by that total. Then every token has the same weight, as in the full batch."
6. "Why byte-level BPE?" — "The 256-byte base vocabulary covers every string, so there is no unknown token, and the merges make frequent text short. Shorter sequences mean less attention cost and a longer effective context."

Typical follow-ups. "What is the memory of AdamW for a 7B model?" (112 GB at 16 bytes per parameter, before activations). "Where does clipping go with accumulation?" (once, after the last micro-batch and after unscale_). "How would you make BPE training fast?" (pair counts with an index from pairs to chunks; update only the chunks a merge touches; a heap for the best pair). "What happens to the decay when the learning rate is warmed up?" (PyTorch multiplies $\lambda$ by $\gamma$, so the decay follows the schedule).

10. Practice #

The notebook for this lesson runs in Colab, in NumPy. Each exercise has tests against a reference:

1. adamw_step: one AdamW step in NumPy, checked against the worked example and a reference run.
2. bias_correction: check numerically that $\mathbb{E}[m_t] = (1 - \beta_1^t)\mu$ for random gradients.
3. accumulate_token_loss: the token-level loss over unequal micro-batches, equal to the full-batch token mean.
4. bf16_round: round float32 to bf16 with bit operations (round to nearest, ties to even), then show that a bf16 accumulator stops at 256.
5. train_bpe, encode and decode: byte-level BPE on a toy corpus, with the CS336 example as a test and a round-trip check.

Timed practice (25–30 minutes each, talking aloud): "implement AdamW as a torch.optim.Optimizer"; "implement clip_grad_norm_ with the global norm"; "write the training step with accumulation and AMP from memory"; "implement BPE training and encoding"; the CS336 assignment 1 tests for the AdamW, learning-rate schedule, gradient clipping and BPE functions.

Say it aloud drill. Without notes, explain in 2 minutes: the order of the training step and what each line protects against. Then explain in 1 minute why bf16 needs no loss scaling.

11. Watch and read #

• Stanford CS336 (Spring 2025), Lecture 2: PyTorch, resource accounting (video on YouTube; embedding is disabled): floating-point formats, FLOP and memory accounting, and the memory of the optimizer state.
• Stanford CS336, assignment 1 handout and tests: §2 (BPE training, the small example, encoding), §4 (AdamW as Algorithm 1, the cosine schedule, gradient clipping), §5 (the training loop).
• Karpathy, build-nanogpt, train_gpt2.py: accumulation, bf16 autocast, clipping and AdamW parameter groups in one file.
• PyTorch, Automatic Mixed Precision examples: gradient clipping and accumulation with GradScaler.
• PyTorch, torch.optim.AdamW: the algorithm box that section 3.1 follows.
• Loshchilov and Hutter (2019), Decoupled Weight Decay Regularization, §2 and Algorithm 2.
• Kingma and Ba (2015), Adam: A Method for Stochastic Optimization, §3 (bias correction).
• Sennrich, Haddow and Birch (2016), Neural Machine Translation of Rare Words with Subword Units: the original BPE for NLP.