DL S050 · Model Architecture Choices — GPT, Llama, Mistral

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 50 / 80 · Module M09 · ~2 hours

The story we're starting with

Open the Llama 3, Mistral 7B, Qwen 2.5, DeepSeek-V3, and Gemma 2 papers side by side. Skim the architecture sections. You will see the same seven decisions, in the same order, with tiny tweaks. Pre-LayerNorm. RMSNorm instead of LayerNorm. SwiGLU feedforward. RoPE positional encoding. Grouped-Query Attention. Tied embeddings. No bias on linears. That's the 2024 recipe. Every serious open model uses ~95% of these; the differences are which tricks they add on top (sliding window in Mistral, MoE in Mixtral/DeepSeek, MLA in DeepSeek-V3).

You will not build vanilla "Attention is All You Need" for your 100M model. That transformer is a museum piece. It has post-LN (unstable), sinusoidal PE (worse than RoPE), no gating in the FFN (leaves 15% of quality on the table), and per-head attention (wastes KV-cache memory). We're going to build the Llama-style decoder — because that's what every checkpoint from 2023 onward uses, and because you can trace every design decision to a real paper with a real ablation.

This session is a decision walkthrough: for each of the seven choices, we look at the alternative, the paper that killed the alternative, the empirical gap, and the code. By the end you can look at any modern LLM's config file and see the shape of the model without running it.

Prerequisites

• S040 (nanoGPT rebuild) and S043 (RoPE + FlashAttention) — you should have a GPT-2-style decoder mentally loaded.
• S021 (BatchNorm + LayerNorm) — for the norm section.

1 · Vanilla GPT-2 vs Llama-style — the seven diffs

Here's the diff, at a glance.

We'll walk each one.

2 · Pre-LN vs Post-LN — the change nobody talks about anymore

Vanilla "Attention is All You Need" placed LayerNorm after each sub-block:

x = x + Attention(x)
x = LayerNorm(x)
x = x + FFN(x)
x = LayerNorm(x)

That's post-LN. It works OK up to ~12 layers, then training becomes brittle — you need careful learning rate warmup, gradient clipping, and even then loss spikes are common past 24 layers.

Pre-LN moves the norm inside each residual branch:

x = x + Attention(LayerNorm(x))
x = x + FFN(LayerNorm(x))

This one change made ~100-layer transformers trainable without any of the exotic warmup dances. Why? Because the residual path stays a pure identity + accumulating add. Gradients flow through the identity path unchanged; the norm sits inside the branch where it can normalize per-branch activations without disturbing the residual.

Every LLM since GPT-2 uses pre-LN. This one is a settled question.

2.1 Code diff

# Post-LN (GPT-1 style — don't use)
def block(x):
    x = layer_norm(x + attention(x))
    x = layer_norm(x + ffn(x))
    return x
 
# Pre-LN (Llama style — use this)
def block(x):
    x = x + attention(layer_norm(x))
    x = x + ffn(layer_norm(x))
    return x

Two lines. One line each per direction. That's the whole difference. Also add a final norm after the last block, before the output projection — otherwise the last block's output isn't normalized.

3 · RMSNorm vs LayerNorm

LayerNorm:

LN(x) = γ · (x - μ) / σ + β

Where μ, σ are computed per token across the feature dim. Two parameters (γ, β) per feature.

RMSNorm (Zhang & Sennrich, 2019):

RMSNorm(x) = γ · x / RMS(x)   where   RMS(x) = sqrt(mean(x²) + ε)

Dropped the mean subtraction. Dropped the bias. That's it.

Empirically: RMSNorm matches LayerNorm on downstream loss to within noise (< 0.1% relative), while being ~15% faster on GPU (skips the mean computation). Also cleaner in bf16 — the mean subtraction is a common source of precision loss.

Adopted by Llama, PaLM, Chinchilla, Gemma. Not universally — GPT-4 seems to still use LN — but the trend is clear.

3.1 Ten-line implementation

import torch.nn as nn, torch
 
class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))  # γ only, no β
 
    def forward(self, x):
        rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).sqrt()
        return self.weight * (x / rms)

That's the whole thing.

4 · RoPE (Rotary Position Embeddings)

