Mixed Precision, Gradient Accumulation, and Distributed Trai

Efficient PyTorch training is usually a memory and communication problem before it is an algorithm problem. Mixed precision reduces the bytes moved through accelerators, gradient accumulation simulates a larger batch when memory is limited, and distributed training splits work across processes. Used together, they let you train larger models, keep hardware busier, and preserve the optimizer behavior you intended.

The outcome of this lesson is practical: you should be able to read a training loop and know which tensors are low precision, when gradients are synchronized, how many examples contribute to one optimizer update, and why a run becomes unstable or slow. This connects directly to the course because efficient training changes experiment design: a model that is impossible to train on one device may become feasible, but only if the numerical and synchronization details remain correct.

How the Mechanisms Fit Together

Mixed precision means some operations run in a smaller floating-point format, usually float16 or bfloat16, while sensitive state remains in float32. Model weights may be stored as float32, matrix multiplications may execute in lower precision, and the optimizer updates float32 parameters. PyTorch automatic mixed precision, often called AMP, uses autocast to choose lower precision for operations that usually tolerate it and float32 for operations that need more range or accuracy.

Float16 has less exponent range than float32, so small gradients can underflow to zero. GradScaler addresses this by multiplying the loss before backpropagation, creating larger intermediate gradients. Before the optimizer step, the scaler unscales gradients back to the correct magnitude, checks for infinities or NaNs, and skips the step if the update is unsafe. Bfloat16 has a wider exponent range, so explicit scaling is often less important, but autocast remains useful because it selects efficient kernels.

Gradient accumulation separates the microbatch from the effective batch. A microbatch is the set of samples that fits in memory for one forward and backward pass. The effective batch is the total number of samples whose gradients are combined before one optimizer step. If the loss reduction is a mean, each microbatch loss must usually be divided by the accumulation count so that the accumulated gradient matches one large batch rather than becoming too large.

DistributedDataParallel, or DDP, runs one Python process per device in the common setup. Each process owns a full model replica and a different shard of the input batch. During backward, DDP groups parameter gradients into buckets and all-reduces them across processes, so every replica receives the averaged gradient before stepping. The optimizer then runs locally but deterministically on equivalent parameters. This is data parallelism, not model parallelism: the model is copied, while the data is split.

API Anatomy

A typical efficient training loop has five moving parts. The model and optimizer are normal PyTorch objects. The dataloader uses a DistributedSampler in distributed runs so ranks do not train on the same samples. The forward pass sits inside torch.autocast. The loss is divided by accumulation_steps. Backward uses either scaler.scale(loss).backward() for scaled AMP or loss.backward() for regular precision. The optimizer step occurs only after the requested number of microbatches.

DDP adds process configuration: RANK identifies a process, WORLD_SIZE is the total number of processes, and LOCAL_RANK maps a process to a device on one node. Launchers such as torchrun populate these environment variables. In a multi-process run, you initialize the process group, move the model to the rank’s device, wrap it in DistributedDataParallel, train, then destroy the process group.

Example 1: Autocast Around the Forward Pass

This example shows the boundary of mixed precision. The model definition is unchanged. Only the forward pass and loss calculation are placed inside autocast. On a CUDA device, eligible operations can use float16 kernels. On CPU in this example, autocast is disabled so the deterministic output remains float32.

import torch
from torch import nn

torch.manual_seed(7)
device_type = "cuda" if torch.cuda.is_available() else "cpu"
model = nn.Linear(3, 1)
x = torch.tensor([[1.0, 2.0, 3.0], [2.0, 0.0, 1.0]])
y = torch.tensor([[1.0], [0.0]])

with torch.autocast(device_type=device_type, dtype=torch.float16, enabled=device_type == "cuda"):
    pred = model(x)
    loss = nn.functional.mse_loss(pred, y)

print(f"device_type={device_type}")
print(f"prediction_dtype={pred.dtype}")
print(f"loss_is_finite={torch.isfinite(loss).item()}")

Expected CPU behavior is device_type=cpu, prediction_dtype=torch.float32, and loss_is_finite=True. On CUDA, the prediction dtype may be float16 for eligible kernels while parameters remain float32. The important design point is that autocast is a scoped policy, not a permanent conversion of the whole model.

Example 2: Accumulation Matches a Larger Batch

The next example compares one full batch with two accumulated microbatches. Dividing each microbatch loss by two keeps the final gradient equivalent to the full batch gradient when the loss uses mean reduction.

import copy
import torch
from torch import nn

