PyTorch Mixed Precision and GPU Training

Memory limits the configurations that fit; throughput limits how much training fits in a budget. AMP can reduce memory and improve throughput when low-precision operations suit the workload. Accumulation fits a larger effective batch into smaller forwards, while activation checkpointing recomputes selected values to reduce saved activations. Their benefits depend on the bottleneck. Run the body blocks in order: the AMP example uses CUDA when available and otherwise runs ordinary float32 on CPU. The CPU path checks control flow, not CUDA AMP behavior or speed.

Where the memory goes

For \(P\) trainable float32 parameters with dense gradients and two float32 Adam moment buffers, these tensors total about \(16P\) bytes: 4 for parameters, 4 for gradients, and 8 for moments. One billion parameters give 16 decimal GB, about 14.9 GiB. This is an accounting estimate when those buffers exist, not an allocation guaranteed before the first forward: Adam usually creates state on its first update, and gradients can be released between updates. It excludes activations, temporary workspaces, other buffers, and allocator overhead.

Saved activations depend on the operations and tensor shapes. Increasing batch size, depth, or sequence length can increase them, but the scaling is not universally linear: materialized attention scores, for example, can grow quadratically with sequence length. Standard autocast leaves float32 model parameters unchanged and reduces precision for eligible operations. Activation checkpointing reduces selected saved intermediates. Neither automatically halves the total training footprint.

import torch

def memory_report():
    if not torch.cuda.is_available():
        return None
    return {"allocated_GiB": torch.cuda.memory_allocated() / 2**30,
            "reserved_GiB": torch.cuda.memory_reserved() / 2**30,
            "peak_allocated_GiB": torch.cuda.max_memory_allocated() / 2**30}

if torch.cuda.is_available():
    torch.cuda.reset_peak_memory_stats()
report = memory_report()

report is a CUDA allocator snapshot, or None without CUDA. Allocated memory tracks live tensor allocations; reserved memory includes allocator-managed blocks available for reuse. The peak is the maximum allocated value since the reset. Reset before the interval being compared, and inspect the report after it. These counters are not the same as total device usage reported by an external tool.

The training loop this code slots into is built in Building a PyTorch Training Loop with nn.Module, Dataset, and DataLoader.

Automatic mixed precision

Autocast selects dtypes per eligible operation and backend. Matrix multiplication often uses a lower precision, while some numerically sensitive operations use float32 or promote their inputs. “All normalization and reductions stay float32” is not a universal rule. Tensor shapes, kernels, hardware, and input-loading cost determine the speedup; compare measured throughput and peak memory against a float32 baseline while also checking validation quality.

Float16 has limited dynamic range, so sufficiently small gradients can round to zero and large values can overflow. GradScaler multiplies the loss by a scale before backward and unscales the resulting gradients before the update. This helps preserve small gradients; it cannot recover values already lost in forward computation or guarantee finite gradients. For a given optimizer it skips an update when its gradients are nonfinite, then adjusts the scale. Per-update schedulers should not advance for that skipped update.

import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader

torch.manual_seed(0)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
use_amp = device.type == "cuda"
model = nn.Linear(8, 3).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.9)
criterion = nn.CrossEntropyLoss()
loader = DataLoader(TensorDataset(torch.randn(10, 8), torch.randint(0, 3, (10,))),
                    batch_size=3, shuffle=False)
if hasattr(torch.amp, "GradScaler"):
    scaler = torch.amp.GradScaler("cuda", enabled=use_amp)
else:
    scaler = torch.cuda.amp.GradScaler(enabled=use_amp)

def train_step(xb, yb):
    model.train()
    xb, yb = xb.to(device), yb.to(device)
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=use_amp):
        loss = criterion(model(xb), yb)
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    old_scale = scaler.get_scale()
    scaler.step(optimizer)
    scaler.update()
    updated = scaler.get_scale() >= old_scale
    if updated:
        scheduler.step()
    return loss.detach(), updated

last_loss, updated = train_step(*next(iter(loader)))
print("finite loss", bool(torch.isfinite(last_loss).item()))
# finite loss True

Current PyTorch exposes torch.amp.GradScaler("cuda"). The code detects that API and falls back to torch.cuda.amp.GradScaler for older installations. torch.autocast controls the forward region; parameters are left in float32. On CPU, both autocasting and scaling are disabled here. The tiny synthetic batch demonstrates the update order, not useful predictive performance.

