PyTorch Automatic Mixed Precision (AMP): A Comprehensive Guide

3.1 Op categorization

Internally, PyTorch maintains three categories of ops relevant to autocast:

1. Allowlist (cast to lower precision): ops that benefit from Tensor Cores and are numerically robust to reduced precision — matmul, mm, bmm, addmm, conv1d/2d/3d, linear, and similar. These execute in fp16/bf16.
2. Denylist (kept in fp32): numerically sensitive reductions and functions — softmax, log_softmax, cross_entropy, nll_loss, layer_norm, batch_norm (in training mode), pow, exp, mse_loss, and other ops prone to overflow or requiring high dynamic range. These are forced to fp32 regardless of the ambient autocast dtype.
3. Widest-input / promote rules: ops like cat, stack, add that combine multiple tensors run in the widest dtype among their inputs, so mixing an fp16 tensor with an fp32 tensor promotes the result to fp32, preventing silent precision loss on concatenation-like ops.

This policy is defined in PyTorch’s aten/src/ATen/autocast_mode.cpp and is not user-configurable per op out of the box, though you can force specific regions to a chosen dtype or exit autocast manually (Section 3.3).

3.2 Basic usage

import torch

device = "cuda"
model = MyModel().to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

for inputs, targets in dataloader:
    inputs, targets = inputs.to(device), targets.to(device)
    optimizer.zero_grad(set_to_none=True)

    with torch.autocast(device_type="cuda", dtype=torch.float16):
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)

    loss.backward()
    optimizer.step()

Note that loss.backward() is deliberately outside the autocast context. Autocast only needs to be active for the forward pass; the backward pass automatically uses the same dtypes that were recorded for each op during the forward pass (autograd replays the forward’s casting decisions), so wrapping backward() in autocast is unnecessary and has no additional effect.

3.3 autocast parameters in depth

torch.autocast(
    device_type="cuda",       # "cuda", "cpu", "xpu", "mps" (support varies by backend)
    dtype=torch.float16,      # or torch.bfloat16
    enabled=True,             # toggle without removing the context manager
    cache_enabled=True,       # cache casted weights across autocast-wrapped calls
)

• device_type: Since PyTorch 1.10+, the unified torch.amp.autocast API requires specifying the device explicitly, replacing the older device-specific torch.cuda.amp.autocast() / torch.cpu.amp.autocast() wrappers (which are now thin aliases). CPU autocast typically only supports bfloat16.
• dtype: fp16 requires a GradScaler (Section 4); bf16 does not, since it has fp32’s dynamic range.
• cache_enabled: Autocast caches the fp16/bf16 cast of fp32 weights within a single autocast region so a weight reused multiple times in the forward pass (e.g., weight tying, recurrent loops) isn’t recast every call. This cache is invalidated at each autocast context entry/exit. Disable it if you’re manually mutating parameters between reuses within a single region, or you’ll get stale casts.
• enabled=False: Useful for a debugging switch that keeps code structure identical between AMP and full-precision runs, or for regions you deliberately want in fp32 (e.g., a numerically fragile custom loss) nested inside a larger autocast block:

with torch.autocast(device_type="cuda", dtype=torch.float16):
    features = backbone(inputs)               # runs in fp16
    with torch.autocast(device_type="cuda", enabled=False):
        # force fp32 for a numerically fragile custom op
        stable_out = fragile_custom_op(features.float())
    logits = head(stable_out)

3.4 Nesting and re-entrancy

Autocast contexts can be nested; the innermost context’s settings take precedence for the ops executed within it, and settings revert to the enclosing context’s values on exit. This is the standard pattern for locally opting an operation out of autocast (as shown above) without disabling it for the whole forward pass.

4.1 Why scaling is necessary

In fp16, gradients with magnitude below ~6e-8 (or below the representable range once you account for the format’s granularity near zero) flush to zero. Empirically, a large fraction of gradient values in deep networks — especially in later training stages or deep layers — fall into this dead zone. GradScaler multiplies the loss by a scale factor S before backpropagation, so gradients are shifted up into fp16’s representable range by the same multiplicative factor (since gradients are linear in the loss via the chain rule). Before the optimizer step, gradients are divided by S to recover their true magnitude.

4.2 The scaling algorithm

scaler = torch.amp.GradScaler(device="cuda")

for inputs, targets in dataloader:
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)

    scaler.scale(loss).backward()      # loss * S, then backward
    scaler.step(optimizer)             # unscale grads, check for inf/nan, conditionally step
    scaler.update()                    # adjust S for next iteration

