Build a GPT from scratch — Large Language Models

Why this matters

Every elite LLM engineer has implemented GPT once. It is the rite of passage — the moment all the abstract theory clicks into "oh, this is just a few hundred lines of PyTorch." After this lesson you will be able to read any LLM codebase (HF Transformers, vLLM, Megatron, torchtune) and follow it.

We will build a decoder-only Transformer with all modern conventions: pre-norm RMSNorm, SwiGLU FFN, RoPE, GQA, weight tying, causal attention via SDPA. Train it on tiny_shakespeare. Sample from it.

This lesson is a single long worked example. Open a notebook and follow along.

Learning objectives

Implement a modern decoder-only Transformer in <300 lines.
Train it from scratch on a tiny corpus.
Sample text from your trained model.
Refactor for clarity, readability, and reproducibility.

1. The plan

text → BPE tokens → embedding → [block × L] → final norm → output projection (tied) → logits

Where each block is:

x = x + Attn(RMSNorm(x))   # GQA + RoPE + causal + SDPA
x = x + FFN(RMSNorm(x))    # SwiGLU

Goal model size for the lab: ~10M params, trains on a single laptop GPU in 30 minutes.

2. Imports and config

python

import math, time, os
from dataclasses import dataclass
import torch, torch.nn as nn
import torch.nn.functional as F

@dataclass
class Config:
    vocab_size: int = 2000          # BPE vocab from Lesson 0.3
    d_model:   int = 256
    n_layers:  int = 6
    n_heads_q: int = 8
    n_heads_kv: int = 4              # GQA
    d_ff:      int = 768             # 3 * d_model for SwiGLU
    block_size: int = 256
    rope_theta: float = 10000.0
    dropout:   float = 0.0

3. The components

RMSNorm

python

class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d))
        self.eps = eps
    def forward(self, x):
        rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return x * rms * self.weight

SwiGLU FFN

python

class SwiGLU(nn.Module):
    def __init__(self, d, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d, d_ff, bias=False)        # gate
        self.w3 = nn.Linear(d, d_ff, bias=False)        # up-proj
        self.w2 = nn.Linear(d_ff, d, bias=False)        # down-proj
    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

RoPE helpers

python

def precompute_rope(d_head, max_T, theta=10000.0, device="cpu"):
    inv = 1.0 / (theta ** (torch.arange(0, d_head, 2, device=device).float() / d_head))
    t = torch.arange(max_T, device=device).float()
    freqs = torch.outer(t, inv)                              # (T, d_head/2)
    return freqs.cos(), freqs.sin()                          # broadcastable

def apply_rope(x, cos, sin):
    # x: (..., T, d_head)
    x1, x2 = x[..., 0::2], x[..., 1::2]
    return torch.stack([x1 * cos - x2 * sin,
                        x1 * sin + x2 * cos], dim=-1).flatten(-2)

Grouped-Query Attention with SDPA

python

class GQA(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.h_q  = cfg.n_heads_q
        self.h_kv = cfg.n_heads_kv
        self.dk   = cfg.d_model // cfg.n_heads_q
        self.W_Q  = nn.Linear(cfg.d_model, self.h_q  * self.dk, bias=False)
        self.W_K  = nn.Linear(cfg.d_model, self.h_kv * self.dk, bias=False)
        self.W_V  = nn.Linear(cfg.d_model, self.h_kv * self.dk, bias=False)
        self.W_O  = nn.Linear(self.h_q * self.dk, cfg.d_model, bias=False)

    def forward(self, x, cos, sin):
        B, T, _ = x.shape
        Q = self.W_Q(x).view(B, T, self.h_q,  self.dk).transpose(1, 2)
        K = self.W_K(x).view(B, T, self.h_kv, self.dk).transpose(1, 2)
        V = self.W_V(x).view(B, T, self.h_kv, self.dk).transpose(1, 2)
        # apply RoPE to Q and K (use first T rows of cos/sin)
        Q = apply_rope(Q, cos[:T, None, :], sin[:T, None, :])
        K = apply_rope(K, cos[:T, None, :], sin[:T, None, :])
        # repeat K, V to match Q heads (so SDPA works without manual broadcast)
        rep = self.h_q // self.h_kv
        K = K.repeat_interleave(rep, dim=1)
        V = V.repeat_interleave(rep, dim=1)
        # SDPA = FlashAttention under the hood
        out = F.scaled_dot_product_attention(Q, K, V, is_causal=True)
        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.W_O(out)

Block

python

class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.norm1 = RMSNorm(cfg.d_model)
        self.attn  = GQA(cfg)
        self.norm2 = RMSNorm(cfg.d_model)
        self.mlp   = SwiGLU(cfg.d_model, cfg.d_ff)
    def forward(self, x, cos, sin):
        x = x + self.attn(self.norm1(x), cos, sin)
        x = x + self.mlp(self.norm2(x))
        return x

Full model

python

class GPT(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.blocks  = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])
        self.norm    = RMSNorm(cfg.d_model)
        # weight tying: output proj shares weights with input embedding
        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        self.lm_head.weight = self.tok_emb.weight
        cos, sin = precompute_rope(cfg.d_model // cfg.n_heads_q,
                                   cfg.block_size, cfg.rope_theta)
        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)
        self.apply(self._init)

    def _init(self, m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, std=0.02)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        x = self.tok_emb(idx)
        for block in self.blocks:
            x = block(x, self.cos, self.sin)
        x = self.norm(x)
        logits = self.lm_head(x)                  # (B, T, V)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)),
                                   targets.view(-1), ignore_index=-1)
        return logits, loss

