GPUs, Mixed Precision & Reproducibility

"Run me twice and get two different answers, and you do not have a result; you have a rumor. Seed me properly and the rumor becomes evidence."

A Random Number Generator With a Strong Sense of Accountability

The training loop of the previous section is correct but, on a CPU, slow. Deep learning is fast because the GPU executes the same matrix operations on thousands of cores at once, and the convolutions of Chapter 19 only become practical there. Three concerns dominate once you leave the toy regime: getting the computation onto the GPU at all, getting more of it per second and per gigabyte through reduced precision, and getting the same result twice so a number means something. None is conceptually deep, but each has a sharp edge that bites the unprepared, and getting all three right is what makes the rest of Part III reproducible on your own hardware. The same matrix-multiply hardware that accelerates a network was already lurking under the optimized filtering of Chapter 3; here we drive it deliberately, and Chapter 21 folds these knobs into a full training recipe.

1. Device Placement Beginner

A tensor lives on a device, cpu or cuda (an NVIDIA GPU; Apple silicon exposes mps). An operation requires all its inputs on the same device, and the model's parameters and the input batch must agree, or PyTorch raises a device-mismatch error. The pattern is to detect the device once, move the model there at the start, and move each batch there inside the loop, exactly the two .to(device) calls already present in Section 18.5. The code shows the canonical detection-and-placement idiom.

# The device idiom: detect the best available accelerator once, move the model's
# parameters and each input batch there so the matrix multiplies run on the GPU,
# and pull only scalars back to the CPU for logging.
import torch

# detect the best available device once, at program start
device = (
    "cuda" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()
    else "cpu"
)
print("training on:", device)

model = torch.nn.Linear(784, 10).to(device)     # move ALL parameters to the device

x = torch.randn(64, 784).to(device)             # the batch must live there too
out = model(x)                                   # both on device -> runs on GPU
print(out.device)                                # cuda:0  (or mps / cpu)

# moving a result back to the CPU for logging or numpy:
acc_value = out.argmax(1).float().mean().cpu().item()

The most common beginner error is forgetting to move the batch, which gives a clear "expected all tensors to be on the same device" message. The second is unnecessary transfers: moving tensors back and forth between CPU and GPU inside the inner loop is slow because the transfer crosses the PCIe bus. Keep data on the GPU for the whole forward and backward pass and only pull scalars (losses, metrics) back to the CPU, with .item(), for logging.

2. Mixed Precision: Speed and Memory for Almost Free Intermediate

By default tensors are 32-bit floats (FP32). Modern GPUs compute far faster in 16-bit (FP16 or the more robust BF16) and store half the bytes, so using 16-bit where it is safe roughly doubles throughput and halves activation memory. Automatic mixed precision (AMP) does this selectively: it runs the matrix multiplies and convolutions in 16-bit (where the speedup lives and the precision loss is tolerable) while keeping numerically sensitive operations like the loss reduction and normalization statistics in FP32.

The one wrinkle is that FP16 has a narrow range, and small gradients can underflow to zero, silently stalling learning. The fix is loss scaling: multiply the loss by a large factor before backward so the gradients land in FP16's representable range, then unscale them before the optimizer step. This is exact rather than a hack, because the gradient is linear in the loss. Scaling the loss by a constant $s$ scales every gradient by the same $s$ (the chain rule just carries the constant through), so dividing the gradients back by $s$ before the step recovers precisely the true update. The only effect is that the in-between numbers were large enough for FP16 to represent. PyTorch automates the entire dance with two objects.

# Automatic mixed precision in one training step: autocast runs the heavy
# matrix work in FP16 while keeping the loss reduction in FP32, and GradScaler
# applies dynamic loss scaling so small FP16 gradients do not underflow to zero.
import torch

device = "cuda"
model = torch.nn.Linear(784, 10).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
scaler = torch.amp.GradScaler(device)          # manages dynamic loss scaling

x = torch.randn(64, 784, device=device)
y = torch.randint(0, 10, (64,), device=device)

optimizer.zero_grad()
with torch.amp.autocast(device_type=device, dtype=torch.float16):
    logits = model(x)                          # matmuls run in FP16
    loss = torch.nn.functional.cross_entropy(logits, y)   # reduction stays FP32

scaler.scale(loss).backward()                  # scale loss -> gradients in FP16 range
scaler.step(optimizer)                         # unscale, then optimizer.step()
scaler.update()                                # adapt the scale factor for next step
print(round(loss.item(), 4))

Folding these four lines into the Section 18.5 loop is the entire change needed to train in mixed precision: wrap the forward pass in autocast, and route the backward and step through the scaler. On a recent GPU the speedup is often 1.5 to 3 times with no measurable accuracy loss, and the halved memory lets you fit larger batches or models. BF16, available on newer hardware, has the same range as FP32 and so does not need the scaler at all, which is why it is increasingly the default for training large models. Figure 18.6.1 shows where the precisions differ.

3. Reproducibility: Making a Number Mean Something Intermediate

Randomness enters training from many sources: weight initialization, data shuffling in the DataLoader (Section 18.4), dropout masks, and augmentation. Run the same script twice without controlling these and you get two different results, which makes it impossible to tell whether a change you made helped or whether you just got luckier on the second run. Reproducibility means pinning every random source so the run is repeatable. The minimum is seeding Python's random, NumPy, and PyTorch (CPU and CUDA); full determinism additionally requires telling cuDNN to use deterministic algorithms and disabling its autotuner. The code packages this into a reusable utility.

# A reusable seeding utility: set_seed pins every random source (Python, NumPy,
# PyTorch CPU and CUDA) and optionally forces deterministic algorithms, while
# seed_worker gives each DataLoader worker its own repeatable seed.
import os, random
import numpy as np
import torch

def set_seed(seed=42, deterministic=True):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)                 # seeds CPU and all CUDA devices
    if deterministic:
        torch.use_deterministic_algorithms(True)   # error on nondeterministic ops
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False      # disable autotuner (it varies)
        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"   # needed for determinism

def seed_worker(worker_id):
    # give each DataLoader worker a derived, deterministic seed
    s = torch.initial_seed() % 2**32
    np.random.seed(s); random.seed(s)

set_seed(42)
g = torch.Generator(); g.manual_seed(42)
# pass worker_init_fn=seed_worker and generator=g to the training DataLoader
print("seeded; runs are now repeatable")

Two honest caveats. Full determinism can be slightly slower (the cuDNN autotuner, disabled above, normally picks the fastest convolution algorithm per shape) and a handful of GPU operations have no deterministic implementation, in which case use_deterministic_algorithms(True) will tell you by raising an error rather than failing silently. The pragmatic stance for research is to always seed (so a result is repeatable on the same machine) and to report results across several seeds (so you can distinguish a real improvement from seed luck), the very practice the experiment-tracking culture of modern vision papers expects. Reproducibility across different GPUs or library versions is harder still and rarely bit-exact; the goal is repeatability on your own setup plus seed-averaged honesty in what you report.