Chapter 23

Training a GPT from Scratch

A complete NumPy training loop, evaluation, and text generation.

A GPT is a decoder-only transformer trained to predict the next token. At production scale this means trillions of tokens and specialized systems; here it means a tiny character model, a short public-domain text embedded in the source file, and enough NumPy to see the whole loop. The goal is not a useful model, but a complete, inspectable training run.

23.1 Data and batches

The model reads characters, not subword tokens. build_vocab sorts the characters in the text, assigns each one an integer ID, and encodes the whole string into an array. A batch samples start positions and returns two arrays: $X$ is a sequence of IDs, and $Y$ is the same sequence shifted one character to the right. The loss is therefore next-character cross-entropy at every position. Character modeling is deliberately inefficient: the model must learn words, spaces, and punctuation from individual symbols. That weakness is useful for teaching because the vocabulary is small, the data can live in the source file, and every tensor shape is easy to print. The code still keeps a validation split, because even a toy model can memorize a tiny excerpt while its held-out loss stops improving.

p(x_1,\ldots,x_T)=\prod_{t=1}^{T}p(x_t\mid x_{<t}) .\tag{23.1}

This is the autoregressive factorization used by neural language models [bengio2003] and GPTs [radford2019]. The causal mask in Chapter 22 enforces the conditioning: position $t$ can read earlier positions but not the target it is asked to predict.

23.2 Model and loss

The code builds an embedding matrix, a list of transformer blocks from Chapter 22, a final RMSNorm, and a tied output head. For tokens $x_{b,t}$ , the embedding lookup creates $\mH_0$ . After $L$ blocks, logits are

\mZ_{b,t} = \operatorname{RMSNorm}(\mH_L)_{b,t}\mE^{\T} .\tag{23.2}

The cross-entropy averages over batch and time:

\mathcal{L}=-\frac{1}{BT}\sum_{b,t}\log\softmax(\mZ_{b,t})_{y_{b,t}} .\tag{23.3}

Weight tying means $\mE$ receives two gradients: sparse additions from the input lookup and a dense classifier gradient from the output head [press2016using]. The implementation accumulates both into the same array before the optimizer step. This tying is not required, but it is a good default for a small model: the vectors used to read characters are also the vectors used to score characters. That reduces parameters and gives the embedding table a learning signal from every predicted position, not only from characters that appeared in the input side of the batch. The rest of the backward pass is ordinary reverse mode through the stack. Cross-entropy gives the logit gradient $(\softmax(\vz)-\vy)/(BT)$ . The tied head sends that gradient into the hidden states and into $\mE$ . The final RMSNorm and each transformer block then run the manual backward functions from the previous chapter. Because the blocks are stored in a Python list, the backward loop simply walks that list in reverse and stores one gradient tree per layer.

23.3 AdamW and warmup

Each training step samples a batch, runs gpt_loss_and_grads, updates the parameters with AdamW, and occasionally estimates validation loss on held-out text. Adam comes from moving averages of gradients and squared gradients [kingma2014adam]; AdamW applies weight decay as a separate shrinkage step [loshchilov2017decoupled]. Chapter 15 covers optimizers and schedules in more detail: Chapter 15. The learning rate warms up linearly so the first updates are not as large as the steady-state updates.

Listing 23.1 AdamW update used in the loop

def adamw_step(params, grads, state, lr, weight_decay=0.01,
               beta1=0.9, beta2=0.999, eps=1e-8):
    state["t"] = state.get("t", 0) + 1
    t = state["t"]
    for path, param, grad in tree_items(params, grads):
        slot = state.setdefault(path, {
            "m": np.zeros_like(param),
            "v": np.zeros_like(param),
        })
        slot["m"] = beta1 * slot["m"] + (1.0 - beta1) * grad
        slot["v"] = beta2 * slot["v"] + (1.0 - beta2) * (grad * grad)
        m_hat = slot["m"] / (1.0 - beta1 ** t)
        v_hat = slot["v"] / (1.0 - beta2 ** t)
        param *= 1.0 - lr * weight_decay
        param -= lr * m_hat / (np.sqrt(v_hat) + eps)

The figure is generated at build time by running a short CPU training job. It should bend downward, but do not over-interpret it: the dataset is tiny, the model is tiny, and a character-level model learns spelling and punctuation long before it learns anything resembling reasoning. The validation curve is noisier than the training-batch curve because it is estimated from a few small held-out batches. That is acceptable here: the figure is a smoke test that optimization is wired correctly, not a benchmark. On a laptop-scale CPU run, this chapter should take seconds to build; increasing widths, layers, or context length quickly turns the same code into a minutes-long experiment.

Figure 23.1 A tiny GPT’s training-batch loss and validation loss during a short NumPy run.

Listing 23.2 Complete tiny-GPT training loop

def train_tiny_gpt(steps=30, seed=23, d_model=24, n_heads=2, n_layers=1,
                   hidden_dim=48, batch_size=8, seq_len=16, base_lr=3e-3):
    rng = np.random.default_rng(seed)
    data, stoi, itos = build_vocab()
    train_data, val_data = split_data(data)
    params = init_gpt_params(len(stoi), d_model, n_heads, n_layers, hidden_dim, rng)
    opt_state = {}
    train_losses, val_points = [], []
    for step in range(1, steps + 1):
        x, y = get_batch(train_data, batch_size, seq_len, rng)
        loss, grads = gpt_loss_and_grads(params, x, y, n_heads)
        lr = warmup_lr(step, base_lr, warmup_steps=5)
        adamw_step(params, grads, opt_state, lr)
        train_losses.append(loss)
        if step == 1 or step == steps or step % 5 == 0:
            val_rng = np.random.default_rng(seed + 10_000 + step)
            val = estimate_loss(params, val_data, n_heads, val_rng, batch_size, seq_len)
            val_points.append((step, val))
    history = {"train": np.array(train_losses), "val": np.array(val_points)}
    return params, history, stoi, itos

