Architectures & Hyperparameters

From 2017 to Now

Four specific changes separate a 2017 Transformer from what nearly every modern open-weight model actually ships with today: where normalization sits relative to each sublayer, what kind of normalization it is, how position is encoded, and what the feed-forward layer's internal activation looks like. Each one is a real, published, independently-motivated fix for a real problem the original design had.

The Original Transformer Block

Pre-LN vs. Post-LN: A Training-Stability Fix

RMSNorm: A Cheaper Normalization

# Original illustrative implementations for this series --
# not copied from any CS336 file or model codebase.

import torch

def layer_norm(x, eps=1e-5):
    mean = x.mean(dim=-1, keepdim=True)
    var = x.var(dim=-1, keepdim=True, unbiased=False)
    return (x - mean) / torch.sqrt(var + eps)

def rms_norm(x, eps=1e-5):
    # No mean subtraction -- only rescale by root-mean-square
    rms = torch.sqrt((x ** 2).mean(dim=-1, keepdim=True) + eps)
    return x / rms

RoPE: Encoding Position by Rotation

SwiGLU: A Gated Feed-Forward Layer

Putting It Together: The Modern Default Block

Stack all four changes on top of the original design and you get the transformer block architecture nearly every current open-weight LLM actually ships: Pre-LN placement, RMSNorm instead of LayerNorm, RoPE instead of learned or sinusoidal absolute position embeddings, and a SwiGLU feed-forward layer instead of a ReLU MLP. None of these four changes depend on each other — each is independently motivated and independently validated in its own paper — but they've converged into a de facto standard combination because they compose cleanly and each solves a real, separate problem: training stability, compute cost, length generalization, and representational flexibility, respectively.

Counting Every Parameter

Every change above can be judged by one practical question: where do the weights actually live, and how many are there? The answer is short. A decoder-only transformer has trainable parameters in only four places — the embedding tables, the attention projections, the feed-forward layers, and the normalization scales — and the whole count follows from five numbers. This section walks the model top to bottom, counts each piece, and then checks the arithmetic against a real PyTorch model.

1. The pipeline: which stages have weights?

Most of the steps between your prompt and the next token have no trainable parameters at all. Green stages learn; grey stages are fixed computation.

1. Input text"The cat sat on the…"0
2. TokenizationBPE from Part 2 — a lookup, not a layer0
3. Token IDs464 · 3797 · 3332 · …0
4. Token embeddingsE ∈ ℝ^V×dV·d
5. Position embeddingslearned P ∈ ℝ^T_max×d; RoPE or sinusoidal: noneT_max·d or 0
6. N × transformer blockattention + MLP + 2 norms, repeated N timesN · (per-block)
7. Final normγ and β (LayerNorm) or γ only (RMSNorm)2d or d
8. LM headd → V logits; often tied to Ed·V + V, or 0 if tied
9. Softmaxlogits → probabilities0
10. Samplinggreedy, top-k, nucleus0

2. Inside one block

This is the classic GPT-2-style block — LayerNorm with scale and shift, biases on every linear layer, a two-matrix GELU MLP — because it's the one the standard counting formula describes. The modern variant comes right after. The numbers on the right use the worked-example sizes below (d = 1,024, dff = 4,096).

Adding the pieces up gives the per-block formula. Splitting attention across h heads changes the shapes of the intermediate tensors, not the total: the four projections are d × d no matter how many heads share them.

3. The whole model

Weight tying — reusing the token-embedding matrix as the output projection — is worth knowing about because at small scales it's a large fraction of the model. It's the single biggest difference between the untied and tied totals below.

4. Worked example

V = 50,000, d = 1,024, d_ff = 4,096, N = 24, T_max = 2,048:

Where the 407M parameters live (untied):

• MLPs (24)201,449,472 · 49.5%
• Attention (24)100,761,600 · 24.8%
• LM head51,250,000 · 12.6%
• Token embeddings51,200,000 · 12.6%
• Position embeddings2,097,152 · 0.5%
• All 49 LayerNorms100,352 · 0.02%

5. What the modern block changes

Each of the four upgrades from earlier in this article moves the count, some up and some down:

That last row is the one to understand. SwiGLU adds a third d × dff matrix (the gate), so to keep the MLP near its old size Llama sets dff = ⅔ · 4d ≈ 2.67d, then rounds up to a multiple of 256 for hardware efficiency. In Meta's reference code that's int(2 * 4d / 3) rounded up to multiple_of = 256: for d = 4,096 it gives 11,008, Llama-7B's actual MLP width, and for our d = 1,024 it gives 2,816. The modern per-block count is 4d² + 3·d·dff + 2d:

