One-sentence summary: model.py is not code that appears from nowhere. It connects the embeddings, positional encoding, causal self-attention, FFN, normalization, residual paths, and output projection whose data flow we have already derived.

📦 Original 2024 demo repository: github.com/waylandzhang/Transformer-from-scratch. That historical demo keeps model, training, and generation in one model.py; this chapter is the corrected standalone teaching version, not a line-for-line mirror.

18.1 Draw the Model's Boundary First

We are implementing a small decoder-only Transformer:

Token IDs [B, T]
    ↓
Token embedding + fixed sinusoidal position encoding
    ↓
N × Pre-Norm Transformer block
    ├── LayerNorm → causal multi-head attention → residual add
    └── LayerNorm → FFN → residual add
    ↓
Final LayerNorm
    ↓
Tied output projection
    ↓
Logits [B, T, vocab_size]

To keep the mechanism visible, we will write Attention ourselves and leave out KV Cache, RoPE, Flash Attention, padding masks, and distributed training. The code supports equal-length batches without PAD tokens. Padded batches also need the key padding mask discussed in Chapter 16.

This is not a component-by-component GPT-2 reproduction. It uses a GPT-style Pre-Norm decoder block while retaining the original Transformer's fixed sinusoidal positions. That is a teaching choice, and we will state each such choice instead of hiding it.

18.2 Let the Configuration Guard Our Shapes

from dataclasses import dataclass

@dataclass
class ModelConfig:
    vocab_size: int
    context_length: int = 256
    d_model: int = 512
    n_layers: int = 6
    n_heads: int = 8
    dropout: float = 0.1

    def __post_init__(self):
        if self.vocab_size <= 0 or self.context_length <= 0:
            raise ValueError("vocab_size and context_length must be positive")
        if self.d_model <= 0 or self.n_layers <= 0 or self.n_heads <= 0:
            raise ValueError("d_model, n_layers, and n_heads must be positive")
        if self.d_model % self.n_heads != 0:
            raise ValueError("d_model must be divisible by n_heads")
        if not 0.0 <= self.dropout < 1.0:
            raise ValueError("dropout must be in [0, 1)")

    @property
    def head_dim(self):
        return self.d_model // self.n_heads

head_dim = d_model // n_heads is valid only when the division is exact. The old code assumed that it was, while its output projection still expected exactly d_model features. An invalid configuration therefore failed much later than it should have.

vocab_size should come from the tokenizer's vocabulary, not from the largest token ID that happened to occur in one training file. The latter makes support for a valid token depend on whether the sample happened to contain it.

18.3 Feed-Forward Network

Every sequence position independently passes through the same FFN:

[B, T, d_model]
       ↓ Linear
[B, T, 4 × d_model]
       ↓ GELU
[B, T, 4 × d_model]
       ↓ Linear + Dropout
[B, T, d_model]

class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        hidden_dim = 4 * config.d_model
        self.net = nn.Sequential(
            nn.Linear(config.d_model, hidden_dim),
            nn.GELU(approximate="tanh"),
            nn.Linear(hidden_dim, config.d_model),
            nn.Dropout(config.dropout),
        )

    def forward(self, x):
        return self.net(x)

We use GELU to stay consistent with the GPT-2 reference in the full-forward-pass chapter. A modern model may instead use SwiGLU and an intermediate width that is not simply 4 × d_model; this chapter first builds the classic two-layer FFN.

18.4 Write One Readable Attention Head First

With a causal mask, the equation is:

\operatorname{Attention}(Q,K,V)=\operatorname{softmax}\left(\frac{QK^\top}{\sqrt{d_h}}+M\right)V

Here $M_{ij}=0$ means visible, while $M_{ij}=-\infty$ means Query position $i$ must not see Key position $j$ .

