Mixed precision, accumulation and activation checkpointing

Why it matters

Lower precision is the default way to make training faster and activations smaller; gradient accumulation and activation checkpointing are the default ways to fit larger effective batches and longer contexts into memory. Each knob changes numerics or cost in specific ways. Interviewers probe whether you know why loss scaling exists, what autocast keeps in fp32, and when accumulation is not equivalent to a larger batch.

Core ideas

Formats trade exponent bits (range) against mantissa bits (precision).

• fp32: 8 exponent bits, 23 mantissa bits.
• fp16: 5 exponent, 10 mantissa. The largest finite value is 65504 and values below about 6e-8 flush to zero, so small gradients underflow and large activations overflow.
• bf16: 8 exponent, 7 mantissa. The same range as fp32 with much coarser precision.
• TF32: a tensor-core mode that keeps the fp32 exponent with a 10-bit mantissa for matmul inputs and accumulates in fp32.
• fp8 (E4M3, E5M2): narrower still; in practice used with per-tensor or finer scaling factors managed by the framework.

The mixed-precision recipe. Keep an fp32 master copy of the weights; run forward and backward mostly in 16-bit; accumulate matmul results in fp32; keep precision-sensitive operations (softmax, normalization, reductions, losses) in fp32. PyTorch autocast makes this per-operator choice from explicit lists. With fp16, multiply the loss by a scale S before backward so small gradients stay representable, unscale before the optimizer step and before gradient clipping, skip the step and lower S when an inf or NaN appears, and raise S periodically (dynamic loss scaling). bf16's fp32-like range usually makes scaling unnecessary.

Gradient accumulation. Run k micro-batches, sum their gradients, step once. To reproduce a large-batch gradient, normalize correctly: divide each micro-batch's mean loss by k when micro-batches are equal-sized, or normalize by the total token count across all micro-batches when sequence lengths vary. With DDP, skip the all-reduce on all but the last micro-batch (no_sync), or you pay k times the communication. The effective batch is micro-batch × k × data-parallel ranks, and schedules count optimizer steps, not micro-batches. Changing the effective batch usually needs learning-rate retuning; linear scaling is a heuristic, not a law. Batch normalization computes statistics per micro-batch, so accumulation does not reproduce large-batch behavior for models that use it.

Activation checkpointing. Store activations only at segment boundaries and recompute the inside of each segment during backward. Recomputing everything costs about one extra forward pass, roughly a third more compute when backward costs about twice the forward. Chen et al. show that checkpointing every √n layers reduces activation memory for n layers to O(√n). Selective checkpointing recomputes only operations that are cheap to redo but expensive to store, such as attention softmax outputs. Random state must be replayed so dropout masks match; PyTorch's checkpoint utility does this by default.

Common misconceptions

• "bf16 is strictly better than fp16." bf16 trades precision for range. It avoids loss scaling, but its 7-bit mantissa makes accumulation and small updates noisier, so reductions and master weights still need fp32.
• "Mixed precision means everything runs in 16-bit." Autocast keeps sensitive operations in fp32, matmuls accumulate in fp32, and the optimizer updates fp32 master weights.
• "Loss scaling changes the optimization." Done correctly, gradients are unscaled before the update, so updates match except for skipped overflow steps. Clipping still-scaled gradients is a bug.
• "Gradient accumulation is identical to a larger batch." Only with correct normalization and no batch-dependent layers; batch norm, dropout randomness and data order differ, and throughput is not that of more devices.
• "Checkpointing reduces all memory." It shrinks activations only; weights, gradients and optimizer state are untouched.

How to approach it

Break things on purpose: train with fp16 and no scaler until gradients underflow or overflow, then add the scaler and watch the scale adapt. Verify that accumulated gradients match large-batch gradients within a tolerance. Measure peak memory and step time with checkpointing on and off. In interviews, connect each choice to range, precision, memory and compute, and say what you would monitor (loss-scale history, inf and NaN counts, gradient norms).