Mixed precision with autocast and GradScaler

Lower precision where it is safe

By default PyTorch computes in float32. Using 16-bit formats halves memory for activations and lets modern NVIDIA GPUs use Tensor Cores, which run matrix multiplications much faster. Automatic mixed precision keeps this safe: inside torch.autocast(device_type="cuda", dtype=...), operations that benefit (matmuls, convolutions) run in 16-bit, while numerically sensitive ones (reductions, softmax, many losses) stay in float32. The weights themselves stay in float32.

There are two 16-bit formats. float16 has a small range, so tiny gradients can underflow to zero. The fix is a gradient scaler, torch.amp.GradScaler("cuda"): it multiplies the loss by a large factor before backward, unscales the gradients before the optimizer step, skips steps whose gradients contain inf or NaN, and adjusts the factor over time. bfloat16 keeps float32's exponent range with fewer mantissa bits, so it usually trains without a scaler; it needs hardware support, such as recent NVIDIA GPUs and many modern CPUs.

Avoid calling .half() on the whole model for training; that loses the float32 safety net autocast provides.

A float16 training step

Only the forward pass and loss sit inside autocast. The backward pass and optimizer step go through the scaler. model, optimizer, loss_fn and train_loader come from earlier topics.

import torch

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

for x, y in train_loader:
    x, y = x.to("cuda"), y.to("cuda")
    optimizer.zero_grad()
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        logits = model(x)
        loss = loss_fn(logits, y)
    scaler.scale(loss).backward()   # scale up so small gradients survive
    scaler.step(optimizer)          # unscale, then step unless inf/NaN
    scaler.update()                 # adjust the scale factor

Rough sums, exact ledger

A shopkeeper adds up the day's rush quickly in round numbers on a notepad, but the account book is kept to the paisa. Autocast is the notepad for fast arithmetic; the float32 weights are the account book. The scaler is like counting in paise instead of rupees, so small amounts do not round away to nothing.

When to choose bfloat16

If your hardware supports bfloat16, try dtype=torch.bfloat16 with no scaler first; it is simpler and rarely underflows. Use float16 plus GradScaler on GPUs without bfloat16 support. Either way, compare validation metrics with a float32 run once to confirm nothing degraded.

त्वरित जाँच: Why is a gradient scaler used with float16 mixed precision?

• Small gradient values can underflow to zero in float16, and scaling the loss keeps them representable
• float16 cannot represent negative numbers
• The optimizer cannot read float16 gradients
• Scaling makes the learning rate larger