Chapter 41 — Project 1: train a language model from scratch

Step 1 — Get data

mkdir -p llm-project && cd llm-project
curl -o input.txt https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
wc -c input.txt          # ~1.1 MB

Or use your own: a book from Project Gutenberg, your blog posts, your company's docs, code from a repo. Using your own text makes this far more interesting — you'll recognise the style it learns.

Step 2 — Train a tokenizer

# tokenizer.py
from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders

tok = Tokenizer(models.BPE())
tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=True)
tok.decoder = decoders.ByteLevel()

trainer = trainers.BpeTrainer(
    vocab_size=4096,                      # small vocab for a small model
    special_tokens=["<|endoftext|>"],
    initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
)
tok.train(["input.txt"], trainer)
tok.save("tokenizer.json")

ids = tok.encode("To be or not to be").ids
print(ids, "->", tok.decode(ids))
print("vocab size:", tok.get_vocab_size())

Why 4096 and not 50,000? A bigger vocabulary means a bigger embedding table and output layer — wasteful for a tiny model on 1 MB of text. Rule of thumb for toys: vocab of a few thousand.

Step 3 — Prepare the data

# prepare.py
import numpy as np
from tokenizers import Tokenizer

tok = Tokenizer.from_file("tokenizer.json")
text = open("input.txt", encoding="utf-8").read()
ids = tok.encode(text).ids
print(f"{len(text)} characters -> {len(ids)} tokens "
      f"({len(text)/len(ids):.2f} chars per token)")

data = np.array(ids, dtype=np.uint16)
n = int(0.9 * len(data))
data[:n].tofile("train.bin")
data[n:].tofile("val.bin")
print("saved train.bin and val.bin")

What "packing" means: we concatenate everything into one long stream of token IDs and later slice random windows from it. No padding, no waste.

Step 4 — The model

# model.py
import math, torch, torch.nn as nn, torch.nn.functional as F
from dataclasses import dataclass

@dataclass
class Config:
    vocab_size: int = 4096
    block_size: int = 256      # context length in tokens
    n_layer:    int = 6
    n_head:     int = 6
    n_embd:     int = 384      # must be divisible by n_head
    dropout:  float = 0.1

class Block(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.ln1  = nn.LayerNorm(c.n_embd)
        self.attn = nn.MultiheadAttention(c.n_embd, c.n_head,
                                          dropout=c.dropout, batch_first=True)
        self.ln2  = nn.LayerNorm(c.n_embd)
        self.ff   = nn.Sequential(
            nn.Linear(c.n_embd, 4 * c.n_embd), nn.GELU(),
            nn.Linear(4 * c.n_embd, c.n_embd), nn.Dropout(c.dropout))

    def forward(self, x, mask):
        h = self.ln1(x)
        a, _ = self.attn(h, h, h, attn_mask=mask, need_weights=False)
        x = x + a                                  # residual
        x = x + self.ff(self.ln2(x))               # residual
        return x

class GPT(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.c = c
        self.tok  = nn.Embedding(c.vocab_size, c.n_embd)
        self.pos  = nn.Embedding(c.block_size, c.n_embd)
        self.drop = nn.Dropout(c.dropout)
        self.blocks = nn.ModuleList([Block(c) for _ in range(c.n_layer)])
        self.lnf  = nn.LayerNorm(c.n_embd)
        self.head = nn.Linear(c.n_embd, c.vocab_size, bias=False)
        self.head.weight = self.tok.weight          # weight tying: saves params

    def forward(self, idx, targets=None):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)
        x = self.drop(self.tok(idx) + self.pos(pos))
        mask = torch.triu(torch.ones(T, T, device=idx.device, dtype=torch.bool), 1)
        for b in self.blocks:
            x = b(x, mask)
        logits = self.head(self.lnf(x))
        if targets is None:
            return logits, None
        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=0.8, top_k=50):
        self.eval()
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.c.block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature
            if top_k:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float("inf")
            probs = F.softmax(logits, dim=-1)
            idx = torch.cat([idx, torch.multinomial(probs, 1)], dim=1)
        return idx

