llama.py

THE EVOLVING TRANSFORMER · 04

LLaMA

From GPT-2 to modern open LLMs: LLaMA swaps learned positional embeddings for RoPE, LayerNorm for RMSNorm, GELU MLP for SwiGLU, and full multi-head attention for grouped-query attention — all without linear biases. Defaults match LLaMA 3 8B, the same size as the architecture diagram on this page.

CLIP
LLaMA (you are here)
Mistral
4096d_model
32layers N
GQA8 KV · 32 Q
8BLLaMA 3
LLaMA 3 8B architecture

LLaMA 3 8B · decoder stack (32× blocks, GQA 32/8, RoPE, Pre-RMSNorm, SwiGLU)

This chapter evolves chapter 02 GPT-2 into LLaMA — the modern decoder-only stack that drops learned positional embeddings for RoPE, swaps LayerNorm for RMSNorm, replaces the GELU MLP with SwiGLU, and uses grouped-query attention instead of full multi-head attention. Hyperparameters follow LLaMA 3 8B from Meta, matching the diagram above. Same teaching pattern: one config object up front, then each layer in order until the full model runs.

We do not walk through a separate LLaMA 2 chapter: the block diagram is the same decoder-only Transformer (RoPE, RMSNorm, SwiGLU, pre-norm residuals). LLaMA 3 is essentially a scaled and refined LLaMA 2 — same backbone, better tokenizer, longer context, and training scale. The code below matches LLaMA 3 because that is what the 8B defaults and this series use going forward.

Each section below is one building block. You get a short explanation, the code for that module, a delta from chapter 02 where it helps, and how it connects to what came before. By the end you have a runnable LLAMA class, not just a diagram.

Delta from GPT-2 (02)

  • Learned wpe removed → RoPE applied in attention
  • LayerNorm → RMSNorm
  • GELU MLP → SwiGLU
  • MHA → GQA (grouped-query attention)
  • Linear biases removed (bias=False everywhere)
  • No embedding dropout

Full comparisons: §08 (vs GPT-2 and LLaMA 2 lineage).

00 · CONFIG

One config object for the whole model

Before any module, we define a single LlamaConfig dataclass for LLaMA 3 8B. Every class below takes config instead of a long list of constructor arguments — same pattern as chapter 02, with fields adjusted for LLaMA's architecture.

LLaMA 3 8B uses d_model=4096, N=32 blocks, H=32 query heads, n_kv_heads=8 (4:1 GQA ratio), vocab=128,256, context length 8,192, and d_ff=14,336 — matching Meta’s Llama-3-8B release (~8.0B parameters with untied embeddings). Unlike GPT-2, there is no dropout field — LLaMA trains without embedding or attention dropout in the reference stack.

dk property: dmodel // H → 4096 // 32 = 128. Multi-head attention splits d_model across H query heads; d_model must be divisible by H.

