Dl Engineering And Capstone
GPUs, Mixed Precision & Scaling Training
Module 1, Chapter 4 introduced GPUs conceptually. This chapter covers the practical techniques for using them efficiently — the difference between "it runs on a
Jr Codex Deep Learning Notes
Level: Advanced Prerequisites: Chapter 1: PyTorch End-to-End Workflow Time to complete: ~25 minutes
Table of Contents
- Revisiting Why GPUs Matter
- Numerical Precision: FP32 vs FP16
- Mixed Precision Training
- Gradient Accumulation
- Multi-GPU Training — a Practical Orientation
- Profiling: Finding the Actual Bottleneck
- Summary & Next Steps
1. Revisiting Why GPUs Matter
Module 1, Chapter 4 introduced GPUs conceptually. This chapter covers the practical techniques for using them efficiently — the difference between "it runs on a GPU" and "it uses the GPU's full capability."
2. Numerical Precision: FP32 vs FP16
Every number in a neural network — weights, activations, gradients — is stored with a specific numerical precision, most commonly 32-bit floating point (FP32) by default.
FP32 vs FP16 (Half Precision)
─────────────────────────────────────────
FP32: 32 bits per number — HIGH precision, but uses
MORE memory and compute per operation
FP16: 16 bits per number — HALF the memory, and
modern GPUs can perform FP16 operations
substantially FASTER — but LOWER precision,
which can cause numerical instability in
some computations (e.g. very small gradient
values may round to exactly zero)
─────────────────────────────────────────
3. Mixed Precision Training
Mixed precision training uses FP16 for most operations (for speed and memory savings) while selectively keeping FP32 for numerically sensitive operations (like the final loss computation), getting most of FP16's speed benefit without FP16's full numerical risk.
import torch
model = model.to("cuda")
optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)
scaler = torch.cuda.amp.GradScaler() # handles the numerical stability concerns automatically
for images, labels in train_loader:
images, labels = images.to("cuda"), labels.to("cuda")
optimizer.zero_grad()
with torch.cuda.amp.autocast(): # automatically uses FP16 where SAFE, FP32 where NEEDED
outputs = model(images)
loss = criterion(outputs, labels)
scaler.scale(loss).backward() # SCALES the loss up before backprop — prevents small
# gradients from rounding to zero in FP16 (Section 2)
scaler.step(optimizer) # unscales gradients before the actual weight update
scaler.update() # adjusts the scaling factor for the NEXT iterationWhy the GradScaler Is Necessary
─────────────────────────────────────────
FP16's limited precision means very SMALL gradient values
(common, especially in early/deep layers — echoing Module
3, Chapter 6's vanishing gradient discussion) can round down
to EXACTLY zero, silently losing the learning signal
entirely. The GradScaler multiplies the loss by a LARGE
factor before backpropagation (making small gradients large
enough to survive FP16's precision limits), then divides
the gradients back down by the same factor before the
actual weight update — solving the precision problem
without changing the underlying MATH.
─────────────────────────────────────────
Mixed precision training commonly provides a 2-3x speedup with minimal or no accuracy loss — it is close to a standard default for training on modern GPUs, particularly for the large architectures in Modules 4-8.
4. Gradient Accumulation
When a desired batch size (Module 2, Chapter 5) doesn't fit in available GPU memory, gradient accumulation simulates a larger batch by accumulating gradients across several smaller "micro-batches" before performing a single weight update.
accumulation_steps = 4 # simulates a batch size 4x LARGER than what actually fits in memory
optimizer.zero_grad()
for i, (images, labels) in enumerate(train_loader):
images, labels = images.to("cuda"), labels.to("cuda")
outputs = model(images)
loss = criterion(outputs, labels) / accumulation_steps # SCALE DOWN — since gradients ACCUMULATE
loss.backward() # gradients ACCUMULATE across iterations — note there's NO zero_grad() here yet
if (i + 1) % accumulation_steps == 0:
optimizer.step() # apply the ACCUMULATED gradient — equivalent to ONE step on a LARGER batch
optimizer.zero_grad() # NOW reset gradients, for the next accumulation cycleWhy Divide the Loss by accumulation_steps
─────────────────────────────────────────
Without dividing, the ACCUMULATED gradient after 4
micro-batches would be 4x LARGER than a genuine batch of
that combined size would produce (since gradients from EACH
micro-batch's backward() call simply ADD together) —
dividing the loss beforehand keeps the EFFECTIVE gradient
magnitude consistent with what a true larger batch would
have produced.
─────────────────────────────────────────
5. Multi-GPU Training — a Practical Orientation
For training across multiple GPUs simultaneously, PyTorch's DistributedDataParallel is the modern standard approach (worth knowing exists, even at an orientation level).
The Core Idea (Data Parallelism)
─────────────────────────────────────────
Each GPU holds a COMPLETE COPY of the model. A large batch
is SPLIT across the available GPUs — each GPU computes a
forward/backward pass on its OWN slice of the batch,
independently and in PARALLEL. Gradients from all GPUs are
then SYNCHRONIZED (averaged) before each weight update,
ensuring every GPU's model copy stays IDENTICAL throughout
training.
─────────────────────────────────────────
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
# A SIMPLIFIED sketch — real multi-GPU setup involves process launching infrastructure
# beyond this orientation's scope
def setup_distributed(rank, world_size):
dist.init_process_group("nccl", rank=rank, world_size=world_size)
model = MyModel().to(rank)
model = DDP(model, device_ids=[rank]) # wraps the model, handling gradient SYNCHRONIZATION automatically
return model6. Profiling: Finding the Actual Bottleneck
Before applying any of Sections 2-5's optimizations, profile first — a common mistake is optimizing GPU compute (mixed precision) when the actual bottleneck is data loading (Chapter 1, Section 3) or something else entirely.
import torch.profiler
with torch.profiler.profile(
activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
record_shapes=True,
) as prof:
for images, labels in train_loader:
images, labels = images.to("cuda"), labels.to("cuda")
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
optimizer.zero_grad()
break # profile just ONE iteration for a quick check
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))What to Look For
─────────────────────────────────────────
GPU utilization consistently LOW: likely a DATA LOADING
bottleneck — increase
num_workers (Ch.1,
Sec.3), or check for
slow preprocessing
A SPECIFIC layer/operation consider whether
dominating GPU time: THAT specific
operation can
be optimized
or replaced
Memory usage near the GPU's reduce batch
LIMIT: size, use
gradient
accumulation
(Section 4)
or mixed
precision
(Section 3)
─────────────────────────────────────────
7. Summary & Next Steps
Key Takeaways
- Mixed precision training uses FP16 for most operations (faster, less memory) while selectively preserving FP32 precision where numerically necessary, commonly providing a 2-3x speedup.
- A GradScaler prevents small gradients from rounding to zero under FP16's limited precision by temporarily scaling the loss up before backpropagation.
- Gradient accumulation simulates a larger batch size than fits in memory by accumulating gradients across several smaller micro-batches before each weight update.
- Always profile before optimizing — a slow data pipeline is a more common real-world bottleneck than raw GPU compute, and different problems require different fixes.
Concept Check
- Why does mixed precision training require a GradScaler rather than simply switching all computations to FP16?
- If you want to simulate a batch size of 128 but only 32 fits in GPU memory, how would you configure gradient accumulation?
- Why should you profile a training pipeline before applying optimizations like mixed precision?
Next Chapter
→ Chapter 3: Debugging & Experiment Tracking
Jr Codex — 1-on-1 Personalized Coaching | Back to Module Index | Back to DL Index