Step 5 — The training loop

# train.py
import math, time, numpy as np, torch
from model import GPT, Config

cfg = Config()
device = "cuda" if torch.cuda.is_available() else "cpu"
print("device:", device)

train_data = np.memmap("train.bin", dtype=np.uint16, mode="r")
val_data   = np.memmap("val.bin",   dtype=np.uint16, mode="r")

batch_size = 32
def get_batch(split):
    d = train_data if split == "train" else val_data
    ix = torch.randint(len(d) - cfg.block_size - 1, (batch_size,))
    x = torch.stack([torch.from_numpy(d[i:i+cfg.block_size].astype(np.int64)) for i in ix])
    y = torch.stack([torch.from_numpy(d[i+1:i+1+cfg.block_size].astype(np.int64)) for i in ix])
    return x.to(device), y.to(device)

model = GPT(cfg).to(device)
n_params = sum(p.numel() for p in model.parameters())
print(f"parameters: {n_params/1e6:.2f}M")

# AdamW, with weight decay only on 2-D weight matrices (not norms/biases)
decay   = [p for p in model.parameters() if p.dim() >= 2]
nodecay = [p for p in model.parameters() if p.dim() <  2]
opt = torch.optim.AdamW([{"params": decay,   "weight_decay": 0.1},
                         {"params": nodecay, "weight_decay": 0.0}],
                        lr=3e-4, betas=(0.9, 0.95))

MAX_STEPS, WARMUP = 5000, 100
def lr_at(step):                              # warmup then cosine decay
    if step < WARMUP:
        return 3e-4 * step / WARMUP
    prog = (step - WARMUP) / (MAX_STEPS - WARMUP)
    return 3e-4 * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * prog)))

@torch.no_grad()
def estimate_loss(n=20):
    model.eval()
    out = {}
    for split in ["train", "val"]:
        losses = [model(*get_batch(split))[1].item() for _ in range(n)]
        out[split] = sum(losses) / len(losses)
    model.train()
    return out

t0 = time.time()
for step in range(MAX_STEPS + 1):
    for g in opt.param_groups:
        g["lr"] = lr_at(step)

    x, y = get_batch("train")
    _, loss = model(x, y)
    opt.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)   # stability
    opt.step()

    if step % 500 == 0:
        l = estimate_loss()
        print(f"step {step:5d} | train {l['train']:.4f} | val {l['val']:.4f} "
              f"| lr {lr_at(step):.2e} | {time.time()-t0:.0f}s")
        torch.save({"model": model.state_dict(), "cfg": cfg}, "ckpt.pt")

print("done. saved ckpt.pt")

Step 6 — Generate

# sample.py
import torch
from tokenizers import Tokenizer
from model import GPT, Config

device = "cuda" if torch.cuda.is_available() else "cpu"
ck = torch.load("ckpt.pt", map_location=device)
model = GPT(ck["cfg"]).to(device); model.load_state_dict(ck["model"]); model.eval()
tok = Tokenizer.from_file("tokenizer.json")

prompt = "KING RICHARD:"
ids = torch.tensor([tok.encode(prompt).ids], device=device)
out = model.generate(ids, max_new_tokens=300, temperature=0.8, top_k=50)
print(tok.decode(out[0].tolist()))

Experiments (do at least four)

1. Train on your own text. Your emails, a favourite author, a codebase. The style transfer is striking.
2. Vary n_layer (2, 4, 8) and plot final val loss. Find where it stops helping on your small dataset.
3. Vary block_size (64, 256, 512). Longer context = better coherence and slower training.
4. Break the causal mask (pass attn_mask=None). Train loss plummets; generation becomes nonsense. Explain why in notes.md.
5. Remove weight tying (delete self.head.weight = self.tok.weight). Count parameters before and after.
6. Remove clip_grad_norm_ and raise the LR to 3e-3. Watch it diverge to nan.
7. Overfit deliberately: use only the first 10 KB of the file. Watch val loss turn upward within a few hundred steps.