The Training Loop, Logging, and Reproducibility

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

0 · Recovery is a state-machine contract

A training loop is a state machine. The forward/backward/update body is short; recovering the next update is the harder contract. Imagine a run interrupted after update 47,200. Loading weights proves that a file can be read. Comparing the next batch, loss, learning rate and parameter update with an uninterrupted run proves much more. This is an illustrative failure scenario, not a claimed personal incident or a seven-line production guarantee.

This lesson builds a complete synthetic CPU experiment with dropout, AdamW, a scheduler, an EMA model, validation, logging and resumable sampling. CUDA, TensorBoard's UI, compiler performance and distributed storage remain explicitly unexecuted references.

Prerequisites

Understand nn.Module, a scalar loss and DataLoader batches. The neighboring dataset lesson explains workers; the precision lesson extends the update with AMP. No downloaded dataset is needed here.

1 · The five-line body (and everything around it)

A pit stop is a useful analogy: reset the gradient buffer, measure the current prediction, differentiate, change parameters, then prepare the rate for the next update. The scheduler is optional and its clock must be declared; ReduceLROnPlateau consumes a validation metric instead of a batch counter.

For a single ordinary update the order is zero_grad → forward/loss → backward → clip if required → optimizer.step → update-clock scheduler.step. Training and evaluation modes are separate from gradient recording. model.eval() changes dropout/BatchNorm behavior; torch.no_grad() disables graph construction for the enclosed operations. Neither substitutes for the other.

Why clear gradients? PyTorch accumulates into leaf .grad across backward calls. Within a graph the chain rule must sum contributions from shared uses, but that does not force every framework to accumulate across separate backward calls: it is an API policy that conveniently supports microbatches and multiple losses. Two freshly constructed identical losses backpropagated without clearing produce twice the gradient. Backpropagating the same freed graph a second time is a different operation and normally errors.

set_to_none=True leaves untouched parameters with no gradient; optimizers can treat None differently from an explicit zero tensor. Inspect this distinction when a branch is conditional rather than declaring every absent gradient a bug.

2 · Reproducibility — the seed sandwich

A seed fixes an RNG's starting position, not all arithmetic or the order in which random values are consumed. Record software versions, device, architecture/config, data identity, preprocessing, sampling policy and precision alongside Python, NumPy and torch RNG state. Explicit NumPy Generator objects have their own state; np.random.seed does not reset them.

For CUDA, disabling cuDNN benchmarking stabilizes algorithm selection and enabling deterministic algorithms requests supported deterministic implementations or errors. It is not a cross-release/hardware guarantee or a fixed 5–20% tax. Decide whether you need statistical agreement or exact same-environment continuation, and test that requirement.

A top-level worker initializer is spawn-picklable; a lambda often is not. This complete CPU-only worker example verifies deterministic order with zero and two workers for a deterministic dataset. Workers do not inevitably reshuffle a plain TensorDataset differently. Worker-local random transforms, prefetch, persistent workers and changing worker counts require additional state/planning for exact mid-epoch resume.

import random
import numpy as np
import torch
from torch.utils.data import DataLoader, TensorDataset
 
def seed_worker(worker_id):
    seed = torch.initial_seed() % 2**32
    np.random.seed(seed)
    random.seed(seed)
 
def first_batch(workers):
    generator = torch.Generator().manual_seed(42)
    dataset = TensorDataset(torch.arange(64))
    loader = DataLoader(dataset, batch_size=8, shuffle=True,
                        num_workers=workers, generator=generator,
                        worker_init_fn=seed_worker)
    return next(iter(loader))[0]
 
if __name__ == "__main__":
    assert torch.equal(first_batch(0), first_batch(0))
    assert torch.equal(first_batch(0), first_batch(2))
    print("deterministic dataset worker-order check passed")

The next program uses its own permutation generator and a saved cursor rather than pretending that saving an epoch number can replay a partially consumed worker queue.

3 · The full loop (annotated)

Run this whole block with PyTorch 2.9 on CPU. It deliberately keeps model/data tiny, uses no services, and serializes only tensors and primitive containers. max_updates simulates stopping at a committed optimizer boundary. dry_run performs one train batch and one validation batch and returns a marked diagnostic checkpoint that cannot be mistaken for a normal continuation.

