GPT-2 from Scratch

Eight chapters of pieces: tokens, embeddings, positions, attention, blocks, the loss, the sampler. The question now is whether they actually fit together into a model that works. Not a toy that runs, but the real GPT-2, bit-for-bit compatible with the weights OpenAI released in 2019.

That compatibility is the point. If our model can load their weights and produce coherent text, then every shape, every mask, every transpose in our code is right, and we know it because a 124-million-parameter oracle told us so. Then we throw the weights away and train our own.

The code in this chapter mirrors code/lumen/gpt2.py and follows the structure of Karpathy's nanoGPT, which in turn follows OpenAI's original TensorFlow code. It needs only torch; the weight-loading step additionally needs transformers to download the checkpoint.

The plan

GPT-2 is a decoder-only transformer. Text comes in as token ids, each id becomes a vector, a position vector is added, the sum passes through twelve identical blocks, a final LayerNorm cleans it up, and a linear layer turns each position's vector into logits over the vocabulary. That is the whole thing.

The config: 124M in five numbers

Every size of GPT-2 is the same code with different numbers. We keep them in a dataclass so that the model can be built from one object and so that "GPT-2 medium" is a one-line change.

import math
from dataclasses import dataclass
import torch
import torch.nn as nn
import torch.nn.functional as F

@dataclass
class GPTConfig:
    vocab_size: int = 50257   # GPT-2 BPE vocabulary (50,000 merges + 256 bytes + <|endoftext|>)
    block_size: int = 1024    # maximum context length T
    n_layer: int = 12
    n_head: int = 12
    n_embd: int = 768         # C, the width of the residual stream
    dropout: float = 0.0
    bias: bool = True         # GPT-2 uses biases in Linear and LayerNorm

The sizes OpenAI released: small (12 layers, 768 wide, 12 heads, 124M), medium (24, 1024, 16, 355M), large (36, 1280, 20, 774M) and XL (48, 1600, 25, 1.5B). In every one the head width is $C/H = 64$ and the MLP hidden width is $4C$. Click around the map below to see where the parameters live.

Embeddings: tokens and positions

The input is a tensor idx of integer token ids with shape $(B, T)$. Two lookup tables turn it into vectors, as covered in The Embedding Layer and Positional Encoding. GPT-2 uses learned absolute positions: a second embedding table with one row per position up to 1024.

self.transformer = nn.ModuleDict(dict(
    wte  = nn.Embedding(config.vocab_size, config.n_embd),   # token embeddings   (V, C)
    wpe  = nn.Embedding(config.block_size, config.n_embd),   # position embeddings (T_max, C)
    drop = nn.Dropout(config.dropout),
    h    = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
    ln_f = nn.LayerNorm(config.n_embd, bias=config.bias),
))

# in forward(idx):
B, T = idx.size()
assert T <= self.config.block_size, f"sequence of length {T} exceeds block_size"
pos = torch.arange(0, T, dtype=torch.long, device=idx.device)   # (T,)
tok_emb = self.transformer.wte(idx)     # (B, T, C)
pos_emb = self.transformer.wpe(pos)     # (T, C), broadcasts over B
x = self.transformer.drop(tok_emb + pos_emb)

The names wte, wpe, h, ln_f are not arbitrary. They are exactly the names in OpenAI's checkpoint, and matching them is what lets us load the weights later with almost no renaming.

CausalSelfAttention

This is the module people get wrong most often, so we go slowly. The attention chapter derived the math; here the job is to implement it with the exact tensor layout GPT-2 uses.

The fused qkv projection

Multi-head attention needs a query, key and value for every token, for every head. Instead of three separate Linear layers we use one Linear of width $3C$ and split its output. This is purely an efficiency trick (one big matmul beats three small ones), but it also dictates the checkpoint layout: c_attn.weight has shape $(3C, C)$ and the first $C$ output columns are the queries for all heads, the next $C$ the keys, the last $C$ the values.

Reshaping to heads

After the split we have q, k, v each of shape $(B, T, C)$. To run 12 heads in parallel we view the $C = 768$ channels as $12 \times 64$ and move the head axis before the sequence axis: $(B, T, H, h_d) \to (B, H, T, h_d)$. Now a batched matrix multiply treats $B \times H$ as independent problems, each attending over a $T \times 64$ matrix.

The causal mask as a registered buffer

The mask is a lower-triangular matrix of ones, shape $(1, 1, T_{max}, T_{max})$, and it never changes. We register it as a buffer, not a parameter: it moves to the GPU with the model and is saved with it, but the optimizer ignores it and it gets no gradient. At forward time we slice out the top-left $T \times T$ corner for the current sequence length.