class AttentionHead(nn.Module):
    def __init__(self, d_model, head_dim, context_length, dropout):
        super().__init__()
        self.head_dim = head_dim
        self.query = nn.Linear(d_model, head_dim, bias=False)
        self.key = nn.Linear(d_model, head_dim, bias=False)
        self.value = nn.Linear(d_model, head_dim, bias=False)
        self.attn_dropout = nn.Dropout(dropout)

        mask = torch.tril(
            torch.ones(context_length, context_length, dtype=torch.bool)
        )
        self.register_buffer("causal_mask", mask, persistent=False)

    def forward(self, x):
        _, T, _ = x.shape
        if T > self.causal_mask.size(0):
            raise ValueError("sequence is longer than context_length")

        q = self.query(x)                          # [B, T, head_dim]
        k = self.key(x)                            # [B, T, head_dim]
        v = self.value(x)                          # [B, T, head_dim]

        scores = q @ k.transpose(-2, -1)           # [B, T, T]
        scores = scores * (self.head_dim ** -0.5)
        allowed = self.causal_mask[:T, :T]
        scores = scores.masked_fill(~allowed, float("-inf"))

        weights = F.softmax(scores, dim=-1)
        weights = self.attn_dropout(weights)
        return weights @ v                         # [B, T, head_dim]

The mask is a boolean buffer: it is not a trainable parameter, but it follows model.to(device). persistent=False also keeps this easily reconstructed triangle out of checkpoints.

After Dropout is applied to Softmax probabilities, an Attention row during training need not sum to exactly one. Inverted Dropout preserves the expectation. model.eval() disables it.

18.5 From Separate Heads to a Fused Implementation

18.5.1 Teaching version: one module per head

class TeachingMultiHeadAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.heads = nn.ModuleList([
            AttentionHead(
                config.d_model,
                config.head_dim,
                config.context_length,
                config.dropout,
            )
            for _ in range(config.n_heads)
        ])
        self.out_proj = nn.Linear(config.d_model, config.d_model, bias=False)
        self.resid_dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        y = torch.cat([head(x) for head in self.heads], dim=-1)
        return self.resid_dropout(self.out_proj(y))

The heads have no mathematical dependence on one another, but this Python list comprehension still invokes separate modules. It is easy to read; it is not the most efficient way to launch the work on a GPU.

18.5.2 Practical version: project once, then reshape

The original Transformer paper defines per-head matrices $W_i^Q,W_i^K,W_i^V$ . It does not require one Python object per head. Concatenating those matrices along their output dimension lets one large matrix multiplication produce every Q, K, and V. This is a common fused implementation that is algebraically equivalent to the paper's equations—not “the paper's source code.”

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_heads = config.n_heads
        self.head_dim = config.head_dim
        self.context_length = config.context_length

        self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False)
        self.out_proj = nn.Linear(config.d_model, config.d_model, bias=False)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)

        mask = torch.tril(torch.ones(
            config.context_length,
            config.context_length,
            dtype=torch.bool,
        )).view(1, 1, config.context_length, config.context_length)
        self.register_buffer("causal_mask", mask, persistent=False)

    def forward(self, x):
        B, T, C = x.shape
        if T > self.context_length:
            raise ValueError("sequence is longer than context_length")

        q, k, v = self.qkv(x).split(C, dim=-1)
        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)

        scores = q @ k.transpose(-2, -1)
        scores = scores * (self.head_dim ** -0.5)
        allowed = self.causal_mask[:, :, :T, :T]
        scores = scores.masked_fill(~allowed, float("-inf"))
        weights = F.softmax(scores, dim=-1)
        weights = self.attn_dropout(weights)

        y = weights @ v                              # [B, H, T, D]
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_dropout(self.out_proj(y))

The shapes move as follows:

x                 [B, T, C]
q, k, v           [B, T, C] each
split heads       [B, H, T, D]   where C = H × D
scores / weights  [B, H, T, T]
weighted values   [B, H, T, D]
merge heads       [B, T, C]

Production code can replace the manual core with PyTorch's scaled_dot_product_attention, allowing the backend to select kernels such as Flash Attention when eligible. We keep the explicit version so the mask, Softmax, and shapes remain visible.

18.6 Transformer Block

Attention, FFN, and two residual paths in a Pre-Norm Transformer block

class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln1 = nn.LayerNorm(config.d_model)
        self.attn = CausalSelfAttention(config)
        self.ln2 = nn.LayerNorm(config.d_model)
        self.ffn = FeedForward(config)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x

This is Pre-Norm: normalization happens before the sublayer, while the residual branch preserves the unnormalized x. GPT-2 uses Pre-Norm with LayerNorm. LLaMA also uses a Pre-Norm topology, but with RMSNorm. Pre-Norm often makes deep networks easier to optimize; that does not make it unconditionally superior to Post-Norm on every metric.

18.7 How Tokens and Positions Enter the Model

18.7.1 Build the position table once