The logger is an ordinary list so the experiment needs no TensorBoard installation. A supplied callback can publish a snapshot after a safe boundary; exceptions before that boundary must resume the previously committed snapshot, not a half-updated state. The bounded implementation is single-process CPU, not a general multiworker/distributed trainer.

import copy
import io
import math
import random
import torch
from torch import nn
import torch.nn.functional as F
 
def pack(state):
    stream = io.BytesIO()
    torch.save(state, stream)
    return stream.getvalue()  # immutable snapshot, not live state_dict aliases
 
def unpack(blob):
    return torch.load(io.BytesIO(blob), map_location="cpu", weights_only=True)
 
def make_model():
    return nn.Sequential(nn.Linear(3, 12), nn.Tanh(), nn.Dropout(.2),
                         nn.Linear(12, 2))
 
def data():
    g = torch.Generator().manual_seed(900)
    x = torch.randn(31, 3, generator=g)
    y = (x[:, 0] + .4*x[:, 1] > 0).long()
    v = torch.randn(13, 3, generator=g)
    z = (v[:, 0] + .4*v[:, 1] > 0).long()
    return x, y, v, z
 
def train(seed=42, resume=None, max_updates=16, epochs=4, patience=10,
          min_delta=0., dry_run=False, should_stop=lambda: False,
          publish=lambda blob: None):
    torch.set_num_threads(1)
    torch.manual_seed(seed)
    random.seed(seed)
    model = make_model()
    ema = copy.deepcopy(model)
    opt = torch.optim.AdamW(model.parameters(), lr=.02, weight_decay=.01)
    # Construct scheduler before restoring optimizer LR/state.
    sched = torch.optim.lr_scheduler.StepLR(opt, step_size=4, gamma=.8)
    sampler = torch.Generator().manual_seed(seed + 1)
    x, y, vx, vy = data()
    config = dict(seed=seed, epochs=epochs, batch=8, ema=.999,
                  patience=patience, min_delta=min_delta, dataset="synthetic-v1")
    epoch, cursor, updates = 0, 0, 0
    perm = torch.randperm(len(x), generator=sampler)
    best, stale, logs, best_weights = math.inf, 0, [], None
    if resume is not None:
        s = unpack(resume)
        assert s["config"] == config and not s["dry_run"]
        model.load_state_dict(s["model"]); ema.load_state_dict(s["ema"])
        sched.load_state_dict(s["sched"]); opt.load_state_dict(s["optim"])
        sampler.set_state(s["sampler"])
        epoch, cursor, updates = s["epoch"], s["cursor"], s["updates"]
        perm, best, stale = s["perm"], s["best"], s["stale"]
        logs, best_weights = s["logs"], s["best_weights"]
        # Restore after model construction, which consumed random numbers.
        torch.set_rng_state(s["torch_rng"]); random.setstate(s["python_rng"])
    def snapshot():
        return pack(dict(model=model.state_dict(), ema=ema.state_dict(),
                         optim=opt.state_dict(), sched=sched.state_dict(),
                         config=config, epoch=epoch, cursor=cursor,
                         perm=perm, updates=updates, best=best, stale=stale,
                         logs=logs, best_weights=best_weights, dry_run=dry_run,
                         sampler=sampler.get_state(), torch_rng=torch.get_rng_state(),
                         python_rng=random.getstate()))
    committed = snapshot()
    while epoch < epochs and updates < max_updates and stale < patience:
        if should_stop():
            break  # only inspect the flag at a committed boundary
        model.train()
        idx = perm[cursor:cursor+8]
        opt.zero_grad(set_to_none=True)
        loss = F.cross_entropy(model(x[idx]), y[idx])
        if not torch.isfinite(loss):
            raise FloatingPointError("nonfinite loss: retain prior checkpoint")
        loss.backward()
        preclip = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.,
                                                error_if_nonfinite=True)
        used_lr = opt.param_groups[0]["lr"]
        opt.step(); sched.step()
        with torch.no_grad():
            for e, p in zip(ema.parameters(), model.parameters()):
                e.lerp_(p, 1-config["ema"])
        updates += 1; cursor += len(idx)
        logs.append(dict(update=updates, indices=idx.tolist(), loss=loss.item(),
                         preclip=preclip.item(), lr=used_lr,
                         next_lr=opt.param_groups[0]["lr"]))
        if cursor == len(x) or dry_run:
            ema.eval()
            with torch.no_grad():
                # Weight summed losses by examples, including the 5-example tail.
                total, n = 0., 0
                for a in range(0, len(vx), 8):
                    total += F.cross_entropy(ema(vx[a:a+8]), vy[a:a+8],
                                             reduction="sum").item()
                    n += len(vx[a:a+8])
                    if dry_run: break
                metric = total/n
            improved = metric < best-min_delta
            if improved:
                best, stale = metric, 0
                best_weights = copy.deepcopy(ema.state_dict())
            else:
                stale += 1
            logs[-1]["val_loss"] = metric
            epoch += 1; cursor = 0
            perm = torch.randperm(len(x), generator=sampler)
        # best and cursor are updated BEFORE both latest/best snapshots.
        committed = snapshot(); publish(committed)
        if dry_run: break
    return committed
 
