DL S017 · The Training Loop, Logging, and Reproducibility

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

The story — the 3 AM checkpoint that saved a PhD

A friend of mine — call her Priya, she'd rather not be named — was 11 days into pretraining a 350M-parameter model for her thesis on a rented A100. At 3:17 AM her SSH tunnel dropped, tmux crashed, the Python process died. She woke up expecting to find rubble. What she found was ckpt.pt at step 47,200, ckpt.pt.best at step 41,800, and a copy of the RNG state. Resume took one command. Her thesis submitted on time. It cost her exactly seven extra lines of code, written the day before. The training loop she'd copied from a tutorial had none of them.

Every experienced ML engineer has this story, with different names attached. The gap between a tutorial training loop and one you'd trust in production is exactly the gap between passing and failing your next deadline. Not model quality — engineering hygiene.

The training loop is where every promise the framework made comes to collect. Autograd works? Great — I need to see the loss go down. nn.Module composes cleanly? Fine — I need to save it and load it two weeks from now on a different machine. DataLoader batches efficiently? Beautiful — I need to know when a worker crashes at 3 AM after 6 hours of training.

There's a version of the training loop you write in a tutorial: 8 lines, prints loss occasionally, celebrates when things work. And there's the version you write when you actually care about the outcome — reproducible seed, checkpoint every N steps, resume-from-checkpoint, log every scalar to TensorBoard, gradient-norm tracking so you notice explosions early, best-model saving, KeyboardInterrupt handling so Ctrl-C doesn't lose an hour of progress.

The gap between the two versions is what separates hobbyists from people whose models actually make it to Monday's meeting. Today we bridge it.

Prerequisites

• Session 015 — nn.Module.
• Session 016 — DataLoader.

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

The heart of every training loop is the same five lines:

optimizer.zero_grad()          # 1. reset accumulated grads
loss = criterion(model(xb), yb)# 2. forward
loss.backward()                # 3. backward
optimizer.step()               # 4. update params
scheduler.step()               # 5. update LR (optional; per-batch OR per-epoch)

Everything else in the loop — reproducibility, logging, checkpointing, monitoring — wraps around these five lines. Learn to see them as a unit.

2 · Reproducibility — the seed sandwich

Two experiments with "the same code" that give different numbers erode all your ability to reason. Set seeds at the very start:

import random, numpy as np, torch
 
