Mixed precision and autocast

What it is

Mixed-precision training runs most ops in float16 or bfloat16 while keeping a few sensitive ones (e.g., reductions, certain norms, loss scaling) in float32. PyTorch's machinery for this is torch.amp / torch.autocast (the "autocast" key in the dispatcher) and torch.amp.GradScaler.

from torch.amp import autocast, GradScaler

scaler = GradScaler()
with autocast(device_type="cuda", dtype=torch.bfloat16):
    out = model(x)
    loss = criterion(out, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

How autocast works

torch.autocast is implemented as a per-device dispatch key (AutocastCUDA, AutocastCPU, AutocastMPS, AutocastXPU). Inside the context manager, the key is added to TLS; the dispatcher then routes ops to the autocast kernel for that key.

The autocast kernel for an op (registered in aten/src/ATen/autocast_mode.cpp) does one of:

• Cast inputs down to the autocast dtype (most matmul/conv/attention) and redispatch.
• Cast inputs up to float32 (loss-sensitive ops like softmax, log_softmax, mse_loss) and redispatch.
• Pass through (most everything else).

A per-op cache prevents redundant casts when the same parameter is used in multiple ops.

The key thing to understand is that autocast is not a Python-level rewrite. It is a dispatch-key kernel registered for every op that needs special handling. Adding new ops or adjusting a casting policy means editing aten/src/ATen/autocast_mode.cpp.

Loss scaling (GradScaler)

float16 underflows easily on small gradients. GradScaler:

1. Multiplies the loss by a large factor before backward.
2. Backward produces scaled gradients (still in float16).
3. Before the optimizer step, gradients are unscaled and inspected for NaNs/Infs.
4. If clean, the optimizer step proceeds; otherwise the step is skipped and the scale is decreased.

bfloat16 has the same dynamic range as float32, so loss scaling is unnecessary — but GradScaler is still safe to use (it'll be a no-op).

Supported dtypes

Dtype	Used for	Notes
float16	Volta+ NVIDIA GPUs	Use GradScaler
bfloat16	Ampere+ NVIDIA, MI200+ AMD, recent Apple Silicon, Intel	No loss scaling needed; preferred for LLM training
float8_e4m3fn / float8_e5m2	Hopper+, MI300+	Used for inference and (experimentally) training

The float8 dtypes are first-class in c10::ScalarType; ops that support them are tagged in aten/src/ATen/native/native_functions.yaml. Float8 training is built on top of torchao (a separate repo) but the dtypes themselves live in PyTorch.

TF32 (Ampere+ matmul)

A separate, simpler knob: torch.backends.cuda.matmul.allow_tf32 = True (the default) lets cuBLAS use TF32 for float32 matmul. This is independent of autocast and is configured per-device via at::globalContext().

Compile-time

torch.compile is fully autocast-aware: the autocast TLS state is captured during tracing and the right casts appear in the FX graph. Inductor takes advantage of fused kernels that perform mixed-precision math directly.

Where to look

Path	Contents
torch/amp/	Public Python API (autocast, GradScaler)
aten/src/ATen/autocast_mode.cpp	Per-op autocast kernels
aten/src/ATen/autocast_mode.h	Casting policy macros
c10/util/Float8_e4m3fn.h, Float8_e5m2.h	Float8 dtype headers
torch/backends/cuda.py, cudnn.py	TF32 toggles

Where to read next

• Systems / Dispatcher — how autocast keys work.
• Features / torch.compile — compile + autocast.