Training Objectives and Architectural Details

Open the GPT-2 you built in the from-scratch chapter next to the config for Llama 3 8B. They are both "transformers". Yet almost every line differs: a different normalization, a different activation, a different way to encode position, a different number of key heads than query heads, no bias terms anywhere, a vocabulary four times larger. Why? Each change was made because something broke, or something was wasteful, at scale.

This chapter has two parts. Part A asks the bigger question first: what should the model be trained to do? There were several serious candidates, and it is worth understanding why one of them won. Part B walks through the modern decoder block one component at a time, each with the problem it solves, the formula, a tiny worked example, and a picture.

Part A: What objective do we train on?

The transformer architecture does not come with a training objective. You choose one. Between 2018 and 2021 there were four serious contenders, and they produced very different models. The choice you make determines what every token in your corpus teaches the model, so it is worth being precise.

Causal language modeling (GPT)

The problem: we want a model that can generate text. The simplest objective that produces a generator is to predict each token from the ones before it. The model reads left to right and at every position outputs a distribution over the next token. The loss is the negative log-probability of the true next token, averaged over all positions.

What just happened: every one of the $T$ positions in the sequence contributes a loss term. A 4,096-token document gives 4,096 training signals from one forward pass. The causal mask in the attention layers is what makes this legal: position $t$ cannot see positions after it, so it cannot cheat.

Masked language modeling (BERT)

The problem BERT was solving was different: understanding, not generation. Devlin et al. (2018) reasoned that a left-to-right model sees only half the context for each prediction, and that for tasks like classification you want each token's representation to depend on both sides. So BERT hides 15% of the tokens and asks the model to fill them in, using bidirectional attention.

Here $M$ is the set of masked positions and $x_{\setminus M}$ means the sequence with those positions replaced by a special [MASK] token. Only masked positions contribute to the loss. That is the catch: from a 512-token sequence you get about 77 training signals, not 512. And the model never learns to generate; it only learns to fill holes.

Prefix LM and seq2seq (T5)

Two hybrids tried to get the best of both. A prefix LM attends bidirectionally over a prefix (the "input") and causally over the rest (the "output"), training loss only on the output part. A seq2seq model like T5 (Raffel et al., 2020) has a separate encoder that reads the input bidirectionally and a decoder that generates the output causally with cross-attention into the encoder. T5's actual pre-training objective was "span corruption": drop out spans of the input and generate them in order.

Why decoder-only won for generation

By 2022 the question was settled empirically. Wang et al. (2022) trained all the variants at matched compute and found that causal decoder-only models pre-trained on plain next-token prediction gave the best zero-shot generalization after adaptation. The reasons are worth internalizing because they explain why nobody revisits the decision.

• Every token is a training signal. Per FLOP, the causal LM sees roughly six times as many supervised predictions as MLM with 15% masking. When compute is the binding constraint, this is decisive.
• It is simple. One stack, one mask, one loss. No encoder, no cross-attention, no span-corruption sampling scheme, no train-test mismatch from [MASK] tokens that never appear at inference.
• It scales cleanly. The KV cache, tensor parallelism, and every inference optimization in this course assume a single causal stack.
• Generation is the product. If what you sell is text coming out of the model, train the model to produce text. Understanding tasks can be phrased as generation ("The sentiment is: positive").

Part B: The modern decoder block

Now the block itself. Here is the same picture drawn twice: the 2019 block (GPT-2) and the 2024 block (Llama 3). Toggle between them and read the table of differences underneath. Then we will justify each row.

RMSNorm: normalization without the mean

The problem LayerNorm solves: activations in a deep residual stack drift in scale, and the next layer's matrix multiplication then sees inputs of wildly different magnitudes from one token to the next. LayerNorm (Ba et al., 2016) fixes this by standardizing each token's vector: subtract the mean, divide by the standard deviation, then apply a learned scale $\gamma$ and shift $\beta$.

Zhang and Sennrich (2019) asked which part of that actually matters. Their answer: the re-scaling. The mean-centering contributes little. So RMSNorm drops the mean and the shift, and divides by the root-mean-square instead.