dff / n_kv_heads: dff = 14336 — not 4 × dmodel (SwiGLU uses a smaller inner dim than GPT-2's GELU FFN). n_kv_heads = 8 gives a 4:1 Q-to-KV ratio (32 query heads share 8 key/value head groups).

All values below are LLaMA 3 8B — verified against Meta’s published config.json (hidden_size, num_hidden_layers, intermediate_size, etc.). Change LlamaConfig fields to explore other sizes from the LLaMA family.

dmodel4096 — hidden_size (LLaMA 3 8B)
N32 — num_hidden_layers
H32 — num_attention_heads
n_kv_heads8 — num_key_value_heads (4:1 GQA)
vocabsize128,256 — vocab_size (tiktoken tokenizer)
max_seq_len8,192 — max_position_embeddings (base 8B)
dff14,336 — intermediate_size (not 4 × dmodel)
norm_eps1e-5 — rms_norm_eps
B2 — batch size for demos only (not in official config)
dk128 — dmodel // H (head dim)
llama.pyimports + LlamaConfig
 1from dataclasses import dataclass
 2import torch
 3import torch.nn as nn
 4import torch.nn.functional as F
 5import math
 6
 7
 8# ============================================================
 9# Config — LLaMA 3 8B from Touvron et al., 2023 / Meta, 2024
10# d_model=4096, N=32, H=32, n_kv_heads=8, vocab=128256, context=8192, d_ff=14336
11# ============================================================
12@dataclass
13class LlamaConfig:
14    dmodel: int = 4096           # LLaMA 3 8B
15    N: int = 32                  # transformer blocks
16    H: int = 32                  # query heads
17    n_kv_heads: int = 8          # GQA — 4:1 ratio (32 Q / 8 KV)
18    vocabsize: int = 128256      # BPE vocabulary size (LLaMA 3)
19    max_seq_len: int = 8192      # context length
20    dff: int = 14336             # SwiGLU inner dim (not 4 × dmodel)
21    norm_eps: float = 1e-5       # RMSNorm epsilon
22    B: int = 2                   # batch size (for demo runs)
23
24    @property
25    def dk(self):
26        # per-head dimension — dmodel must be divisible by H
27        return self.dmodel // self.H
28
29
30config = LlamaConfig()  # LLaMA 3 8B — shared by every module below
01 · EMBEDDINGS

Token embedding only — no learned position table

LLaMA keeps wte — a token embedding table mapping IDs to d_model vectors — but drops GPT-2's wpe entirely. Position is not added at the input; instead, RoPE (section 02) rotates Q and K inside attention. There is also no dropout on the embedding output — LLaMA's reference stack omits the embedding regularisation GPT-2 applies.

Chapter 02 summed wte(x) + wpe(pos) then dropped the result. Here the forward pass is simply wte(x) — one lookup, straight into the block stack.

llama.pywte only
 1# inside LLAMA.__init__ — token embedding only (no wpe):
 2self.wte = nn.Embedding(config.vocabsize, config.dmodel)  # wte(x): (B, S, dmodel)
 3
 4# inside LLAMA.forward — no positional sum, no dropout:
 5_, S = x.shape  # x: (B, S)
 6x = self.wte(x)  # (B, S, dmodel) — no wpe; RoPE handles position in attention
▼ Show original version

GPT-2 in chapter 02: wte + wpe summed, then dropped (excerpt from gpt2.py).

gpt2.pywte + wpe
 1# inside GPT.__init__ — two learned embedding tables:
 2self.wte = nn.Embedding(config.vocabsize, config.dmodel)    # wte(x): (B, S, dmodel)
 3self.wpe = nn.Embedding(config.max_seq_len, config.dmodel)  # wpe(pos): (S, dmodel)
 4
 5# inside GPT.forward — sum and drop:
 6_, S = x.shape                                              # x: (B, S)
 7pos = torch.arange(0, S, device=x.device)                   # (S,)
 8x = self.drop(self.wte(x) + self.wpe(pos))                  # (B, S, dmodel) + (S, dmodel) → (B, S, dmodel)
02 · RoPE

Rotary positional embeddings in attention

RoPE (Rotary Positional Embedding) encodes position by rotating Q and K vectors in each head's 2D subspaces. Frequencies follow a geometric schedule from base θ=10,000 — same idea as sinusoidal PE, but applied as a rotation matrix rather than an additive table.

Cos/sin tables are precomputed for all positions up to max_seq_len and registered as buffers. In forward, even/odd dimensions are paired, rotated, and stacked back — no learned parameters, and the rotation composes naturally for relative position.

llama.pyRotaryPositionalEmbedding
57class RotaryPositionalEmbedding(nn.Module):
58    def __init__(self, config, base: float = 10000.0) -> None:
59        super().__init__()
60        self.dk = config.dk
61        freqs = 1.0 / (
62            base ** (torch.arange(0, self.dk, 2).float() / self.dk)
63        )  # (dk//2,)
64        pos = torch.arange(0, config.max_seq_len).float()  # (max_seq_len,)
65        angles = torch.outer(pos, freqs)  # (max_seq_len, dk//2)
66        self.register_buffer("cos_cached", torch.cos(angles))  # (max_seq_len, dk//2)
67        self.register_buffer("sin_cached", torch.sin(angles))  # (max_seq_len, dk//2)
68
69    def forward(self, x):
70        B, H, S, dk = x.shape  # x: (B, H, S, dk)
71        x_even = x[..., 0::2]  # (B, H, S, dk//2)
72        x_odd = x[..., 1::2]  # (B, H, S, dk//2)
73        cos = self.cos_cached[:S].unsqueeze(0).unsqueeze(0)  # (1, 1, S, dk//2)
74        sin = self.sin_cached[:S].unsqueeze(0).unsqueeze(0)  # (1, 1, S, dk//2)
75        rotated_even = x_even * cos - x_odd * sin  # (B, H, S, dk//2)
76        rotated_odd = x_even * sin + x_odd * cos  # (B, H, S, dk//2)
77        rotated = torch.stack([rotated_even, rotated_odd], dim=-1)  # (B, H, S, dk//2, 2)
78        return rotated.view(B, H, S, dk)  # (B, H, S, dk)
03 · RMSNORM

Root-mean-square normalization

Chapter 01 built LayerNormalization from scratch; chapter 02 switched to PyTorch's nn.LayerNorm. LLaMA goes further with RMSNorm — same scale-and-shift idea, but without centering (no mean subtraction). It divides by the root-mean-square of the features plus norm_eps, then multiplies by a learnable gamma vector.

RMSNorm is cheaper than LayerNorm (one fewer reduction) and works well with pre-norm stacks at large scale. We implement it explicitly here so you see the math, not just an import.

llama.pyRMSNorm
33class RMSNorm(nn.Module):
34    def __init__(self, config) -> None:
35        super().__init__()
36        self.norm_eps = config.norm_eps
37        self.gamma = nn.Parameter(torch.ones(config.dmodel))  # (dmodel,)
38
39    def forward(self, x):
40        # x: (B, S, dmodel)
41        rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.norm_eps)  # (B, S, 1)
42        return x * rms * self.gamma  # (B, S, dmodel)
▼ Show original version

