Mixed-precision training

Deep dive

How it works

Formats (eps = 2^−mantissa bits):

In: fp32 master weights, a batch · Out: updated fp32 weights · State: master weights, optimizer state, and for fp16 the loss scale · Cost: same FLOPs, 16-bit GEMMs run 2x (bf16) to 4x (fp8) faster than tf32.

Per step:

1. Cast fp32 master weights to bf16 (or keep a bf16 copy and refresh it after each update).
2. Forward under autocast. GEMM/conv in bf16 with fp32 accumulators. Softmax, log, exp, norms, losses and reductions in fp32.
3. Backward. Gradients of bf16 tensors are bf16. Accumulate and all-reduce them in fp32 when gradient accumulation spans many micro-batches.
4. fp16 only: scale the loss by S before backward and unscale grads by 1/S. If any grad is inf/NaN, skip the step and halve S.
5. Optimizer step in fp32 on the master weights.

Why the master copy matters: bf16 values near 1 are spaced 2^−7 ≈ 0.8% apart, so a relative update below about 0.4% (half a spacing) rounds away. An Adam step of lr = 3e-4 on a weight of size about 0.02 moves it by about 1.5%, which survives. Late in training (lr 3e-5, 0.15% of the weight), or for large weights, most updates are rounded away entirely.

Code

The numerics of a mixed-precision step on plain slices: bf16 storage with round-to-nearest-even, fp32 accumulation, the fp32 master copy, and dynamic loss scaling.

fn f32_to_bf16(x: f32) -> u16 {
    let b = x.to_bits();
    if x.is_nan() { return ((b >> 16) | 0x40) as u16; }      // keep NaN a (quiet) NaN
    let round = 0x7FFF + ((b >> 16) & 1);                      // +0.5 ulp, ties to even
    ((b + round) >> 16) as u16                                 // overflow rounds to ±inf correctly
}
fn bf16_to_f32(h: u16) -> f32 { f32::from_bits((h as u32) << 16) }

/// Dot product as a tensor core does it: bf16 inputs, fp32 accumulator.
fn dot_bf16(a: &[u16], b: &[u16]) -> f32 {
    a.iter().zip(b).map(|(&x, &y)| bf16_to_f32(x) * bf16_to_f32(y)).sum()
}

/// fp16 dynamic loss scaling, with torch.amp.GradScaler's defaults.
pub struct LossScaler { pub scale: f32, good_steps: u32, growth_interval: u32 }

impl LossScaler {
    pub fn new() -> Self { Self { scale: 65536.0, good_steps: 0, growth_interval: 2000 } }
    /// `grads` came from backward(loss · scale). Returns false = skip this optimizer step.
    pub fn unscale(&mut self, grads: &mut [f32]) -> bool {
        if grads.iter().any(|g| !g.is_finite()) {
            self.scale *= 0.5;                                  // overflow: back off, drop the step
            self.good_steps = 0;
            return false;
        }
        let inv = 1.0 / self.scale;
        grads.iter_mut().for_each(|g| *g *= inv);              // unscale BEFORE clipping
        self.good_steps += 1;
        if self.good_steps == self.growth_interval { self.scale *= 2.0; self.good_steps = 0; }
        true
    }
}

fn main() {
    // 1. bf16 spacing near 1.0 is 2^-7: a relative change below 2^-8 rounds away.
    assert_eq!(bf16_to_f32(f32_to_bf16(1.0 + 1.0 / 512.0)), 1.0);
    assert_eq!(bf16_to_f32(f32_to_bf16(1.0 + 3.0 / 512.0)), 1.0 + 1.0 / 128.0);
    assert_eq!(bf16_to_f32(f32_to_bf16(f32::MAX)), f32::INFINITY);

    // 2. Why master weights: 1000 updates of -3e-5 on w = 0.02 (each 0.15% of |w|).
    let (mut w_bf16, mut w_master) = (f32_to_bf16(0.02), 0.02f32);
    for _ in 0..1000 {
        w_bf16 = f32_to_bf16(bf16_to_f32(w_bf16) - 3e-5);      // every update rounds back to 0.02
        w_master -= 3e-5;
    }
    println!("bf16-only weight {:.4}, fp32 master {:.4}", bf16_to_f32(w_bf16), w_master);

    // 3. Why fp32 accumulators: a bf16 running sum of 4096 ones stalls at 256.
    let one = f32_to_bf16(1.0);
    let acc_bf16 = (0..4096).fold(f32_to_bf16(0.0), |s, _| f32_to_bf16(bf16_to_f32(s) + bf16_to_f32(one)));
    let ones = vec![one; 4096];
    println!("bf16 accumulator {}, fp32 accumulator {}", bf16_to_f32(acc_bf16), dot_bf16(&ones, &ones));

    // 4. Loss scaling: an inf in the scaled grads halves the scale and skips the step.
    let mut sc = LossScaler::new();
    assert!(!sc.unscale(&mut [1.0, f32::INFINITY]) && sc.scale == 32768.0);
    let mut g = [32768.0f32 * 1e-3, -32768.0 * 2e-6];
    assert!(sc.unscale(&mut g) && (g[0] - 1e-3).abs() < 1e-9);
}

