One-sentence summary: Training repeats “sample a batch → compute loss → backpropagate → update parameters,” but a train.py worth trusting must also protect data boundaries, evaluation mode, learning-rate state, gradient norms, and checkpoint recovery.

📦 Original 2024 demo repository: github.com/waylandzhang/Transformer-from-scratch. The historical demo keeps training inside one model.py; this chapter pairs with the audited standalone model.py from Chapter 18.

19.1 A Random Model Is Not an Empty Model

A newly created model already has random parameters and can produce logits. Those parameters simply have not learned stable patterns from the corpus.

For a more memorable English example, imagine a catalog entry about an early instrument maker:

Token stream: [The, hurdy, -, gurdy, maker, tuned, the, tangent, ...]

Input x:      [The, hurdy, -, gurdy, maker, tuned, the]
Target y:     [hurdy, -, gurdy, maker, tuned, the, tangent]

At every position, the target is the next token, not necessarily the next word. “Hurdy-gurdy” itself may occupy several tokens; that is useful here because it stops the diagram from pretending tokens and words are the same thing.

If a model assigned exactly uniform probability to V tokens, its cross-entropy would be $\ln V$. That is a useful reference line, not a guaranteed initializer test. A random network need not emit uniform logits—especially with the tied weights and embedding scaling used in Chapter 18.

19.2 Keep Model and Training Configuration Separate

ModelConfig describes tensor shapes. Training needs a different configuration:

from dataclasses import dataclass

@dataclass
class TrainConfig:
    batch_size: int = 8
    total_steps: int = 500
    eval_interval: int = 50
    eval_batches: int = 10
    peak_lr: float = 3e-4
    min_lr: float = 3e-5
    warmup_steps: int = 50
    weight_decay: float = 0.1
    grad_clip: float = 1.0
    seed: int = 1337

This separation matters at recovery time: model shape must not change because batch size or evaluation frequency changed, and a device string does not belong in the architecture.

These values are candidates for the book's small model, not a universal large-model recipe. Chapter 17 explains why learning rate must be validated with the actual batch, data, and budget.

19.3 Respect the Data's Original Structure

19.3.1 A CSV is not an unstructured text file

The complete script accepts a prepared UTF-8 text file through --data. A plain text transcription of hurdy-gurdy tunings can go straight in. If the source is a CSV catalog and only one column is prose, extract that column with a CSV parser first; do not quietly train on headers, accession numbers, and unrelated fields:

import csv

def csv_column_to_text(path, column):
    with path.open("r", encoding="utf-8", newline="") as handle:
        reader = csv.DictReader(handle)
        if column not in (reader.fieldnames or []):
            raise ValueError(f"missing CSV column: {column}")
        rows = [
            row[column].strip()
            for row in reader
            if row[column].strip()
        ]
    if not rows:
        raise ValueError("the selected CSV column contains no text")
    return "\n".join(rows)

19.3.2 The tokenizer vocabulary is not the largest observed ID

tokenizer_name = "cl100k_base"
tokenizer = tiktoken.get_encoding(tokenizer_name)
token_ids = tokenizer.encode(text, disallowed_special=())
tokens = torch.tensor(token_ids, dtype=torch.long)  # keep corpus on CPU

model_config = ModelConfig(
    vocab_size=tokenizer.n_vocab,
    context_length=128,
    d_model=80,
    n_layers=6,
    n_heads=4,
    dropout=0.1,
)

max(token_ids) + 1 describes IDs observed in this sample; it does not describe every valid tokenizer ID. Use tokenizer.n_vocab. Keep the corpus tensor on CPU and move only the current batch, rather than spending accelerator memory on the entire dataset.

19.3.3 Split first, then sample windows inside each side

def split_tokens(tokens, train_fraction, context_length):
    split = int(len(tokens) * train_fraction)
    train_data = tokens[:split]
    valid_data = tokens[split:]
    if min(len(train_data), len(valid_data)) <= context_length:
        raise ValueError("each split needs more than context_length tokens")
    return train_data, valid_data

Each training window now stays inside train, and each validation window stays inside valid. Duplicate catalog entries can still leak across the boundary. A real dataset may need entity-level deduplication or grouped splitting before tokenization.

19.4 get_batch(): Where One-Off Errors Hide

def get_batch(data, batch_size, context_length, device, generator):
    max_start = len(data) - context_length
    if max_start <= 0:
        raise ValueError("data is too short for context_length")
    starts = torch.randint(
        0, max_start, (batch_size,), generator=generator
    )
    offsets = torch.arange(context_length)
    x = data[starts[:, None] + offsets]
    y = data[starts[:, None] + offsets + 1]
    return x.to(device), y.to(device)