GPT-2 in chapter 02 used PyTorch's built-in LayerNorm (mean + variance normalization with learnable weight and bias).

gpt2.pyLayerNorm (PyTorch built-in)
 1# Drop-in replacement for the educational LayerNormalization class: same math,
 2# fused kernels, and per-feature learnable weight/bias in modern PyTorch.
 3ln = nn.LayerNorm(config.dmodel)  # (B, S, dmodel) → (B, S, dmodel)
04 · SWIGLU

SwiGLU feed-forward — gate, up, down

LLaMA replaces GPT-2's two-linear GELU MLP with SwiGLU: three matrices (w1 gate, w2 up, w3 down) and the activation silu(w1(x)) * w2(x). The inner dimension dff is set explicitly in config (14,336 for 8B) — not derived as 4 × d_model.

All three linears use bias=False, consistent with LLaMA's bias-free design. No dropout inside the FFN either.

llama.pySwiGLUFeedForward
45class SwiGLUFeedForward(nn.Module):
46    def __init__(self, config) -> None:
47        super().__init__()
48        self.w1 = nn.Linear(config.dmodel, config.dff, bias=False)  # (B,S,dmodel) → (B,S,dff) gate
49        self.w2 = nn.Linear(config.dmodel, config.dff, bias=False)  # (B,S,dmodel) → (B,S,dff) up
50        self.w3 = nn.Linear(config.dff, config.dmodel, bias=False)  # (B,S,dff) → (B,S,dmodel) down
51
52    def forward(self, x):
53        # x: (B, S, dmodel) → silu(w1) * w2 → w3 → (B, S, dmodel)
54        return self.w3(F.silu(self.w1(x)) * self.w2(x))
▼ Show original version

