DL S041 · Training Your Decoder on Tiny Shakespeare

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 41 / 80 · Module M07 · ~2 hours

The story

Yesterday you built a GPT. Today you turn it on and watch it learn.

Tiny Shakespeare — 1.1 MB of the Bard's plays, ~40k lines — is the "hello world" of language modelling. Small enough to overfit on a laptop, complex enough to produce something that fools a casual reader for half a paragraph. The training run takes 15 minutes on an M2 Mac or 3 minutes on an A100. By the end you'll have generated text like this:

DUKE OF YORK: Now, by my faith, thou art a knave and knavery. Bear him thence to my heart, and let him be A gentleman, that hath a great occasion.

That's not Shakespeare, but it's not gibberish either. Every noun is a noun. Sentences almost parse. The characters do dialogue with speaker attribution. Nothing in this session is task-specific — we don't tell the model about verbs or characters or Shakespeare. It just predicts the next character from the previous 256. Everything else emerges.

This session covers: the char-level tokenizer, the on-the-fly dataloader (no torch.utils.data.DataLoader, just a batch sampler), the training loop with mixed precision, the eval loop for val loss, learning-rate scheduling with warmup + cosine, and sampling.

Prerequisites

• S040 — the model.
• S018 — training loop pattern.
• S014 — cross-entropy loss.

1 · The tokenizer

Char-level. Simplest possible.

with open('input.txt', 'r', encoding='utf-8') as f:
    text = f.read()
chars = sorted(set(text))
vocab_size = len(chars)          # ~65 for Tiny Shakespeare
 
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: ''.join(itos[i] for i in ids)

That's it. Five lines, no library. Each character becomes an integer, each integer maps back. Vocab size 65 (letters, digits, punctuation, newline). No subword magic. Real LLMs use BPE (S044) — for our toy, char-level trains fast and generalises poorly, which is fine for teaching.

1.1 · Encode the corpus

import torch
data = torch.tensor(encode(text), dtype=torch.long)  # (~1.1M,) tensor of ints
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]

Standard 90/10 split. No shuffling — you want the val split to be text the model hasn't seen. In a bigger corpus you'd hold out entire documents.

2 · The dataloader (all 8 lines of it)

Tiny Shakespeare is small enough to sit in GPU memory. So we skip PyTorch's Dataset/DataLoader and just sample random windows:

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

x is (B, T) of token IDs. y is (B, T) shifted by one — the "next token" targets. Each of the B · T positions in x gets a next-token label in y. That's B · T supervised examples per batch, which is why LM training is so data-efficient per step.

3 · Cross-entropy loss — the LM objective

For each position, the model outputs a distribution over the vocabulary. The target is the ID of the true next token. Cross-entropy asks: "How much probability did you assign to the correct answer?" and penalises assigning too little.

In code, PyTorch handles it in one line:

logits, loss = model(x, targets=y)   # from S040's forward

where inside the model:

loss = F.cross_entropy(logits.view(-1, vocab_size), targets.view(-1))

The .view(-1, vocab_size) flattens (B, T, V) → (B*T, V) and (B, T) → (B*T,) so cross-entropy treats each position as one example. F.cross_entropy combines log-softmax and NLL loss for numerical stability.

3.1 · Perplexity — the interpretable metric

Perplexity is "the effective number of tokens the model is guessing between at each position". A model with PPL = 5 is as uncertain as if it had to pick uniformly among 5 tokens. A perfect model has PPL = 1. A random model on Shakespeare has PPL = 65 (uniform over the 65-char vocab).

Our target: loss ≈ 1.48, so PPL ≈ e^1.48 ≈ 4.4. The model is picking from ~4 plausible next chars at each position on average. Not bad for 15 minutes of training on 1 MB.

4 · The training loop

import torch
from torch.nn import functional as F
 
# hyperparameters (Karpathy's nanoGPT Tiny Shakespeare config)
batch_size = 64
block_size = 256
max_iters = 5000
eval_interval = 500
eval_iters = 200
learning_rate = 3e-4
device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
config = GPTConfig(
    block_size=block_size, vocab_size=vocab_size,
    n_layer=6, n_head=6, n_embd=384, dropout=0.2)
model = GPT(config).to(device)
optimizer = model.configure_optimizers(
    weight_decay=0.1, learning_rate=learning_rate, betas=(0.9, 0.95))

Now the loop:

@torch.no_grad()
def estimate_loss():
    model.eval()
    out = {}
    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = get_batch(split, block_size, batch_size, device)
            _, loss = model(X, targets=Y)
            losses[k] = loss.item()
        out[split] = losses.mean().item()
    model.train()
    return out
 
for step in range(max_iters):
    if step % eval_interval == 0:
        losses = estimate_loss()
        print(f"step {step}: train {losses['train']:.4f}, val {losses['val']:.4f}")
 
    xb, yb = get_batch('train', block_size, batch_size, device)
    _, loss = model(xb, targets=yb)
 
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

Forty lines including the imports. That is a language model trainer.

4.1 · Learning rate schedule — warmup + cosine

The above uses a fixed LR. For better results, add warmup + cosine decay:

import math
warmup_iters = 100
lr_decay_iters = max_iters
min_lr = learning_rate / 10
 
def get_lr(it):
    if it < warmup_iters:
        return learning_rate * it / warmup_iters
    if it > lr_decay_iters:
        return min_lr
    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return min_lr + coeff * (learning_rate - min_lr)
 
