PyTorch Training skill for Codex

Done well means: the loop overfits a tiny batch before scaling up, every epoch logs train and validation metrics, a crash at hour 9 resumes from the last checkpoint with identical optimizer and scheduler state, and multi-GPU runs give the same validation curve as one GPU.

Before changing anything

1. Read AGENTS.md, the existing train script, config files and the lockfile. Record torch.__version__, CUDA version and GPU (nvidia-smi).
2. Find how the project launches training (Makefile, torchrun, Lightning, Accelerate). Stay with the existing framework unless asked to change it.
3. Run one step end to end on a tiny subset before any long run.

Installing CUDA wheels needs network and the right index URL for the CUDA version; ask before installing in a Codex sandbox.

Sanity ladder (do in order)

- [ ] 1. Shapes and dtypes printed for one batch; loss computes
- [ ] 2. Initial loss matches theory (classification: about ln(num_classes))
- [ ] 3. Overfit one batch to near-zero loss in a few hundred steps
- [ ] 4. Train on 1% of data; validation metric moves
- [ ] 5. Full run with checkpoints, logging and early stopping

If step 3 fails, the bug is in the model, loss, labels or optimizer, not the data size.

Reference loop

import math, os, random
import numpy as np
import torch
from torch.utils.data import DataLoader

def seed_all(seed: int) -> None:
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)

seed_all(42)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
use_bf16 = device.type == "cuda" and torch.cuda.is_bf16_supported()
amp_dtype = torch.bfloat16 if use_bf16 else torch.float16
scaler = torch.amp.GradScaler("cuda", enabled=(device.type == "cuda" and not use_bf16))

train_dl = DataLoader(train_ds, batch_size=64, shuffle=True, num_workers=8,
                      pin_memory=True, persistent_workers=True, drop_last=True)
val_dl = DataLoader(val_ds, batch_size=128, num_workers=4, pin_memory=True)

model = Model().to(device)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
accum = 4
total_steps = math.ceil(len(train_dl) / accum) * epochs
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=3e-4, total_steps=total_steps, pct_start=0.05)

for epoch in range(start_epoch, epochs):
    model.train()
    for i, (x, y) in enumerate(train_dl):
        x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
        with torch.autocast(device_type=device.type, dtype=amp_dtype, enabled=device.type == "cuda"):
            loss = loss_fn(model(x), y) / accum
        scaler.scale(loss).backward()
        if (i + 1) % accum == 0:
            scaler.unscale_(opt)
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            scaler.step(opt); scaler.update()
            opt.zero_grad(set_to_none=True)
            sched.step()

    model.eval()
    with torch.inference_mode():
        val_loss = sum(loss_fn(model(x.to(device)), y.to(device)).item() for x, y in val_dl) / len(val_dl)
    save_checkpoint(epoch, val_loss)

With bf16 the scaler is disabled and its calls are pass-throughs, so one code path serves both.

Checkpoints that actually resume

def save_checkpoint(epoch, val_loss, path="ckpt/last.pt"):
    state = {
        "model": model.state_dict(), "optimizer": opt.state_dict(),
        "scheduler": sched.state_dict(), "scaler": scaler.state_dict(),
        "epoch": epoch, "val_loss": val_loss,
        "rng": {"torch": torch.get_rng_state(), "cuda": torch.cuda.get_rng_state_all(),
                "numpy": np.random.get_state(), "python": random.getstate()},
        "config": cfg,
    }
    tmp = path + ".tmp"
    torch.save(state, tmp)
    os.replace(tmp, path)          # atomic: a crash mid-save never corrupts last.pt

ckpt = torch.load("ckpt/last.pt", map_location=device, weights_only=False)  # trusted file only
model.load_state_dict(ckpt["model"]); opt.load_state_dict(ckpt["optimizer"])
sched.load_state_dict(ckpt["scheduler"]); scaler.load_state_dict(ckpt["scaler"])
start_epoch = ckpt["epoch"] + 1

• torch.load defaults to weights_only=True since 2.6; full checkpoints with RNG state need weights_only=False, which unpickles code: only for files you produced.
• Keep last.pt plus best.pt (by validation metric); write model.module.state_dict() under DDP or strip the module. / _orig_mod. prefixes when loading.

Multi-GPU with DDP

import os, torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

dist.init_process_group("nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
model = DDP(Model().cuda(local_rank), device_ids=[local_rank])
sampler = DistributedSampler(train_ds, shuffle=True)
train_dl = DataLoader(train_ds, batch_size=64, sampler=sampler, num_workers=8, pin_memory=True)

for epoch in range(epochs):
    sampler.set_epoch(epoch)          # otherwise every epoch sees the same order
    ...
if dist.get_rank() == 0:
    save_checkpoint(...)              # only rank 0 writes
dist.barrier()
dist.destroy_process_group()

torchrun --standalone --nproc_per_node=4 train.py --config cfg.yaml
torchrun --nnodes=2 --nproc_per_node=8 --rdzv_backend=c10d --rdzv_endpoint=$HEAD:29500 train.py

• Batch size is per GPU: global batch = per-GPU x GPUs x accumulation. Scale LR accordingly or keep global batch constant when comparing runs.
• Under gradient accumulation, wrap non-final micro-steps in model.no_sync() to skip redundant all-reduces.
• Metrics computed per rank must be all-reduced before logging, or rank 0 reports a shard.
• Model too big for one GPU: FSDP2 (torch.distributed.fsdp.fully_shard) instead of DDP.

Speed

Bottleneck signal	Fix
GPU utilization low, CPU busy	More num_workers, pin_memory=True, preprocess offline, faster decoding
GPU busy but slow	bf16 autocast, torch.compile(model), channels_last for CNNs, larger batch
Many small kernels	torch.compile; avoid Python loops over tensor elements
Sync stalls	Remove .item() / .cpu() inside the step; log every N steps

Turn on TF32 for float32 matmuls on Ampere+: torch.backends.cuda.matmul.allow_tf32 = True (or torch.set_float32_matmul_precision("high")). Profile before guessing: torch.profiler.profile(...) around 20 steps.

Debugging

Symptom	Check
Loss NaN or inf	fp16 overflow (switch to bf16), LR too high, log of zero, bad labels; torch.autograd.set_detect_anomaly(True) for a few steps
Loss flat from step 0	optimizer.zero_grad missing, frozen params, wrong loss target dtype, LR far too low
Train improves, val does not	Leakage-free split? Augmentations too strong? Overfitting: weight decay, dropout, more data
Val better than train	Dropout/augmentation only in train mode: expected; check model.eval() is called
CUDA OOM	Smaller batch + accumulation, activation checkpointing, bf16, free tensors held for logging
Results differ every run	Seeds, torch.use_deterministic_algorithms(True) for audits (slower), worker seeding

Output contract

**Run:** resnet50-defects-v4 | torch 2.x, CUDA 12.x, 4x A10G DDP, bf16, compile on
**Global batch:** 256 (64 x 4 GPUs x 1) | **LR:** OneCycle max 3e-4 | **Epochs:** 30 (early stop at 24)
**Best val:** F1 0.912 at epoch 21 (ckpt/best.pt) | **Throughput:** 1,840 img/s
**Resume tested:** killed at epoch 5, resumed, curve matched within noise

Log runs with $ml-experiment-tracking; for LLM adapter training use $lora-fine-tuning.