THE EVOLVING TRANSFORMER · 99
Other chapters teach what each architecture is. This chapter teaches how any of them learns — dataset, batches, loss, backward pass, validation, and generation in one reusable script.
Every architecture chapter in this series ends with forward(x) → logits of shape (B, S, vocab). That is the inference path: given token IDs, predict a distribution over the vocabulary at every position. Training adds one more ingredient: a target sequence y and a loss that measures how wrong those predictions are.
The objective is causal language modeling: at position i, the model may only attend to tokens 0…i (already enforced by the causal mask in each decoder), and we ask it to predict token y[i] — the token that actually comes next in the text. We train on every position in parallel inside one forward pass, then average the cross-entropy across all B × S positions.
This chapter is deliberately model-agnostic. train.py owns the loop; dataset.py owns the data; the only hook that names a specific architecture is build_model(vocab_size). The default example imports GPT-2 from Stop 02 so you reuse the same gpt2.py you already studied.
We train on TinyShakespeare at character level (~65 unique symbols) so a machine finishes 500 steps in minutes. Production stacks use BPE or SentencePiece with vocab sizes in the tens of thousands; the loop does not change — only the tokenizer and vocab_size passed into build_model.
Think of training as a pipeline with two halves. The data half turns raw text into batches (x, y) on your device. The optimization half runs the model, computes loss, backpropagates, and updates weights. Evaluation and generation sit at the end to answer: “did anything actually learn?”
Files in this folder: dataset.py, train.py, test.py. Weights for the default run live in 02-gpt-2/gpt2.py and are imported at runtime via sys.path.
The same five-line inner step trains GPT-2, LLaMA, Mistral, or any other decoder in this series. Swapping models means changing build_model() and the config object — not rewriting the loop.
All training hyperparameters live as module-level constants in train.py. That keeps the script easy to read and tweak without argparse for every knob. Only the post-training --prompt flag is exposed on the command line.
MAX_SEQ_LEN = 256 is smaller than GPT-2’s paper context (1024) on purpose: shorter windows mean more samples per epoch and faster iterations on CPU/MPS. The model’s GPTConfig.max_seq_len must match so positional embeddings and causal masks agree with the batch shape.
N_STEPS500 optimizer steps (not epochs)EVAL_INTERVALLog train + val loss every 50 stepsEVAL_STEPS20 val batches per evaluationBATCH_SIZE16 sequences per stepLR3e-4 — AdamW learning rateMAX_SEQ_LEN256 tokens per window1N_STEPS = 500
2EVAL_INTERVAL = 50
3EVAL_STEPS = 20
4BATCH_SIZE = 16
5LR = 3e-4
6MAX_SEQ_LEN = 256
7
8if torch.cuda.is_available():
9 DEVICE = "cuda"
10elif torch.backends.mps.is_available():
11 DEVICE = "mps"
12else:
13 DEVICE = "cpu"Device selection is automatic: NVIDIA GPU if present, else Apple Silicon mps, else CPU. Every batch is moved with x.to(DEVICE) and y.to(DEVICE) right before the forward pass.
Karpathy’s TinyShakespeare is a ~1MB slice of Shakespeare plays — small enough to download in seconds, rich enough to see loss drop. download_dataset() checks for tinyshakespeare.txt beside the script; if missing, it pulls from GitHub:
https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
Keeping the URL in DATA_URL documents the source in code and lets train.py print it when you start a run.
1DATA_URL = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
2DATA_PATH = "tinyshakespeare.txt"
3
4def download_dataset():
5 if not os.path.exists(DATA_PATH):
6 print(f"Downloading TinyShakespeare to {DATA_PATH}...")
7 urllib.request.urlretrieve(DATA_URL, DATA_PATH)
8 print("Done.")
9 else:
10 print(f"Dataset already exists at {DATA_PATH}")CharTokenizer is the simplest possible tokenizer: one integer ID per unique character in the training text. Shakespeare’s charset includes letters, punctuation, and newlines — typically ~65 symbols, far smaller than GPT-2’s 50,257 BPE merges.
encode maps a string to a list of ints; decode maps back. There is no <unk> token because every character in the file appeared when we built the vocab. At inference time, any character outside that set would crash — another reason production systems use subword tokenizers with explicit unknown handling.
1class CharTokenizer:
2 def __init__(self, text: str):
3 chars = sorted(set(text))
4 self.vocab_size = len(chars)
5 self.char_to_id = {c: i for i, c in enumerate(chars)}
6 self.id_to_char = {i: c for i, c in enumerate(chars)}
7
8 def encode(self, text: str) -> list[int]:
9 return [self.char_to_id[c] for c in text]
10
11 def decode(self, ids: list[int]) -> str:
12 return "".join(self.id_to_char[i] for i in ids)After tokenization, the full play is one long tensor of IDs. Training never shuffles characters inside a window — only which starting index each batch sample uses.
Language modeling needs pairs: input x and target y shifted by one token. For window length S = max_seq_len, index idx takes a contiguous slice of length S + 1 from the token stream:
x = chunk[:-1] — tokens the model is allowed to see (length S)y = chunk[1:] — the next character at each position (length S)Concrete example with S = 4 and text "ROMEO" encoded as [17, 24, 22, 8, 20]:
At position 0 the model sees 17 ('R') and should predict 24 ('O'). At position 3 it sees 8 ('E') and should predict 20 ('O'). Causal attention in the transformer ensures position i never peeks at future tokens in x.
1class CharDataset(Dataset):
2 def __init__(self, token_ids: list[int], max_seq_len: int):
3 self.data = torch.tensor(token_ids, dtype=torch.long)
4 self.max_seq_len = max_seq_len
5
6 def __len__(self):
7 return len(self.data) - self.max_seq_len
8
9 def __getitem__(self, idx):
10 chunk = self.data[idx : idx + self.max_seq_len + 1]
11 return chunk[:-1], chunk[1:]__len__ is len(tokens) - S because the last valid start index must still fit S + 1 characters. With ~1M characters and S = 256, you get on the order of a million training windows — far more than 500 optimizer steps, which is why the training loop restarts the DataLoader iterator when an epoch ends.
get_dataloaders() ties everything together: download, read file, build tokenizer, encode full text, split, wrap in DataLoader.
The split is chronological: the first 90% of tokens train, the last 10% validate. We do not shuffle tokens before splitting — that would leak future plot into training windows. Shuffling happens only across window start indices in the train loader (shuffle=True).
Both loaders use drop_last=True so every batch has exactly BATCH_SIZE sequences — no partial batch edge cases in the demo loop.
1def get_dataloaders(max_seq_len: int, batch_size: int, val_fraction: float = 0.1):
2 download_dataset()
3 with open(DATA_PATH, "r", encoding="utf-8") as f:
4 text = f.read()
5
6 tokenizer = CharTokenizer(text)
7 token_ids = tokenizer.encode(text)
8
9 split = int(len(token_ids) * (1 - val_fraction))
10 train_dataset = CharDataset(token_ids[:split], max_seq_len)
11 val_dataset = CharDataset(token_ids[split:], max_seq_len)
12
13 train_loader = DataLoader(
14 train_dataset, batch_size=batch_size, shuffle=True, drop_last=True
15 )
16 val_loader = DataLoader(
17 val_dataset, batch_size=batch_size, shuffle=False, drop_last=True
18 )
19 return train_loader, val_loader, tokenizer| Tensor | Shape | Meaning |
|---|---|---|
x | (B, S) | Input token IDs — here B=16, S=256 |
y | (B, S) | Target token IDs, shifted by 1 |
logits | (B, S, vocab) | Unnormalized scores per position |
The training script adds the sibling folder 02-gpt-2 to sys.path and imports GPT and GPTConfig. That way there is a single source of truth for the architecture — fixes in the GPT-2 chapter automatically apply here.
build_model(vocab_size) is the only function you replace to train another decoder. Pass the tokenizer’s vocab size (here ~65, not 50,257) so the embedding table matches the data.
1_GPT2_DIR = Path(__file__).resolve().parent.parent / "02-gpt-2"
2if str(_GPT2_DIR) not in sys.path:
3 sys.path.insert(0, str(_GPT2_DIR))
4from gpt2 import GPT, GPTConfig
5
6def build_model(vocab_size: int):
7 config = GPTConfig(vocabsize=vocab_size, max_seq_len=MAX_SEQ_LEN)
8 return GPT(config)nn.CrossEntropyLoss expects class scores of shape (N, C) and targets of shape (N,). Here N is the number of positions we score at once, and C is the number of classes — for us, C = vocab_size (one logit per possible next token, ~65 on TinyShakespeare). Our model returns (B, S, vocab), so we flatten batch and time: N = B × S, and the last dimension is already C.
GPT returns raw logits, not log-softmax. CrossEntropyLoss applies log_softmax internally for numerical stability. Do not pass probabilities or log-probabilities unless you switch to NLLLoss.
Per-step loss is the average negative log-likelihood across all positions in the batch. Lower loss means the model assigns higher probability to the true next characters.
1optimizer = AdamW(model.parameters(), lr=LR)
2criterion = nn.CrossEntropyLoss()
3
4# inside the training loop:
5logits = model(x) # (B, S, vocab)
6loss = criterion(
7 logits.view(-1, logits.size(-1)),
8 y.view(-1),
9)AdamW is Adam with decoupled weight decay — the default choice for transformer LMs since GPT-2. For this tiny demo we use a single learning rate for all parameters; larger systems use warmup, cosine decay, and sometimes separate LR for embeddings.
One optimizer step is always the same five operations, regardless of whether the model has RoPE, GQA, or MoE:
Gradient clipping prevents occasional huge updates when loss spikes early in training. Norm 1.0 is a common default for small LMs; large-scale training often uses different thresholds or adaptive clipping.
1optimizer.zero_grad()
2logits = model(x)
3loss = criterion(logits.view(-1, logits.size(-1)), y.view(-1))
4loss.backward()
5nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
6optimizer.step()The outer loop counts optimizer steps, not full epochs. With ~1M windows and batch 16, one epoch is tens of thousands of steps — overkill for a teaching demo. N_STEPS = 500 is enough to watch loss fall and generate recognizable Shakespeare-ish gibberish.
We hold a manual iterator train_iter = iter(train_loader) and call next(train_iter) each step. When the epoch exhausts, Python raises StopIteration; we catch it, rebuild the iterator, and fetch the next batch. That pattern streams data indefinitely without nesting epoch loops.
Every EVAL_INTERVAL steps we average the last 50 training losses (running sum) and call get_val_loss() for a sanity check on held-out text.
1train_iter = iter(train_loader)
2running_loss = 0.0
3
4for step in range(1, N_STEPS + 1):
5 try:
6 x, y = next(train_iter)
7 except StopIteration:
8 train_iter = iter(train_loader)
9 x, y = next(train_iter)
10
11 x, y = x.to(DEVICE), y.to(DEVICE)
12 # ... zero_grad, forward, backward, step ...
13 running_loss += loss.item()
14
15 if step % EVAL_INTERVAL == 0:
16 avg_train_loss = running_loss / EVAL_INTERVAL
17 val_loss = get_val_loss(model, val_loader, criterion)
18 print(f"step {step:4d}/{N_STEPS} | train: {avg_train_loss:.4f} | val: {val_loss:.4f}")
19 running_loss = 0.0Typical console output shows train and val loss in the same ballpark early on; if val diverges far above train, you are overfitting — less of an issue on this tiny corpus with a fresh model, but critical on real data.
get_val_loss sets model.eval() so dropout is disabled, wraps inference in torch.no_grad(), and averages loss over at most EVAL_STEPS validation batches. It restores model.train() before returning so the next training step still gets dropout.
We cap validation batches because the val loader can be large; 20 batches × 16 × 256 positions is already plenty for a stable estimate during a short demo run.
1def get_val_loss(model, val_loader, criterion, n_steps=EVAL_STEPS):
2 model.eval()
3 total_loss = 0.0
4 steps = 0
5 with torch.no_grad():
6 for x, y in val_loader:
7 x, y = x.to(DEVICE), y.to(DEVICE)
8 logits = model(x)
9 loss = criterion(logits.view(-1, logits.size(-1)), y.view(-1))
10 total_loss += loss.item()
11 steps += 1
12 if steps >= n_steps:
13 break
14 model.train()
15 return total_loss / stepsSummaryWriter writes scalars to runs/training-loop/ for TensorBoard. After training, training_losses.json stores the step / train / val series locally — useful if you plot curves yourself. Neither file is required to understand the chapter; they are artifacts of a real run.
Training optimizes teacher-forced next-token prediction: at every position the model sees the true prefix from x. Generation is different: we feed the model its own predictions autoregressively.
generate() encodes a prompt (default "ROMEO:"), then repeats:
MAX_SEQ_LEN tokens (same as training window)temperature (higher = more random), softmax, sample one IDAfter ~500 steps you should see Elizabethan-flavored character soup — not coherent acts, but loss should have dropped from ~4 (random) toward ~2 or below if training ran correctly. That gap is the proof the weights moved.
1@torch.no_grad()
2def generate(model, tokenizer, prompt="ROMEO:", max_new_tokens=200, temperature=0.8):
3 model.eval()
4 token_ids = tokenizer.encode(prompt)
5 x = torch.tensor([token_ids], dtype=torch.long, device=DEVICE)
6
7 for _ in range(max_new_tokens):
8 x_cond = x[:, -MAX_SEQ_LEN:]
9 logits = model(x_cond)
10 logits = logits[:, -1, :] / temperature
11 probs = torch.softmax(logits, dim=-1)
12 next_token = torch.multinomial(probs, num_samples=1)
13 x = torch.cat([x, next_token], dim=1)
14
15 return tokenizer.decode(x[0].tolist())This sampling loop is the same one you would use for LLaMA, Mistral, or any decoder in the series after swapping build_model. Training changes weights; generation code does not.
1cd the-evolving-transformer/99-the-training-loop
2python3 train.py
3python3 train.py --prompt "First Citizen:"test.py smoke-tests download, batch shapes, one forward/backward step, and that build_model returns logits of the right shape — run it before a long train if you change imports or configs.