Mixed Precision Training › Code Implementation
Language AI HandbookAll chapters Mixed Precision Training › Code ImplementationOne section · Sections in this chapter7 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let us implement mixed precision training in PyTorch, showing both FP16 with dynamic loss scaling and BF16 without scaling. We will train a small transformer-like model on a synthetic task to measure memory usage, throughput, and training stability.
Setup and Imports
import torch
# Check hardware support
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
has_fp16 = torch.cuda.is_available()
has_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()Device: cpu FP16 support: False BF16 support: False
A Simple Transformer Block
We define a small transformer block that exercises the key components: multi-head attention and a feed-forward network. This model is deliberately small so it can run on a CPU if no GPU is available, but the mixed precision patterns are identical for large models.
import torch.nn as nn
class TransformerBlock(nn.Module):
def __init__(self, d_model=256, n_heads=8, d_ff=1024, dropout=0.1):
super().__init__()
self.attention = nn.MultiheadAttention(
d_model, n_heads, dropout=dropout, batch_first=True
)
self.ff = nn.Sequential(
nn.Linear(d_model, d_ff),
nn.GELU(),
nn.Linear(d_ff, d_model),
)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
attn_out, _ = self.attention(x, x, x)
x = self.norm1(x + self.dropout(attn_out))
ff_out = self.ff(x)
x = self.norm2(x + self.dropout(ff_out))
return x
class SimpleTransformer(nn.Module):
def __init__(
self,
vocab_size=1000,
d_model=256,
n_layers=4,
n_heads=8,
d_ff=1024,
seq_len=64,
):
super().__init__()
self.embedding = nn.Embedding(vocab_size, d_model)
self.pos_embedding = nn.Embedding(seq_len, d_model)
self.layers = nn.ModuleList(
[TransformerBlock(d_model, n_heads, d_ff) for _ in range(n_layers)]
)
self.norm = nn.LayerNorm(d_model)
self.head = nn.Linear(d_model, vocab_size)
self.seq_len = seq_len
def forward(self, x):
B, T = x.shape
positions = torch.arange(T, device=x.device).unsqueeze(0).expand(B, T)
h = self.embedding(x) + self.pos_embedding(positions)
for layer in self.layers:
h = layer(h)
h = self.norm(h)
return self.head(h)Model parameters: 921,064 FP32 model size: 3.7 MB FP16 model size: 1.8 MB
The model prints its parameter count and the theoretical size difference between FP32 and FP16 representations. In practice, the activation memory savings are even more significant at scale, but parameter size gives a sense of the baseline compression.
FP16 Training with Dynamic Loss Scaling
The standard PyTorch mixed precision training loop uses autocast for the forward pass and GradScaler for loss scaling. We demonstrate this with a complete training step.
from torch.amp import autocast
def train_step_fp16(model, optimizer, scaler, x, y):
"""Single training step with FP16 AMP and dynamic loss scaling."""
optimizer.zero_grad()
# Forward pass in FP16
with autocast(device_type=device.type, dtype=torch.float16):
logits = model(x)
# logits: (B, T, vocab_size); y: (B, T)
loss = nn.CrossEntropyLoss()(logits.view(-1, vocab_size), y.view(-1))
# Scale loss, backward, unscale, clip, step
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
return loss.item(), grad_norm.item(), scaler.get_scale()import torch.optim as optim
from torch.amp import GradScaler
# Run a short training loop and collect metrics
optimizer_fp16 = optim.AdamW(model.parameters(), lr=1e-4)
scaler = GradScaler(device=device.type)
losses_fp16 = []
scales = []
overflow_steps = []
n_steps = 30
rng = np.random.default_rng(42)
for step in range(n_steps):
x = torch.randint(0, vocab_size, (batch_size, seq_len), device=device)
y = torch.randint(0, vocab_size, (batch_size, seq_len), device=device)
prev_scale = scaler.get_scale()
loss, grad_norm, current_scale = train_step_fp16(
model, optimizer_fp16, scaler, x, y
)
losses_fp16.append(loss)
scales.append(current_scale)
if current_scale < prev_scale:
overflow_steps.append(step)FP16 training completed: 30 steps Final loss: 7.0555 Final loss scale: 65536 Overflow events (scale reductions): 0 Mean loss (last 20 steps): 7.0747
The training loop runs cleanly. The GradScaler monitors gradient overflow at each step and reduces the scale only when it detects non-finite gradients. A short, well-behaved run may report no reductions at all, especially on the CPU execution path used for this notebook. That absence is a property of this smoke test, not evidence that dynamic scaling is unnecessary on production FP16 workloads.
BF16 Training without Loss Scaling
BF16 training simplifies the loop: no GradScaler, no overflow checking, no scale management.
def train_step_bf16(model, optimizer, x, y):
"""Single training step with BF16 AMP. No loss scaling needed."""
optimizer.zero_grad()
# Forward pass in BF16
with autocast(device_type=device.type, dtype=torch.bfloat16):
logits = model(x)
loss = nn.CrossEntropyLoss()(logits.view(-1, vocab_size), y.view(-1))
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
return loss.item()# Reset model and run BF16 training
model_bf16 = SimpleTransformer(
vocab_size=vocab_size, d_model=d_model, n_layers=n_layers, seq_len=seq_len
)
model_bf16 = model_bf16.to(device)
optimizer_bf16 = optim.AdamW(model_bf16.parameters(), lr=1e-4)
losses_bf16 = []
for step in range(n_steps):
x = torch.randint(0, vocab_size, (batch_size, seq_len), device=device)
y = torch.randint(0, vocab_size, (batch_size, seq_len), device=device)
if has_bf16:
loss = train_step_bf16(model_bf16, optimizer_bf16, x, y)
else:
# Fallback to FP32 on hardware without BF16
optimizer_bf16.zero_grad()
logits = model_bf16(x)
loss = nn.CrossEntropyLoss()(logits.view(-1, vocab_size), y.view(-1))
loss.backward()
torch.nn.utils.clip_grad_norm_(model_bf16.parameters(), max_norm=1.0)
optimizer_bf16.step()
loss = loss.item()
losses_bf16.append(loss)BF16 training completed: 30 steps Final loss: 7.0786 Mean loss (last 20 steps): 7.0671 No loss scaling required
The BF16 loop is shorter and easier to reason about. There is no scale factor to monitor, no overflow detection, and no conditional step-skipping logic. The tradeoff is that you need compatible hardware (Ampere GPUs or newer, or TPUs).
Visualizing Training Loss
Let us visualize the loss values for both precision modes. The targets are freshly sampled random tokens at every step, so this 30-step run is a numerical-stability smoke test rather than a convergence experiment. On hardware without BF16 support, the second trace uses FP32.
Visualizing Dynamic Loss Scale Evolution
The recorded scale can remain flat when the short runtime produces no overflows, as it does on this notebook's CPU path. To make the controller mechanics visible on every platform, the next chart uses a deterministic illustrative schedule with four overflow decisions and a shortened five-step growth interval. PyTorch's production default growth interval is much longer at 2,000 successful steps.
Measuring the Throughput Benefit
The actual speedup from mixed precision depends on the GPU model, the operation sizes, and the model architecture. Let us measure time per step for FP32 vs FP16 on this model.
import time
def benchmark_step(
model, use_amp, amp_dtype, use_scaler, n_warmup=3, n_bench=5
):
"""Benchmark a training step and return mean time in milliseconds."""
optimizer = optim.AdamW(model.parameters(), lr=1e-4)
scaler = GradScaler(device=device.type) if use_scaler else None
# Warmup
for _ in range(n_warmup):
x = torch.randint(0, vocab_size, (batch_size, seq_len), device=device)
y = torch.randint(0, vocab_size, (batch_size, seq_len), device=device)
optimizer.zero_grad()
if use_amp:
with autocast(device_type=device.type, dtype=amp_dtype):
logits = model(x)
loss = nn.CrossEntropyLoss()(
logits.view(-1, vocab_size), y.view(-1)
)
if scaler:
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(optimizer)
scaler.update()
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
else:
logits = model(x)
loss = nn.CrossEntropyLoss()(
logits.view(-1, vocab_size), y.view(-1)
)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
if device.type == "cuda":
torch.cuda.synchronize()
# Benchmark
times = []
for _ in range(n_bench):
x = torch.randint(0, vocab_size, (batch_size, seq_len), device=device)
y = torch.randint(0, vocab_size, (batch_size, seq_len), device=device)
optimizer.zero_grad()
start = time.perf_counter()
if use_amp:
with autocast(device_type=device.type, dtype=amp_dtype):
logits = model(x)
loss = nn.CrossEntropyLoss()(
logits.view(-1, vocab_size), y.view(-1)
)
if scaler:
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(optimizer)
scaler.update()
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
else:
logits = model(x)
loss = nn.CrossEntropyLoss()(
logits.view(-1, vocab_size), y.view(-1)
)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
if device.type == "cuda":
torch.cuda.synchronize()
end = time.perf_counter()
times.append((end - start) * 1000)
return float(np.mean(times)), float(np.std(times))# Create fresh models for benchmarking
model_fp32_bench = SimpleTransformer(
vocab_size=vocab_size, d_model=d_model, n_layers=n_layers, seq_len=seq_len
).to(device)
model_fp16_bench = SimpleTransformer(
vocab_size=vocab_size, d_model=d_model, n_layers=n_layers, seq_len=seq_len
).to(device)
t_fp32_mean, t_fp32_std = benchmark_step(
model_fp32_bench, use_amp=False, amp_dtype=None, use_scaler=False
)
t_fp16_mean, t_fp16_std = benchmark_step(
model_fp16_bench, use_amp=True, amp_dtype=torch.float16, use_scaler=True
)
speedup = t_fp32_mean / t_fp16_meanFP32 step time: 21.30 +/- 0.27 ms FP16 step time: 346.84 +/- 1.37 ms Speedup (FP16 vs FP32): 0.06x Note: On CPU, speedup may be < 1x because FP16 tensor cores are GPU-only. On A100/H100 GPUs, typical speedups are 1.5-3x for transformer workloads.
The speedup on CPU will be small or negative because CPUs do not have dedicated FP16 hardware. On a modern GPU with tensor cores (A100, H100, RTX 30/40 series), the speedup for transformer matrix multiplications is typically 1.5x to 3x, depending on the batch and sequence dimensions. Memory savings are consistent regardless of hardware.
Checking Gradient Dtypes at Runtime
During debugging, it is useful to inspect the dtypes flowing through the model during a mixed precision forward pass.
def inspect_dtypes(model, x, amp_dtype):
"""Run one forward pass and print the dtype of each module's output."""
dtypes = {}
def hook(name):
def _hook(module, input, output):
if isinstance(output, torch.Tensor):
dtypes[name] = str(output.dtype)
return _hook
hooks = []
for name, module in model.named_modules():
if isinstance(module, (nn.Linear, nn.LayerNorm, nn.MultiheadAttention)):
hooks.append(module.register_forward_hook(hook(name)))
with autocast(device_type=device.type, dtype=amp_dtype):
_ = model(x)
for h in hooks:
h.remove()
return dtypesDtype inspection during FP16 autocast: -------------------------------------------------- 0.norm1 -> torch.float32 ff.0 -> torch.float16 ff.2 -> torch.float16 0.norm2 -> torch.float32 1.norm1 -> torch.float32 1.norm2 -> torch.float32 norm -> torch.float32 head -> torch.float16
The dtype inspection confirms that nn.Linear layers (which host matrix multiplications) run in FP16 within the autocast context, while nn.LayerNorm stays in FP32. This is the autocast dispatch table in action. The pattern aligns exactly with the theory: operations sensitive to precision stay in FP32, while compute-heavy projection operations use FP16.
Key Parameters
The key parameters for mixed precision training are:
dtypeinautocast:torch.float16for FP16 (requires loss scaling on older hardware),torch.bfloat16for BF16 (preferred on Ampere GPUs and newer).init_scaleinGradScaler: The initial loss scale value. Default is 65536. Lower values reduce early-training overflow at the cost of reduced gradient precision.growth_factorinGradScaler: How much to multiply the scale after a clean stretch. Default is 2.0.backoff_factorinGradScaler: How much to reduce the scale after overflow. Default is 0.5.growth_intervalinGradScaler: Number of consecutive steps without overflow before scaling up. Default is 2000. Lower values make the scale more aggressive.max_norminclip_grad_norm_: The gradient clipping threshold. Apply this after callingscaler.unscale_().