What just happened: we removed one reduction (the mean), one subtraction, and one learned vector ($\beta$). That is 10 to 15 percent less time in the normalization kernel, which matters because norms are memory-bound and run twice per layer. Quality, measured on the models that have tried both, is the same. Every major model since Llama 1 uses RMSNorm.

Pre-norm: where the normalization goes

The original transformer (Vaswani et al., 2017) put the norm after each residual addition: $x \leftarrow \text{Norm}(x + \text{Sublayer}(x))$. The problem: with post-norm, gradients at initialization are large in the top layers and tiny in the bottom, so training needs a careful warmup and is fragile in deep stacks (Xiong et al., 2020).

Pre-norm moves the norm inside the residual branch: $x \leftarrow x + \text{Sublayer}(\text{Norm}(x))$. Now the residual stream is a clean highway from input to output; every block only adds to it. GPT-2 already used pre-norm, and so does every model since. (A few recent models add a second norm on the branch output, "sandwich" or "peri" norm, for extra stability; Gemma 2 does this.)

SwiGLU: the gated MLP

The problem with the classic MLP: it is a fixed nonlinearity applied to a projection, $\text{GELU}(xW_1)W_2$. Every hidden unit fires based on one linear view of the input. Shazeer (2020) tried a small change borrowed from LSTMs: compute two projections, and let one of them gate the other by elementwise multiplication.

What just happened: $xW$ passes through Swish (also called SiLU) and becomes a soft gate; $xV$ is the "content"; their product is the hidden activation, which $W_2$ projects back down. GeGLU is the same with GELU as the gate. The intuition for why it helps: a multiplicative interaction lets the MLP compute things like "feature A and feature B" in one layer, which a single nonlinearity cannot. Shazeer's paper offers no theory, only the honest line that "we offer no explanation... and attribute their success, as all else, to divine benevolence", but the perplexity improvements have held up at every scale since, and PaLM, Llama, Mistral and Gemma all use a gated MLP.

The 8/3 hidden multiplier

The gated MLP has three matrices instead of two. If you kept the hidden width at $4d$ the MLP would cost 50% more parameters. To keep the parameter count matched, papers shrink the hidden width. Standard MLP: $W_1$ is $d \times 4d$ and $W_2$ is $4d \times d$, total $8d^2$. Gated MLP with hidden width $h$: $W$, $V$ are $d \times h$ and $W_2$ is $h \times d$, total $3dh$. Setting $3dh = 8d^2$ gives:

Llama 1 and 2 use exactly this, rounded up to a multiple of 256 (for $d = 4096$: $10{,}923 \to 11{,}008$). Llama 3 8B multiplies by a further 1.3 and rounds to a multiple of 1024, giving $14{,}336 = 3.5d$; the extra width is spent on purpose because the MLP is where knowledge lives and the 128k vocabulary freed up budget elsewhere.

RoPE: a recap and a pointer

The problem with GPT-2's learned absolute position embeddings: they are a lookup table of 1,024 vectors. Position 1,025 does not exist. And "how far apart are these two tokens" has to be inferred from two absolute vectors, which is a clumsy way to encode the thing attention actually cares about.

Rotary position embedding (Su et al., 2021) instead rotates each query and key vector by an angle proportional to its position, in a set of 2-d planes with different frequencies. Because a dot product between two rotated vectors depends only on the difference of their angles, the attention score depends only on relative position. It is applied inside every attention layer, not once at the input, and it adds no parameters. You built it in the positional encoding chapter; here the only new fact is the base frequency: Llama 3 raised $\theta$ from 10,000 to 500,000 so that the lowest-frequency planes rotate slowly enough to distinguish positions up to 128k apart.

Grouped-query attention: shrinking the KV cache

The problem does not show up in training at all. It shows up at inference. To generate token $t+1$ the model needs the keys and values of all previous tokens in every layer, and recomputing them every step would be absurd, so they are cached. The KV cache per sequence is:

