Python / PyTorch Fundamentals Interview Questions
What is mixed precision training in PyTorch and how do you implement it with torch.cuda.amp?
Mixed precision training runs most operations in FP16 (or BF16) for speed while keeping a master copy of weights in FP32 for numerical stability. Modern GPUs (Volta and later) have dedicated hardware (Tensor Cores) that make FP16 matrix multiplication significantly faster than FP32.
import torch import torch.nn as nn from torch.cuda.amp import autocast, GradScaler model = nn.Linear(1024, 512).cuda() optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3) scaler = GradScaler() # manages loss scaling to prevent FP16 underflow x = torch.randn(256, 1024).cuda() y = torch.randn(256, 512).cuda() for step in range(100): optimizer.zero_grad() # autocast: automatically runs eligible ops in FP16/BF16 with autocast(device_type="cuda", dtype=torch.float16): y_hat = model(x) # matmul runs in FP16 â faster! loss = nn.MSELoss()(y_hat, y) # Loss scaling: inflate loss before backward to prevent small # gradients from underflowing to zero in FP16's limited range scaler.scale(loss).backward() scaler.unscale_(optimizer) # restore original gradient magnitudes nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) # skips the step if grads are inf/NaN scaler.update() # adjusts scale factor for next iteration # BFloat16: no GradScaler needed (same exponent range as FP32) with autocast(device_type="cuda", dtype=torch.bfloat16): y_hat = model(x) # no underflow risk â scaling unnecessary
More Related questions...