Complexity and constants

• FLOPs unchanged. Matmul throughput rises 2x (bf16 vs tf32) or 4x (fp8 vs tf32). End-to-end speedups are about 1.5-3x, because non-GEMM ops are memory-bound and gain only from halved bytes.
• Model state with Adam: 2 (bf16 weights) + 2 (bf16 grads) + 4 + 4 + 4 (fp32 master, m, v) = 16 B/param. Mixed precision does not reduce this. It reduces activations (2 B instead of 4) and bandwidth.
• fp8 on H100 is reported at about 1.3-1.5x end-to-end over bf16 at LLM scale, less than the 2x peak because attention, norms and communication stay 16-bit.

Efficiency

• Cost: H100 SXM dense peaks (NVIDIA datasheet): fp32 67, tf32 495, bf16 989, fp8 1979 TFLOPS. Model state stays 16 B/param with Adam; activations take 2 B instead of 4. Bound: tensor-core compute for GEMMs, memory bandwidth for everything else.
• Levers:
  • bf16 or fp16 autocast with fp32 master weights — nearly 2x less memory (Micikevicius et al.), end-to-end 1.5-3x (order of) — all GPU training
  • fp8 with per-tensor or block scales — peak GEMM rate 2x bf16; end-to-end order of 1.3-1.5x at LLM scale — Hopper and newer
  • TF32 matmul precision for fp32 code — 495 instead of 67 TFLOPS — legacy fp32 models
  • keep the overflow flag on the device — no host sync per step — fp16 loss scaling
  • fp32 reduce dtype for gradients — avoids accumulating in bf16 — large data-parallel worlds
• Naive → standard → state of the art:
  • fp32 on CUDA cores · 67 TFLOPS peak
  • bf16 autocast with fp32 master · 989 TFLOPS peak
  • fp8 with block scales, MXFP8 on Blackwell · 1979 TFLOPS peak on H100

Implementation notes

• Accumulate in fp32 inside GEMMs: mma instructions take bf16 inputs and fp32 accumulators. Never sum long dot products in bf16; after about 256 equal terms, new additions vanish.
• fp8 recipes. Transformer Engine uses per-tensor scales from an amax history ("delayed scaling"), E4M3 forward and E5M2 for grads. DeepSeek-V3 uses 1×128 activation tiles and 128×128 weight blocks with their own scales, and promotes partial sums to fp32 every 128 elements of K. Blackwell MXFP8 is hardware block scaling (32 elements share one E8M0 scale).
• Keep the residual stream and the logits/loss in fp32 or bf16 with fp32 softmax. Attention logits are the first thing to overflow in fp16.
• On GPU, keep the overflow flag on the device and skip the update conditionally, as PyTorch's GradScaler does. Reading it on the host costs a sync every step.
• With ZeRO/FSDP, choose param_dtype=bf16 for the gathered weights and reduce_dtype=fp32 for gradient reduce-scatter when the data-parallel world is large.

Pitfalls

• Clipping before unscaling clips against a norm 65536x too large, so it never triggers.
• Pure-bf16 weights without a master copy: training stalls late as updates round to zero. Use an fp32 master, Kahan-compensated updates, or stochastic rounding.
• Custom CUDA ops not registered with autocast receive bf16 inputs and accumulate in bf16.
• Overflow in fp16 shows up as a scale that keeps halving to tiny values (below about 1). The model itself has a problem; loss scaling cannot fix it.
• Comparing against fp32 references with tight tolerances. Expect relative differences of order 1e-2 for bf16 activations.
• PyTorch keeps TF32 off for matmuls by default. Plain fp32 code runs at 67 instead of 495 TFLOPS on H100 unless you set the matmul precision.

Variants and alternatives

In production

PyTorch AMP (autocast, GradScaler), Megatron-LM (bf16 + fp32 gradient accumulation), NVIDIA Transformer Engine (fp8), DeepSeek-V3 (fp8 at 671B MoE scale), JAX on TPU (bf16 is TPU-native), Whisper speech models (FP16 with dynamic loss scaling).

Implementations

Sources

• Micikevicius et al., Mixed Precision Training, ICLR 2018. paper
• Kalamkar et al., A Study of BFLOAT16 for Deep Learning Training, 2019. paper
• Micikevicius et al., FP8 Formats for Deep Learning, 2022. paper
• Rouhani et al., Microscaling Data Formats for Deep Learning, 2023. paper
• DeepSeek-AI, DeepSeek-V3 Technical Report, 2024. paper · code
• Radford et al., Robust Speech Recognition via Large-Scale Weak Supervision (Whisper), 2022 (FP16 with dynamic loss scaling). paper