Lesson 44: Precision and Parallel Training¶
Every network in Lessons 31-43 trained on a small, CPU-friendly toy problem in default 32-bit floating point (fp32), on a single process. Real training runs rarely look like that: models are trained in reduced precision to go faster and use less memory, and split across many GPUs because no single device holds a modern model plus its training data. This lesson closes out Part 3 with the engineering underneath how training actually happens at scale, building each mechanism in miniature and checking it against the real thing: floating-point precision and why reduced precision needs help to work at all, mixed-precision training's specific fix for that, and the two ways of splitting work across devices, data parallelism and model parallelism.
import torch
import torch.nn as nn
import torch.nn.functional as F
import copy
import numpy as np
import matplotlib.pyplot as plt
Floating-point precision: fp32 vs. fp16¶
A 32-bit float (fp32) and a 16-bit float (fp16) both use the same idea -- sign, exponent, mantissa -- but fp16 spends far fewer bits on the exponent, so it simply cannot represent very small or very large magnitudes. torch.finfo reports the smallest positive normal value each format can hold: about 1.2e-38 for fp32, but only 6.1e-5 for fp16 (fp16 subnormals stretch a little further, down to about 6e-8, at the cost of losing precision entirely). Below that floor, a value doesn't round to something small -- it rounds to exactly zero.
small_values = torch.tensor([1e-2, 1e-3, 1e-4, 1e-5, 1e-6, 1e-7, 1e-8, 1e-9], dtype=torch.float32)
small_fp16 = small_values.half()
print(f'{"fp32 value":>12} {"cast to fp16":>14} {"underflowed to 0?":>19}')
for v32, v16 in zip(small_values, small_fp16):
print(f'{v32.item():>12.1e} {v16.item():>14.1e} {str(v16.item() == 0):>19}')
Values down to 1e-6 survive the cast intact; 1e-7 survives but loses precision (rounds to 1.2e-7); 1e-8 and smaller vanish to exactly 0.0. That's not a rounding error in the usual sense -- it's total information loss, and it matters enormously for training because gradients routinely take on exactly these magnitudes, especially in deep networks. Lessons 33 and 37 built and studied the vanishing-gradient problem in fp32, where early-layer gradients shrink but stay nonzero. In fp16, "shrink" can mean "disappear."
class DeepSigmoidNet(nn.Module):
def __init__(self, depth=12, width=16, in_dim=4):
super().__init__()
layers = []
d = in_dim
for _ in range(depth):
layers += [nn.Linear(d, width), nn.Sigmoid()]
d = width
layers += [nn.Linear(d, 1)]
self.net = nn.Sequential(*layers)
def forward(self, x):
return self.net(x)
torch.manual_seed(0)
model = DeepSigmoidNet(depth=12)
X = torch.randn(32, 4)
y = torch.randn(32, 1)
model.zero_grad()
loss = F.mse_loss(model(X), y)
loss.backward()
first_layer_grad = model.net[0].weight.grad
print(f'first-layer gradient magnitude: mean={first_layer_grad.abs().mean():.2e}, max={first_layer_grad.abs().max():.2e}')
print(f'fraction exactly zero in fp32: {(first_layer_grad == 0).float().mean():.1%}')
print(f'fraction exactly zero after casting to fp16: {(first_layer_grad.half() == 0).float().mean():.1%}')
A 12-layer sigmoid network (deliberately deep and derivative-shrinking, the same vanishing-gradient setup Lesson 37 diagnosed) produces first-layer gradients around 1e-11 -- comfortably nonzero in fp32, but 100% zero the instant they're cast to fp16. Train this network directly in fp16 and the first layer would never update at all; the optimizer would faithfully apply a learning rate to a gradient of exactly zero, forever. This is the actual problem mixed-precision training has to solve, not a hypothetical one.
Mixed-precision training: the recipe¶
Training entirely in fp16 hits the wall demonstrated above; training entirely in fp32 gives up the speed and memory savings that were the whole point. Mixed-precision training (Micikevicius et al., 2017) keeps both: the forward and backward pass run in fp16 for speed, but two safeguards protect against exactly the failure shown above:
- fp32 master weights: a full-precision copy of every weight is kept off to the side. Updates apply to the fp32 copy (where small updates don't get rounded away), and an fp16 copy is re-derived from it before each forward pass.
- Loss scaling: multiply the loss by a large constant before calling
.backward(). Since gradients are linear in the loss, every gradient scales up by the same factor -- pushing values that would have underflowed back into fp16's representable range. The update step divides the gradient back down by that same factor before it's applied, so the math is otherwise unaffected.
torch.manual_seed(0)
model = DeepSigmoidNet(depth=12)
X = torch.randn(32, 4)
y = torch.randn(32, 1)
print(f'{"loss scale":>12} {"fraction of fp16 grads still zero":>34}')
for scale in [1.0, 65536.0, 2**20, 2**24, 2**28]:
model.zero_grad()
loss = F.mse_loss(model(X), y)
(loss * scale).backward() # every gradient in the graph is scaled by `scale`, exactly
grad_fp16 = model.net[0].weight.grad.half()
print(f'{scale:>12.0f} {float((grad_fp16 == 0).float().mean()):>33.1%}')
With no scaling (scale=1), every first-layer gradient underflows, exactly as before. Scaling by 65536 (2**16, a common real-world default) already rescues most of them, but over a third are still small enough to vanish. Only at 2**28 does the last gradient clear fp16's floor. This is precisely why real implementations (PyTorch's torch.cuda.amp.GradScaler) use dynamic loss scaling rather than a fixed constant: start with some scale, and increase it further whenever a step completes without producing an inf/NaN (a sign the scale could safely be larger), or back off whenever one appears (a sign the scale overflowed something). One more detail worth naming: bfloat16 (bf16), an increasingly common alternative to fp16, sidesteps loss scaling entirely -- torch.finfo shows it has the same exponent range as fp32 (smallest_normal around 1.2e-38 for both), so it can't underflow the way fp16 does. What it gives up instead is precision (eps around 0.008 for bf16 versus 0.001 for fp16) -- a different trade-off along the same speed-versus-representable-range axis, not a free lunch.
Data parallelism: correctness proven in miniature¶
Precision controls how much work one device does per step; parallelism controls how many devices share that work. Data parallelism (the standard recipe behind large-batch training, e.g. Goyal et al., 2017) splits a batch into shards, runs an identical copy of the model on each shard on its own device, and averages the resulting gradients before the update step. On one CPU there's no speedup to show, but the correctness of that recipe -- that averaging per-shard gradients gives the same answer as computing the gradient on the whole batch at once -- is exactly checkable, because MSE loss is itself an average, and the gradient of an average of per-shard averages is the average of the per-shard gradients.
torch.manual_seed(1)
model = nn.Sequential(nn.Linear(8, 16), nn.ReLU(), nn.Linear(16, 1))
X = torch.randn(64, 8)
y = torch.randn(64, 1)
# "one device": gradient over the whole batch at once
model.zero_grad()
F.mse_loss(model(X), y).backward()
grad_single_device = [p.grad.clone() for p in model.parameters()]
# "data parallel, 4 shards": an identical model copy per shard, each computing its own
# gradient independently; only the gradients get communicated and averaged, never the data
n_shards = 4
shard_size = len(X) // n_shards
per_shard_grads = []
for s in range(n_shards):
shard_model = copy.deepcopy(model)
shard_model.zero_grad()
xs, ys = X[s * shard_size:(s + 1) * shard_size], y[s * shard_size:(s + 1) * shard_size]
F.mse_loss(shard_model(xs), ys).backward()
per_shard_grads.append([p.grad.clone() for p in shard_model.parameters()])
averaged_grad = [sum(shard[i] for shard in per_shard_grads) / n_shards for i in range(len(grad_single_device))]
print(f'{"parameter":>10} {"max abs diff: single-device vs. averaged-shard grad":>52}')
for i, (g_single, g_avg) in enumerate(zip(grad_single_device, averaged_grad)):
print(f'{i:>10} {(g_single - g_avg).abs().max().item():>52.2e}')
Every parameter's averaged-shard gradient matches the single-device gradient to within floating-point roundoff (~1e-8, not a real discrepancy). This is the whole correctness argument for data parallelism in one demonstration: splitting the batch and averaging afterward is mathematically the same computation as never splitting it at all. What this toy version doesn't -- can't -- show is the actual cost real data-parallel training pays: the all-reduce, the communication step that sums/averages gradients across every device before any of them can proceed to the update. At small scale that communication is nearly free; at hundreds of GPUs it can dominate the step time, which is why large-scale training spends real engineering effort on overlapping communication with computation rather than doing them one after another.
Model parallelism: shipping activations across a boundary¶
Data parallelism copies the whole model onto every device and splits the data. Model parallelism (Shoeybi et al., 2019) does the opposite: split the model itself across devices, because it's too large to fit on one. In the simplest form (pipeline parallelism), each device holds a contiguous block of layers; the forward pass computes one block's output, hands that activation across the device boundary, and the next device picks up from there. The backward pass runs the same handoff in reverse: each device needs only the gradient of the loss with respect to the activation it sent, not anything about what happens on the other side of the boundary.
class TwoStageNet(nn.Module):
def __init__(self):
super().__init__()
self.stage1 = nn.Sequential(nn.Linear(8, 16), nn.ReLU(), nn.Linear(16, 16), nn.ReLU())
self.stage2 = nn.Sequential(nn.Linear(16, 16), nn.ReLU(), nn.Linear(16, 1))
def forward(self, x):
return self.stage2(self.stage1(x))
torch.manual_seed(2)
whole_model = TwoStageNet()
X = torch.randn(8, 8)
y = torch.randn(8, 1)
# "one device": run the whole model, one backward pass, end to end
F.mse_loss(whole_model(X), y).backward()
grad_whole = [p.grad.clone() for p in whole_model.parameters()]
# "two devices": stage1 on device A, stage2 on device B, with the activation manually
# shipped across the boundary in each direction
torch.manual_seed(2)
split_model = TwoStageNet()
split_model.zero_grad()
activation_out = split_model.stage1(X) # entirely on "device A"
activation_in = activation_out.detach().requires_grad_(True) # crosses the boundary
prediction = split_model.stage2(activation_in) # entirely on "device B"
F.mse_loss(prediction, y).backward() # backprop stops at the boundary
activation_grad = activation_in.grad # shipped back across the boundary
activation_out.backward(activation_grad) # resumes backprop into stage1
grad_split = [p.grad.clone() for p in split_model.parameters()]
print(f'{"parameter":>10} {"max abs diff: whole model vs. split across 2 stages":>52}')
for i, (g_whole, g_split) in enumerate(zip(grad_whole, grad_split)):
print(f'{i:>10} {(g_whole - g_split).abs().max().item():>52.2e}')
Every gradient matches exactly -- 0.0 difference, not just roundoff, because the split-and-reattach mechanism (detach() at the boundary going forward, then manually resuming .backward() with the shipped-back gradient) is mathematically identical to autograd tracing straight through an unsplit model; it just makes the "device boundary" explicit as a place activations get detached and gradients get reattached by hand instead of automatically. Real pipeline parallelism adds one more concern this two-device, one-batch toy sidesteps entirely: while device B works on the tail of the pipeline for one batch, device A sits idle unless a second batch is already flowing in behind it -- the "pipeline bubble" that real implementations fight with careful scheduling (splitting each batch into smaller microbatches so every stage stays busy). Data parallelism's bottleneck is communication (the all-reduce above); model parallelism's is idle time.
Two more practical tricks¶
Two more techniques round out the practical picture without needing their own full demonstration, because both are natural extensions of mechanisms already built above:
- Gradient accumulation simulates a larger batch than fits in memory by running several small forward/backward passes without an optimizer step in between, summing their gradients, and only then updating -- mathematically the same averaging idea as the data-parallel shards above, just spread across time on one device instead of across devices at once.
- Gradient checkpointing trades compute for memory: instead of storing every intermediate activation for the backward pass (what autograd normally does), it stores only a handful of checkpoints and recomputes the activations in between during the backward pass. More forward-pass compute, much less memory -- useful when a model's activations, not its weights, are what doesn't fit.
Exercises¶
- In the fp16 underflow demo, change
DeepSigmoidNet'sdepthfrom 12 to 4. Does the first-layer gradient still underflow completely in fp16? Relate the answer back to Lesson 37's discussion of why depth and sigmoid activations together cause vanishing gradients. - The loss-scaling table stops at
scale=2**28. Keep increasing the scale (try2**32,2**40) -- at what point does the scaled gradient itself start to risk overflowing fp16'smaxof65504instead of underflowing? What does that imply about why dynamic loss scaling needs to search for a scale, rather than always using the largest one possible? - Change
n_shardsin the data-parallelism demo from 4 to 3, with a batch size that doesn't divide evenly (try 65). Does the averaged-shard gradient still match the single-device gradient exactly, or does the mismatch reveal something about how real data-parallel training handles uneven shard sizes? - Add a third stage to
TwoStageNet(threenn.Sequentialblocks instead of two) and extend the manual forward/backward handoff to three devices. Does the gradient still match a whole-model backward pass exactly?