Under the hood, each call does the following:

1. scaler.scale(loss): returns loss * S, where S starts at init_scale (default 2**16).
2. .backward(): backpropagates the scaled loss; every gradient in .grad is implicitly scaled by S due to linearity.
3. scaler.step(optimizer):
   • Internally calls scaler.unscale_(optimizer) if you haven’t already called it manually (divides all gradients by S in place).
   • Inspects the unscaled gradients for inf/nan values (a cheap torch.isfinite-style reduction across all parameter grads registered with that optimizer).
   • If any non-finite values are found, the optimizer step is skipped entirely for this iteration (the corrupted update is discarded), and this is reported to update().
   • Otherwise, optimizer.step() proceeds normally with the correctly-scaled gradients.
4. scaler.update():
   • If the last step was skipped (inf/nan detected), S is multiplied by backoff_factor (default 0.5) — scale down, since we overflowed.
   • If growth_interval (default 2000) consecutive steps have passed without an overflow, S is multiplied by growth_factor (default 2.0) — cautiously scale up to use more of fp16’s range.

This is a dynamic scheme: it hunts for the largest scale factor that doesn’t cause overflow, adapting automatically over the course of training without user tuning in the common case.

4.3 Constructor parameters

scaler = torch.amp.GradScaler(
    device="cuda",
    init_scale=2.0**16,
    growth_factor=2.0,
    backoff_factor=0.5,
    growth_interval=2000,
    enabled=True,
)

• init_scale: Start conservatively high; the algorithm will back off quickly if it’s too aggressive, so there’s little cost to starting high, but starting too low wastes early-training gradient signal.
• growth_interval: Lower values adapt scale faster but risk oscillating around the overflow boundary, wasting steps. 2000 is a reasonable default for most vision/NLP workloads; for very unstable models (some RL, GANs), consider tightening this or reducing growth_factor.
• enabled=False: Lets you keep an identical code path for bf16 or fp32 runs — scaler.scale(loss) becomes a no-op passthrough, step() calls optimizer.step() directly, and update() does nothing. This is the idiomatic way to write dtype-agnostic training loops (see Section 9.1).

4.4 Gradient clipping under AMP

A common correctness bug: calling torch.nn.utils.clip_grad_norm_ directly on .grad tensors that are still in scaled space clips against the wrong magnitude (scaled by S, not the true gradient norm). You must explicitly unscale first:

scaler.scale(loss).backward()
scaler.unscale_(optimizer)                          # bring grads back to true scale
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)                               # will NOT re-unscale, since unscale_() already ran
scaler.update()

GradScaler tracks, per optimizer, whether unscale_() has already been called this iteration and will not double-unscale. Calling unscale_() more than once per step() for the same optimizer without an intervening update() raises a RuntimeError.

4.5 Gradient accumulation under AMP

When accumulating gradients over multiple micro-batches before stepping, scale and accumulate every micro-batch, but only call step()/update() on the batch boundary:

scaler = torch.amp.GradScaler(device="cuda")
accum_steps = 4

for i, (inputs, targets) in enumerate(dataloader):
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        outputs = model(inputs)
        loss = loss_fn(outputs, targets) / accum_steps

    scaler.scale(loss).backward()

    if (i + 1) % accum_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad(set_to_none=True)

Because gradients accumulate additively in .grad across the micro-batches and all of them were scaled by the same S (it only changes on update()), the accumulated gradient is consistently scaled and a single unscale_/step at the boundary is correct.

4.6 Multiple models / optimizers in one scaler

A single GradScaler instance can drive multiple optimizers (e.g., separate generator/discriminator optimizers in a GAN), since it tracks per-optimizer state internally:

scaler = torch.amp.GradScaler(device="cuda")

scaler.scale(loss_d).backward()
scaler.step(optimizer_d)

scaler.scale(loss_g).backward()
scaler.step(optimizer_g)

scaler.update()   # one update call after all steps for this iteration

Call update() once per iteration, after all step() calls, not after each one — otherwise the scale factor updates based on partial information and the shared S diverges between optimizers’ actual needs.

5.1 The problem

autocast’s dispatch-level interception works by patching registered ATen ops. A torch.autograd.Function subclass with a custom forward/backward is opaque to this mechanism by default — the ops inside it don’t automatically participate in the allowlist/denylist casting policy, and worse, autograd’s replay of forward-pass dtypes during backward doesn’t apply because the function is a single opaque node in the graph.