torch.randint excludes its upper bound. The largest sampled start is therefore len(data) - context_length - 1, leaving exactly one extra token for the target.

A dedicated torch.Generator makes training-batch order independent of evaluation sampling, and its state can be checkpointed.

19.5 Loss Takes Logits, Not Precomputed Probabilities

Chapter 18's model computes:

loss = F.cross_entropy(
    logits.reshape(-1, logits.size(-1)),
    targets.reshape(-1),
)

cross_entropy combines a numerically stable LogSoftmax with negative log likelihood. Its input is raw logits and integer token targets.

A falling training loss proves improvement on this objective. It does not by itself prove factuality, safety, readable generations, or even improving validation loss.

19.6 Evaluation Must Not Contaminate Training State

@torch.inference_mode()
def estimate_loss(model, train_data, valid_data, config, device):
    was_training = model.training
    model.eval()
    try:
        result = {}
        eval_generator = torch.Generator().manual_seed(config.seed + 1)
        for name, data in (("train", train_data), ("valid", valid_data)):
            losses = []
            for _ in range(config.eval_batches):
                x, y = get_batch(
                    data,
                    config.batch_size,
                    model.config.context_length,
                    device,
                    eval_generator,
                )
                _, loss = model(x, y)
                losses.append(loss.detach().cpu())
            result[name] = torch.stack(losses).mean().item()
        return result
    finally:
        model.train(was_training)

inference_mode() avoids autograd state, eval() disables this model's Dropout, and finally restores the mode that the caller actually had. The book model has no BatchNorm, so BatchNorm is not the explanation for the mode switch here.

The fixed evaluation generator samples the same train and validation windows each time, making comparisons less noisy without consuming the training-batch generator.

19.7 AdamW Groups and the Chapter 17 Schedule

def configure_optimizer(model, config):
    decay, no_decay = [], []
    for _, parameter in model.named_parameters():
        if not parameter.requires_grad:
            continue
        (decay if parameter.ndim >= 2 else no_decay).append(parameter)

    groups = [
        {"params": decay, "weight_decay": config.weight_decay},
        {"params": no_decay, "weight_decay": 0.0},
    ]
    return torch.optim.AdamW(
        groups,
        lr=config.peak_lr,
        betas=(0.9, 0.95),
    )

This common grouping decays matrix-like embedding and linear weights while leaving one-dimensional biases and Norm parameters out. The tied embedding/output parameter appears once in named_parameters().

AdamW is not “Adam plus L2.” It decouples weight decay from Adam's moment-normalized loss-gradient update; it also does not discover a perfect private learning rate for every parameter.

The schedule uses the exact boundary convention from Chapter 17:

def lr_at_step(step, total_steps, warmup_steps, peak_lr, min_lr):
    if not 1 <= step <= total_steps:
        raise ValueError("step must be in [1, total_steps]")
    if not 0 < warmup_steps < total_steps:
        raise ValueError("warmup_steps must be in (0, total_steps)")
    if step <= warmup_steps:
        return peak_lr * step / warmup_steps
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
    return min_lr + (peak_lr - min_lr) * cosine

19.8 Update First, Then Report That Step

Each iteration sets the current LR, samples a batch, clears gradients, computes loss, backpropagates, clips the global norm, and calls optimizer.step(). Evaluation for step=500 happens after update 500.

The old code evaluated at index 499, performed one more update, and then saved a different model state than the printed metrics described.

clip_grad_norm_ returns the total norm before clipping. Logging it reveals spikes; it does not mean every printed norm must be at most 1.0.

19.9 Saving Weights Is Not the Same as Resuming Training

An honest recovery checkpoint needs model parameters, optimizer moments, completed step count, both configurations, tokenizer identity, a corpus fingerprint, the batch generator state, and PyTorch RNG state.

The complete script writes a temporary file and atomically replaces the destination. On load it uses weights_only=True and rejects mismatched configuration or corpus hashes. Still load only checkpoints you created or otherwise trust.

Saving RNG state cannot promise bit-for-bit equality across different GPUs, PyTorch versions, or nondeterministic kernels. It preserves the training state; it cannot make unlike environments mathematically identical.

19.10 Complete train.py

import argparse
import hashlib
import math
from dataclasses import asdict, dataclass
from pathlib import Path

import tiktoken
import torch

from model import Model, ModelConfig