23.4 Sampling

After training, generation repeats the same forward pass. Encode the prompt, run the model on the recent context, take the last-position logits, divide by temperature, optionally keep only the largest $k$ logits, softmax, and sample. Lower temperature sharpens the distribution; higher temperature makes unlikely characters easier to draw. Top-k sampling sets all but the largest $k$ logits to $-\infty$ before softmax; it is a simple way to avoid sampling from a long low-probability tail, related to later work on text degeneration [holtzman2019curious]. Sampling is seeded in the tests, so the same probabilities produce the same string. Without a seed, generation is intentionally random: two runs can diverge after the first sampled character because that character becomes part of the next context. This feedback loop is why poor sampling settings can make text drift even when the one-step validation loss looks reasonable.

Listing 23.3 Sampling with temperature and top-k

def sample_text(params, prompt, stoi, itos, n_heads, steps, seed=0,
                temperature=1.0, top_k=None, max_context=64):
    rng = np.random.default_rng(seed)
    ids = [stoi[ch] for ch in prompt]
    for _ in range(steps):
        context = np.array([ids[-max_context:]], dtype=np.int64)
        logits = gpt_logits(params, context, n_heads)[0, -1] / temperature
        if top_k is not None and top_k < logits.size:
            keep = np.argpartition(logits, -top_k)[-top_k:]
            masked = np.full_like(logits, -np.inf)
            masked[keep] = logits[keep]
            logits = masked
        probs = softmax(logits)
        ids.append(int(sample_categorical(probs[None, :], rng)[0]))
    return "".join(itos[i] for i in ids)

In practice

The same loop scales conceptually to token-level GPTs: bigger datasets, larger batches, more layers, and distributed matrix multiplies. The small code follows the shape of minimalist GPT implementations such as nanoGPT [karpathy2023nanogpt], but it is intentionally slower because it keeps every backward pass visible. Treat its samples as debugging artifacts, not evidence of model quality. A real training run reports validation loss on data not used for updates, plus downstream evaluations appropriate to the intended use. The honest takeaway is scale, not magic: this chapter trains a tiny character model for a few steps; useful GPTs train token models for vast numbers of updates. The equations and data flow are the same, but the engineering constraints are completely different.

Key equations

p(x_1,\ldots,x_T)=\prod_t p(x_t\mid x_{<t})

\mZ_{b,t}=\vh_{b,t}\mE^{\T}

\mathcal{L}=-\frac{1}{BT}\sum_{b,t}\log\softmax(\mZ_{b,t})_{y_{b,t}}

\eta_s=\eta_{\max}\min(1, s/S_{\text{warmup}})

p_i=\softmax(z_i/\tau) \quad \text{after optional top-}k\text{ masking}

23.5 Teach it

The one-sentence version. Train a GPT by showing it many prefixes, asking it to predict the next token at every position, backpropagating cross-entropy, and sampling from the last-position logits.

An analogy. It is autocomplete with a strict blindfold: while practicing each character, the model may read only the characters to its left.

At the board.

Write a text string, then two rows: input characters and the same row shifted left as targets.
Draw embedding, transformer blocks, final norm, and tied head.
Write cross-entropy over batch and time.
Show sampling: logits, temperature, optional top-k mask, softmax, draw.

Misconceptions to address.

"The tiny model is a real chatbot." It is a character toy trained on a tiny excerpt.
"Validation loss is optional." Without held-out text, you cannot see memorization.
"Sampling is training." Sampling uses the trained weights; it does not update them.

Check for understanding. Why does a tied embedding matrix receive gradients from both the input lookup and the output classifier?

23.6 Exercises

Exercise 23.1 ★ Shifted targets

Given an encoded character sequence $d_0,d_1,\ldots$ , write the input and target rows for a context starting at index $i$ with length $T$ . Why does this create $T$ supervised examples from one slice?

Solution in Appendix D

Exercise 23.2 ★★ Tied embedding gradient

Derive the two contributions to the gradient of the tied embedding matrix: one from the output head $\vh\mE^{\T}$ and one from the input lookup.

Solution in Appendix D

Exercise 23.3 ★★ Warmup and AdamW

Explain why the code applies warmup to the learning rate and weight decay directly to parameters, not by adding $\lambda\theta$ to the Adam gradient.

Solution in Appendix D

Exercise 23.4 ★★★ Train and sample

Run a few optimization steps on one fixed tiny batch and check that the loss falls. Then sample with a fixed random seed and verify that the generated string is deterministic.

Solution in Appendix D

References

[radford2019] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever. Language models are unsupervised multitask learners. OpenAI technical report, 2019.
[bengio2003] Y. Bengio, R. Ducharme, P. Vincent, and C. Jauvin. A neural probabilistic language model. Journal of Machine Learning Research 3, 1137–1155, 2003.
[karpathy2023nanogpt] A. Karpathy. nanoGPT. Source code, 2023. https://github.com/karpathy/nanoGPT
[holtzman2019curious] A. Holtzman et al. The Curious Case of Neural Text Degeneration. 2019. arXiv:1904.09751
[kingma2014adam] D. P. Kingma and J. Ba. Adam: A Method for Stochastic Optimization. 2014. arXiv:1412.6980
[loshchilov2017decoupled] I. Loshchilov and F. Hutter. Decoupled Weight Decay Regularization. 2017. arXiv:1711.05101
[press2016using] O. Press and L. Wolf. Using the Output Embedding to Improve Language Models. 2016. arXiv:1608.05859