whole = unpack(train())
part = train(max_updates=3)
continued = unpack(train(resume=part))
assert whole["logs"] == continued["logs"]
for group in ["model", "ema"]:
    for key in whole[group]:
        assert torch.equal(whole[group][key], continued[group][key])
assert whole["best"] == min(r["val_loss"] for r in whole["logs"] if "val_loss" in r)
assert unpack(train(seed=43))["logs"] != whole["logs"]
assert unpack(train(dry_run=True))["updates"] == 1
assert unpack(train(patience=1, min_delta=1e6))["epoch"] == 2
assert unpack(train(should_stop=lambda: True))["updates"] == 0
# Negative control: resetting optimizer moments changes the resumed trajectory.
bad = unpack(part); bad["optim"]["state"] = {}
wrong = unpack(train(resume=pack(bad)))
assert wrong["logs"] != whole["logs"]
# Two separately constructed graphs accumulate, by API contract.
p = torch.tensor(2., requires_grad=True)
(p*p).backward(); first = p.grad.clone(); (p*p).backward()
assert p.grad == 2*first
print("exact CPU continuation, EMA, early stop, dry-run and negative controls passed")

A saved permutation plus cursor points to the next unconsumed sample. Advancing it before a successful update or reconstructing it from only an epoch skips/repeats training. The whole-run comparison includes sampled dropout through the global torch RNG, sampler state and the next learning rate. It demonstrates this program's same-runtime contract, not all future PyTorch versions.

4 · Checkpointing — save state, not just weights

Weights-only checkpoints are useful for inference and intentional optimizer restarts. They are insufficient for continuing the same Adam trajectory. Save parameters and buffers, optimizer moments/step counts, scheduler/scaler state if present, EMA state, update/epoch counters, best metric and early-stop patience, RNGs actually used, data cursor/sampler state and immutable configuration/version identifiers. A tokenizer/config belongs in a language-model bundle; this numerical example has no tokenizer. Metrics that drive stopping or scheduling are control state, not disposable logs.

FP32 Adam with two FP32 moments has approximately three parameter-sized persistent tensors, not automatically four: an extra master-weight copy depends on the precision implementation. Gradients, buffers, sharding metadata and storage encoding change the complete checkpoint footprint. Estimate bytes from actual tensors.

4.1 Restricted loading and atomic publication

PyTorch 2.6 changed the default to weights_only=True when pickle_module is not explicitly supplied; use the explicit flag with the 2.9 contract. Restricted loading reduces arbitrary-code-execution exposure but does not prove a file harmless or protect against every denial-of-service/corruption case. Do not switch to unrestricted pickle merely to make an unknown download load. The program saves primitive RNG tuples/tensors rather than a lambda or whole optimizer object.

This independent local helper writes, flushes, verifies and renames a snapshot. It retains three periodic snapshots and a separate best file. These are local filesystem semantics; a cloud object store or sharded checkpoint requires its own manifest/commit protocol. An atomic rename is not itself proof of crash durability of the directory entry.

from pathlib import Path
import os
import tempfile
 
def atomic_checkpoint(destination, blob):
    destination = Path(destination)
    fd, temp = tempfile.mkstemp(prefix=destination.name+".", dir=destination.parent)
    try:
        with os.fdopen(fd, "wb") as f:
            f.write(blob); f.flush(); os.fsync(f.fileno())
        unpack(Path(temp).read_bytes())
        os.replace(temp, destination)
    finally:
        if os.path.exists(temp): os.unlink(temp)
 