We covered RoPE in depth in S043. Recap: instead of adding positional info once at the input embedding, RoPE rotates the query and key vectors at each attention layer by a position-dependent angle. Two nice properties:

1. Relative position emerges for free — the attention score q_i · k_j becomes a function of (i - j) only (after the rotation).
2. Extrapolation friendly — because the rotation angles are fixed by position, you can linearly interpolate them at inference to handle contexts longer than what you trained on. RoPE + linear interpolation → 4× context extension for near-zero degradation. Sinusoidal PE has no equivalent trick.

Every 2023+ open model uses RoPE. Not negotiable for a modern build.

We use S043's apply_rope(q, k, freqs) here — you already have the code.

5 · Grouped-Query Attention (GQA)

This one is the modern KV-cache saver.

5.1 The setup

Standard Multi-Head Attention (MHA): n_heads query heads, n_heads key heads, n_heads value heads. All independent. Total params in the QKV projection = 3 · n_heads · head_dim · d_model = 3 · d_model².

Multi-Query Attention (MQA, Shazeer 2019): n_heads queries, but only 1 shared key and 1 shared value across all heads. QKV params = (n_heads + 2) · head_dim · d_model ≈ d_model². Save 66% of KV memory.

MQA turned out to be too aggressive — it hurt quality by ~1-2% loss.

Grouped-Query Attention (GQA, Ainslie et al., 2023): the compromise. n_heads queries but only n_kv_heads keys and values (where n_kv_heads divides n_heads). Each group of n_heads / n_kv_heads query heads shares one KV pair.

At n_kv_heads = 8 vs n_kv_heads = n_heads, quality drops <0.5% but KV cache shrinks 4-8×. That directly translates to 4-8× more concurrent users on the same GPU at inference time. Deployed models all use GQA now.

5.2 What changes in code

Instead of one QKV projection producing [B, T, 3 · n_heads · head_dim], you have:

• Q projection: [B, T, n_heads · head_dim]
• KV projection: [B, T, 2 · n_kv_heads · head_dim]

In the attention math, you repeat each KV head n_heads / n_kv_heads times so they broadcast against the queries:

# k, v shape: [B, T, n_kv_heads, head_dim]
# expand to [B, T, n_heads, head_dim] by repeating each KV head n_rep times
n_rep = n_heads // n_kv_heads
k = k.repeat_interleave(n_rep, dim=2)
v = v.repeat_interleave(n_rep, dim=2)
# now proceed as vanilla attention

That's the whole change. repeat_interleave is a memory-view op in modern PyTorch; no actual copy at runtime with flash-attention.

6 · SwiGLU FFN — the "silent 1-point improvement"

Vanilla transformer FFN:

FFN(x) = W2 · GELU(W1 · x + b1) + b2

Two projections, one nonlinearity. Hidden dim = 4 × d_model.

SwiGLU FFN (Shazeer 2020):

FFN(x) = W2 · (Swish(W_g · x) ⊙ (W_v · x))

Three projections (W_g, W_v, W_2), and instead of a single nonlinearity, a gate: element-wise multiply between the Swish-activated gate and the value projection.

Extra params for the same "width" — so real implementations shrink the hidden dim to ⅔ · 4 = 8/3 of d_model to hold total params constant. Llama uses hidden = int(2/3 · 4 · d_model) rounded to a multiple of 256.

6.1 Twenty-line implementation

import torch.nn.functional as F
 