5.2 Making custom Functions autocast-compatible

PyTorch provides decorators to explicitly declare a custom function’s casting behavior:

from torch.amp import custom_fwd, custom_bwd

class MyCustomOp(torch.autograd.Function):
    @staticmethod
    @custom_fwd(device_type="cuda", cast_inputs=torch.float16)
    def forward(ctx, x, weight):
        ctx.save_for_backward(x, weight)
        return some_low_level_kernel(x, weight)

    @staticmethod
    @custom_bwd(device_type="cuda")
    def backward(ctx, grad_output):
        x, weight = ctx.saved_tensors
        # grad_output arrives in the same dtype forward() computed in
        return compute_grad_x(grad_output, weight), compute_grad_w(grad_output, x)

• @custom_fwd(cast_inputs=torch.float16) forces all floating-point inputs to forward to be cast to fp16 before the function executes, regardless of ambient autocast state, then disables further autocasting inside the function body (so internal ops run in the requested dtype deterministically).
• If you omit cast_inputs, the function runs in whatever dtype its inputs happen to be — you’re opting out of automatic casting but the function will still execute correctly if the caller passes already-cast tensors.
• @custom_bwd ensures the backward pass replays with the same effectively-disabled-autocast semantics as the forward, keeping dtype behavior symmetric.

5.3 Registering ops with autocast’s policy (advanced)

If you’re authoring a C++/CUDA extension and want a custom op to participate in the same allowlist/promote/denylist machinery as built-in ops (rather than hand-casting in Python), PyTorch exposes torch.library-based registration hooks (m.impl(..., c10::DispatchKey::Autocast...) on the C++ side, or the torch._C._nn registration APIs) to attach an autocast rule to a custom dispatcher key. This is rarely necessary for typical model code — reach for custom_fwd/custom_bwd first — and is primarily relevant if you’re shipping a custom op library that should feel like a first-class autocast citizen for downstream users.

6.1 Basic pattern

AMP composes with DDP with essentially no special-casing: each process runs its own local autocast region and its own GradScaler. DDP’s gradient all-reduce happens on the (still-scaled) .grad tensors during backward(), which is safe because all-reduce is a linear (averaging) operation and every rank uses the same scale factor S for a given iteration — the averaged gradient is scaled by the same S as any individual rank’s, so unscaling downstream is still correct.

model = MyModel().to(local_rank)
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])
scaler = torch.amp.GradScaler(device="cuda")

for inputs, targets in dataloader:
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)

    scaler.scale(loss).backward()   # DDP all-reduces scaled grads during this call
    scaler.step(optimizer)
    scaler.update()

6.2 The subtle inf/nan divergence risk

Because scaler.step() decides locally, per-rank whether to skip the optimizer step based on that rank’s gradients being finite, and because scaler.update() adjusts S locally based on that rank’s skip/no-skip outcome, there’s a theoretical risk of scale factor divergence across ranks if different ranks observe different overflow behavior (e.g., due to non-deterministic ops or rank-specific data causing localized overflow). In practice, since DDP all-reduces gradients before the inf/nan check effectively sees the post-all-reduce values (the check happens on the same all-reduced .grad tensors every rank has after backward() returns), all ranks observe the same averaged gradient tensor and therefore the same finite/non-finite verdict, keeping S synchronized across ranks without extra effort. This only breaks down if you do something unusual like per-rank gradient modification between backward() and step().

6.3 find_unused_parameters and autocast caching interaction

If your model has conditional branches (common in mixture-of-experts or dynamic architectures) combined with cache_enabled=True autocast weight caching, be aware that DDP’s bucketing and the autocast cache both add bookkeeping around parameter reuse; if you hit unexpected staleness or DDP “unused parameter” errors under AMP, temporarily disabling cache_enabled is a useful debugging step.

7.1 FSDP’s native MixedPrecision policy

Rather than (or in addition to) wrapping the forward pass in torch.autocast, FSDP exposes a MixedPrecision config that controls the dtype of parameters, gradient reduction, and buffers directly at the sharding level:

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision

mp_policy = MixedPrecision(
    param_dtype=torch.bfloat16,     # dtype params are cast to for forward/backward compute
    reduce_dtype=torch.bfloat16,    # dtype used for gradient all-reduce/reduce-scatter
    buffer_dtype=torch.bfloat16,    # dtype for buffers (e.g., BatchNorm running stats)
)