Run the forward and loss inside autocast, then backward outside it. Before clipping, unscale once so that the threshold refers to the actual gradient norm. The example has one optimizer: a decrease in scale after update() indicates its update was skipped, so the scheduler advances only when the scale did not decrease. This convention needs reconsideration for multiple optimizers or a custom scaler. Clipping limits gradients; it does not by itself bound AdamW’s final parameter update.

Bfloat16 has float32’s exponent width and a much wider range than float16, but fewer fraction bits. It is less prone to underflow and usually does not need loss scaling; it can still underflow, overflow, or lose important precision. On CUDA, check support with torch.cuda.is_bf16_supported() and validate the model before choosing it. In PyTorch 2.10, this check includes emulation by default; torch.cuda.is_bf16_supported(including_emulation=False) excludes it. A successful check or a completed bfloat16 operation does not establish native acceleration, training quality, or a speed advantage. See the PyTorch support-check API.

Gradient accumulation

Gradient accumulation holds parameters fixed across several micro-batches, adds their gradients, and updates once. For an example-wise mean objective, each example must receive the same weight. Dividing every micro-batch mean by a fixed count is correct only when the count and batch sizes match the intended window. A short final window needs an update too.

import torch.nn.functional as F

def accumulated_epoch(model, loader, optimizer, scheduler, accum_steps):
    if not isinstance(accum_steps, int) or accum_steps < 1:
        raise ValueError("accum_steps must be a positive integer")
    model.train()
    optimizer.zero_grad(set_to_none=True)
    window_n, window_batches = 0, 0
    sizes = []
    for i, (xb, yb) in enumerate(loader):
        loss_sum = F.cross_entropy(model(xb), yb, reduction="sum")
        loss_sum.backward()
        window_n += len(yb)
        window_batches += 1
        if window_batches == accum_steps or i + 1 == len(loader):
            for p in model.parameters():
                if p.grad is not None:
                    p.grad.div_(window_n)
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0, error_if_nonfinite=True)
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad(set_to_none=True)
            sizes.append(window_n)
            window_n, window_batches = 0, 0
    if not sizes:
        raise ValueError("loader is empty")
    return sizes

torch.manual_seed(2)
accum_model = nn.Linear(8, 3)
accum_opt = torch.optim.SGD(accum_model.parameters(), lr=0.01)
accum_scheduler = torch.optim.lr_scheduler.StepLR(accum_opt, step_size=10)
sizes = accumulated_epoch(accum_model, loader, accum_opt, accum_scheduler, accum_steps=3)
print("examples per update", sizes)
print("updates", len(sizes))
# examples per update [9, 1]
# updates 2

The code above uses summed, unweighted cross-entropy with no ignored labels. At the update boundary it divides accumulated gradients by the actual number of examples, then clips and steps. Ten examples loaded in batches of 3 form windows of 9 and 1 with accum_steps=3. Class weighting, ignored targets, and token-based losses require a denominator matched to their objective. This example keeps both the model and batches on CPU and uses float32 to isolate accumulation from AMP. Moving it to an accelerator also requires moving each batch to the model’s device.

For a separable loss and matching forward computations, this produces the same mathematical gradient as the combined window. Floating-point reduction order can introduce small differences. BatchNorm uses different statistics per micro-batch; dropout masks and other batch-dependent or stochastic computations can also differ. Accumulation does not guarantee exact equivalence merely because BatchNorm is absent. Tune the training recipe for the resulting effective batch size.

A scheduler defined per optimizer update advances once per completed accumulation window. Its total update budget is epochs * ceil(len(loader) / accum_steps) for this fixed-length loader, unless updates are skipped. To combine accumulation with float16 scaling, keep the scale unchanged within a window, unscale only after its last backward, normalize and clip once, then step and update the scaler. Do not mix scaled and unscaled contributions in one gradient buffer.

Activation checkpointing

Activation checkpointing saves selected boundaries and recomputes needed intermediates during backward. In an idealized chain of \(L\) equal-cost layers, splitting into roughly \(\sqrt L\) segments can reduce the activation-storage term to \(O(\sqrt L)\) by balancing saved boundaries against one recomputed segment. Checkpointing every block individually, as below, still stores block boundaries and does not by itself establish that bound. Measure the actual memory and recomputation cost for the chosen partition.