GPT-2's MLP in chapter 02: two linears with GELU and dropout between them.

gpt2.pyMLP
35class MLP(nn.Module):
36    def __init__(self, config):
37        super().__init__()
38        self.w1 = nn.Linear(config.dmodel, config.dff)  # (B, S, dmodel) → (B, S, dff)
39        self.w2 = nn.Linear(config.dff, config.dmodel)  # (B, S, dff) → (B, S, dmodel)
40        self.dropout = nn.Dropout(config.dropout)
41
42    def forward(self, x):
43        # x: (B, S, dmodel) → w1 → gelu → dropout → w2 → (B, S, dmodel)
44        return self.w2(self.dropout(F.gelu(self.w1(x))))  # (B, S, dmodel)
05 · GQA

Grouped-query attention with RoPE

GQA (Grouped-Query Attention) uses H query heads but only n_kv_heads key/value head groups. After projecting K and V, repeat_kv broadcasts each KV head to serve multiple Q heads — here a 4:1 ratio (32 Q / 8 KV).

This is the main architectural fork from LLaMA 2 7B / 13B, which still use full MHA (n_kv_heads = H). Only LLaMA 2 70B used GQA; LLaMA 3 enables GQA on every size, which cuts KV-cache memory and speeds inference.

RoPE is applied to Q and K after projection. The causal mask and scaled dot-product attention follow the same pattern as GPT-2, but projections are bias-free and there is no attention dropout.

Attention(Q, K, V) = softmax(QKT / √dk) V
llama.pyGQABlock
81class GQABlock(nn.Module):
82    def __init__(self, config) -> None:
83        super().__init__()
84        self.dk = config.dk
85        self.H = config.H
86        self.n_kv_heads = config.n_kv_heads
87        self.n_rep = self.H // config.n_kv_heads
88        self.wq = nn.Linear(config.dmodel, config.H * self.dk, bias=False)  # dmodel → H*dk
89        self.wk = nn.Linear(config.dmodel, config.n_kv_heads * self.dk, bias=False)
90        self.wv = nn.Linear(config.dmodel, config.n_kv_heads * self.dk, bias=False)
91        self.wo = nn.Linear(config.dmodel, config.dmodel, bias=False)  # dmodel → dmodel
92
93    @staticmethod
94    def repeat_kv(x, n_rep):
95        if n_rep == 1:
96            return x
97        B, n_kv_heads, S, dk = x.shape
98        x = x.unsqueeze(2).expand(B, n_kv_heads, n_rep, S, dk)
99        return x.contiguous().view(B, n_kv_heads * n_rep, S, dk)  # (B, H, S, dk)
100
101    def forward(self, x, rope, mask):
102        B, S, _ = x.shape  # x: (B, S, dmodel)
103        query = self.wq(x).view(B, S, self.H, self.dk).transpose(1, 2)  # (B, H, S, dk)
104        key = self.wk(x).view(B, S, self.n_kv_heads, self.dk).transpose(1, 2)  # (B, n_kv, S, dk)
105        value = self.wv(x).view(B, S, self.n_kv_heads, self.dk).transpose(1, 2)  # (B, n_kv, S, dk)
106        query = rope(query)  # (B, H, S, dk)
107        key = rope(key)  # (B, n_kv, S, dk)
108        key = GQABlock.repeat_kv(key, self.n_rep)  # (B, H, S, dk)
109        value = GQABlock.repeat_kv(value, self.n_rep)  # (B, H, S, dk)
110        attn_score = query @ key.transpose(-1, -2) / math.sqrt(self.dk)  # (B, H, S, S)
111        attn_score = attn_score.masked_fill(mask[:S, :S] == 0, float("-inf"))
112        attn_score = F.softmax(attn_score, dim=-1)  # (B, H, S, S)
113        x = attn_score @ value  # (B, H, S, dk)
114        x = x.transpose(1, 2).contiguous().view(B, S, self.H * self.dk)  # (B, S, dmodel)
115        return self.wo(x)  # (B, S, dmodel)
▼ Show original version

