Build a transformer from scratch (TinyGPT)

Attention Is All You Need is the paper title you should know. The replica you will ship is TinyGPT: embeddings, causal self-attention, MLP, residuals, layer norm, and a tied or untied language-model head.

Do not wave your hands at multi-head attention. Write the shapes. B batch, T time, C width, H heads, D = C / H head dim. Q, K, V are produced by one linear map and split. Scores are QKT / sqrt(D). The causal mask forbids looking forward. Softmax is over keys. The output is a weighted sum of V, then a projection back to C.

Coding: CausalSelfAttention

This is the starter block. Type it. Do not paste it into a repo you never run.

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class CausalSelfAttention(nn.Module):
    def __init__(self, n_embd: int, n_head: int, block_size: int, dropout: float = 0.0):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.n_embd = n_embd
        self.c_attn = nn.Linear(n_embd, 3 * n_embd)
        self.c_proj = nn.Linear(n_embd, n_embd)
        self.resid_dropout = nn.Dropout(dropout)
        self.register_buffer(
            "bias",
            torch.tril(torch.ones(block_size, block_size)).view(1, 1, block_size, block_size),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.size()
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float("-inf"))
        att = F.softmax(att, dim=-1)
        y = att @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_dropout(self.c_proj(y))

Then wrap it in a block: norm, attention, residual, norm, MLP, residual. MLP is usually 4C wide with GELU or SiLU. Positional information can be a learned table of length block_size. That is enough to train on Shakespeare or a tiny code dump.

What you must be able to say out loud

• Why divide by sqrt(D): variance of the dot product grows with dim if you do not.
• Why the mask is lower-triangular: next-token training cannot see the answer token.
• Why KV cache exists later: at decode you reuse K and V for past tokens. You will meet this again in week 9.
• Why FlashAttention matters later: the attention matrix is T x T per head. Materialising it is the naive replica. Fusion is the production replica.

Assignment

• Train TinyGPT on a small text file until train loss drops. Log tokens per second. No claimed SOTA.
• Print parameter count from sum(p.numel() for p in model.parameters()).
• Break the mask (set it to ones) and show that the model cheats on a next-token batch. That is the causal test.

Opinion checkpoint

Core project 2: TinyGPT that trains on a tiny corpus.

Next: Scaling laws, data, and synthetic data.