model = FSDP(model, mixed_precision=mp_policy, device_id=local_rank)

This casts the sharded parameters themselves to param_dtype for compute, communicates gradients in reduce_dtype, and keeps an fp32 master copy of parameters (by default) for the optimizer step — conceptually similar to what autocast + fp32 master weights achieves manually, but implemented natively in the sharding/communication path rather than at op dispatch time.

7.2 Combining FSDP with torch.autocast

You can still layer torch.autocast on top of FSDP’s MixedPrecision (or instead of it), but the two mechanisms address different things: FSDP’s policy governs what dtype sharded parameters and gradient communication use, while autocast governs what dtype individual ops execute in during the forward pass, independent of sharding. A common pattern for bf16 training on large models is to rely on FSDP’s MixedPrecision alone (bf16 needs no GradScaler), skipping torch.autocast entirely, since FSDP already ensures compute happens in the requested lower-precision dtype for the sharded modules.

For fp16 with FSDP, you still need loss scaling, but the standard torch.amp.GradScaler was originally written for non-sharded parameters. Since PyTorch 1.12+, torch.distributed.fsdp.sharded_grad_scaler.ShardedGradScaler provides an FSDP-aware equivalent that correctly handles the inf/nan check across sharded gradient shards (which live on different ranks and can’t be checked with a naive local reduction the way DDP’s fully-replicated gradients can):

from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

scaler = ShardedGradScaler()

Its scale/step/update API mirrors torch.amp.GradScaler but performs the necessary cross-rank communication to determine whether any shard anywhere saw an inf/nan, since with FSDP no single rank holds the complete gradient for a global finiteness check.

7.3 Practical recommendation

For large-model training (the primary use case FSDP targets), bf16 + FSDP MixedPrecision (no scaler needed) is by far the more common and lower-friction setup than fp16 + ShardedGradScaler, precisely because bf16 sidesteps the scaling machinery’s added complexity in a sharded, multi-rank setting. Reach for fp16 + ShardedGradScaler primarily when targeting hardware where bf16 throughput is not competitive with fp16 (some older Tensor Core generations) or when you need fp16’s extra mantissa precision relative to bf16 for a numerically delicate model.

8.1 Measuring wall-clock speedup correctly

Naive timing pitfalls to avoid:

• CUDA is asynchronous. A bare time.time() around a forward/backward call measures kernel launch time, not execution time, unless you call torch.cuda.synchronize() immediately before and after the timed region.
• Warm-up matters. The first several iterations pay for CUDA context initialization, cuDNN algorithm autotuning (especially with torch.backends.cudnn.benchmark=True), and JIT/kernel caching. Discard at least 5-10 warm-up iterations before timing.
• Use CUDA events for GPU-side timing, which avoids CPU-GPU synchronization overhead contaminating the measurement:

import torch

def benchmark(fn, warmup=10, iters=50):
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()

    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)

    start.record()
    for _ in range(iters):
        fn()
    end.record()
    torch.cuda.synchronize()

    return start.elapsed_time(end) / iters   # milliseconds per iteration

Run this once with autocast enabled and once with it disabled (same model, same batch, same seed) to get a clean speedup ratio.

8.2 Measuring memory savings

torch.cuda.reset_peak_memory_stats()
# ... run one or more training iterations ...
peak_bytes = torch.cuda.max_memory_allocated()
print(f"Peak memory: {peak_bytes / 1e9:.2f} GB")

Compare max_memory_allocated() between fp32-only and AMP runs at matched batch size to quantify the activation memory savings, or hold memory fixed and increase batch size under AMP until you hit the same peak memory as the fp32 baseline to quantify the effective batch size gain.

8.3 Using torch.profiler to see the dtype breakdown

from torch.profiler import profile, ProfilerActivity, record_function

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
    with_stack=False,
) as prof:
    with record_function("train_step"):
        with torch.autocast(device_type="cuda", dtype=torch.float16):
            outputs = model(inputs)
            loss = loss_fn(outputs, targets)
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))

The resulting table lets you confirm that the expected ops (matmuls, convs) are actually dispatching to fp16 Tensor Core kernels (look for kernel names containing half or Tensor Core-specific kernel naming, and check the Input Shapes/dtype columns) rather than silently falling back to fp32 execution paths — a useful sanity check when speedup is smaller than expected, since it can reveal that a supposedly-eligible op isn’t actually being cast (often due to being outside the autocast region, or receiving inputs of an unexpected dtype that triggers a promotion rule).

