PyTorch II: real training pipelines

The data pipeline

A map-style Dataset implements len and getitem(i); an iterable-style one yields samples (for streams too large to index, such as web-scale text). The DataLoader wraps it with batching, shuffling, parallel workers and memory pinning:

from torch.utils.data import Dataset, DataLoader

class ImageFolderDS(Dataset):
    def __init__(self, paths, labels, transform): self.p, self.y, self.t = paths, labels, transform
    def __len__(self): return len(self.p)
    def __getitem__(self, i): return self.t(load_image(self.p[i])), self.y[i]

loader = DataLoader(ds, batch_size=128, shuffle=True, num_workers=8,
                    pin_memory=True, persistent_workers=True, drop_last=True)

The rules that matter: do decoding and augmentation inside getitem so the workers (separate processes) do it in parallel with GPU compute; pin_memory=True plus .to(device, non_blocking=True) overlaps host-to-device copies; drop_last=True keeps batch shapes static (which torch.compile and BatchNorm appreciate); a custom collate_fn pads variable-length sequences (6.4 Recurrent neural networks and BPTT); samplers handle class-balanced or distributed sampling. The single most common performance bug is a GPU idling at 30% utilization while 2 workers decode JPEGs — measure it (below) before optimizing the model.

Models as modules

nn.Module is a tree. Parameters registered in init are found by parameters(), moved by .to(), saved by state_dict(); submodules nest; nn.Sequential, nn.ModuleList and nn.ModuleDict are the containers (a plain Python list of layers is not registered — a classic bug). model.train()/model.eval() toggle dropout and BatchNorm behaviour (4.6 Initialization and normalization, 4.7 Regularization in deep nets); torch.no_grad() or the faster torch.inference_mode() disables graph building during evaluation. Keep the forward pass pure (no I/O, no device-specific code) so the same module runs on CPU, GPU, under compile and under DDP.

class Block(nn.Module):
    def __init__(self, d, p=0.1):
        super().__init__()
        self.norm = nn.LayerNorm(d)
        self.mlp = nn.Sequential(nn.Linear(d, 4*d), nn.GELU(), nn.Linear(4*d, d), nn.Dropout(p))
    def forward(self, x): return x + self.mlp(self.norm(x))        # pre-norm residual block

class Net(nn.Module):
    def __init__(self, d_in, d, n_blocks, n_out):
        super().__init__()
        self.inp = nn.Linear(d_in, d)
        self.blocks = nn.ModuleList([Block(d) for _ in range(n_blocks)])
        self.out = nn.Linear(d, n_out)
    def forward(self, x):
        x = self.inp(x)
        for b in self.blocks: x = b(x)
        return self.out(x)                                            # logits

The training template

Everything from 4.5 Optimizers and learning-rate schedules and 4.8 Training dynamics and debugging in one loop, written to be resumable:

import torch, math, time
from dataclasses import dataclass, asdict

@dataclass
class Cfg:
    lr: float = 3e-4; wd: float = 0.1; warmup: int = 500; steps: int = 20_000
    bs: int = 128; clip: float = 1.0; amp_dtype: torch.dtype = torch.bfloat16; eval_every: int = 500

def lr_at(step, cfg):                      # linear warmup + cosine to 10%
    if step < cfg.warmup: return cfg.lr * step / cfg.warmup
    t = (step - cfg.warmup) / max(1, cfg.steps - cfg.warmup)
    return cfg.lr * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * t)))

def train(model, train_loader, val_loader, cfg, device="cuda", ckpt_path="ckpt.pt", resume=False):
    model.to(device)
    decay = [p for n, p in model.named_parameters() if p.dim() >= 2]
    no_decay = [p for n, p in model.named_parameters() if p.dim() < 2]       # biases, norms
    opt = torch.optim.AdamW([{"params": decay, "weight_decay": cfg.wd},
                             {"params": no_decay, "weight_decay": 0.0}], lr=cfg.lr, betas=(0.9, 0.95))
    scaler = torch.amp.GradScaler(enabled=(cfg.amp_dtype == torch.float16))
    step = 0
    if resume:
        ck = torch.load(ckpt_path, map_location=device)
        model.load_state_dict(ck["model"]); opt.load_state_dict(ck["opt"])
        scaler.load_state_dict(ck["scaler"]); step = ck["step"]; torch.set_rng_state(ck["rng"])
    model.train(); it = iter(train_loader); t0 = time.time()
    while step < cfg.steps:
        try: x, y = next(it)
        except StopIteration: it = iter(train_loader); x, y = next(it)
        x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
        for g in opt.param_groups: g["lr"] = lr_at(step, cfg)
        with torch.autocast(device_type="cuda", dtype=cfg.amp_dtype):
            loss = torch.nn.functional.cross_entropy(model(x), y)
        opt.zero_grad(set_to_none=True)
        scaler.scale(loss).backward()
        scaler.unscale_(opt)
        gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.clip)
        scaler.step(opt); scaler.update(); step += 1
        if step % 50 == 0:
            print(f"step {step} loss {loss.item():.4f} gnorm {gnorm:.2f} lr {lr_at(step, cfg):.2e} "
                  f"{50 * cfg.bs / (time.time() - t0):.0f} ex/s"); t0 = time.time()
        if step % cfg.eval_every == 0:
            val = evaluate(model, val_loader, device); model.train()
            torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "scaler": scaler.state_dict(),
                        "step": step, "rng": torch.get_rng_state(), "cfg": asdict(cfg), "val": val}, ckpt_path)