from torch.utils.checkpoint import checkpoint
from copy import deepcopy

class CheckpointedBlocks(nn.Module):
    def __init__(self, blocks):
        super().__init__()
        self.blocks = nn.ModuleList(blocks)
    def forward(self, x):
        for block in self.blocks:
            x = checkpoint(block, x, use_reentrant=False, preserve_rng_state=True)
        return x

torch.manual_seed(3)
plain = nn.Sequential(nn.Sequential(nn.Linear(8, 8), nn.GELU()), nn.Linear(8, 3))
checked = CheckpointedBlocks(deepcopy(list(plain.children())))
x = torch.randn(4, 8)
a, b = plain(x), checked(x)
a.square().mean().backward()
b.square().mean().backward()
print("outputs agree", torch.allclose(a, b))
print("gradients agree", all(torch.allclose(p.grad, q.grad) for p, q in
      zip(plain.parameters(), checked.parameters())))
# outputs agree True
# gradients agree True

Checkpointing can make an otherwise too-large model, sequence, or batch fit, but recomputation costs time. PyTorch normally preserves relevant PyTorch RNG state so dropout can match the original computation; this comes from preserve_rng_state=True, not simply from choosing the non-reentrant variant. Python/NumPy random calls, device changes inside a block, or mutable state need separate care. BatchNorm running statistics can be updated again during recomputation. The runnable comparison uses stateless blocks and checks gradients as well as outputs; it does not measure a memory saving.

Finding the real bottleneck

GPU utilization from nvidia-smi is a sampled activity measure, not a diagnosis of input starvation or a measure of how efficiently active kernels use the hardware. Low utilization can reflect input loading, host launch overhead, synchronization, small kernels, communication, or short alternating phases. A threshold such as 80% cannot identify the cause.

Compare warmed-up step time with real inputs against a repeated, already-on-device batch of the same shape. A large reduction implicates loading, preprocessing, transfers, or their interaction with compute; it does not identify which one by itself. Repeating a batch also removes input variability and may improve caching. AMP can still help the compute portion even when input work is significant, though the overall gain may be limited.

from torch.profiler import profile, ProfilerActivity

def profile_training(loader, train_step, device, optimizer):
    for _ in range(2):
        for xb, yb in loader:
            train_step(xb, yb)
    if device.type == "cuda":
        torch.cuda.synchronize(device)
    optimizer.zero_grad(set_to_none=True)
    activities = [ProfilerActivity.CPU]
    if device.type == "cuda":
        activities.append(ProfilerActivity.CUDA)
    with profile(activities=activities, record_shapes=True, profile_memory=True) as prof:
        for _ in range(3):
            for xb, yb in loader:
                train_step(xb, yb)
                prof.step()
        if device.type == "cuda":
            torch.cuda.synchronize(device)
    return prof

prof = profile_training(loader, train_step, device, optimizer)
events = prof.key_averages()
print("operations recorded", len(events) > 0)
timing_key = "self_cuda_time_total" if device.type == "cuda" else "self_cpu_time_total"
# Inspect table in a notebook, or print(table) in a script.
table = events.table(sort_by=timing_key, row_limit=15)
# operations recorded True

The profiling function warms up before measuring three complete loader passes. The profiler records CPU activity and adds CUDA only when used. train_step includes transfer and update work; the iterator also lets the trace capture main-process loading. A real worker process needs separate investigation. Inspect the returned table and timeline: CPU and GPU execution can overlap, so totals should not simply be added or compared to declare one side the bottleneck. Recording shapes and memory adds overhead; use a lighter run for final throughput measurement. The profiler may warn that a memory block was allocated before recording began, so its deallocation cannot be fully accounted for. Its memory-event table is not a complete inventory of live memory; use allocator counters for a separately defined peak-memory interval.

For a wall-clock benchmark, warm up, synchronize CUDA before starting the timer, run several steps, and synchronize again before stopping. Otherwise the timer can measure only asynchronous work submission. CUDA events can measure a chosen device interval; wall-clock time including loading answers an end-to-end question. Report the batch size, precision, hardware, and measurement interval. The first steps may include initialization, autotuning, allocation, or compilation, which should be reported separately when startup cost matters.