Scaled dot-product and the output projection

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.n_head, self.n_embd = config.n_head, config.n_embd
        # one Linear produces q, k and v for all heads at once: (C) -> (3C)
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
        # output projection back into the residual stream
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        self.c_proj.LUMEN_SCALE_INIT = 1          # flag: scale this layer's init (see below)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        # causal mask: ones on and below the diagonal. A buffer, not a parameter.
        self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
                                          .view(1, 1, config.block_size, config.block_size))

    def forward(self, x):
        B, T, C = x.size()
        hd = C // self.n_head
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)          # each (B, T, C)
        q = q.view(B, T, self.n_head, hd).transpose(1, 2)           # (B, H, T, hd)
        k = k.view(B, T, self.n_head, hd).transpose(1, 2)
        v = v.view(B, T, self.n_head, hd).transpose(1, 2)
        att = (q @ k.transpose(-2, -1)) / math.sqrt(hd)             # (B, H, T, T)
        att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float("-inf"))
        att = F.softmax(att, dim=-1)                                # rows sum to 1, future is 0
        att = self.attn_dropout(att)
        y = att @ v                                                 # (B, H, T, hd)
        y = y.transpose(1, 2).contiguous().view(B, T, C)            # back to (B, T, C)
        return self.resid_dropout(self.c_proj(y))

Three lines carry all the danger. The view then transpose must be in that order (view first, while the channels are still contiguous). The mask is applied with $-\infty$ before softmax, so masked positions get exactly zero weight, not merely small weight. And after attention the transpose must be followed by contiguous() before the view back to $(B, T, C)$, or PyTorch will refuse.

The MLP and GPT-2's GELU

The second half of each block is a two-layer feed-forward network applied to every position independently: expand from $C$ to $4C$, apply a nonlinearity, project back to $C$. See Layers of Understanding for why the expansion helps.

One detail matters for weight compatibility. GPT-2 uses the GELU activation (Hendrycks and Gimpel, 2016), but the tanh approximation of it, because the exact error-function version was slow in TensorFlow in 2019:

The difference from exact GELU is tiny (below $10^{-3}$ everywhere) but it is not zero, and if you want to reproduce OpenAI's logits to several decimals you need the same function. PyTorch exposes it as nn.GELU(approximate="tanh").

class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc   = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)   # (C) -> (4C)
        self.gelu   = nn.GELU(approximate="tanh")                                      # GPT-2's variant
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)   # (4C) -> (C)
        self.c_proj.LUMEN_SCALE_INIT = 1
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        return self.dropout(self.c_proj(self.gelu(self.c_fc(x))))

The Block: pre-LN and residuals

A block glues attention and MLP together with two rules. First, residual connections: each sublayer's output is added to its input rather than replacing it, so the residual stream $x$ carries information straight through all 12 layers and gradients flow back just as directly. Second, pre-LN: LayerNorm is applied to the input of each sublayer, not to its output (Xiong et al., 2020 explain why this trains more stably than the original post-LN transformer).

class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd, bias=config.bias)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd, bias=config.bias)
        self.mlp  = MLP(config)

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))   # communicate: tokens exchange information
        x = x + self.mlp(self.ln_2(x))    # compute: each token thinks on its own
        return x

The full model

Forward: logits, and loss when targets are given

The model's forward takes token ids and optionally targets. With targets it returns the mean cross-entropy over all $B \times T$ positions, exactly as in Learning to Predict. Without them it returns logits, and we can skip the expensive LM head on all but the last position during generation.

class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.transformer = nn.ModuleDict(dict(
            wte  = nn.Embedding(config.vocab_size, config.n_embd),
            wpe  = nn.Embedding(config.block_size, config.n_embd),
            drop = nn.Dropout(config.dropout),
            h    = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f = nn.LayerNorm(config.n_embd, bias=config.bias),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.transformer.wte.weight = self.lm_head.weight        # weight tying (see below)
        self.apply(self._init_weights)
        for name, p in self.named_parameters():                   # residual-projection scaling
            if name.endswith("c_proj.weight"):
                nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layer))

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.size()
        pos = torch.arange(0, T, dtype=torch.long, device=idx.device)
        x = self.transformer.drop(self.transformer.wte(idx) + self.transformer.wpe(pos))
        for block in self.transformer.h:
            x = block(x)
        x = self.transformer.ln_f(x)                              # (B, T, C)
        if targets is not None:
            logits = self.lm_head(x)                              # (B, T, V)
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
            return logits, loss
        logits = self.lm_head(x[:, [-1], :])                      # only the last position: (B, 1, V)
        return logits, None