The totals end up within about 1% of each other. That's by design: the modern block isn't meant to be bigger, it's meant to spend the same budget better.

6. Check it yourself in PyTorch

Formulas are easy to get subtly wrong, so here is the same count done by PyTorch itself: build both blocks out of real nn.Linear, nn.LayerNorm and nn.RMSNorm layers and sum their .parameters(). Building on the meta device allocates no memory, so this runs instantly on a laptop. Needs PyTorch 2.4+ (for nn.RMSNorm); install with pip install torch.

import torch
import torch.nn as nn


class ClassicBlock(nn.Module):
    """GPT-2-style block: LayerNorm (gamma+beta), biased Q/K/V/O, GELU MLP."""
    def __init__(self, d, d_ff):
        super().__init__()
        self.ln1 = nn.LayerNorm(d)
        self.q, self.k, self.v, self.o = (nn.Linear(d, d) for _ in range(4))
        self.ln2 = nn.LayerNorm(d)
        self.up, self.down = nn.Linear(d, d_ff), nn.Linear(d_ff, d)


class ModernBlock(nn.Module):
    """Llama-style block: RMSNorm (gamma only), no biases, SwiGLU MLP, RoPE (no weights)."""
    def __init__(self, d, d_ff):
        super().__init__()
        self.norm1 = nn.RMSNorm(d)
        self.q, self.k, self.v, self.o = (nn.Linear(d, d, bias=False) for _ in range(4))
        self.norm2 = nn.RMSNorm(d)
        self.gate = nn.Linear(d, d_ff, bias=False)
        self.up = nn.Linear(d, d_ff, bias=False)
        self.down = nn.Linear(d_ff, d, bias=False)


class LM(nn.Module):
    def __init__(self, block, V, d, d_ff, N, T_max=None, final_norm=nn.LayerNorm, tied=False, head_bias=True):
        super().__init__()
        self.tok = nn.Embedding(V, d)
        self.pos = nn.Embedding(T_max, d) if T_max else None   # None = RoPE: nothing to learn
        self.blocks = nn.ModuleList(block(d, d_ff) for _ in range(N))
        self.norm = final_norm(d)
        self.head = nn.Linear(d, V, bias=head_bias)
        if tied:
            self.head.weight = self.tok.weight               # shared tensor, counted once


def count(model):
    return sum(p.numel() for p in model.parameters())      # .parameters() de-duplicates tied weights


V, d, d_ff, N, T_max = 50_000, 1_024, 4_096, 24, 2_048

with torch.device("meta"):                                  # shapes only, no memory allocated
    block = ClassicBlock(d, d_ff)
    untied = LM(ClassicBlock, V, d, d_ff, N, T_max)
    tied = LM(ClassicBlock, V, d, d_ff, N, T_max, tied=True)
    modern = LM(ModernBlock, V, d, 2_816, N, final_norm=nn.RMSNorm, tied=True, head_bias=False)
    modern_block = ModernBlock(d, 2_816)

print(f"classic block:     {count(block):>12,}")
print(f"classic untied LM: {count(untied):>12,}")
print(f"classic tied LM:   {count(tied):>12,}")
print(f"modern block:      {count(modern_block):>12,}")
print(f"modern tied LM:    {count(modern):>12,}")

# Executed output (PyTorch 2.14):
# classic block:       12,596,224
# classic untied LM:  406,858,576
# classic tied LM:    355,658,576
# modern block:        12,847,104
# modern tied LM:     359,531,520

Every number matches the hand calculation. The one subtle line is weight tying: assigning self.head.weight = self.tok.weight makes both modules point at the same tensor, and model.parameters() yields each tensor once, so the shared matrix is counted once — exactly the 51.2M drop between the untied and tied totals.

Hyperparameter Choices

Beyond the structural choices above, the same lecture covers the numeric knobs that define a specific model's shape: number of layers, model dimension (d_model), number of attention heads, feed-forward hidden dimension (typically some multiple of d_model, often smaller for gated variants like SwiGLU since the gating branch effectively adds a third weight matrix), vocabulary size (carried over directly from Part 2's tokenizer), and context length. None of these have one universally correct value — they're chosen jointly with the training compute budget available, which is exactly the resource-accounting question Part 3 formalizes. The parameter count from the section above plugs straight into it: training compute is roughly 6 × parameters × tokens, so the 407M-parameter example trained on 20B tokens costs about 6 × 0.41B × 20B ≈ 4.9 × 10¹⁹ FLOPs.

⚠️ Confidence Notes and Gaps

🔗 Full Reference List