GPT-2's CausalSelfAttention in chapter 02: full MHA — every head gets its own K and V (excerpt from gpt2.py).

gpt2.pyCausalSelfAttention (excerpt)
47class CausalSelfAttention(nn.Module):
48    def __init__(self, config) -> None:
49        super().__init__()
50        self.dmodel = config.dmodel
51        self.H = config.H
52        self.dk = config.dk
53        self.wq = nn.Linear(config.dmodel, config.dmodel)  # dmodel → dmodel
54        self.wk = nn.Linear(config.dmodel, config.dmodel)  # dmodel → dmodel
55        self.wv = nn.Linear(config.dmodel, config.dmodel)  # dmodel → dmodel
56        self.wo = nn.Linear(config.dmodel, config.dmodel)  # dmodel → dmodel
57        self.attn_dropout = nn.Dropout(config.dropout)
58        self.resid_dropout = nn.Dropout(config.dropout)
59
60    def attention(self, q, k, v, S, mask):
61        # q, k, v: (B, H, S, dk)
62        attn_score = q @ k.transpose(-1, -2) / math.sqrt(self.dk)  # (B, H, S, S)
63        attn_score = attn_score.masked_fill(mask[:S, :S] == 0, -1e9)  # mask: (S, S)
64        attn_score = attn_score.softmax(dim=-1)  # (B, H, S, S)
65        attn_score = self.attn_dropout(attn_score)  # (B, H, S, S)
66        return attn_score @ v  # (B, H, S, dk)
67
68    def forward(self, x, mask):
69        B, S, _ = x.shape  # x: (B, S, dmodel)
70        # project and split into H heads — same Wq, Wk, Wv pattern as chapter 01
71        query = self.wq(x).view(B, S, self.H, self.dk).transpose(1, 2)  # (B, H, S, dk)
72        key = self.wk(x).view(B, S, self.H, self.dk).transpose(1, 2)  # (B, H, S, dk)
73        value = self.wv(x).view(B, S, self.H, self.dk).transpose(1, 2)  # (B, H, S, dk)
74        x = self.attention(query, key, value, S, mask)  # (B, H, S, dk)
75        # (B, H, S, dk) -> (B, S, H, dk) -> (B, S, dmodel)
76        x = x.transpose(1, 2).contiguous().view(B, S, self.dmodel)  # (B, S, dmodel)
77        return self.resid_dropout(self.wo(x))  # (B, S, dmodel)
06 · BLOCK

RMSNorm + GQA + RMSNorm + SwiGLU

A LLaMA block is Block(config): two RMSNorms (rms_1, rms_2), one GQABlock, one SwiGLUFeedForward. Pre-norm is done inline in forward — same pattern as GPT-2, but norms are RMSNorm and attention receives the shared rope module.

Each sublayer gets a residual add. The block signature is (x, rope, mask) because RoPE lives outside the block and is passed through from the model.

llama.pyBlock
118class Block(nn.Module):
119    def __init__(self, config) -> None:
120        super().__init__()
121        self.gqa = GQABlock(config)
122        self.ff = SwiGLUFeedForward(config)
123        self.rms_1 = RMSNorm(config)  # pre-norm before attention
124        self.rms_2 = RMSNorm(config)  # pre-norm before FFN
125
126    def forward(self, x, rope, mask):
127        # x: (B, S, dmodel)
128        x = x + self.gqa(self.rms_1(x), rope, mask)  # (B, S, dmodel)
129        x = x + self.ff(self.rms_2(x))  # (B, S, dmodel)
130        return x  # (B, S, dmodel)
07 · ASSEMBLE LLAMA

Token embed, RoPE, N blocks, final norm, lm_head

