The full model and its training loop

Shapes, loss and the cost of one training step

Chapter 9 of 21 · II. The transformer · about 35 minutes. Builds on: 3. Backpropagation and optimizers, 7. The transformer block, 8. Position and rotary embeddings.

Skip this chapter if

you can trace tensor shapes through a full decoder-only transformer, write the training loop, and explain the $6\cdot N\cdot D$ compute rule and the memory cost of training (about 16 or more bytes per parameter).

Run a whole forward pass

Run a whole forward pass from prompt to next-token probabilities. The guided tour walks through it. Open in the Model Explorer.

Assembling the pieces

A decoder-only language model, the design used by essentially every current chat model, is the following pipeline. Let $B$ be the batch size, $T$ the sequence length, $d$ the width, $V$ the vocabulary and $L$ the number of blocks.

token ids              [B, T]              integers in [0, V)
embedding lookup       [B, T, d]           row of E for each id
L × block              [B, T, d]           attention + MLP with residuals (Chapters 6–8)
final norm             [B, T, d]
unembedding (· Wout)   [B, T, V]           logits for every position
softmax + loss         scalar              cross-entropy against the next tokens

Training uses teacher forcing: the targets are simply the input sequence shifted left by one, so position $t$ is trained to predict token $t + 1$, and the causal mask guarantees it cannot peek. In pseudocode (PyTorch-style, illustrative rather than a tuned implementation):

def forward(ids):                       # ids: [B, T]
    x = E[ids]                          # embedding, [B, T, d]
    for blk in blocks:                  # L pre-norm blocks
        x = x + attn(norm1(x), causal=True)
        x = x + mlp(norm2(x))
    return norm_f(x) @ W_out            # logits, [B, T, V]

for ids in data_loader:                 # ids: [B, T+1]
    logits = forward(ids[:, :-1])
    loss = cross_entropy(logits, ids[:, 1:])   # average over B*T positions
    loss.backward()
    clip_grad_norm_(params, 1.0)
    optimizer.step(); scheduler.step(); optimizer.zero_grad()

Documents are concatenated and cut into fixed-length sequences (“packing”), with end-of-text tokens between them; some pipelines also prevent attention from crossing document boundaries, while others (including the one described for DeepSeek-V3) do not .

How much compute does training take?

A matrix multiplication of a $1\times a$ vector by an $a\times b$ matrix costs $a\cdot b$ multiplies and about as many adds, so about $2\cdot a\cdot b$ FLOPs, which is 2 FLOPs per parameter. Summing over all weights, the forward pass costs about $2\cdot N$ FLOPs per token, where $N$ is the number of parameters that are multiplied per token. The backward pass computes two products of the same size per matrix (gradient w.r.t. the weights and w.r.t. the input), so it costs about $4\cdot N$. Total:

$$\text{training FLOPs}\;\approx\;6\,N\,D$$

This ignores the attention-score term, which grows with sequence length and is small at moderate $T$. Experiment:

For example, 405 billion parameters on 15.6 trillion tokens is $6 \times 4.05\times 10^{11} \times 1.56\times 10^{13} \approx 3.8\times 10^{25}$ FLOPs, which matches the order of magnitude Llama 3 reports .

How much memory does training take?

Per parameter, with Adam and mixed precision (Chapter 15): the working weights in 16-bit (2 bytes), gradients (2 bytes), and optimizer state in 32-bit: a master copy of the weights (4), first moment (4), second moment (4). That is about 16 bytes per parameter before any activations. A 70-billion-parameter model therefore needs over a terabyte just for these states, far beyond one GPU, which is why training uses sharding and parallelism (Chapter 16).

Auxiliary objectives

The next-token loss can be supplemented. Multi-token prediction adds heads that predict tokens further ahead from the same hidden state, giving a denser training signal; DeepSeek-V3 uses a single extra predicted token (depth 1) and can reuse that module to speed up inference .

Measuring progress

Validation loss (or perplexity, Chapter 4) on held-out text is the cleanest pretraining signal because it is smooth and cheap. Downstream benchmarks measure what people care about but are noisier, can saturate, and can leak into training data. Neither tells you whether a model is secure; that requires targeted evaluation (Chapters 20 and 21).

Check yourself

Why is the backward pass about twice the forward pass?

For each weight matrix, the forward pass does one product (activations × weights). The backward pass does two of the same size: one for the gradient of the weights and one to propagate the gradient to the layer’s input.

A mixture-of-experts model has 671B total and 37B active parameters. Which one goes in 6ND?

The active count, since only those weights are multiplied per token: about $6 \times 37\times 10^{9} \times 14.8\times 10^{12} \approx 3.3\times 10^{24}$ FLOPs for 14.8T tokens.

Roughly how much GPU memory do the parameters, gradients and Adam state for an 8B model need?

About 16 bytes × $8\times 10^{9} \approx 128$ GB, before activations.

Security link

The training objective is to reproduce the data. Nothing in the loss distinguishes instructions from facts or knowledge from secrets. Whatever is in the corpus, including private data and attacker-planted text, can be learned, and sufficiently repeated text is memorized and can be extracted later (Chapter 20).

Papers

• The original architecture (with an encoder; modern chat models keep only the decoder).
• GPT-3: shows what happens when this recipe is scaled.
• LLaMA: a clear description of a modern open dense recipe.
• Llama 3: architecture, data and infrastructure at scale.
• Multi-token prediction.