The old version rebuilt the entire sinusoidal table on every forward() call and stored a device string in its configuration. That repeats work and can disagree with the model after model.to(...).

def sinusoidal_positions(max_length, d_model):
    position = torch.arange(max_length, dtype=torch.float32).unsqueeze(1)
    frequency = torch.exp(
        torch.arange(0, d_model, 2, dtype=torch.float32)
        * (-math.log(10000.0) / d_model)
    )
    angles = position * frequency

    table = torch.zeros(max_length, d_model)
    table[:, 0::2] = torch.sin(angles)
    table[:, 1::2] = torch.cos(angles[:, :table[:, 1::2].shape[1]])
    return table.unsqueeze(0)  # [1, context_length, d_model]

The final slice also keeps an odd d_model from causing a shape mismatch in the cosine assignment. Common configurations are even, but a helper need not depend on that silently.

self.register_buffer(
    "position_encoding",
    sinusoidal_positions(config.context_length, config.d_model),
    persistent=False,
)

The table now follows the model across devices and dtypes; there is no self.device to go stale.

18.7.2 Weight tying

self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
self.apply(self._init_weights)
self.lm_head.weight = self.token_embedding.weight

The last line makes input embeddings and output projection refer to one actual Parameter. It does not copy the current values into a second matrix.

PyTorch's documentation states that nn.Embedding initializes weights from $\mathcal{N}(0,1)$ by default. Combined with $\sqrt{d_{model}}$ input scaling and a tied output head, that default produces unusually large initial logits. We therefore initialize the model before tying the weights:

@staticmethod
def _init_weights(module):
    if isinstance(module, (nn.Linear, nn.Embedding)):
        nn.init.normal_(module.weight, mean=0.0, std=0.02)
    if isinstance(module, nn.Linear) and module.bias is not None:
        nn.init.zeros_(module.bias)

0.02 is this teaching model's stated starting point, not an initialization law for every architecture. Deeper GPT recipes may additionally scale selected residual projections with layer count.

18.7.3 Forward pass and loss

def forward(self, idx, targets=None):
    if idx.ndim != 2:
        raise ValueError("idx must have shape [batch, time]")
    _, T = idx.shape
    if T == 0:
        raise ValueError("sequence must contain at least one token")
    if T > self.config.context_length:
        raise ValueError("sequence is longer than context_length")
    if targets is not None and targets.shape != idx.shape:
        raise ValueError("targets must have the same shape as idx")

    token = self.token_embedding(idx) * math.sqrt(self.config.d_model)
    position = self.position_encoding[:, :T, :]
    x = self.embedding_dropout(token + position)

    for block in self.blocks:
        x = block(x)
    x = self.final_norm(x)
    logits = self.lm_head(x)

    loss = None
    if targets is not None:
        loss = F.cross_entropy(
            logits.reshape(-1, logits.size(-1)),
            targets.reshape(-1),
        )
    return logits, loss

The original Transformer scales token embeddings by $\sqrt{d_{model}}$ before adding sinusoidal positions; this model explicitly follows that convention.

F.cross_entropy takes raw logits, not manually Softmax-normalized probabilities, plus integer token targets. Chapter 19 must prepare the one-token shift first: idx=[x0,...,xT-1] and targets=[x1,...,xT].

18.8 Generation Needs Explicit Boundaries Too

@torch.no_grad()
def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
    if idx.ndim != 2 or idx.size(1) == 0:
        raise ValueError("idx must have shape [batch, time] with time > 0")
    if max_new_tokens < 0:
        raise ValueError("max_new_tokens must be non-negative")
    if temperature < 0:
        raise ValueError("temperature must be non-negative")
    if top_k is not None and top_k <= 0:
        raise ValueError("top_k must be positive")

    for _ in range(max_new_tokens):
        idx_crop = idx[:, -self.config.context_length:]
        logits, _ = self(idx_crop)
        next_logits = logits[:, -1, :]

        if temperature == 0:
            idx_next = torch.argmax(next_logits, dim=-1, keepdim=True)
        else:
            next_logits = next_logits / temperature
            if top_k is None:
                probs = F.softmax(next_logits, dim=-1)
                idx_next = torch.multinomial(probs, num_samples=1)
            else:
                k = min(top_k, next_logits.size(-1))
                top_logits, top_indices = torch.topk(next_logits, k, dim=-1)
                top_probs = F.softmax(top_logits, dim=-1)
                choice = torch.multinomial(top_probs, num_samples=1)
                idx_next = top_indices.gather(-1, choice)

        idx = torch.cat((idx, idx_next), dim=1)
    return idx