A measured example on one T4

A separate Kaggle validation notebook ran on one Tesla T4 using Python 3.12.13 and PyTorch 2.10.0+cu128. Two GPUs were available, but all tests used cuda:0. The functional checks completed eight updates each in float32, float16, and bfloat16, with parameters remaining float32. An intentional overflow left both the parameter and scheduler unchanged while the loss scale fell from 65,536 to 32,768. Accumulation windows of 9 and 1 examples matched the combined-window reference with a maximum parameter difference of 0.0. Checkpointed and ordinary computations also matched in the dropout test (probability 0.2), with a maximum gradient difference of 0.0. These checks cover those small test cases; they do not establish long-run convergence or validation quality.

The performance test used eight width-1,024 Linear–GELU blocks, a ten-class head, batch size 256, and AdamW. Dropout was disabled. Each condition started from the same initialization and synthetic batch, already on the GPU. After five warmup updates, 20 complete steps were timed with CUDA synchronization at the interval boundaries. Each condition ran three times in shuffled order. The table reports median step time and the median of the three peak allocated-memory measurements. Each peak covers the full measured interval, including backward and optimizer updates; it is not activation memory alone.

Condition	Median ms per step	Peak allocated MiB	Successful / attempted updates
Float32	6.61	179.617	60 / 60
Float16 AMP	7.16	179.613	60 / 60
Bfloat16 AMP	8.43	179.612	60 / 60
Float16 AMP with checkpointing	12.70	179.613	60 / 60

Float16 took about 8% longer per step than float32 here; bfloat16 took about 28% longer. Checkpointing added about 77% to the float16 step time without reducing its recorded peak. The four peak values differ by less than 0.005 MiB. This workload therefore showed no meaningful whole-step peak-memory saving. A reduction in saved activations can be hidden if the overall peak occurs elsewhere, but these measurements do not identify the peak’s source. Measuring forward, backward, and optimizer phases separately would help test that explanation. Other model shapes and batch sizes need their own comparison.

The notebook also collected a trace from a smaller width-256, four-block model with batch size 64. That trace confirms CUDA kernels were recorded, but cannot explain the larger benchmark’s timing by itself. Its summary showed an AdamW GPU annotation of 1.375 ms as 130.95% of a 1.050 ms kernel-and-copy total. The raw trace shows that the annotation spans and the kernel events are different measurements. Treat that percentage as a reason to inspect the timeline, not as GPU utilization or an exclusive share of runtime.

Scaling this across several GPUs is covered in Distributed Training with PyTorch DDP.

Choose a change from the measurement

Observed constraint	Candidate change	What to verify
Slow data preparation or transfer	Worker count, caching, pinned memory and nonblocking CUDA transfers	End-to-end time, host memory, multiprocessing behavior, and overlap
Gradient-buffer overhead	zero_grad(set_to_none=True)	Unused parameters with None gradients may be skipped by the optimizer
Convolution kernel choice	cudnn.benchmark=True	Warmup cost, changing shapes, and reproducibility requirements
Eligible dense compute or activation memory	AMP	Kernel support, measured memory, convergence, and validation quality
Framework overhead or fusion opportunities	torch.compile	Compilation, graph breaks, recompilation, and steady-state time
Desired batch does not fit	Gradient accumulation	Correct scaling and remainder handling; changed update count and throughput
Saved activations dominate memory	Activation checkpointing	Recomputed state, gradients, peak memory, and elapsed time

Change one factor at a time on a representative workload and retain the baseline measurement. More workers can add overhead or exhaust host resources, and pinned memory is not free. Performance improvements should preserve the intended learning objective and acceptable validation quality; a faster iteration alone is not enough.

The inference-time cost of this, and how to reduce it, is covered in LLM Inference Optimization: KV Cache, FlashAttention, and Quantization.

References

• Chen, Xu, Zhang, and Guestrin (2016). Training Deep Nets with Sublinear Memory Cost. arXiv preprint.
• Micikevicius et al. (2018). Mixed Precision Training.
• PyTorch. AMP.
• PyTorch. scaling and accumulation.
• PyTorch. activation checkpointing.
• PyTorch. profiling.