From GPT-2 to Llama: four changes and the reason for each

Part 6 ended on GPT-2’s five decisions and the loop that serves it. Open a model from this decade and the block is still Part 5’s: two stages, each reading a normalised copy of the bus and adding its result back. Four things inside it changed.

I read modeling_llama.py from transformers 5.17 next to nanoGPT’s model.py, and loaded SmolLM2-135M, a model with Llama’s architecture at GPT-2 small’s size. Then I undid each change on its trained weights and measured what broke.

Same block, four edits

LlamaDecoderLayer.forward, modeling_llama.py:305-323, condensed:

residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)            # RMSNorm, not LayerNorm
hidden_states, _ = self.self_attn(hidden_states, position_embeddings=...)
hidden_states = residual + hidden_states                       # the bus, as in GPT-2
residual = hidden_states
hidden_states = self.mlp(self.post_attention_layernorm(hidden_states))
hidden_states = residual + hidden_states

This is Part 5’s four lines with more names; the plus signs are where they were, and so is pre-LN, the norm applied before each stage rather than after. What changed sits inside the calls: how positions enter attention, what the norm computes, the shape of the feed-forward, and how many keys and values attention keeps. (Llama also drops the bias vectors from its linear projections; nothing below depends on them.)

The two models make a fair pair. GPT-2 small is 124 million parameters in 12 blocks of width 768. SmolLM2-135M is 135 million in 30 blocks of width 576. None of the four changes is about size; you could make either model with or without them.

Positions become a rotation

GPT-2 adds a learned row to each token, one row per position, and the table ends at 1,024. Llama has no position table. It rotates the query and key instead, at modeling_llama.py:158-159:

q_embed = (q * cos) + (rotate_half(q) * sin)
k_embed = (k * cos) + (rotate_half(k) * sin)

This is rotary position embedding, RoPE. In Part 4, attention scores a pair of tokens by comparing a query vector with a key vector. Each query and key is split into heads of 64 numbers. Each head’s vector is read as 32 pairs, each pair a point on a circle, and each point is turned like a clock hand by an angle proportional to the position. rotate_half is how the code writes that turn without a matrix. Every pair has its own speed. One config number, the rotation base, sets all of them; SmolLM2’s is 100,000. In SmolLM2 the fastest pair goes round once every 6.3 positions and the slowest once every 438,000.

The rotation is not decoration. Rotate a query by position m and a key by position n, and their dot product, the score from Part 4, depends only on m − n. I took a random query and key through SmolLM2’s rotation. At positions 10 and 3 the score was −13.4724; at 1,010 and 1,003 it was −13.4723. One step closer, at 10 and 4, it was −10.786.

The whole model inherits that. I shifted every position id in a 512-token text by 500, then by 5,000. SmolLM2’s loss on it, its average surprise at the next token, was 3.381 unshifted and 3.381 after each shift. GPT-2 given the same shift of 500 went from 3.85 to 8.97, because every token now reads a different row of its table. Removing the rotation altogether, by setting every position to 0, took SmolLM2 from 3.12 to 6.28 on the full 1,024-token text. Using 10,000 as the rotation base, the value in the original RoPE paper and in Llama 2, instead of SmolLM2’s 100,000 cost 3%.

Drag the key’s position and watch the score move, then shift both and watch it stay. Pick pair 17, the slowest of the three, and nearby positions barely differ. The score here is one pair’s share, between −1 and 1.

No table also means no last row. Llama 3.1 8B’s config records an original length of 8,192 positions and a maximum of 131,072. It gets there by stretching the rotation. Its heads have 128 numbers, so 64 pairs: the 29 slowest are slowed by a factor of 8, the 29 fastest left alone, and the six between blended. Meta also pre-trained the family at growing lengths up to 128K, in six stages for the 405B. There is no assert to hit: past the length a model was trained at, quality is a measurement, not a guarantee.

The norm stops centring

LayerNorm subtracts each row’s mean, divides by its standard deviation, then scales and shifts. LlamaRMSNorm, modeling_llama.py:62-67, skips the mean and the shift, and divides by the root mean square instead: square each number, average, take the root. If you’ve standardised a column in a dataframe, you’ve subtracted and divided; RMSNorm keeps only the division:

variance = hidden_states.pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
return self.weight * hidden_states.to(input_dtype)

The bet is that the mean is small anyway, and in SmolLM2 it is. In a typical norm layer, a typical row’s mean is 2% of its root mean square; in the worst layer, 4.5%. Swap full LayerNorm back in with the same learned scale and the loss moves from 3.124 to 3.175, under 2%.

Be fair about the reason given for it. The RMSNorm paper sells it as cheaper: one statistic instead of two, no shift. On my laptop’s CPU, torch’s F.layer_norm took 50 to 75 microseconds on a [1, 1024, 576] tensor across runs, and F.rms_norm 110 to 135. The profiler says why: on CPU, rms_norm still runs as a chain of small operations, square, average, add, root and two multiplies, while layer_norm is a single tuned kernel, one compiled routine for the whole job.