Three boundaries matter:

temperature == 0 explicitly takes the greedy branch instead of dividing by zero.
Sampling directly inside the K returned logits keeps exactly K indices even when the boundary score is tied.
@torch.no_grad() disables gradient recording, not Dropout. The caller must still run model.eval() before generation.

This teaching loop recomputes the visible window at every step; it has no KV Cache. Once the oldest token is cropped, fixed absolute positions restart from zero within the window. That is this implementation's explicit sliding-window policy, not every serving system's default.

18.9 Complete `model.py`

import math
from dataclasses import dataclass

import torch
import torch.nn as nn
from torch.nn import functional as F


@dataclass
class ModelConfig:
    vocab_size: int
    context_length: int = 256
    d_model: int = 512
    n_layers: int = 6
    n_heads: int = 8
    dropout: float = 0.1

    def __post_init__(self):
        if self.vocab_size <= 0 or self.context_length <= 0:
            raise ValueError("vocab_size and context_length must be positive")
        if self.d_model <= 0 or self.n_layers <= 0 or self.n_heads <= 0:
            raise ValueError("d_model, n_layers, and n_heads must be positive")
        if self.d_model % self.n_heads != 0:
            raise ValueError("d_model must be divisible by n_heads")
        if not 0.0 <= self.dropout < 1.0:
            raise ValueError("dropout must be in [0, 1)")

    @property
    def head_dim(self):
        return self.d_model // self.n_heads


def sinusoidal_positions(max_length, d_model):
    position = torch.arange(max_length, dtype=torch.float32).unsqueeze(1)
    frequency = torch.exp(
        torch.arange(0, d_model, 2, dtype=torch.float32)
        * (-math.log(10000.0) / d_model)
    )
    angles = position * frequency
    table = torch.zeros(max_length, d_model)
    table[:, 0::2] = torch.sin(angles)
    table[:, 1::2] = torch.cos(angles[:, :table[:, 1::2].shape[1]])
    return table.unsqueeze(0)


class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        hidden_dim = 4 * config.d_model
        self.net = nn.Sequential(
            nn.Linear(config.d_model, hidden_dim),
            nn.GELU(approximate="tanh"),
            nn.Linear(hidden_dim, config.d_model),
            nn.Dropout(config.dropout),
        )

    def forward(self, x):
        return self.net(x)


class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_heads = config.n_heads
        self.head_dim = config.head_dim
        self.context_length = config.context_length
        self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False)
        self.out_proj = nn.Linear(config.d_model, config.d_model, bias=False)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)

        mask = torch.tril(torch.ones(
            config.context_length,
            config.context_length,
            dtype=torch.bool,
        )).view(1, 1, config.context_length, config.context_length)
        self.register_buffer("causal_mask", mask, persistent=False)

    def forward(self, x):
        B, T, C = x.shape
        if T > self.context_length:
            raise ValueError("sequence is longer than context_length")

        q, k, v = self.qkv(x).split(C, dim=-1)
        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)

        scores = q @ k.transpose(-2, -1)
        scores = scores * (self.head_dim ** -0.5)
        allowed = self.causal_mask[:, :, :T, :T]
        scores = scores.masked_fill(~allowed, float("-inf"))
        weights = F.softmax(scores, dim=-1)
        weights = self.attn_dropout(weights)

        y = weights @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_dropout(self.out_proj(y))


class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln1 = nn.LayerNorm(config.d_model)
        self.attn = CausalSelfAttention(config)
        self.ln2 = nn.LayerNorm(config.d_model)
        self.ffn = FeedForward(config)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x