torch.manual_seed(3)
base = nn.Linear(2, 1)
full_batch_model = copy.deepcopy(base)
accum_model = copy.deepcopy(base)

x = torch.tensor([[1.0, 0.0], [0.0, 1.0], [2.0, 1.0], [1.0, 2.0]])
y = torch.tensor([[1.0], [0.5], [2.0], [1.5]])

full_loss = nn.functional.mse_loss(full_batch_model(x), y)
full_loss.backward()

for start in range(0, 4, 2):
    micro_x = x[start:start + 2]
    micro_y = y[start:start + 2]
    micro_loss = nn.functional.mse_loss(accum_model(micro_x), micro_y) / 2
    micro_loss.backward()

max_grad_diff = max(
    (left.grad - right.grad).abs().max().item()
    for left, right in zip(full_batch_model.parameters(), accum_model.parameters())
)
print(f"max_grad_diff={max_grad_diff:.8f}")

The expected output is max_grad_diff=0.00000000, aside from tiny floating-point variation on some hardware. If you remove / 2, the accumulated gradient becomes twice as large and the optimizer behaves as if the learning rate changed. This is one of the most common accumulation bugs.

Example 3: DDP Process Shape

This small script runs as a single process by default and becomes distributed when launched with a multi-process launcher. It illustrates the control flow without hiding the fact that DDP is process based.

import os
import torch
import torch.distributed as dist
from torch import nn


def main() -> None:
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    rank = int(os.environ.get("RANK", "0"))

    if world_size > 1:
        dist.init_process_group(backend="gloo")

    model = nn.Linear(2, 1)
    wrapped = model
    if world_size > 1:
        wrapped = nn.parallel.DistributedDataParallel(model)

    print(f"rank={rank} world_size={world_size} wrapper={type(wrapped).__name__}")

    if world_size > 1:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()

Run normally, the deterministic output is rank=0 world_size=1 wrapper=Linear. Launched with multiple ranks, each process prints its own rank and the wrapper becomes DistributedDataParallel. In GPU training you would also set the CUDA device from LOCAL_RANK and move the model and batches to that device before wrapping.

Example 4: AMP and Accumulation Together

This training step combines the two local techniques. The scaler touches backward and step, while accumulation controls when the optimizer is allowed to update parameters.

import torch
from torch import nn