8.4 What speedup to expect

Realistic expectations, highly hardware- and model-dependent:

• Compute-bound models (large matmul/conv-dominated architectures, large batch sizes) on Tensor Core-equipped GPUs (V100 and later): commonly 1.5-3x throughput improvement.
• Memory-bandwidth-bound models (small batch, many elementwise/normalization ops relative to matmuls): smaller speedup, since the denylist ops (softmax, layer norm, etc.) still run in fp32 and elementwise op throughput isn’t Tensor-Core-accelerated the same way.
• Memory capacity gains are often the more valuable win in practice — the freed-up memory lets you increase batch size or model size, which itself often improves both throughput and, for large-batch-sensitive optimizers, training dynamics.

9.1 Writing dtype-agnostic training code

A robust pattern that lets you switch between fp32/fp16/bf16 via configuration without branching your training loop:

use_amp = True
amp_dtype = torch.bfloat16   # or torch.float16, or None for fp32

scaler = torch.amp.GradScaler(device="cuda", enabled=(use_amp and amp_dtype == torch.float16))

for inputs, targets in dataloader:
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=use_amp):
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Since GradScaler(enabled=False) makes scale()/step()/update() no-ops that fall through to plain backward()/optimizer.step(), this same code path is correct for bf16 (scaler disabled, autocast active) and fp32 (scaler disabled, autocast disabled) without any if branches in the hot loop.

9.2 Loss becomes NaN immediately

Checklist, roughly in order of likelihood:

1. Scale started too high for this model. Try a lower init_scale (e.g., 2**12) — some architectures (RNNs with long sequences, certain GAN losses) overflow even at moderate scales before the backoff mechanism has a chance to correct.
2. A custom loss or op outside the denylist is overflowing in fp16. Force it to fp32 explicitly with a nested autocast(enabled=False) block (Section 3.3), casting inputs with .float() first.
3. Gradient clipping applied before unscale_() — clipping scaled gradients against an unscaled threshold can mask real explosions or trigger false-positive clipping that destabilizes training. Verify the ordering in Section 4.4.
4. Data-dependent overflow — e.g., an input containing inf/nan upstream of the model, unrelated to AMP at all. Sanity-check by running one iteration in full fp32 with torch.autograd.set_detect_anomaly(True).
5. Mixing bf16 assumptions into fp16 code — e.g., initializing GradScaler but ambiently running the model in a bfloat16-dtyped tensor pipeline outside of autocast (bf16 tensors created directly, not via autocast casting), which desyncs the scaler’s assumptions about what dtype is actually in play.

9.3 “No speedup observed”

• Model is memory-bandwidth-bound rather than compute-bound (Section 8.4) — check with the profiler whether GPU utilization was already near 100% in fp32; if so, AMP has limited room to help.
• cudnn.benchmark not enabled — set torch.backends.cudnn.benchmark = True for convolutional models with fixed input shapes, so cuDNN can autotune fp16-optimized algorithms.
• GPU generation lacks Tensor Cores or has poor fp16 throughput relative to fp32 (pre-Volta NVIDIA GPUs) — AMP provides little to no benefit and mainly serves as a memory-savings tool there.
• Small batch size or small model — kernel launch overhead and non-Tensor-Core ops dominate; the fp16 matmul speedup is a small fraction of total step time.
• Autocast region doesn’t actually wrap the compute-heavy ops (e.g., wrapping only the loss computation and not the forward pass) — verify region boundaries.

9.4 In-place operations and autocast

In-place ops (add_, mul_, etc.) on tensors that came out of an autocast region can produce dtype-mismatch errors or silently defeat the casting policy, since in-place ops don’t go through the same dispatch-time casting decision in all cases. Prefer out-of-place ops within autocast regions unless you’ve specifically verified the in-place variant is autocast-safe for your PyTorch version.

9.5 Reproducibility

AMP introduces additional nondeterminism beyond the usual CUDA nondeterminism concerns: dynamic loss scaling means the exact scale factor trajectory (and hence which steps get skipped) can depend on timing-sensitive floating point behavior. For strict reproducibility across runs, log the scaler’s state_dict() (see Section 10) alongside model checkpoints, and be aware that bit-exact reproducibility of AMP training across different GPU architectures is generally not achievable even with fixed seeds, due to differing Tensor Core rounding behavior.