Training
at Scale
The practical engineering that separates a working prototype from a production training run: mixed precision, distributed training across multiple GPUs, debugging the inevitable failures, and tracking experiments so you can reproduce what worked.
Mixed Precision Training
By default, PyTorch uses 32-bit floating point (FP32) for all computations. Mixed precision training uses 16-bit floating point (FP16 or BF16) for most operations, falling back to FP32 only where precision is critical. This roughly halves memory usage and can give a 2-3× speedup on modern GPUs with Tensor Cores — at essentially no cost to final model accuracy when done correctly.
Memory savings
FP16/BF16 uses half the memory of FP32. This means larger batch sizes, larger models, or longer sequences fit in the same GPU memory — often the binding constraint in deep learning.
Speed
Modern GPUs (NVIDIA Volta and later) have Tensor Cores specifically optimized for FP16/BF16 matrix operations — 2-8× faster than FP32 for the same operation, depending on hardware generation.
FP16 risk: underflow
FP16 has a much smaller representable range than FP32. Small gradients can underflow to zero, stalling training. Loss scaling (multiply loss by a large constant before backward, divide gradients after) addresses this.
BF16: the safer choice
BFloat16 has the same exponent range as FP32 (no underflow risk) but less mantissa precision. Modern GPUs (A100+) support it natively. If available, BF16 is generally preferred over FP16 for its stability.
import torch
from torch.cuda.amp import autocast, GradScaler
model = model.to('cuda')
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
scaler = GradScaler() # handles loss scaling for FP16 stability
for X_batch, y_batch in train_loader:
X_batch, y_batch = X_batch.cuda(), y_batch.cuda()
optimizer.zero_grad()
# autocast: automatically casts ops to FP16/BF16 where safe,
# keeps FP32 for numerically sensitive ops (e.g. softmax, loss)
with autocast(dtype=torch.bfloat16): # or torch.float16
outputs = model(X_batch)
loss = criterion(outputs, y_batch)
# GradScaler scales the loss up before backward (prevents underflow),
# then unscales gradients before the optimizer step
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
# Note: with BF16, GradScaler is often unnecessary (no underflow risk)
# Simplified BF16-only loop:
for X_batch, y_batch in train_loader:
optimizer.zero_grad()
with autocast(dtype=torch.bfloat16):
loss = criterion(model(X_batch), y_batch)
loss.backward() # no scaler needed for BF16
optimizer.step()Distributed Training Concepts: Data vs. Model Parallelism
When a single GPU isn't enough — either the model doesn't fit in memory, or training would take too long — you distribute the work across multiple GPUs (or multiple machines, each with multiple GPUs). There are two fundamentally different strategies, often combined.
| Strategy | What's split | When to use | Communication cost |
|---|---|---|---|
| Data Parallelism | Batch split across GPUs; each has a full model copy | Model fits on 1 GPU, want faster training | Sync gradients each step |
| Model Parallelism | Model layers split across GPUs | Model too large for 1 GPU | Pass activations between GPUs |
| Tensor Parallelism | Individual layers (matrices) split across GPUs | Very large layers (LLMs) | High — frequent all-reduce ops |
| Pipeline Parallelism | Sequential model stages across GPUs, micro-batched | Very deep models, multi-node | Lower — staged activation passing |
Each GPU holds an identical copy of the model. A large batch is split across GPUs — each computes forward and backward passes on its slice independently. Gradients are then averaged (all-reduced) across all GPUs before the parameter update, so every copy of the model stays synchronized.
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
def setup(rank, world_size):
dist.init_process_group("nccl", rank=rank, world_size=world_size)
torch.cuda.set_device(rank)
def train(rank, world_size):
setup(rank, world_size)
model = MyModel().to(rank)
model = DDP(model, device_ids=[rank]) # wraps model — handles gradient sync
# DistributedSampler ensures each GPU sees a different data slice
sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank)
loader = DataLoader(train_dataset, batch_size=32, sampler=sampler)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
for epoch in range(10):
sampler.set_epoch(epoch) # reshuffles differently each epoch
for X, y in loader:
X, y = X.to(rank), y.to(rank)
optimizer.zero_grad()
loss = criterion(model(X), y)
loss.backward() # DDP automatically all-reduces gradients here
optimizer.step()
# Launch with: torchrun --nproc_per_node=4 train_script.py
# This automatically sets rank, world_size, and spawns 4 processesDebugging a Model That Won't Train
Every deep learning practitioner faces this: loss is flat, loss is NaN, accuracy is stuck at random-guess level, or the model trains but doesn't generalize. This lesson is a systematic debugging checklist — the difference between hours of confused guessing and minutes of targeted diagnosis.
- Check learning rate — too high is the most common cause. Try reducing by 10×.
- Check for division by zero or log(0) in custom loss functions — add epsilon (1e-8).
- Check input data for NaN/Inf values before they even reach the model.
- Enable gradient clipping:
clip_grad_norm_(params, max_norm=1.0). - If using mixed precision, verify GradScaler is configured correctly.
- Verify you're calling
optimizer.zero_grad()before each backward pass. - Verify
loss.backward()andoptimizer.step()are both being called. - Check the learning rate isn't absurdly small (try 1e-3 as a sane default to test).
- Sanity check: can the model overfit a tiny dataset (e.g. 10 examples)? If not, there's a bug, not a generalization issue.
- Verify labels are correctly aligned with inputs — shuffled/mismatched labels is a common silent bug.
- Check loss function matches the task (e.g. not using MSE for classification).
- Add regularization: dropout, weight decay, data augmentation (Module 4).
- Reduce model capacity if dataset is small relative to parameter count.
- Verify train/val split has no data leakage (duplicate or near-duplicate samples in both).
- Get more training data, or use transfer learning if applicable.
- Verify the model and data are actually on GPU, not silently running on CPU.
- Check DataLoader
num_workers— too low causes CPU data-loading to bottleneck the GPU. - Enable mixed precision training (Lesson 11.1).
- Profile with
torch.profilerto find the actual bottleneck rather than guessing.
import torch
# The single most valuable debugging step: can your model memorize
# 5-10 examples perfectly? If not, you have a bug — not a data problem.
tiny_X = X_train[:8]
tiny_y = y_train[:8]
model = MyModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = torch.nn.CrossEntropyLoss()
for step in range(200):
optimizer.zero_grad()
out = model(tiny_X)
loss = criterion(out, tiny_y)
loss.backward()
optimizer.step()
if step % 20 == 0:
acc = (out.argmax(1) == tiny_y).float().mean()
print(f"step {step}: loss={loss.item():.4f}, acc={acc.item():.2f}")
# Expected: loss → ~0, accuracy → 1.0 within ~100 steps
# If this fails, the bug is in your model/loss/optimizer setup —
# not your data, regularization, or hyperparametersExperiment Tracking (Weights & Biases / MLflow Basics)
After your tenth training run with slightly different hyperparameters, you will not remember which configuration produced which result unless you tracked it. Experiment tracking tools log metrics, hyperparameters, and artifacts automatically, turning "I think run 7 was the good one?" into a searchable, comparable record.
What to track
Hyperparameters (lr, batch size, architecture), metrics per step/epoch (train/val loss, accuracy), system metrics (GPU utilization, memory), model checkpoints, and example predictions/visualizations.
Comparing runs
Tools like Weights & Biases provide dashboards to overlay loss curves from multiple runs, filter by hyperparameter values, and sort by final metric — essential for hyperparameter tuning and understanding what actually matters.
Reproducibility
Logging the exact code version (git commit), random seed, and full configuration means you can reproduce any past result exactly — critical for debugging regressions and for scientific rigor.
Free tools
Weights & Biases (free tier for individuals), MLflow (fully open source, self-hosted), TensorBoard (built into PyTorch, simplest option for local use). All are more than sufficient for learning and small projects.
import wandb
import torch
# ── Initialize a run — logs config automatically ───────────────────────
wandb.init(
project="cifar10-resnet",
config={
"learning_rate": 1e-3,
"batch_size": 128,
"epochs": 100,
"architecture": "SmallResNet",
"optimizer": "AdamW"
}
)
config = wandb.config # access hyperparameters from config
model = SmallResNet()
optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)
for epoch in range(config.epochs):
# ── training step ──
train_loss, train_acc = 0.42, 0.85 # placeholder — your real metrics
val_loss, val_acc = 0.51, 0.81
# Log metrics — automatically creates plots in the dashboard
wandb.log({
"epoch": epoch,
"train/loss": train_loss,
"train/accuracy": train_acc,
"val/loss": val_loss,
"val/accuracy": val_acc,
"lr": optimizer.param_groups[0]['lr']
})
# Log the final model as an artifact (versioned, downloadable)
torch.save(model.state_dict(), "model.pt")
artifact = wandb.Artifact("trained-model", type="model")
artifact.add_file("model.pt")
wandb.log_artifact(artifact)
wandb.finish() # always close the run cleanly
# Quick alternative — TensorBoard (no account needed, runs locally)
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter("runs/experiment_1")
writer.add_scalar("Loss/train", train_loss, epoch)
writer.add_scalar("Loss/val", val_loss, epoch)
# View with: tensorboard --logdir=runs