The leading 2 is for K and V. In standard multi-head attention (MHA) $n_{kv}$ equals the number of query heads, and the cache is enormous. Multi-query attention (MQA; Shazeer, 2019) shares one K and one V head across all query heads, dividing the cache by $n_{\text{heads}}$, but costs some quality. Grouped-query attention (GQA; Ainslie et al., 2023) is the compromise: put query heads into $g$ groups and give each group its own K/V head.

No bias terms

GPT-2 has a bias vector on every linear layer and every LayerNorm. Llama has none. The problem they solved by removing them is subtle: PaLM (Chowdhery et al., 2022) reported that dropping biases improved training stability for large models, and nobody has found a quality cost. Biases are a negligible fraction of parameters, so the saving is not about size; it is one fewer thing that can drift. The RMSNorm scale $\gamma$ stays.

Untied embeddings

GPT-2 uses the same matrix to embed input tokens and to produce output logits ("weight tying", Press and Wolf, 2017). For a small model this saves a large fraction of parameters and acts as a regularizer. The problem at scale: the two jobs are different. The input embedding wants to represent "what this token means"; the output projection wants to represent "what predicts this token next". Llama 3 8B unties them, spending an extra $128{,}256 \times 4{,}096 \approx 525$M parameters, about 6.5% of the model, on a separate output head. Whether this is worth it depends on the model: Gemma keeps tying because its huge 256k vocabulary would otherwise dominate the parameter count at small sizes.

QK-norm

The problem: attention logits are $q \cdot k / \sqrt{d}$, and if the norms of $q$ and $k$ grow during training, the logits grow, the softmax saturates, and gradients vanish or explode. This is a leading cause of the loss spikes discussed in the next chapter. QK-norm (Henry et al., 2020; used at scale by Dehghani et al., 2023 in ViT-22B) applies a normalization to $q$ and $k$ separately before the dot product, so the logit magnitude is bounded by the learned scale rather than by whatever the projections drift to. OLMo 2 and Gemma 3 adopt it. The cost is two extra small norms per layer.

Attention logit soft-capping

A blunter tool for the same problem: pass each logit through a scaled $\tanh$ so it can never exceed a cap.

For $|z| \ll c$ the function is nearly the identity; for large $|z|$ it saturates at $\pm c$. Gemma 2 (Gemma Team, 2024) uses $c = 50$ on attention logits and $c = 30$ on the final output logits. With $c = 50$, a logit of 10 becomes $50\tanh(0.2) = 9.87$ (almost unchanged) and a logit of 200 becomes $50\tanh(4) = 49.97$ (clamped). The downside: fused attention kernels like FlashAttention need explicit support for the extra nonlinearity, which delayed adoption.

Sliding-window attention

The problem: full attention over a 128k context costs $O(T^2)$ per layer and a KV cache of 128k entries per layer. Sliding-window attention (Beltagy et al., 2020) restricts each token to attend to only the previous $w$ tokens. Information from further back still arrives, but indirectly: after $\ell$ layers the receptive field is $\ell \cdot w$. Mistral 7B (Jiang et al., 2023) used $w = 4096$ in every layer. Gemma 2 alternates local (window 4096) and global layers so that some layers still see everything. This caps the KV cache of the local layers at $w$ entries regardless of context length.

Larger vocabularies

GPT-2 had 50,257 tokens; Llama 1 and 2 had 32,000; Llama 3 jumped to 128,256; Gemma uses 256,000. The problem: with a small vocabulary, non-English text and code fragment into many tokens. Llama 3 reports that its tokenizer needs about 15% fewer tokens than Llama 2's for the same English text, and much fewer for other languages. Fewer tokens per document means each training token carries more information and each generated token gets you further, so at fixed compute the model effectively sees more text. Tao et al. (2024) fit scaling laws for vocabulary size and find the optimal vocabulary grows with model size, which argues that older models were too small in this dimension. The costs are a bigger embedding table, a bigger output softmax, and rarer tokens that each get fewer training examples.

Context length and how it is grown

Why not train at 128k context from the start? Because attention cost grows with the square of the sequence length within a batch, and because most training documents are short; padding or packing them into 128k windows would waste compute on cross-document attention. So the standard recipe (Llama 3, and most others) trains at a modest length (8,192 tokens for Llama 3) for almost all of the tokens, then runs a long-context stage at the end.