@dataclass
class TrainConfig:
    batch_size: int = 8
    total_steps: int = 500
    eval_interval: int = 50
    eval_batches: int = 10
    peak_lr: float = 3e-4
    min_lr: float = 3e-5
    warmup_steps: int = 50
    weight_decay: float = 0.1
    grad_clip: float = 1.0
    seed: int = 1337

    def __post_init__(self):
        if min(self.batch_size, self.total_steps, self.eval_interval,
               self.eval_batches, self.warmup_steps) <= 0:
            raise ValueError("step and batch settings must be positive")
        if self.warmup_steps >= self.total_steps:
            raise ValueError("warmup_steps must be smaller than total_steps")
        if not 0.0 <= self.min_lr <= self.peak_lr:
            raise ValueError("learning rates must satisfy 0 <= min_lr <= peak_lr")
        if self.weight_decay < 0 or self.grad_clip <= 0:
            raise ValueError("weight_decay and grad_clip are invalid")


def select_device():
    if torch.cuda.is_available():
        return torch.device("cuda")
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def split_tokens(tokens, train_fraction, context_length):
    if not 0.0 < train_fraction < 1.0:
        raise ValueError("train_fraction must be between 0 and 1")
    split = int(len(tokens) * train_fraction)
    train_data = tokens[:split]
    valid_data = tokens[split:]
    if min(len(train_data), len(valid_data)) <= context_length:
        raise ValueError("each split needs more than context_length tokens")
    return train_data, valid_data


def get_batch(data, batch_size, context_length, device, generator):
    max_start = len(data) - context_length
    if max_start <= 0:
        raise ValueError("data is too short for context_length")
    starts = torch.randint(
        0, max_start, (batch_size,), generator=generator
    )
    offsets = torch.arange(context_length)
    x = data[starts[:, None] + offsets]
    y = data[starts[:, None] + offsets + 1]
    return x.to(device), y.to(device)


@torch.inference_mode()
def estimate_loss(model, train_data, valid_data, config, device):
    was_training = model.training
    model.eval()
    try:
        result = {}
        eval_generator = torch.Generator().manual_seed(config.seed + 1)
        for name, data in (("train", train_data), ("valid", valid_data)):
            losses = []
            for _ in range(config.eval_batches):
                x, y = get_batch(
                    data,
                    config.batch_size,
                    model.config.context_length,
                    device,
                    eval_generator,
                )
                _, loss = model(x, y)
                losses.append(loss.detach().cpu())
            result[name] = torch.stack(losses).mean().item()
        return result
    finally:
        model.train(was_training)


def configure_optimizer(model, config):
    decay, no_decay = [], []
    for _, parameter in model.named_parameters():
        if not parameter.requires_grad:
            continue
        (decay if parameter.ndim >= 2 else no_decay).append(parameter)
    groups = [
        {"params": decay, "weight_decay": config.weight_decay},
        {"params": no_decay, "weight_decay": 0.0},
    ]
    return torch.optim.AdamW(
        groups,
        lr=config.peak_lr,
        betas=(0.9, 0.95),
    )


def lr_at_step(step, total_steps, warmup_steps, peak_lr, min_lr):
    if not 1 <= step <= total_steps:
        raise ValueError("step must be inside the training range")
    if not 0 < warmup_steps < total_steps:
        raise ValueError("warmup_steps must be inside the training range")
    if step <= warmup_steps:
        return peak_lr * step / warmup_steps
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
    return min_lr + (peak_lr - min_lr) * cosine


def save_checkpoint(path, model, optimizer, step, model_config,
                    train_config, tokenizer_name, corpus_sha256,
                    train_generator):
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    cuda_rng = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None
    mps_rng = (
        torch.mps.get_rng_state()
        if hasattr(torch, "mps") and torch.backends.mps.is_available()
        else None
    )
    torch.save({
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "step": step,
        "model_config": asdict(model_config),
        "train_config": asdict(train_config),
        "tokenizer_name": tokenizer_name,
        "corpus_sha256": corpus_sha256,
        "torch_rng_state": torch.get_rng_state(),
        "cuda_rng_state": cuda_rng,
        "mps_rng_state": mps_rng,
        "train_generator_state": train_generator.get_state(),
    }, temporary)
    temporary.replace(path)


def load_checkpoint(path, model, optimizer, model_config,
                    train_config, tokenizer_name, corpus_sha256,
                    train_generator, device):
    checkpoint = torch.load(path, map_location="cpu", weights_only=True)
    if checkpoint["model_config"] != asdict(model_config):
        raise ValueError("checkpoint model_config does not match")
    if checkpoint["train_config"] != asdict(train_config):
        raise ValueError("checkpoint train_config does not match")
    if checkpoint["tokenizer_name"] != tokenizer_name:
        raise ValueError("checkpoint tokenizer does not match")
    if checkpoint["corpus_sha256"] != corpus_sha256:
        raise ValueError("checkpoint corpus does not match")
    model.load_state_dict(checkpoint["model_state_dict"])
    optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    torch.set_rng_state(checkpoint["torch_rng_state"])
    train_generator.set_state(checkpoint["train_generator_state"])
    if device.type == "cuda" and checkpoint["cuda_rng_state"] is not None:
        torch.cuda.set_rng_state_all(checkpoint["cuda_rng_state"])
    if device.type == "mps" and checkpoint["mps_rng_state"] is not None:
        torch.mps.set_rng_state(checkpoint["mps_rng_state"])
    return int(checkpoint["step"])


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data", type=Path, required=True)
    parser.add_argument("--output", type=Path,
                        default=Path("model/checkpoint.pt"))
    parser.add_argument("--resume", type=Path)
    return parser.parse_args()