# inside the loop, before optimizer.step():
lr = get_lr(step)
for pg in optimizer.param_groups:
    pg['lr'] = lr

Linear warmup (100 steps of ramp) prevents early-training instability. Cosine decay (to 10% of peak by end) squeezes out the last bit of loss reduction. Standard for transformer training.

5 · Reading the loss curve

Expected output:

If val starts rising while train drops → severe overfit; add more dropout or stop early. If both plateau at ~2.0 → your LR is too small, or block_size too small, or n_layer too small.

6 · Generate!

context = torch.zeros((1, 1), dtype=torch.long, device=device)   # start with token 0 (\n)
out = model.generate(context, max_new_tokens=500, temperature=0.8, top_k=40)
print(decode(out[0].tolist()))

Sample output (yours will differ — random sampling):

 
BUCKINGHAM: Come, cousin; I'll to your uncle Warwick.
Farewell, sweet lord, and let us pray to God
That we may fight and win the crown at last.
 
QUEEN MARGARET: O, God's own name, thou art a proud usurper!
Thy father slew my father, and thou dost
Continue the war upon our righteous king.

Not real Shakespeare. But: line breaks in the right places, speaker attribution, capitalisation, punctuation, iambic-ish rhythm, invented but plausible dialogue. From a model with 10 million parameters trained for 15 minutes on 1 MB of text. That's what next-token prediction with attention gets you.

7 · Pitfalls

8 · What we're NOT doing (and why)

• No wandb / tensorboard logging. Simple print. Real projects use tools; teaching uses print.
• No checkpointing. 15-minute training run, we can just re-run. See S055 for how to do this properly.
• No distributed training. One GPU. See M09 for multi-GPU.
• No bf16 mixed precision. Would give 2× speedup on modern hardware; adds ~10 lines of torch.autocast boilerplate. Optional for now.

Real nanoGPT train.py has all these. Read it after finishing this session to see the production versions.

9 · Modern-2025 twist — Tiny Shakespeare in the age of frontier models

Why still train char-level on Tiny Shakespeare in 2026, when you can download Llama 3.1 8B and it already "knows" Shakespeare far better than any toy model ever will? Because the shape of the training curve, the failure modes, and the intuitions transfer directly. Every scaling-law paper of the last three years — Chinchilla (Hoffmann et al., 2022), the Kaplan revisit (2024), the Muennighoff "overtrained" paper — uses the same recipe as your Tiny Shakespeare run, just scaled up 10 million×.

What has changed at frontier scale as of 2026 that you should be aware of:

• Chinchilla-optimal token budget: for a model with N params, train on ~20N tokens (Hoffmann 2022). Your Tiny Shakespeare setup: N=10M, tokens≈200M (5000 steps × 64 × 256) — you're actually training at ~20× Chinchilla, which is why the model overfits so happily.
• Overtraining is now standard. Llama 3 8B trained on 15T tokens, ~1875× Chinchilla-optimal. The Muennighoff 2023 paper showed inference-cost dominance flips the compute-optimal frontier: over-train smaller models to get better inference-cost-per-quality.
• Muon optimizer (Jordan et al., 2024) is the first optimizer to seriously challenge AdamW for LM pretraining. modded-nanoGPT's speedrun uses it. Try swapping it in and you'll see the loss curve drop ~15% faster.
• Data mixing (S045–S046) matters more than architecture. DeepSeek-V3 attributes most of its edge over Llama 3 to data curation.
• 1.58-bit BitNet (Ma et al., Microsoft, Feb 2024) trains a Tiny-Shakespeare-scale model with ternary weights and matches FP16 loss. If you have a spare hour, port your S040 model to BitNet and watch it work.

Further reading:

• Chinchilla paper (Hoffmann et al., 2022) — the token-budget formula.
• Muennighoff et al., "Scaling Data-Constrained Language Models" (2023) — why over-training is now standard.
• Karpathy's llm.c training log — same loss-curve intuitions at GPT-2 scale.

# Split before you train, and evaluate both every N steps.
n = int(0.9 * len(data))
train_data, val_data = data[:n], data[n:]

for step in range(max_steps):
    ...
    if step % eval_interval == 0:
        print(step, estimate_loss(train_data), estimate_loss(val_data))

# The diagnostic that settles it:
sample = generate(model, n_tokens=300)
for k in (20, 40, 80):
    spans = [sample[i:i+k] for i in range(len(sample) - k)]
    hits = sum(s in raw_text for s in spans)
    print(k, hits / len(spans))
# Long verbatim spans copied from the source = memorisation,
# regardless of how good the training curve looks.

Retention scaffold

One-line summary (write it in your own words): _______________________________________________

Spaced review: re-read §5 (loss curve interpretation) and §9 (Chinchilla ratios) in 24 hours. On day 7, re-run the training with a different hyperparameter and predict the loss curve before running.

Next session (S042): we just trained a model. Now let's make INFERENCE 100× faster with a KV cache. Naive generation re-computes attention over all past tokens at every step; the KV cache is the fix. Every LLM serving stack in the world runs on this trick.

Sticky note (keep on your desk): loss(step 0) ≈ ln(vocab). PPL = exp(loss). Warmup + cosine. Sample at T≈0.8, top-k=40. Overfit if tokens/params ≫ 20 (Chinchilla).

Previous: ← DL S040 · Next: DL S042 →