def save_generation(directory, step, blob, improved=False):
    directory = Path(directory)
    atomic_checkpoint(directory/"latest.pt", blob)
    atomic_checkpoint(directory/f"periodic-{step:09d}.pt", blob)
    if improved: atomic_checkpoint(directory/"best.pt", blob)
    for old in sorted(directory.glob("periodic-*.pt"))[:-3]: old.unlink()
 
# Test only inside a temporary directory; no existing checkpoints are touched.
with tempfile.TemporaryDirectory(prefix="batch6-dl-training-") as d:
    for step in range(5): save_generation(d, step, part, improved=(step == 1))
    assert len(list(Path(d).glob("periodic-*.pt"))) == 3
    assert (Path(d)/"best.pt").exists()
    assert unpack((Path(d)/"latest.pt").read_bytes())["updates"] == 3
    # Wire the actual trainer to durable local publication, then resume its bytes.
    train(max_updates=3, publish=lambda blob:atomic_checkpoint(Path(d)/"resume.pt",blob))
    resumed_disk=unpack(train(resume=(Path(d)/"resume.pt").read_bytes()))
    assert resumed_disk["logs"] == whole["logs"]
print("atomic local publication, on-disk continuation and rolling retention passed")

4.2 EMA — a second model, not a guaranteed gain

With decay $d$, update $\theta_{EMA}\leftarrow d\theta_{EMA}+(1-d)\theta$ after each successful optimizer step. For old EMA 2, new weight 4 and $d=.9$, the result is 2.2. Initialize from model weights and save the EMA state and update count; do not reset it at resume. The complete loop uses .999, a slow average whose appropriateness must be validated. PyTorch's AveragedModel and get_ema_multi_avg_fn are alternative implementations. Models with buffers need an explicit policy: copy/update buffers or recompute BatchNorm statistics; blindly averaging integer counters is wrong. EMA is useful in several training settings, including diffusion, but no universal accuracy increment or claim about every released checkpoint follows.

5 · TensorBoard — make the update clock visible

Log the unscaled training loss, example/token denominator, validation loss, accuracy or perplexity as appropriate, pre-clip norm, actual parameter-update norm and every parameter group's used learning rate. Log validation at its completed update count to align curves. A pre-step LR and get_last_lr() after stepping refer to different updates. Never compare dropout-active, augmented training loss with clean evaluation loss without identifying the mode difference.

For an environment where TensorBoard is already installed, use SummaryWriter(log_dir=...), call add_scalar(tag, scalar, global_step), optional add_histogram for selected parameter/gradient distributions, and close() in finally. add_graph(model, sample_input) traces a particular invocation; it is not proof of all dynamic execution paths. A local UI command is tensorboard --logdir=runs --host=127.0.0.1 --port=6006. This lesson does not install or launch it. The executed logger demonstrates the same scalar/update contract without claiming event-file or browser verification.

5.1 Team experiment tracking

W&B and other trackers can aggregate runs, configs and artifacts. They are optional integrations, not a verified industry default or a drop-in semantic identity with TensorBoard. Before enabling one, inspect what metadata, code, environment information and artifacts it uploads, obtain authorization and avoid logging secrets or private examples. Record the git revision and dependency/config hashes explicitly rather than assuming a service captured them. No account, upload, API call or integration run is part of this lesson.

6 · Gradient monitoring — interpret a norm, do not worship a threshold

clip_grad_norm_ returns the concatenated pre-clip norm. With gradients 3 and 4 it reports 5; clipping to 1 produces approximately .6 and .8. The API uses a stabilizing epsilon, so exact equality to 1 is not promised. Values depend on loss reduction, parameterization, model size and data: 1–10 is not a universal healthy range. A spike suggests investigation; zero may mean convergence, a masked/frozen branch or saturation. NaN is not repaired by clipping. Use finite checks and preserve the offending data/config safely.

6.1 Compilation changes the execution boundary

torch.compile(model) can optimize captured regions; loss/metrics outside the wrapped callable remain outside that capture. First-call compilation, graph breaks, guards and recompilation affect measured time. Variable shapes are one trigger, not the only one; dynamic=True is not a guarantee that all recompiles vanish. .item(), printing and data-dependent Python branches can break capture. Compare eager and compiled outputs/gradients first; report cold compile and warmed steady-state measurements separately. No 1.3–2.5× gain, exact compile latency or GPU parity is asserted here.

7 · Handling interruptions at a safe boundary