LLAMA owns everything. wte maps token IDs to vectors — no wpe, no dropout. A single RotaryPositionalEmbedding is shared across all blocks. N Blocks via nn.ModuleList, each taking (x, rope, mask). Final RMSNorm then lm_head projects to vocab.

Two design choices carry over from GPT-2: (1) weight tying — lm_head.weight = wte.weight; (2) a pre-registered causal mask sliced per sequence at runtime. With the §00 defaults, that is the full LLaMA 3 8B stack at 4096 / 32 / 32 / 14336.

llama.pyLLAMA
133class LLAMA(nn.Module):
134    def __init__(self, config) -> None:
135        super().__init__()
136        self.wte = nn.Embedding(config.vocabsize, config.dmodel)  # wte(x): (B, S, dmodel)
137        self.rope = RotaryPositionalEmbedding(config)
138        self.blocks = nn.ModuleList([Block(config) for _ in range(config.N)])
139        self.norm = RMSNorm(config)  # (B, S, dmodel) → (B, S, dmodel)
140        self.lm_head = nn.Linear(config.dmodel, config.vocabsize, bias=False)  # (B, S, dmodel) → (B, S, vocabsize)
141        self.lm_head.weight = self.wte.weight  # weight tying
142        self.register_buffer(
143            "mask", torch.tril(torch.ones(config.max_seq_len, config.max_seq_len))
144        )  # (max_seq_len, max_seq_len)
145
146    def forward(self, x):
147        _, S = x.shape  # x: (B, S)
148        x = self.wte(x)  # (B, S, dmodel) — no wpe; RoPE handles position in attention
149        for block in self.blocks:
150            x = block(x, self.rope, self.mask)  # (B, S, dmodel)
151        x = self.norm(x)  # (B, S, dmodel)
152        return self.lm_head(x)  # (B, S, vocabsize)
08 · COMPARISON

GPT-2 Small (02) vs LLaMA 3 8B (03)

Canonical comparison vs the previous chapter in this series (code-level detail after you have read the stack).

AxisGPT-2 Small (02)LLaMA 3 8B (03)
BackboneDecoder-only · learned wpe · LayerNorm · GELU MLPDecoder-only · RoPE · RMSNorm · SwiGLU · GQA
Position encodingLearned wpe, summed at inputRoPE on Q/K inside attention (no wpe)
AttentionFull MHA (12 heads on Small)GQA — n_kv_heads=8, H=32 (4:1)
FFN activationGELU two-linear MLP (dff = 4 × dmodel)SwiGLU three-linear (silu(w1) * w2, explicit dff)
NormalizationPre-norm LayerNormPre-norm RMSNorm
Linear biasesBiases on attention and MLPbias=False everywhere
Dropout0.1 on embeddings, attention, MLPNone in reference stack
Context / vocab1,024 · 50,257 BPE8,192 · 128,256 (tiktoken)

LLaMA 2 vs LLaMA 3 8B (same backbone, config deltas)

We do not have a separate LLaMA 2 chapter — the block diagram is the same. What changes in the numbers you would put in LlamaConfig:

AxisLLaMA 2LLaMA 3 8B (this chapter)
BackboneRoPE · RMSNorm · SwiGLUSame
Attention (7B / 8B class)7B / 13B: full MHA8B: GQA on every size
Attention (70B)GQAGQA
Vocabulary~32K (SentencePiece)128K — stronger multilingual / code
Context (base)4,096 tokens8,192 (3.1+ extends with RoPE scaling)

Training scale (~2T vs ~15T+ tokens) and data quality explain much of LLaMA 3’s capability jump — not a new layer type. For LLaMA 2 7B you would set n_kv_heads = H (32) for MHA and use the smaller vocab / 4K context in config.

09 · REFERENCES

Papers & technical sources

Primary reports and references for this chapter. Read these for full equations, training details, and official hyperparameters.

Share this post

Comments