class Model(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.embedding_dropout = nn.Dropout(config.dropout)
        self.blocks = nn.ModuleList([
            TransformerBlock(config) for _ in range(config.n_layers)
        ])
        self.final_norm = nn.LayerNorm(config.d_model)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.apply(self._init_weights)
        self.lm_head.weight = self.token_embedding.weight

        self.register_buffer(
            "position_encoding",
            sinusoidal_positions(config.context_length, config.d_model),
            persistent=False,
        )

    @staticmethod
    def _init_weights(module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
        if isinstance(module, nn.Linear) and module.bias is not None:
            nn.init.zeros_(module.bias)

    def forward(self, idx, targets=None):
        if idx.ndim != 2:
            raise ValueError("idx must have shape [batch, time]")
        _, T = idx.shape
        if T == 0:
            raise ValueError("sequence must contain at least one token")
        if T > self.config.context_length:
            raise ValueError("sequence is longer than context_length")
        if targets is not None and targets.shape != idx.shape:
            raise ValueError("targets must have the same shape as idx")

        token = self.token_embedding(idx) * math.sqrt(self.config.d_model)
        position = self.position_encoding[:, :T, :]
        x = self.embedding_dropout(token + position)

        for block in self.blocks:
            x = block(x)
        x = self.final_norm(x)
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.reshape(-1, logits.size(-1)),
                targets.reshape(-1),
            )
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        if idx.ndim != 2 or idx.size(1) == 0:
            raise ValueError("idx must have shape [batch, time] with time > 0")
        if max_new_tokens < 0:
            raise ValueError("max_new_tokens must be non-negative")
        if temperature < 0:
            raise ValueError("temperature must be non-negative")
        if top_k is not None and top_k <= 0:
            raise ValueError("top_k must be positive")

        for _ in range(max_new_tokens):
            idx_crop = idx[:, -self.config.context_length:]
            logits, _ = self(idx_crop)
            next_logits = logits[:, -1, :]

            if temperature == 0:
                idx_next = torch.argmax(next_logits, dim=-1, keepdim=True)
            else:
                next_logits = next_logits / temperature
                if top_k is None:
                    probs = F.softmax(next_logits, dim=-1)
                    idx_next = torch.multinomial(probs, num_samples=1)
                else:
                    k = min(top_k, next_logits.size(-1))
                    top_logits, top_indices = torch.topk(next_logits, k, dim=-1)
                    top_probs = F.softmax(top_logits, dim=-1)
                    choice = torch.multinomial(top_probs, num_samples=1)
                    idx_next = top_indices.gather(-1, choice)

            idx = torch.cat((idx, idx_next), dim=1)
        return idx

18.10 Verify Semantics, Not Just “It Runs”

torch.manual_seed(7)
config = ModelConfig(
    vocab_size=32,
    context_length=8,
    d_model=12,
    n_layers=2,
    n_heads=3,
    dropout=0.0,
)
model = Model(config).eval()

idx = torch.tensor([[1, 2, 3, 4]])
targets = torch.tensor([[2, 3, 4, 5]])

with torch.inference_mode():
    logits, loss = model(idx, targets)
assert logits.shape == (1, 4, 32)
assert loss.ndim == 0

# Changing future tokens must not alter logits at the first two positions.
future_changed = torch.tensor([[1, 2, 9, 10]])
with torch.inference_mode():
    a, _ = model(idx)
    b, _ = model(future_changed)
torch.testing.assert_close(a[:, :2], b[:, :2])

# Three greedy tokens should increase the length by three.
generated = model.generate(idx[:, :2], 3, temperature=0)
assert generated.shape == (1, 5)

These assertions verify the output contract, causal behavior, and generation length. Merely checking that no exception occurred is not enough.

18.11 Parameter Count: Tied and Untied Are Far Apart

For d_model=512, n_heads=8, n_layers=6, and vocab_size=50,000:

Component	Parameters
Tied token embedding / output	25,600,000
Attention, 6 layers	6,291,456
FFN, 6 layers including biases	12,598,272
13 LayerNorms	13,312
Sinusoidal positions	0
Total	44,503,040

Without tying, the output matrix adds another 512 × 50,000 = 25,600,000, for 70,103,040 total parameters. The old edition counted the output as independent and then hid bias and Norm assumptions inside “about 70 million.” The assumptions and exact totals are now explicit.

Chapter Checklist

After this chapter, you should be able to:

Validate shape constraints before a matrix multiplication fails
Write one Attention head with a boolean causal mask
Explain why separate heads and fused QKV are algebraically equivalent but execute differently
Implement a Pre-Norm block, fixed position buffer, and actual weight tying
Handle temperature=0, Top-K ties, and model.eval() correctly
Test causal invariance so the model cannot silently look ahead

See You in the Next Chapter

The model can now run a forward pass and generate token IDs from random parameters, but it has not learned from data.

In Chapter 19 we will write train.py: prepare shifted inputs and targets, split training and validation data, backpropagate, and let the optimizer actually change the parameters.