Llama 3 grew context in six stages from 8k to 128k, on about 800B tokens, at each stage waiting until short-context evaluations recovered and needle-in-a-haystack retrieval worked at the new length. RoPE's base frequency is raised to match (500,000 for Llama 3). This is cheap relative to the main run and gives you most of the long-context capability. The advanced objectives chapter covers the mechanics.

The config sheet: Llama 3 8B, with the arithmetic

Here is the actual configuration. Then we count every parameter by hand, because being able to do this is the difference between reading a config file and understanding it.

Counting the parameters

Per layer, attention has four matrices. $W_Q$ and $W_O$ are full $d \times d$; $W_K$ and $W_V$ are $d \times (n_{kv} \cdot d_{\text{head}}) = 4096 \times 1024$.

Attention per layer: $16.78\text{M} \times 2 + 4.19\text{M} \times 2 = 41{,}943{,}040$. The gated MLP has three matrices of $4096 \times 14336$:

Two RMSNorms add $2 \times 4096 = 8{,}192$. Per layer total: $41{,}943{,}040 + 176{,}160{,}768 + 8{,}192 = 218{,}112{,}000$. Times 32 layers: $6{,}979{,}584{,}000$.

Outside the layers: the input embedding $128{,}256 \times 4096 = 525{,}336{,}576$, the untied output head of the same size, and a final RMSNorm of 4,096.

That is 8.03B, matching the reported size. Two observations fall out. The MLP is 81% of each layer; that is where the knowledge lives. And the two embedding matrices are 13% of the whole model, which is why the vocabulary size decision is not free.

Companion code

Everything above is implemented in code/lumen/block.py: RMSNorm, SwiGLU, and a TransformerBlock that can be configured as GPT-2-style or Llama-style. The attention module in code/lumen/attention.py takes n_kv_heads and repeats K and V across groups. The heart of it is short enough to show here.

import torch, torch.nn as nn, torch.nn.functional as F

class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d))   # gamma only, no beta

    def forward(self, x):
        rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return x * rms * self.weight

class SwiGLU(nn.Module):
    def __init__(self, d, hidden):
        super().__init__()
        self.w_gate = nn.Linear(d, hidden, bias=False)
        self.w_up   = nn.Linear(d, hidden, bias=False)
        self.w_down = nn.Linear(hidden, d, bias=False)

    def forward(self, x):
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))

def swiglu_hidden(d, multiplier=1.0, multiple_of=256):
    h = int(2 * (4 * d) / 3 * multiplier)              # 8/3 * d, times an optional factor
    return multiple_of * ((h + multiple_of - 1) // multiple_of)

print(swiglu_hidden(4096), swiglu_hidden(4096, 1.3, 1024))

The first number is Llama 2 7B's hidden size; the second is Llama 3 8B's. The helper reproduces both from the rule.

Practice

Further reading

• Wang et al. (2022). What Language Model Architecture and Pretraining Objective Work Best for Zero-Shot Generalization? The controlled comparison of objectives.
• Zhang and Sennrich (2019). Root Mean Square Layer Normalization. RMSNorm.
• Xiong et al. (2020). On Layer Normalization in the Transformer Architecture. Why pre-norm trains more easily.
• Shazeer (2020). GLU Variants Improve Transformer. SwiGLU and GeGLU; short and worth reading in full.
• Su et al. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. RoPE.
• Ainslie et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. Grouped-query attention.
• Shazeer (2019). Fast Transformer Decoding: One Write-Head is All You Need. Multi-query attention and the KV-cache argument.
• Henry et al. (2020). Query-Key Normalization for Transformers. QK-norm.
• Gemma Team (2024). Gemma 2: Improving Open Language Models at a Practical Size. Logit soft-capping and alternating local/global attention.
• Jiang et al. (2023). Mistral 7B. Sliding-window attention in a production model.
• Tao et al. (2024). Scaling Laws with Vocabulary. Larger models want larger vocabularies.
• Grattafiori et al. (2024). The Llama 3 Herd of Models. The config sheet's source.