def seed_everything(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
 
seed_everything(42)

Two things worth flagging:

• cudnn.deterministic = True forces cuDNN to pick deterministic algorithms. Costs ~5-20% training speed. Worth it for research, tune off for production.
• Even with all this, DataLoader with num_workers > 0 shuffles differently unless you also set a generator and a worker_init_fn. For strict determinism:

g = torch.Generator().manual_seed(42)
loader = DataLoader(ds, batch_size=128, shuffle=True, generator=g,
                    worker_init_fn=lambda id: np.random.seed(42 + id))

Reproducibility is a spectrum, not a switch. Get "same seed → same loss curve to 3 decimals" and stop chasing bit-for-bit. It's not worth the last 5% of engineering effort.

3 · The full loop (annotated)

Read this end to end, then we'll unpack each concern.

import torch
from torch.utils.tensorboard import SummaryWriter
 
def train(model, train_loader, val_loader, epochs, device,
          lr=1e-3, ckpt_path="ckpt.pt", log_dir="runs/exp1"):
 
    model.to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    criterion = torch.nn.CrossEntropyLoss()
    writer    = SummaryWriter(log_dir)
 
    start_epoch, best_val = 0, float("inf")
 
    # Resume if a checkpoint exists
    import os
    if os.path.exists(ckpt_path):
        ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
        model.load_state_dict(ckpt["model"])
        optimizer.load_state_dict(ckpt["optim"])
        start_epoch = ckpt["epoch"] + 1
        best_val    = ckpt["best_val"]
        print(f"resumed from epoch {start_epoch}, best_val={best_val:.4f}")
 
    try:
        for epoch in range(start_epoch, epochs):
            # ---- train ----
            model.train()
            for step, (xb, yb) in enumerate(train_loader):
                xb, yb = xb.to(device, non_blocking=True), yb.to(device, non_blocking=True)
                optimizer.zero_grad()
                logits = model(xb)
                loss = criterion(logits, yb)
                loss.backward()
                grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()
 
                if step % 100 == 0:
                    gs = epoch * len(train_loader) + step
                    writer.add_scalar("train/loss", loss.item(), gs)
                    writer.add_scalar("train/grad_norm", grad_norm.item(), gs)
 
            # ---- validate ----
            model.eval()
            val_loss, n = 0.0, 0
            with torch.no_grad():
                for xb, yb in val_loader:
                    xb, yb = xb.to(device), yb.to(device)
                    val_loss += criterion(model(xb), yb).item() * xb.size(0)
                    n += xb.size(0)
            val_loss /= n
            writer.add_scalar("val/loss", val_loss, epoch)
            print(f"epoch {epoch}  val_loss {val_loss:.4f}")
 
            # ---- checkpoint ----
            state = {
                "epoch": epoch,
                "model": model.state_dict(),
                "optim": optimizer.state_dict(),
                "best_val": best_val,
            }
            torch.save(state, ckpt_path)              # every-epoch ckpt
            if val_loss < best_val:
                best_val = val_loss
                torch.save(state, ckpt_path + ".best")
 
    except KeyboardInterrupt:
        print("interrupted; saving current state")
        torch.save({"epoch": epoch, "model": model.state_dict(),
                    "optim": optimizer.state_dict(), "best_val": best_val},
                   ckpt_path + ".interrupted")
 
    writer.close()

This is ~50 lines of dense code. Every piece is load-bearing. Below, one section per concern.

4 · Checkpointing — save state, not just weights

A checkpoint that only saves model.state_dict() is almost useless for resume. You'll resume with a freshly-initialized optimizer, meaning momentum buffers reset, Adam's m and v reset, learning-rate schedule position resets. Your next epoch's dynamics differ from what would have happened without the interruption.

Save all four things:

Save on:

• Every epoch (or every N steps for long jobs)
• Best-so-far validation metric (separate file, don't overwrite)
• On KeyboardInterrupt (see §7)

Naming convention I use: ckpt.pt = latest, ckpt.pt.best = best val, ckpt-epoch{N}.pt = periodic snapshots (keep last 3).

4.1 weights_only=True — the 2024 security default flip

Starting in PyTorch 2.4 (July 2024), torch.load prints a FutureWarning if you don't pass weights_only, and in 2.6 (Jan 2025) the default flipped to weights_only=True. Reason: an attacker who ships you a .pt file can execute arbitrary code on your machine when you torch.load it — pickle is Turing-complete. weights_only=True uses a restricted unpickler that only permits tensors and a few safe types. For downloaded checkpoints from the internet, always weights_only=True. For your own checkpoints that store a dict of tensors + integers + floats, it Just Works. For checkpoints that store custom classes (a scheduler with a lambda, a full optimizer object with custom state), you may need weights_only=False — be sure you trust the file. See PyTorch security policy 2024.

4.2 EMA — the trick every diffusion paper uses

Exponential Moving Average of weights: maintain a shadow copy θ_ema = 0.999 · θ_ema + 0.001 · θ after every step, and use θ_ema for evaluation and final release. This smooths out the noise from stochastic updates and typically gives 0.3–1.0% higher validation accuracy on ImageNet, and dramatically better sample quality in diffusion models (this is why every Stable Diffusion / SDXL / FLUX checkpoint you download is the EMA weights, not the raw training weights). PyTorch ships it as torch.optim.swa_utils.AveragedModel(model, multi_avg_fn=torch.optim.swa_utils.get_ema_multi_avg_fn(0.999)). One line, big win. See Karras et al. 2024, "Analyzing and Improving the Training Dynamics of Diffusion Models" for the state of the art on EMA schedules.

5 · TensorBoard — the debugger you didn't know you needed

TensorBoard turns time-series data into interactive plots in a browser. You start it once:

tensorboard --logdir=runs --port=6006

...and open localhost:6006. Every writer.add_scalar call shows up as a line on a plot. Cheap and life-changing.

Minimum set of scalars to log:

• train/loss (every ~100 steps)
• val/loss (per epoch)
• train/grad_norm (every ~100 steps — see §6)
• train/lr (per step, if scheduler is active)
• Any validation metric you care about (accuracy, F1, perplexity)

Bonus, high-leverage:

• writer.add_histogram("params/fc1_weight", model.fc1.weight, epoch) — layer weight distributions per epoch. If one layer's weights are exploding, this makes it obvious.
• writer.add_graph(model, sample_input) — visualizes the module tree.

There are fancier tools (Weights & Biases, MLflow, Neptune, Comet). Learn TensorBoard first — it's free, local, and the vocabulary translates 1:1.

6 · Gradient monitoring — the canary in the coal mine

Every backward pass produces gradients. If any gradient's L2 norm goes to infinity, your training is one step away from NaN. If they all go to zero, you're not learning. Watching gradient norms is the cheapest way to notice both.

grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
writer.add_scalar("train/grad_norm", grad_norm.item(), global_step)

clip_grad_norm_ computes the total gradient norm across all parameters AND clips it in place to max_norm. The return value is the pre-clip norm — that's what you log. Then you can:

• See if norms are stable (~O(1)–O(10) is healthy for most models)
• Notice a sudden spike (usually a bad batch — corrupted example, huge input, etc.)
• Notice slow drift up (usually LR too high — reduce)
• Notice collapse to ~0 (dead network — check initialization, activations, LR too low)

Session 024 covers clipping in depth. For now, know it's a common line in real training loops even before you understand why.

6.1 torch.compile and the training loop

The simplest wrapping is:

model = model.to(device)
model = torch.compile(model)          # trace on first forward

Expect the first batch to be slow (5–120 seconds while TorchDynamo traces and TorchInductor compiles kernels). Subsequent batches are 1.3–2.5× faster on A100/H100. If your loss/metrics are computed in Python outside model(...), they don't get compiled — that's fine and by design.

Gotchas to know now (details in later modules):

• Recompilation storms if input shapes vary batch-to-batch. Use dynamic=True or pad to fixed shapes.
• Graph breaks on print(x), .item(), and data-dependent if. Debug with torch._dynamo.explain(model)(sample_input).
• AMP interaction: as of 2.5, torch.compile composes cleanly with autocast. Wrap the model, then wrap the forward in autocast. (S018 details.)

6.2 Weights & Biases — when TensorBoard isn't enough

For solo projects, TensorBoard is fine. For teams, comparing dozens of runs, or logging from cluster jobs to a central dashboard, W&B (pip install wandb) is the industry default in 2025. Swap SummaryWriter for wandb.init(...) and writer.add_scalar for wandb.log. The rest of the loop is identical. W&B also auto-captures git hash, GPU util, memory, and env vars — metadata that helps future-you figure out what you actually ran. HuggingFace Trainer, Lightning, and TorchTitan all ship W&B integrations by default.

7 · Handling KeyboardInterrupt — the "please save when I hit Ctrl-C" pattern

Real training runs get killed. Servers restart. You realize the LR is wrong. Ctrl-C should never lose an epoch's worth of progress.

Wrap the entire training loop in try/except KeyboardInterrupt and save on the way out:

try:
    for epoch in range(start_epoch, epochs):
        ...
except KeyboardInterrupt:
    print("caught SIGINT; saving")
    save_checkpoint(...)
finally:
    writer.close()

For cluster jobs (Slurm sending SIGTERM), install a signal handler:

import signal
_shutdown = False
def _handler(sig, frame):
    global _shutdown; _shutdown = True
signal.signal(signal.SIGTERM, _handler)
 
# ... inside your inner loop ...
if _shutdown:
    save_checkpoint(...); break

Small effort, huge peace of mind.

8 · War stories

if val_loss < best_val:
    best_val = val_loss
    torch.save(model.state_dict(), "best.pt")

9 · Diagram — the loop with all the machinery

10 · 🧠 Retention scaffold

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

Spaced review: re-read §4 (checkpoint contents) and §6 (grad monitoring) in 24 hours. Full session on day 7.

Next session (S018): GPU + mixed precision. autocast, GradScaler, bf16 vs fp16 vs fp8, and the two AMP gotchas that cause 90% of "loss went NaN" reports.

Sticky note (keep on your desk): Body = zero → forward → backward → step. Checkpoint = 4 items (model, optim, sched, epoch). Every long loop needs try/except KeyboardInterrupt.

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