class SwiGLU(nn.Module):
    def __init__(self, d_model, hidden=None):
        super().__init__()
        if hidden is None:
            hidden = int(2/3 * 4 * d_model)
            # Round to multiple of 256 for GPU efficiency
            hidden = ((hidden + 255) // 256) * 256
        self.w_gate = nn.Linear(d_model, hidden, bias=False)
        self.w_val  = nn.Linear(d_model, hidden, bias=False)
        self.w_out  = nn.Linear(hidden, d_model, bias=False)
 
    def forward(self, x):
        gate = F.silu(self.w_gate(x))       # SiLU == Swish(β=1)
        val  = self.w_val(x)
        return self.w_out(gate * val)

Compared to GELU-FFN with the same total params, SwiGLU gets ~1 point better validation loss on WikiText and ~1 point better on downstream. Cost: one extra matmul at every layer, so training FLOPs go up ~5%. Payback: worth it.

6.2 Why gating works

Handwave: element-wise multiplying by a gate lets the network learn which channels to attenuate per token — a form of soft dynamic routing. Ordinary FFNs can't do this; they can only apply a static nonlinearity. Gating unlocks a whole slice of computation that ReLU/GELU can't reach.

Sadly there's no crisp theory — this is one of those "empirical wins the field can't quite explain" results.

7 · No bias on linears + tied embeddings

Two small wins:

• No bias: Llama drops the b term from every nn.Linear. Empirically <0.1% quality loss. Saves 0.1% of params. Also removes a common source of unnormalized activations that mess up bf16. Do it.

• Tied embeddings: Share the weight matrix between the input token embedding and the output logit projection. Empirically zero quality loss (sometimes small gain for small models). Saves vocab_size × d_model params — for our 100M model with vocab 32k and d=768, that's 24.5M params saved. Almost 25% of the model. Free.

self.embed = nn.Embedding(vocab_size, d_model)
self.head  = nn.Linear(d_model, vocab_size, bias=False)
self.head.weight = self.embed.weight   # tie!

One line.

8 · Putting it all together — the config

Here's a Llama-style 100M-param config you'd actually use:

CONFIG_100M = dict(
    vocab_size = 32_000,       # from your S045 tokenizer
    d_model    = 768,
    n_layers   = 12,
    n_heads    = 12,
    n_kv_heads = 4,             # GQA: 3× KV savings
    hidden_dim = int(2/3 * 4 * 768),  # SwiGLU sizing
    max_seq_len = 2048,
    rope_theta  = 10_000.0,
    norm_eps    = 1e-6,
    tie_embeddings = True,
    dropout     = 0.0,          # no dropout for pretraining
)

Parameter count breakdown:

Right in the "call it 100M" ballpark.

8.1 A tiny reference-implementation skeleton

class LlamaLikeBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.norm1 = RMSNorm(cfg["d_model"])
        self.attn  = GroupedQueryAttention(cfg)
        self.norm2 = RMSNorm(cfg["d_model"])
        self.ffn   = SwiGLU(cfg["d_model"])
 
    def forward(self, x, freqs):
        x = x + self.attn(self.norm1(x), freqs)   # pre-LN + RoPE inside attn
        x = x + self.ffn(self.norm2(x))
        return x

Every block: two pre-normed residuals. RoPE-flavored GQA. SwiGLU FFN. That's a modern LLM block.

9 · Diagram — the anatomy

10 · What we did NOT choose (and why)

• Sliding-window attention (Mistral): great for long context on inference, adds complexity, unnecessary for our 2k-context 100M model.
• Mixture of Experts (Mixtral, DeepSeek): amortizes compute across many small FFNs. Awesome, but the training and infra are 3× more complex. Cover in S077.
• Multi-head Latent Attention (DeepSeek-V3): compresses KV further. Cutting-edge, unproven for small models. Wait a year.
• Sparse attention (Longformer, Big Bird): a solved problem replaced by RoPE-scaling for long context.

11 · The 2025 receipts — what actually shipped this year

The "seven diffs" list above is the 2023–2024 consensus. Since then, three architectural moves stopped being research curiosities and started shipping in frontier open models. If you're reading arxiv weekly, none of this is news; if you're not, this is the delta from "Llama-2 clone" to "you understand DeepSeek-V3."

11.1 · Multi-Head Latent Attention (MLA) — DeepSeek's KV-cache demolition

DeepSeek-V2 (May 2024) and DeepSeek-V3 (Dec 2024) replaced Multi-Head Attention with Multi-Head Latent Attention. The idea in one sentence: instead of caching K and V per head, cache a small compressed latent vector $c^{KV}_t \in \mathbb{R}^{d_c}$ per token, and reconstruct K and V on the fly with a learned up-projection. For DeepSeek-V3, $d_c = 512$ vs a naive KV footprint that would be $n_h \cdot d_h = 128 \cdot 128 = 16384$ per token. That's a 93.3% KV cache reduction, and the reported quality is better than MHA, not worse.

Why does that work at all? Because K and V for a given token are heavily redundant across heads — they're all views of the same residual-stream vector $x_t$ through different linear maps. Compressing to a shared latent and reprojecting per-head is a form of low-rank factorization that also acts as a mild regularizer. The DeepSeek-V3 tech report (arxiv 2412.19437, §2.1) has the full derivation; it takes 20 minutes and it's worth it.

The catch: MLA is fiddly to implement correctly with RoPE, because rotary needs to be applied in the pre-compression space of the query and a decoupled key. DeepSeek's solution is the "decoupled RoPE" trick — a small chunk of the query/key dims carry rotary, the rest carry content. Read section 2.1.2 of the tech report before you try it. For your 100M model: don't. MLA is a 2026-scale win. Note it, plan to add it in your v2 model.

11.2 · Mixture-of-Experts, but this time it works — DeepSeek's DeepSeekMoE + auxiliary-loss-free balancing

MoE has been "the future" since Shazeer 2017. Mixtral (Dec 2023) proved it could ship. But the two chronic MoE headaches were routing collapse (one expert gets all the traffic) and the auxiliary balancing loss that traded quality for stability. DeepSeek-V3 killed both:

• Fine-grained experts: instead of 8 fat experts × top-2 (Mixtral), DeepSeek-V3 uses 256 small routed experts + 1 shared expert × top-8. More experts × smaller-per-expert = better specialization and better load smoothing at the same active-params budget.
• Auxiliary-loss-free load balancing: each expert has a per-expert bias term that's nudged up when the expert is under-used and down when it's over-used. No gradient contribution — just a running-mean nudge. It kills routing collapse without the quality tax of the standard α · KL(routing || uniform) loss.

The receipts: DeepSeek-V3 is 671B total params, 37B active per token, trained on 14.8T tokens for 2.788M H800-hours (their own $5.576M cost estimate — see tech report §5.4). Compare to Llama-3 405B dense, trained for ~30M H100-hours. DeepSeek got GPT-4-class quality for ~1/10 the compute. That is the biggest architectural story of 2024–2025 and every open-weights lab is copying pieces of it in 2025 releases (Qwen 2.5-Max MoE, Llama 4 Scout/Maverick MoE variants, Mistral's mystery-model iterations).

Lesson for the 100M builder: don't use MoE at 100M. Below ~7B active params, the routing overhead eats the win. But understand it now — your career in this field will involve a lot of MoE.

11.3 · Multi-Token Prediction (MTP) — cheap self-distillation during pretrain

Also from DeepSeek-V3 (and the Gloeckle et al. Meta paper, arxiv 2404.19737, "Better & Faster Large Language Models via Multi-token Prediction"): instead of only predicting the next token, add k = 2 or 4 shallow prediction heads that predict tokens $t+2, t+3, \ldots$ from the same hidden state. During training, each head contributes to the loss with a small weight. At inference, you either discard them (they were free auxiliary supervision) or use them for speculative decoding (2× throughput, no draft model needed).

DeepSeek reports MTP tokens are accepted ~85–90% of the time as speculative draft tokens on their production endpoints — that's why V3 is so cheap to serve. Meta's paper reports MTP models trained on 200B+ tokens generalize better on downstream reasoning tasks (HumanEval, MBPP) than next-token-only baselines at the same param count. This one is cheap enough to consider for your 100M v2: one extra nn.Linear(d_model, vocab) head, a k=2 future-token loss with weight 0.5, and you're in the game.

11.4 · Long-context tricks that shipped in 2025

Context windows in 2025 are absurd: Llama 4 Scout ships with 10M tokens of context, Gemini 2.5 with 2M, Claude 4.5 with 1M. Almost all of this progress came from positional-encoding scaling rather than architecture changes:

• YaRN (Peng et al. 2023, arxiv 2309.00071): frequency-domain interpolation of RoPE. Fine-tune on ~1B tokens at 4× the training context, get quality out to that 4×. Every open long-context model in 2024 used this.
• LongRoPE (Ding et al., Microsoft, ICML 2024, arxiv 2402.13753): evolutionary search over per-dimension RoPE scale factors. Extended Llama-2 to 2M tokens with minimal fine-tuning.
• LongRoPE2 (2025): the follow-up. Applied to Phi-3-mini-128k and Llama-3-8B, achieves >99% recall on "needle-in-a-haystack" at 1M context with only 800B extra fine-tune tokens.
• Ring Attention (Liu et al., Berkeley 2023, arxiv 2310.01889): the trick that unlocked 10M — shard the sequence dimension across GPUs, ring-communicate KV blocks. Compute is unchanged, memory per GPU drops linearly with world size. Llama 4 Scout's 10M window is Ring Attention + a NoPE-style position-tagging change.

Lesson for the 100M builder: your model has 2k context. That's fine. When you graduate to 8k+, use YaRN. When you graduate to 100k+, use YaRN + a partial fine-tune. When you graduate to 1M+, hire someone who's done it before, or read the LongRoPE2 paper five times.

11.5 · What's still on the frontier (and probably won't ship in 2026)

• State-space models / Mamba-2 / Jamba / Griffin hybrids — impressive perplexity numbers, no clear win on in-context learning or long tail. The 2024 Together AI "Zamba" and AI21's Jamba-1.5 are the best hybrids I've seen; both are worth reading.
• 1.58-bit BitNet (Ma et al., Microsoft, 2024, arxiv 2402.17764): ternary weights $\{-1, 0, +1\}$ from scratch. If it scales to 100B params it changes everything. As of mid-2025 the largest confirmed BitNet is a 3B model. Watch, don't build yet.
• Titans / "neural memory" architectures (Google, 2025): learned key-value memories that persist across sequences. Novel, unproven, exciting.

Rule of thumb: anything that doesn't have an open-weights checkpoint you can download and inspect is research, not engineering. Copy what shipped, not what got tweeted.

Further reading (11.x)

• DeepSeek-V3 Tech Report — arxiv:2412.19437 — the single most valuable pretraining doc of 2024
• DeepSeek-V2 (MLA introduction) — arxiv:2405.04434
• Multi-Token Prediction (Meta) — arxiv:2404.19737
• Llama 4 model card (Meta, 2025) — ai.meta.com/blog/llama-4-multimodal-intelligence
• LongRoPE2 (Microsoft) — arxiv:2502.20082
• Ring Attention — arxiv:2310.01889
• BitNet b1.58 — arxiv:2402.17764

12 · War stories

13 · 🧠 Retention scaffold

One-line summary (write it in your own words): _______________________________

Spaced review: re-read §1 (the 7 diffs table) and §8 (CONFIG_100M) in 24 hours. Revisit the full session on day 7, focusing on §11 (2025 receipts).

Next session (S051): DDP vs FSDP — the two flavors of distributed training, when each becomes mandatory, and how much wall-clock each really costs.

Sticky note (keep on your desk): "Pre-LN + RMSNorm + RoPE + GQA(4:1) + SwiGLU(8d/3) + tied embeddings + no bias. That's the 2024 base. Add MLA + fine-grained MoE + MTP for 2025."

Legacy recall drills (kept for spaced-review continuity)

Stretch

Take the CONFIG_100M dict and vary one parameter at a time: (a) n_kv_heads from 1 to 12; (b) hidden_dim from 4×d down to 8/3·d; (c) n_layers at 8/12/16 with d_model scaled to hold params ~constant. Predict which changes affect training FLOPs the most, which affect inference FLOPs, and which affect KV cache size.

In your own words

"The 2024 LLM architecture is defined by (7 things). Each was adopted because ____ paper showed ____ gain."

Spaced-review pointer

• S040 (nanoGPT) — the baseline you're mutating.
• S043 (RoPE) — you're using its apply_rope unchanged.
• S042 (KV cache) — GQA's memory savings show up when you re-implement the cache with GQA.

Next-session teaser

We have an architecture. But how do we train it across multiple GPUs? Session 051 is DDP vs FSDP — the two flavors of distributed training, when to use which, and how much comms overhead each costs. Even if you never touch >1 GPU, understanding these frames the ceiling on any single-GPU run.

Bring back tomorrow

• Table: the 7-diff comparison table (post-LN → pre-LN, LN → RMSNorm, ...)
• Config: CONFIG_100M — you should be able to write it from memory by tomorrow.
• Skeleton: the 8-line LlamaLikeBlock.

Previous: ← DL S049 · Next: DL S051 →