Weight tying

The token embedding wte.weight has shape $(V, C)$. The LM head is a Linear from $C$ to $V$, whose weight also has shape $(V, C)$. GPT-2 uses one tensor for both (Press and Wolf, 2017). The line self.transformer.wte.weight = self.lm_head.weight makes the two modules share a single parameter. It saves $38.6$M parameters, roughly 30% of the small model, and it encodes a sensible prior: the vector used to recognize a token on the way in should be the vector used to predict it on the way out.

Initialization: 0.02, and the residual scaling

All weights start as Gaussians with standard deviation 0.02, biases at zero. That number is from OpenAI's code and is close to the $1/\sqrt{C}$ Xavier-style scale for $C = 768$ ($1/\sqrt{768} = 0.036$).

There is one refinement. Each block adds two things into the residual stream (its attention output and its MLP output), so by the last layer the stream has received $2N = 24$ contributions. If each has the same variance, the stream's variance grows like $2N$, and its standard deviation like $\sqrt{2N}$. To keep the stream at a sane scale, the GPT-2 paper scales the init of the projections that write into the stream, the two c_proj layers, by $1/\sqrt{2N}$. For $N = 12$ that is $0.02/\sqrt{24} = 0.0041$.

Counting parameters: 124,439,808

Before loading anything, count. If your count matches OpenAI's to the last digit, your shapes are right. The number quoted everywhere is "124M"; the exact figure with tied weights is $124{,}439{,}808$. Here is the arithmetic.

A few things to notice. Each block is about $12C^2$ in weights ($4C^2$ for attention, $8C^2$ for the MLP), and $12 \times 768^2 = 7{,}077{,}888$, so the biases and LayerNorms add only ten thousand per block. The embedding table is 31% of the whole model at this size; for GPT-2 XL it is 5%. That is why parameter counts for small models are often quoted "non-embedding".

model = GPT(GPTConfig())
n = sum(p.numel() for p in model.parameters())
print(n)                                      # counts the tied tensor once

Loading OpenAI's pretrained weights

Now the oracle. Hugging Face hosts OpenAI's checkpoint; transformers.GPT2LMHeadModel.from_pretrained("gpt2") downloads it. Its state dict has the same key names as ours (transformer.wte.weight, transformer.h.0.attn.c_attn.weight, …) because we chose our names to match. Two things still need care.

The Conv1D transpose gotcha

OpenAI's TensorFlow code implemented the linear layers as a module called Conv1D, which stores its weight as $(\text{in}, \text{out})$. PyTorch's nn.Linear stores $(\text{out}, \text{in})$. So the four Linear weights inside each block are the transpose of what we need, and must be flipped on the way in. The embeddings and LayerNorms are not affected. The exact keys to transpose are:

transposed = ["attn.c_attn.weight", "attn.c_proj.weight", "mlp.c_fc.weight", "mlp.c_proj.weight"]

Keys to skip

The HF checkpoint also contains .attn.bias (the causal mask) and, in older versions, .attn.masked_bias. They are buffers, not weights; skip them on their side, and skip our own .attn.bias buffer on ours.

@classmethod
def from_pretrained(cls, model_type="gpt2"):
    from transformers import GPT2LMHeadModel
    cfg_args = {
        "gpt2":        dict(n_layer=12, n_head=12, n_embd=768),    # 124M
        "gpt2-medium": dict(n_layer=24, n_head=16, n_embd=1024),   # 355M
        "gpt2-large":  dict(n_layer=36, n_head=20, n_embd=1280),   # 774M
        "gpt2-xl":     dict(n_layer=48, n_head=25, n_embd=1600),   # 1558M
    }[model_type]
    model = cls(GPTConfig(vocab_size=50257, block_size=1024, bias=True, **cfg_args))
    sd = model.state_dict()
    keys = [k for k in sd if not k.endswith(".attn.bias")]             # drop our mask buffer

    hf = GPT2LMHeadModel.from_pretrained(model_type)
    sd_hf = hf.state_dict()
    keys_hf = [k for k in sd_hf if not k.endswith(".attn.masked_bias") and not k.endswith(".attn.bias")]
    assert len(keys_hf) == len(keys), f"{len(keys_hf)} vs {len(keys)}"

    transposed = ["attn.c_attn.weight", "attn.c_proj.weight", "mlp.c_fc.weight", "mlp.c_proj.weight"]
    with torch.no_grad():
        for k in keys_hf:
            if any(k.endswith(t) for t in transposed):
                assert sd_hf[k].shape[::-1] == sd[k].shape, k       # Conv1D stores (in, out)
                sd[k].copy_(sd_hf[k].t())
            else:
                assert sd_hf[k].shape == sd[k].shape, k
                sd[k].copy_(sd_hf[k])
    return model