On a single thread the order flips: layer_norm about 155 microseconds, rms_norm about 130. Less arithmetic isn’t less time until someone writes the kernel, and I didn’t measure a GPU.

The feed-forward grows a gate

GPT-2’s feed-forward is two matrices with GELU between them. Llama’s is three, modeling_llama.py:175:

down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

The input goes through two projections side by side. One passes through SiLU, a smooth activation close to GELU, and then multiplies the other element by element: a gate. This is SwiGLU, from a 2020 paper that tested it and reported better results without offering a theory for why.

Three matrices cost more than two at the same width, so the width shrinks. With d the model width, a plain feed-forward of width 4d holds 8d² numbers, biases aside. Three matrices of width f hold 3·d·f, which matches at f = 8d/3. SmolLM2 uses exactly that: 1,536 is 8/3 of 576. Llama 2 7B’s 11,008 is 8/3 of 4,096 rounded up to a multiple of 256.

At 8d/3 the gate changes nothing in the budget: with full attention heads, the feed-forward still holds two thirds of a block. SmolLM2’s 75% comes from sharing keys and values across query heads, covered next. Llama 3.1 8B moves the split from both sides. Its feed-forward is wider, f = 3.5d, and its attention is smaller because it shares keys and values. Together they give the 81% quoted in Part 5: 176 million of each block’s 218 million parameters. The wider feed-forward alone, with full attention heads, gives 72%.

The gate carries weight. Replace SiLU with GELU in SmolLM2 and the loss goes from 3.12 to 3.79. Drop the gate’s multiplication entirely and it goes to 17.7, well past 10.8, the loss of a model that gives each of 49,152 tokens equal odds.

The feed-forward is also where mixture of experts plugs in: many feed-forwards per block and a router, a small layer that picks a few of them for each token. Llama 3 stayed dense; Llama 4 and other families didn’t.

Keys and values get shared

GPT-2 has as many key and value heads as query heads. SmolLM2 has 9 query heads and 3 key-value heads; each key and value head serves a group of three queries. Think of one index serving three different queries: same data, different questions. In the plain PyTorch path, repeat_kv, modeling_llama.py:179-189, copies the three key and value heads out to nine before the arithmetic; torch’s fused attention kernel can skip the copy when there is no padding mask. This is grouped-query attention, GQA. Llama 2 used it only in its two largest sizes, 70B and an unreleased 34B. Every Llama 3 size has it.

At inference the sharing is exact. I built a copy of SmolLM2 with nine full key and value heads, each a duplicate of its group’s, and ran both. The logits matched to the last bit, and the copy had 147.8 million parameters instead of 134.5. What GQA gives up, it gives up in training, with fewer distinct keys and values to learn.

The saving is in Part 6’s cache. SmolLM2 stores 11,520 numbers per token, which I read back from transformers’ own cache after 256 tokens. With full heads it would store 34,560. GPT-2 small, with 12 blocks to SmolLM2’s 30, stores 18,432.

At Llama 3.1 8B’s size the ratio decides what fits. It has 32 query heads and 8 key-value heads, so each token costs 131 KB in bf16, the 16-bit format from Part 3. At its full 131,072 positions, one sequence’s cache is 17.2 GB, next to 16.1 GB of weights. With full heads it would be 68.7 GB.

Start from GPT-2 small and switch the three changes on one at a time; the feed-forward share barely moves. Then drag the key-value heads down to 1 and watch the cache and the share.

What the ML engineer watches

Four habits, and none of them needs a GPU.

Recompute the cache for every model. Parameter count does not predict it: the cache is 2 × blocks × key-value heads × head width numbers per token, and the key-value heads are a config field. SmolLM2 is larger than GPT-2 small and stores three eighths less per token.

Port the architecture before the weights. Every change above is part of the contract the weights were trained under. The activation, the position scheme with its rotation base and scaling, and the key-value head count are config fields; the norm type comes with the architecture name. Get the head count wrong and loading fails, because the shapes don’t match. Get the count right and the grouping wrong, tiling the shared heads instead of interleaving them, and nothing raises: the loss went past 11. The same silence covers the rest. A bias-free LayerNorm instead of RMSNorm cost 1.6%, GELU instead of SiLU 21%, and the wrong rotation base 3%. Rotation pairs laid out side by side instead of in halves took the loss to 7.24. If you port one by hand, compare your logits against a reference implementation before you trust it.

Test the lengths you serve. Rotation removes the hard limit and keeps the question. If you run a model past the length it was trained at, you are the one who has to measure it there, on your own data.

Count active parameters separately. With mixture of experts, Qwen3-30B-A3B holds 30.5 billion parameters and uses 3.35 billion per token; 95% of the weights are experts. Memory follows the first number; arithmetic per token follows the second, and you pay for both.

None of this is modelling work either. A Llama block is a GPT-2 block with four parts swapped, and each swap is a line of config or code you can read, a cost you can compute, and a mismatch you can test for before a single request.