@torch.inference_mode()
def evaluate(model, loader, device):
    model.eval(); tot, correct, n = 0.0, 0, 0
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        logits = model(x)
        tot += torch.nn.functional.cross_entropy(logits, y, reduction="sum").item()
        correct += (logits.argmax(-1) == y).sum().item(); n += y.numel()
    return {"loss": tot / n, "acc": correct / n}

Points worth noticing: the scheduler is applied per step and computed from the step count, so resuming is exact; weight decay skips 1-D parameters (4.5 Optimizers and learning-rate schedules); the gradient norm is measured after unscaling and before clipping (4.8 Training dynamics and debugging); set_to_none=True saves memory; evaluation accumulates sums rather than averaging per-batch means (unequal last batches bias the mean); and the checkpoint holds everything needed to continue the same trajectory. Add logging to Weights & Biases or MLflow (2.5 Tooling: Git, GPUs, experiments and the Hugging Face Hub) by replacing the print.

Speed: compile and profile

model = torch.compile(model) traces the forward and backward passes, fuses element-wise operations into single kernels and removes Python overhead; typical speedups are 1.3–2× for transformer-style models at the cost of a slow first iteration and recompilation whenever shapes change (keep them static: drop_last, padded sequence lengths). Use TF32 matmuls on Ampere+ GPUs (torch.set_float32_matmul_precision("high")) for a free ~3× on fp32 matmuls. The profiler tells you where time goes:

from torch.profiler import profile, ProfilerActivity
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    for _ in range(5): train_step()
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))

If the top rows are data-loader or .item() synchronizations rather than matmuls, fix the pipeline first; nvidia-smi utilization and the examples-per-second number in the loop are the quick checks.

More than one GPU

DistributedDataParallel (DDP) launches one process per GPU (torchrun --nproc_per_node=8 train.py), gives each a full copy of the model and a different shard of the data (DistributedSampler), and all-reduces gradients after every backward so all copies stay identical. The effective batch is per-GPU batch × world size, so scale the learning rate accordingly (4.5 Optimizers and learning-rate schedules), and only rank 0 should log and save. DDP is the default for models that fit on one GPU.

When the model (plus gradients and Adam states — 16 bytes per parameter in mixed precision, 4.12 Compute, hardware and scale basics) does not fit, FSDP / ZeRO shard parameters, gradients and optimizer states across GPUs and gather them just-in-time for each layer; tensor and pipeline parallelism split individual layers or groups of layers across devices. Those are the concern of 4.12 Compute, hardware and scale basics and 7.9 Scaling laws and compute-optimal training; as a learner you will mostly use DDP, FSDP via a library, or a single GPU with gradient accumulation (sum gradients over $k$ micro-batches before stepping to emulate a $k\times$ larger batch).

Lightning and Accelerate. PyTorch Lightning structures the loop (training_step, configure_optimizers) and handles devices, precision, checkpointing and logging; Hugging Face Accelerate keeps your loop and adds device placement, AMP and distributed launching with a few lines. Use them when you need multi-GPU or many experiments quickly; write the raw loop at least once so you understand what they do, and keep it for research code where you need full control.

Testing ML code

ML bugs are silent, so test what can be tested: shapes (assert logits.shape == (B, K)), the overfit-one-batch test as a unit test (4.8 Training dynamics and debugging), determinism (two runs with the same seed give the same loss), gradient flow (every parameter with requires_grad has a non-None, finite gradient after one step), data invariants (labels in range, no NaNs, splits disjoint), and checkpoint round-trips (save → load → identical evaluation). A ten-minute test suite catches most of the bugs in 4.8 Training dynamics and debugging's checklist before a GPU is ever rented.

Where this shows up later

• Fine-tuning a pretrained CNN with this loop — 5.4 Transfer learning and training CNNs in practice, 5.10 Project: train and deploy an image classifier; padding, collation and sequence batching — 6.4 Recurrent neural networks and BPTT, 6.8 Project: character-level language model.
• The nanoGPT training script is this template with a GPT model — 7.8 Build GPT from scratch (nanoGPT walkthrough); FSDP, tensor and pipeline parallelism at scale — 4.12 Compute, hardware and scale basics, 7.9 Scaling laws and compute-optimal training.
• Hugging Face Trainer/TRL and Unsloth wrap this loop for LLM fine-tuning — 8.2 Fine-tuning and PEFT, 8.13 Project: fine-tune a small LLM with LoRA and ship a chat app; serving and inference loops — 7.10 Efficient inference: KV cache, GQA, FlashAttention, speculative decoding, quantization.