The two asserts are the important lines. If any shape does not line up, one of your modules is built wrong, and the key name tells you which.

A generation loop

Generation is the sampling loop from Learning to Predict, with one GPT-2-specific detail: the model has position embeddings only for 1024 slots, so if the running sequence grows longer than that we must crop the context to the last 1024 tokens before each forward pass.

@torch.no_grad()
def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
    for _ in range(max_new_tokens):
        idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
        logits, _ = self(idx_cond)                       # (B, 1, V): only the last position
        logits = logits[:, -1, :] / temperature          # (B, V)
        if top_k is not None:
            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits[logits < v[:, [-1]]] = float("-inf")  # drop everything below the k-th best
        probs = F.softmax(logits, dim=-1)
        idx_next = torch.multinomial(probs, num_samples=1)   # (B, 1)
        idx = torch.cat((idx, idx_next), dim=1)
    return idx

This recomputes the whole prefix at every step, which is $O(T^2)$ work per token. It is fine for a demo; the KV cache that fixes it is the subject of Inference: KV Cache, Batching, Speculative Decoding and code/lumen/kv_cache.py.

The sanity test

Two checks tell you the model is right. First, loss on ordinary English should be low: the model reads a paragraph as its own target (shifted) and reports the mean cross-entropy. For pretrained GPT-2 small on a clean paragraph, expect roughly 3 to 4 nats (perplexity 20 to 55). A randomly initialized model gives 10.8. Second, generated text should be grammatical and stay on topic for a few sentences.

import tiktoken
enc = tiktoken.get_encoding("gpt2")
model = GPT.from_pretrained("gpt2").eval()

text = ("The Industrial Revolution began in Britain in the late eighteenth century, "
        "driven by new machinery, cheap coal, and a growing network of canals and railways.")
ids = torch.tensor([enc.encode(text)])
_, loss = model(ids[:, :-1], ids[:, 1:])
print(f"loss {loss.item():.2f}  perplexity {loss.exp().item():.1f}")

prompt = torch.tensor([enc.encode("The Industrial Revolution began")])
out = model.generate(prompt, max_new_tokens=40, temperature=0.8, top_k=40)
print(enc.decode(out[0].tolist()))

Your exact numbers will differ (the loss depends on the paragraph; the text depends on the seed), but the pattern should not: a loss in the low single digits and readable text. If the loss is 10 or above, or the text is word salad, go to the debugging checklist below.

Training from scratch on a tiny corpus

Loading weights proved the architecture. Now discard them and learn from data, using the same recipe at small scale that the large runs use. The reference implementation is code/lumen/train.py with data from code/lumen/data.py; the essentials follow.

The data loader with the shift-by-one

Tokenize the whole corpus once into a single long array of ids. A batch is then $B$ random windows of length $T+1$; the first $T$ tokens are the input and the last $T$ are the target.

class TokenLoader:
    def __init__(self, tokens, B, T, device="cpu", seed=0):
        self.tokens = torch.tensor(tokens, dtype=torch.long)
        self.B, self.T, self.device = B, T, device
        self.g = torch.Generator().manual_seed(seed)

    def next_batch(self):
        ix = torch.randint(0, len(self.tokens) - self.T - 1, (self.B,), generator=self.g)
        x = torch.stack([self.tokens[i     : i + self.T]     for i in ix])   # (B, T)
        y = torch.stack([self.tokens[i + 1 : i + self.T + 1] for i in ix])   # (B, T), shifted by one
        return x.to(self.device), y.to(self.device)

AdamW, and weight decay on matrices only

The optimizer is AdamW (Loshchilov and Hutter, 2019) with the settings the GPT-3 paper reports for its 125M model, which nanoGPT and most reproductions reuse: learning rate $6 \times 10^{-4}$, betas $(0.9, 0.95)$, weight decay $0.1$, epsilon $10^{-8}$. The unusual $\beta_2 = 0.95$ (instead of the default $0.999$) makes the second-moment estimate adapt faster, which helps stability with large batches.

