Training Loop Anatomy

Six things every serious training loop has

Config — a dataclass or YAML / TOML file that captures every hyperparameter. Saved into every checkpoint.
Seed — set at the top, applied to everything.
Mixed precision — autocast + GradScaler (or BF16, no scaler).
Gradient clipping — between backward and optimizer step.
Learning-rate schedule — warmup + cosine for transformers, OneCycle for fine-tuning, step decay for old-school CNNs.
Validation + best-checkpoint saving — every epoch (or every N steps), eval on a held-out set, save if val metric improved.

Tip: If you start every project from a 'blank training loop' template that already includes these six pieces, you save yourself the same week of debugging on every project.

The gradient accumulation pattern

For models too big to fit a desired batch size in VRAM: skip opt.zero_grad() for several forward/backward passes, then step + zero. The result is mathematically equivalent to a single bigger batch. Adjust scaler.scale(loss / accum_steps).backward() to keep the loss magnitude consistent.

Where frameworks help

PyTorch Lightning, Hugging Face Trainer, and Accelerate handle the boilerplate (mixed precision, distributed, checkpointing, logging) so you can focus on the model and data. Trade-off: less code, less visibility into what's happening. Write your own loop at least once before adopting a framework.

Principle: Your first training loop should be hand-written. Your tenth should probably use a framework. The middle ground is where you learn the most.

Code

A complete production-quality training loop·python

import torch, json
from torch import nn, optim
from torch.cuda.amp import autocast, GradScaler
from torch.nn.utils import clip_grad_norm_

def train(model, train_loader, val_loader, cfg, device):
    model = model.to(device)
    opt = optim.AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.wd)
    sch = warmup_cosine(opt, cfg.warmup_steps, cfg.total_steps)
    scaler = GradScaler()
    loss_fn = nn.CrossEntropyLoss()

    best_val = 0.0
    step = 0
    for epoch in range(cfg.epochs):
        model.train()
        for xb, yb in train_loader:
            xb, yb = xb.to(device), yb.to(device)
            opt.zero_grad()
            with autocast(dtype=torch.bfloat16):
                logits = model(xb)
                loss = loss_fn(logits, yb)
            scaler.scale(loss).backward()
            scaler.unscale_(opt)
            clip_grad_norm_(model.parameters(), max_norm=cfg.clip)
            scaler.step(opt); scaler.update(); sch.step()
            step += 1

        val_acc = evaluate(model, val_loader, device)
        if val_acc > best_val:
            best_val = val_acc
            torch.save({"model": model.state_dict(), "config": cfg.__dict__,
                        "val_acc": val_acc, "step": step}, "best.pt")
        print(f"epoch={epoch} step={step} val={val_acc:.4f}")
    return best_val

Six things every serious training loop has

The gradient accumulation pattern

Where frameworks help

Code

External links

Exercise

Progress

Comments 0