That is the whole model. ~120 lines.

4. Data loader

python

import numpy as np
data = np.fromfile("tinyshakespeare.bin", dtype=np.uint16)        # token ids you tokenized in 0.3
split = int(0.95 * len(data)); train, val = data[:split], data[split:]

def get_batch(split_arr, block_size, batch_size, device):
    ix = np.random.randint(0, len(split_arr) - block_size - 1, size=batch_size)
    x = torch.tensor(np.stack([split_arr[i:i+block_size] for i in ix]).astype(np.int64))
    y = torch.tensor(np.stack([split_arr[i+1:i+1+block_size] for i in ix]).astype(np.int64))
    return x.to(device), y.to(device)

Save your tokenized corpus once (Lesson 0.3) as a uint16 array → fast random batches.

5. Training loop

python

device = "cuda" if torch.cuda.is_available() else "cpu"
cfg = Config()
m = GPT(cfg).to(device)
print(sum(p.numel() for p in m.parameters()) / 1e6, "M params")

opt = torch.optim.AdamW(m.parameters(), lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1)
ITERS = 5000
for it in range(ITERS):
    x, y = get_batch(train, cfg.block_size, 64, device)
    with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
        _, loss = m(x, y)
    opt.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0)
    opt.step()
    if it % 100 == 0:
        print(f"step {it} loss {loss.item():.4f}")

This trains a ~10M-parameter GPT in ~10-30 minutes on a single GPU. Final loss should be around 1.5-1.9 (perplexity ~5-7).

6. Sampling

python

@torch.no_grad()
def generate(model, idx, max_new=200, temperature=1.0, top_k=None):
    model.eval()
    for _ in range(max_new):
        idx_cond = idx[:, -cfg.block_size:]
        logits, _ = model(idx_cond)
        logits = logits[:, -1] / temperature
        if top_k is not None:
            v, _ = torch.topk(logits, top_k)
            logits[logits < v[:, -1:]] = -float("inf")
        probs = F.softmax(logits, dim=-1)
        nxt = torch.multinomial(probs, 1)
        idx = torch.cat([idx, nxt], dim=1)
    return idx

start = torch.tensor([[your_bos_id]], device=device)
out = generate(m, start, max_new=200, temperature=0.8, top_k=50)
print(decode(out[0].tolist()))

Even at 10M params, your model should produce vaguely Shakespearean text after enough training.

7. Refactor checklist (the engineer's view)

Before declaring success:

Print parameter count and break down by component (embed, attn, FFN). Make sure FFN is the largest.
Verify each block's input/output shape with print statements.
Confirm loss.item() decreases monotonically after warmup.
Run validation loss every 200 steps; plot train/val.
Save a checkpoint and reload — verify identical outputs.
Compare your model's logits on a known prompt with HuggingFace's gpt2 (just to make sure the sampling pipeline is correct).
(Bonus) Compile with torch.compile(m) and measure speedup.

Hands-on lab (1 full day)

build_gpt.ipynb:

Implement everything above. Train on tiny_shakespeare to PPL ≤ 8.
Sample 500 tokens at temperatures 0.5, 0.8, 1.0. Discuss qualitative differences.
Increase n_layers to 12 and d_model to 384 (~30M params). Retrain. Compare PPL.
Replace SwiGLU with vanilla Linear -> GELU -> Linear. Compare PPL.
Replace RoPE with absolute learned embeddings. Compare PPL.
Add a KV cache to your generate() function. Verify identical outputs and ~10× speedup.
Bonus: train on enwik8 (100MB Wikipedia). Compare to literature numbers (~1.0 BPB for ~50M-param GPT).

Common pitfalls

Forgetting is_causal=True in SDPA → silent leakage; loss looks like it converges but the model "memorises."
Forgetting targets.view(-1) and logits.view(-1, V) in the cross-entropy reshape.
Computing RoPE with the wrong dtype on GPU → slow autocast issues.
Returning (logits, loss) but unpacking only logits somewhere — the model trains but you ignore loss.
Mixing block_size between train and inference — your buffers (cos/sin) need to cover both.

Self-check

Where in your model does positional information enter?
How many distinct projection matrices live in one transformer block?
Why is weight tying useful?
What is the expected ratio of FFN params to attention params?
What is the minimum loss achievable on tiny_shakespeare with vocab 2k? (Hint: think about data entropy.)

References

Karpathy, nanoGPT — the canonical reference.
Karpathy, Let's build GPT in code, spelled out — 2-hour video; watch it.
HuggingFace modeling_llama.py — the production version of what you just built.
Touvron et al. (2024), "The Llama 3 Herd of Models" — architecture appendix.
Eleuther AI's GPT-NeoX — research-quality Megatron-style implementation.