Weight decay should apply only to the 2-d tensors: the weight matrices and the embeddings. Biases and LayerNorm gains are 1-d and are excluded. Decaying a LayerNorm gain toward zero would fight the normalization; decaying a bias just adds noise. The split is done with two parameter groups:

def configure_optimizer(model, lr=6e-4, weight_decay=0.1, betas=(0.9, 0.95)):
    params = [p for p in model.parameters() if p.requires_grad]
    decay    = [p for p in params if p.dim() >= 2]    # matrices and embeddings
    no_decay = [p for p in params if p.dim() < 2]     # biases, LayerNorm weights
    groups = [{"params": decay, "weight_decay": weight_decay},
              {"params": no_decay, "weight_decay": 0.0}]
    return torch.optim.AdamW(groups, lr=lr, betas=betas, eps=1e-8)

Warmup and cosine decay

The learning rate is not constant. It ramps linearly from zero over a warmup period (Adam's moment estimates are garbage for the first few hundred steps, and a full-size step on garbage can wreck the init), then follows a cosine curve down to 10% of the peak (Loshchilov and Hutter, 2017).

At $t = t_w$ the cosine is $\cos 0 = 1$, giving $\eta_{\max}$; at $t = t_{\max}$ it is $\cos\pi = -1$, giving $\eta_{\min}$. Halfway between, the rate is the midpoint.

def lr_at(step, max_lr=6e-4, min_lr=6e-5, warmup=100, max_steps=2000):
    if step < warmup:
        return max_lr * (step + 1) / warmup
    if step > max_steps:
        return min_lr
    progress = (step - warmup) / max(1, max_steps - warmup)
    return min_lr + 0.5 * (max_lr - min_lr) * (1.0 + math.cos(math.pi * progress))

Gradient clipping

Before the optimizer step, the global norm of all gradients is computed, and if it exceeds 1.0 every gradient is scaled down so the norm is exactly 1.0. This caps the damage a single bad batch can do. The returned norm is worth logging: it usually sits well below 1 and a sudden jump is the earliest warning of a coming loss spike.

The loop

model = GPT(GPTConfig(block_size=256)).to(device)     # small context for a small corpus
opt = configure_optimizer(model)
loader = TokenLoader(train_tokens, B=8, T=256, device=device)

for step in range(2000):
    for g in opt.param_groups:
        g["lr"] = lr_at(step)
    x, y = loader.next_batch()
    _, loss = model(x, y)
    opt.zero_grad(set_to_none=True)
    loss.backward()
    norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()
    if step % 100 == 0:
        print(f"step {step:4d}  loss {loss.item():.3f}  grad-norm {norm:.2f}  lr {lr_at(step):.2e}")

What loss to expect

The numbers above are typical of a run on a corpus of roughly a million tokens (a few novels, or the "tiny Shakespeare" file) with a 124M model. The loss starts at $\ln 50257 \approx 10.8$ (10.9 with a bit of init noise), drops below 7 in the first hundred steps as the model learns token frequencies, and grinds down toward 4 over a couple of thousand steps. On such a small corpus the model will then begin to overfit: training loss keeps falling but held-out loss turns upward, because 124M parameters can memorize a megabyte of text. That is expected and is not a bug; it is the signal that you need more data, not more steps.

For comparison, the real GPT-2 small reaches about 3.3 nats on held-out web text after seeing tens of billions of tokens. You will not get there on a laptop. But you will see the same curve shape, and that is the point of this exercise.

Measuring held-out loss

Training loss on a tiny corpus lies to you, because the model starts memorizing. Keep the last 10% of the token array aside and never train on it. Every few hundred steps, evaluate the model on a few dozen random windows from that held-out slice, in eval mode with gradients off, and average. This is the number that tells you whether the model is learning language or learning the file.

@torch.no_grad()
def estimate_loss(model, loader, iters=40):
    model.eval()
    total = 0.0
    for _ in range(iters):
        x, y = loader.next_batch()
        _, loss = model(x, y)
        total += loss.item()
    model.train()
    return total / iters          # average the losses; exponentiate afterwards for perplexity

Plot training and validation loss on the same axes. On a small corpus they track each other for the first few hundred steps, then separate: training keeps improving, validation flattens and turns up. The step where validation is lowest is the model you keep.

How to scale up

Three changes take this loop from a laptop to a real run, and all three are in code/lumen/train.py. First, mixed precision: wrap the forward pass in torch.autocast(device_type="cuda", dtype=torch.bfloat16) so the matmuls run in 16-bit while the weights and optimizer state stay in 32-bit (Micikevicius et al., 2018). This roughly doubles or triples throughput on modern GPUs, and bfloat16 needs no loss scaling because it keeps float32's exponent range.

Second, gradient accumulation. GPT-2 was trained with about half a million tokens per optimizer step ($512$ sequences $\times$ $1024$ tokens), far more than fits in one GPU's memory. The fix is to split the big batch into micro-batches, run forward and backward on each, and only call optimizer.step() after all of them. Gradients add up across backward() calls, so the sum over micro-batches equals the gradient of the big batch, provided each micro-batch loss is divided by the number of micro-batches so the mean is a mean and not a sum.

grad_accum = 32                     # 32 micro-batches of 8×256 = 65,536 tokens per step
for step in range(max_steps):
    for g in opt.param_groups:
        g["lr"] = lr_at(step)
    opt.zero_grad(set_to_none=True)
    for micro in range(grad_accum):
        x, y = loader.next_batch()
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            _, loss = model(x, y)
        (loss / grad_accum).backward()          # scale so the accumulated gradient is a mean
    norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()

Third, multiple GPUs via DistributedDataParallel: each process holds a full copy of the model, works on its own micro-batches, and the gradients are averaged across processes before the step. Eight GPUs each doing 4 micro-batches is the same computation as one GPU doing 32. The full treatment, including data pipelines, tokenizing at scale and evaluation, is in Pre-Training.

Debugging checklist

Things go wrong. Here is the order in which to check them, from cheapest to most expensive.

1. Shapes. Push a $(2, 5)$ batch through and print the shape after every module, exactly as in the tracer above. The residual stream must be $(2, 5, 768)$ everywhere; the attention scores $(2, 12, 5, 5)$; the logits $(2, 5, 50257)$.
2. Loss at initialization. With random weights the loss must be $\ln 50257 \approx 10.82$, give or take 0.1. Much lower means the model can see the future: check that the mask is lower-triangular (torch.tril, not triu) and that the targets are shifted by exactly one. Much higher means the init is too large.
3. Mask direction. Print att[0, 0] for a short sequence after softmax: every entry above the diagonal must be exactly 0, and each row must sum to 1.
4. Overfit a single batch. Train on the same $(x, y)$ batch repeatedly. The loss must go to nearly zero within a few hundred steps. If it plateaus, gradients are not reaching some parameters (a detached tensor, a missing residual, a forgotten requires_grad).
5. Parameter count. 124,439,808 with tied weights. Off by 38,597,376? The tie is missing. Off by 12 × 1,048,576? You counted the mask buffers.
6. Pretrained sanity. Load OpenAI's weights; loss on a paragraph should be 3 to 4. If it is 6 to 8 the model is working but something is subtly wrong: usually the GELU variant, the Conv1D transpose, or LayerNorm epsilon (GPT-2 uses $10^{-5}$, PyTorch's default).
7. Gradient norm. Log it every step. A healthy run sits below 1 after warmup. If it explodes, lower the learning rate or lengthen warmup; if it is exactly 0, the loss is not connected to the parameters.
8. Loss turns to nan. Almost always a learning rate too high, a missing warmup, or an $-\infty$ mask applied to an entire row (which makes softmax divide by zero). Check that the diagonal of the mask is unmasked.

Practice

Further reading

• Radford et al. (2019). Language Models are Unsupervised Multitask Learners. The GPT-2 paper, including the $1/\sqrt{N}$ residual-scaling note and the four model sizes.
• Karpathy. nanoGPT. The cleanest reference implementation; this chapter's code follows its structure.
• Brown et al. (2020). Language Models are Few-Shot Learners. Table 2.1 lists the learning rates and batch sizes per model size that most reproductions use.
• Hendrycks and Gimpel (2016). Gaussian Error Linear Units (GELUs). The activation, including the tanh approximation.
• Xiong et al. (2020). On Layer Normalization in the Transformer Architecture. Why pre-LN trains more stably than post-LN.
• Press and Wolf (2017). Using the Output Embedding to Improve Language Models. Weight tying.
• Loshchilov and Hutter (2019). Decoupled Weight Decay Regularization. AdamW.
• Loshchilov and Hutter (2017). SGDR: Stochastic Gradient Descent with Warm Restarts. The cosine schedule.
• Micikevicius et al. (2018). Mixed Precision Training. The technique behind autocast.
• Vaswani et al. (2017). Attention Is All You Need. The original architecture that GPT-2 simplifies to a decoder stack.