torch.manual_seed(11)
device_type = "cuda" if torch.cuda.is_available() else "cpu"
model = nn.Linear(4, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
scaler = torch.cuda.amp.GradScaler(enabled=device_type == "cuda")
accumulation_steps = 2
optimizer_steps = 0

for microbatch in range(4):
    x = torch.ones(3, 4) * (microbatch + 1)
    y = torch.ones(3, 1)
    with torch.autocast(device_type=device_type, dtype=torch.float16, enabled=device_type == "cuda"):
        loss = nn.functional.mse_loss(model(x), y) / accumulation_steps
    scaler.scale(loss).backward()

    if (microbatch + 1) % accumulation_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad(set_to_none=True)
        optimizer_steps += 1

print(f"optimizer_steps={optimizer_steps}")

The expected output is optimizer_steps=2. Four microbatches produced two optimizer updates because the effective batch is two microbatches wide. In a DDP loop, you can also use model.no_sync() on non-step microbatches to avoid unnecessary all-reduce work, then allow synchronization on the final microbatch before the optimizer step.

Design Choices and Trade-offs

Choose float16 AMP when accelerator tensor cores strongly favor it and your model tolerates the smaller exponent range. Choose bfloat16 when supported and when gradient range matters more than mantissa precision. Keep normalization, reductions, and loss computation under PyTorch’s autocast policy unless you have measured a reason to override it.

Choose accumulation when the desired effective batch does not fit in memory. It saves activation memory compared with one large batch, but it does not reduce total compute. It can also change training dynamics if you forget to adjust learning rate schedules, gradient clipping, or logging intervals around optimizer steps rather than microbatches.

Choose DDP when one process cannot provide enough throughput. DDP usually scales well because each rank performs forward and backward locally, then communicates gradients. Its costs are memory replication, communication bandwidth, launcher complexity, and debugging complexity. Small models or slow input pipelines may get slower when distributed because communication and process overhead dominate.

Failure Modes and Troubleshooting

NaN loss after enabling AMP. Symptom: loss becomes NaN or the scaler repeatedly skips optimizer steps. Cause: overflow in lower precision, unstable loss, or gradients that are too large before unscale. Diagnose by printing scaler scale changes, checking torch.isfinite(loss), and temporarily disabling autocast around the suspected operation. Correct by lowering the learning rate, clipping after unscale, using bfloat16 if available, or keeping the unstable operation in float32.

Accumulation trains differently from the same large batch. Symptom: loss curves diverge after switching to microbatches. Cause: missing loss division, stepping the scheduler every microbatch, clipping before all gradients are accumulated, or batch normalization seeing smaller per-rank microbatch statistics. Diagnose by comparing gradients with a tiny deterministic case like Example 2. Correct the loss scale, step the scheduler with the optimizer, clip immediately before the step, and reconsider normalization for very small microbatches.

DDP hangs during backward. Symptom: all ranks stop with no Python exception, often at the first backward pass. Cause: ranks executed different collective operations, one rank crashed earlier, unused parameters differ across ranks, or the dataloader produced uneven work. Diagnose by checking every rank’s last log line, setting a distributed debug level, verifying sampler lengths, and running with one batch. Correct rank-dependent branches, use a distributed sampler, call set_epoch each epoch, and configure unused-parameter handling only when the model genuinely has conditional parameters.

Distributed speedup is poor. Symptom: more devices do not reduce epoch time. Cause: tiny batches, CPU dataloader bottlenecks, excessive synchronization, slow storage, or gradients too small to amortize all-reduce overhead. Diagnose dataloader time separately from step time and profile communication. Correct with larger per-rank batches, pinned memory, more dataloader workers, fewer unnecessary metric reductions, and no_sync during accumulation.

Reliability and Performance Implications

These techniques improve performance only when the measurement target is clear. Track samples per second, optimizer steps per second, validation quality, peak memory, skipped AMP steps, and communication time. Save checkpoints from a single authoritative rank and include model, optimizer, scheduler, scaler, epoch, global step, and sampler state where applicable. Without scaler and optimizer state, resuming can silently change the next updates.

Reliability also depends on identical control flow across ranks. Validation, logging, and checkpointing can be rank-specific, but collectives must be called by all participating ranks in the same order. Seed initialization should account for rank so data augmentation differs where intended, while model initialization must be synchronized or performed before DDP guarantees parameter consistency.

Hands-on Lab

Prerequisites: a Python environment with PyTorch installed. A CUDA GPU is useful for observing AMP speed and memory effects, but the structure can be verified on CPU. For distributed verification on one machine, use two CPU ranks with the gloo backend or two GPU ranks with the appropriate backend for your hardware.

Create a small classifier or regression script with a deterministic dataset, one model, one optimizer, and a baseline full-precision training loop. Record final loss after a few optimizer steps.
Add autocast and GradScaler. Verify that loss remains finite and that validation quality is close to the full-precision baseline.
Add accumulation_steps=2. Divide each microbatch loss by two and print both microbatch count and optimizer step count. Verify that the optimizer steps half as often as microbatches.
Convert the script for DDP. Initialize the process group from environment variables, wrap the model, use DistributedSampler, and save checkpoints only on rank zero.
Launch a single-rank run, then a two-rank run. Verify that each rank receives different samples, all ranks complete the same number of optimizer steps, and validation is computed once or reduced deliberately.
Cleanup: destroy the process group at process exit, delete temporary checkpoints, and restore the baseline script configuration if you are comparing experiment results.

Verification is successful when the run finishes without hangs, the number of optimizer updates matches the configured effective batch, no AMP step is skipped indefinitely, and a resumed checkpoint continues from the same global step with finite loss.

Assessment Exercises

A loop accumulates four microbatches but does not divide the loss. Explain how this changes the gradient and give two possible symptoms in the training curve.
You enable DDP and see every rank training on the same examples. Which component is probably missing, and how would you verify the fix?
A model is faster with AMP for large batches but slower for small batches. Explain why kernel launch overhead, tensor core utilization, and dataloader time might produce that result.
Design a checkpoint dictionary for AMP plus DDP training. Which state must be saved to resume without changing optimizer behavior?
In an accumulation loop with DDP, where would you place gradient clipping and why?

Summary

Mixed precision changes numeric execution, gradient accumulation changes the relationship between microbatches and optimizer updates, and DDP changes how gradients are synchronized across processes. The techniques are powerful because they attack different limits: arithmetic throughput, memory capacity, and single-device throughput. The cost is that loss scaling, update frequency, sampler behavior, and collective communication must be explicit. Treat the training loop as a precise mechanism, measure it with small deterministic examples first, then scale only after the numerical and distributed behavior is understood.

Mixed Precision, Gradient Accumulation, and Distributed Training