Tensors & Neural Plumbing advanced 8 min read 5 flashcards

Accumulation Precision and Mixed-Dtype Plumbing

Storing weights in 16 bits is safe because almost nothing is actually computed in 16 bits — matmuls accumulate in fp32, reductions run in fp32, and the master copy is fp32, and every documented failure of low-precision training is one of these rules being broken.

"We train in bf16" is a statement about storage, not about arithmetic, and treating it as a statement about arithmetic is how people produce mixed-precision training runs that diverge for no visible reason. A bf16 training step involves at least three different numeric formats, and which operation gets which one is not a detail — it is the reason the technique works.

Why accumulation is the crux

Adding a small number to a large one in low precision loses the small number entirely. In bf16, with roughly 8 bits of mantissa, adding 1 to 512 is representable; keep adding 1 to a running sum and eventually each addition rounds to no change at all, and the sum stops growing while the true value keeps increasing. This is stagnation, and a matmul over a reduction dimension of 4096 is precisely a long chain of such additions.

Hardware solves it rather than software: tensor cores multiply in low precision and accumulate in fp32, so the products are 16-bit and the running sum is 32-bit. This is why a bf16 matmul is accurate enough to train with, and it is also why a hand-written kernel that accumulates in bf16 to save registers will produce a model that trains worse for reasons no hyperparameter sweep will fix.

bf16 versus fp16, and what each gives up

Both are 16 bits and they spend those bits oppositely. fp16 has 5 exponent bits and 10 mantissa bits: more precision, and a dynamic range so narrow that gradients underflow to zero, which is why fp16 training needs loss scaling — multiply the loss by a large constant before backward, divide it out of the gradients afterwards — plus machinery to detect overflow and skip the step. bf16 has 8 exponent bits and 7 mantissa bits, matching fp32's range exactly, so it underflows where fp32 underflows and loss scaling becomes unnecessary. The trade is fewer mantissa bits, which is tolerable precisely because accumulation happens in fp32.

The fp32 islands

Several operations stay in fp32 regardless of the nominal training precision, and each for a specific reason:

  • The master weights. Optimiser updates are typically much smaller than the weights themselves, so applying them to a 16-bit copy stagnates: the update rounds away and the weight never moves. A fp32 master copy is updated, and a 16-bit copy is cast from it for the forward pass.
  • Normalisation statistics. LayerNorm and RMSNorm compute a mean and variance over the hidden dimension — a reduction, therefore vulnerable to exactly the accumulation problem above.
  • Softmax. Exponentials produce large dynamic range, and the max-subtraction trick that keeps them stable relies on precision that bf16 does not have to spare.
  • Loss and logits. Cross-entropy over a large vocabulary is another long reduction, and it sits at the point where errors propagate into every gradient.

The optimiser state is its own budget line and often the dominant one: Adam keeps two fp32 moments per parameter, so optimiser state alone can exceed the weights several times over (see optimiser state memory).

When it breaks

  • Custom kernels are where the rules get broken. A fused kernel that keeps a reduction in bf16 to save registers is fast and subtly wrong, and the symptom is a slightly worse final loss rather than a crash. Audit the accumulator dtype of anything hand-written or newly adopted.
  • fp8 moves the problem rather than solving it. Sub-16-bit formats need per-tensor or finer-grained scaling factors to place values inside a much narrower representable range, and choosing and updating those scales is a live source of instability.
  • Casting order changes results. (a.to(bf16) * b.to(bf16)).sum() and (a * b).sum().to(bf16) are different computations; framework autocast rules decide which you get, and they are not always the one you assumed.
  • Determinism is a separate cost. Reduction order is not fixed across runs on parallel hardware, and floating-point addition is not associative, so bitwise-identical reruns require deterministic kernels that are usually slower (see numerical computation gotchas).
  • Gradient all-reduce precision is an easy silent regression. Compressing gradients to 16 bits for the collective saves bandwidth and adds an accumulation error across thousands of ranks; whether it is safe depends on your world size, not on whether it worked at eight GPUs.
Check yourself

5 flashcards for this concept

Click a card to reveal the answer.

Drill the whole track