def main():
    args = parse_args()
    train_config = TrainConfig()
    tokenizer_name = "cl100k_base"
    tokenizer = tiktoken.get_encoding(tokenizer_name)

    text = args.data.read_text(encoding="utf-8")
    corpus_sha256 = hashlib.sha256(text.encode("utf-8")).hexdigest()
    token_ids = tokenizer.encode(text, disallowed_special=())
    tokens = torch.tensor(token_ids, dtype=torch.long)  # keep corpus on CPU

    model_config = ModelConfig(
        vocab_size=tokenizer.n_vocab,
        context_length=128,
        d_model=80,
        n_layers=6,
        n_heads=4,
        dropout=0.1,
    )
    train_data, valid_data = split_tokens(
        tokens, train_fraction=0.9,
        context_length=model_config.context_length,
    )

    device = select_device()
    torch.manual_seed(train_config.seed)
    if device.type == "cuda":
        torch.cuda.manual_seed_all(train_config.seed)

    model = Model(model_config).to(device)
    optimizer = configure_optimizer(model, train_config)
    train_generator = torch.Generator().manual_seed(train_config.seed)
    start_step = 0

    if args.resume is not None:
        start_step = load_checkpoint(
            args.resume, model, optimizer, model_config,
            train_config, tokenizer_name, corpus_sha256,
            train_generator, device,
        )

    model.train()
    if start_step == 0:
        metrics = estimate_loss(
            model, train_data, valid_data, train_config, device
        )
        print(f"step=0 train={metrics['train']:.4f} "
              f"valid={metrics['valid']:.4f}")

    for step in range(start_step + 1, train_config.total_steps + 1):
        lr = lr_at_step(
            step,
            train_config.total_steps,
            train_config.warmup_steps,
            train_config.peak_lr,
            train_config.min_lr,
        )
        for group in optimizer.param_groups:
            group["lr"] = lr

        x, y = get_batch(
            train_data,
            train_config.batch_size,
            model_config.context_length,
            device,
            train_generator,
        )
        optimizer.zero_grad(set_to_none=True)
        _, loss = model(x, y)
        loss.backward()
        grad_norm = torch.nn.utils.clip_grad_norm_(
            model.parameters(), train_config.grad_clip
        )
        optimizer.step()

        should_evaluate = (
            step % train_config.eval_interval == 0
            or step == train_config.total_steps
        )
        if should_evaluate:
            metrics = estimate_loss(
                model, train_data, valid_data, train_config, device
            )
            print(
                f"step={step} lr={lr:.2e} grad={grad_norm.item():.3f} "
                f"train={metrics['train']:.4f} valid={metrics['valid']:.4f}"
            )
            save_checkpoint(
                args.output, model, optimizer, step,
                model_config, train_config, tokenizer_name, corpus_sha256,
                train_generator,
            )


if __name__ == "__main__":
    main()

19.11 How to Tell Whether Training Is Healthy

Do not invent a pretty curve that must fall from 10.8 to 2.8. Check evidence instead:

• Step-zero loss is finite.
• Parameters actually change after optimizer.step().
• Training loss trends down over a window; it need not fall at every step.
• Validation loss also improves.
• Gradient norm does not repeatedly spike.
• A resumed run restores step, LR, optimizer state, and batch RNG.

Training loss falling while validation loss rises is the classic widening generalization gap. Training loss itself falling and then rising may instead point to LR, data order, checkpoint recovery, or numerical trouble.

Chapter Checklist

• Explain why targets are inputs shifted by one token
• Select a CSV text column with a real parser
• Keep the corpus on CPU and move only each batch
• Evaluate without contaminating training mode or RNG
• Use AdamW groups, warmup + cosine, and gradient clipping
• Save and validate enough state to resume training

See You in the Next Chapter

We now have more than a weight file: the checkpoint knows its completed step, optimizer state, configuration, and corpus identity.

Chapter 20 builds inference.py: reconstruct the model from a trusted checkpoint, enter evaluation mode, encode a prompt, sample within explicit boundaries, and decode the result.