A signal handler should set a flag; the complete loop checks it between committed updates. A raw KeyboardInterrupt may arrive halfway through an optimizer step or a checkpoint write. Saving those live tensors and marking the whole epoch complete is not resume-safe. Prefer the previous committed generation, or design an explicit transactional step with rollback. SIGKILL, power loss and storage failure cannot be made safe by except KeyboardInterrupt.

For SIGTERM/SIGINT, install a small signal.signal handler that sets stop_requested=True, supply should_stop=lambda: stop_requested, and close any real writer in finally. Do not perform tensor serialization inside the signal handler. A distributed job also needs rank-consistent shutdown and all required shards; a rank-zero export alone is not a complete restart.

Choose checkpoint cadence from measured write cost and tolerated lost work: epoch boundaries are simple but can be hours apart; every N completed updates bounds replay; asynchronous staging reduces a stall only if copied state is consistent and the write actually finishes. Keep the best model separate from latest and rolling generations. Do not delete the last validated generation before the next one is readable.

8 · Failure cases and worked answers

1. Only weights restored: inference may work, but Adam moments and LR position restart. The negative control resets optimizer state and diverges from the uninterrupted log. Restore the bundle, not just tensors.
2. Same seed, different run: compare RNG consumption, sampler/data order, augmentation, versions, device and deterministic-kernel settings. Worker count alone is not proof of changed shuffled indices; the executable TensorDataset counterexample demonstrates this.
3. Best model never saves: initializing a minimized loss to zero prevents positive losses from improving it. Start at positive infinity; update best before serializing. state_dict() held in memory aliases model tensors unless copied/serialized.
4. Stopped in epoch 2, resumed at epoch 3: this silently discards the rest of epoch 2. Save a next-sample cursor plus permutation, or explicitly restart a whole epoch with the earlier state and accept replay.
5. Validation averages batch means: an 8-example and 5-example batch should contribute 8/13 and 5/13, not 1/2 each. Sum per-example losses and divide by 13.

10 · Retention scaffold — answers and extension contract

Q1 — What lives in a resume-safe checkpoint? The original four categories are model, optimizer, scheduler and position. Exact continuation adds every stateful dependency: scaler/EMA, RNGs, sampler/cursor, best/patience and config. Omitting moments changes the next adaptive update even with identical weights.

Q2 — What does deterministic cuDNN buy and cost? It constrains algorithm selection where supported, potentially reducing speed. It does not promise a fixed slowdown or cross-device equality. Request deterministic algorithms, pin the environment, then test the required tolerance.

Q3 — Which fingerprints do norms catch? Exploding/nonfinite gradients and vanishing/disconnected gradients are candidate causes, not unique diagnoses. Log pre-clip distributions per layer and inspect actual updates and data before changing LR.

Q4 — What is the load default? weights_only=True is the PyTorch 2.6+ default absent an explicit pickle module. The restricted unpickler narrows execution risk; untrusted checkpoints still require provenance and safe handling.

Q5 — What does EMA buy? It smooths the parameter trajectory and can improve evaluation. Validate decay and buffer handling; diffusion usage is not a theorem that all diffusion models require EMA or every checkpoint contains it.

Q6 — Rolling last three, early stopping, dry run and EMA? The full code supplies all four. save_generation deletes only old periodic files, never best.pt; the minimum-improvement/patience state is saved; dry-run state is marked non-resumable; evaluation uses the saved EMA. The tests exercise exact continuation after three updates, stopped-before-start, a one-batch dry run and early stop at epoch two. A CLI can expose these parameters without changing the state contract.

Sources and execution boundary

Primary contracts: PyTorch 2.9 reproducibility, serialization, optimizers and schedulers, EMA/SWA source v2.9.0, and TensorBoard writer source v2.9.0. Exact passages and raw hashes are in the batch review evidence. CPU synthetic programs are executed; no cloud/GPU, event UI, compiled model, external tracker or production durability claim is made.

Optional pre-read resources

These inherited resource links are optional, not the primary evidence for this review. Video identity, duration and current availability are not certified.

• A Recipe for Training Neural Networks
• TensorBoard: TensorFlow's visualization toolkit
• A Recipe for Training Neural Networks
• Saving and loading models

Previous: ← DL S016 · Datasets, DataLoaders · Next: DL S018 · GPU + Mixed Precision →