1.4 Numerical Computing, Floating Point & Precision¶

Numbers in a computer are not mathematical reals — they are finite approximations governed by precise rules laid out in the IEEE 754 standard. For most software this distinction barely matters. For deep learning it is existential: the choice of floating-point format is a first-class design decision that affects training stability, convergence speed, hardware utilization, and memory budget all at once. This chapter gives you a complete mental model of how floating-point works, where it breaks, and why the community settled on bf16 as the workhorse of LLM training — along with the algorithms that keep numerics stable when the format alone is not enough.

Familiarity with the concepts here is assumed throughout the rest of the book. Mixed Precision, bf16 & FP8 Training applies these ideas to the full training loop, and FlashAttention I: IO-Awareness & The Online Softmax builds numerically stable attention kernels on top of the logsumexp identity we derive here.

IEEE 754: The Foundation¶

The IEEE 754 standard (1985, revised 2008 and 2019) defines how a real number is represented, rounded, and how arithmetic operations behave. Every mainstream ML accelerator — NVIDIA GPUs, Google TPUs, AMD MI-series — implements IEEE 754 (or a deliberate approximation of it, as we will see with TF32 and FP8).

The bit layout¶

Every IEEE 754 number is laid out as three fields:

The leading 1. is implicit for normal numbers, buying a free bit of precision. Special bit patterns encode \(\pm\infty\) (all exponent bits set, mantissa zero) and NaN (all exponent bits set, non-zero mantissa).

The critical intuition: exponent bits control dynamic range, mantissa bits control precision. More exponent bits → you can represent very large and very small numbers without overflow/underflow. More mantissa bits → you can represent nearby numbers that differ by a tiny amount.

The Floating-Point Zoo: fp32, tf32, fp16, bf16, fp8¶

The last decade has produced a proliferation of numeric formats as hardware designers chase the sweet spot between precision (which costs mantissa bits) and performance (which rewards narrower types). Here is the full comparison:

* TF32 packs into a 32-bit register; only 19 bits carry information.

The “min normal” column is the smallest normal magnitude; subnormals extend a little lower (fp8 E4M3 down to \(2^{-9} \approx 1.95\times10^{-3}\), E5M2 down to \(2^{-16} \approx 1.5\times10^{-5}\)). Note also that the E4M3 variant hardware actually implements is E4M3FN (“finite”): it spends the all-ones exponent on ordinary numbers rather than infinities, which is how it reaches a max of 448 instead of 240. PyTorch exposes it as torch.float8_e4m3fn; you can confirm every row of this table yourself with torch.finfo(dtype).

FP32  [s|eeeeeeee|mmmmmmmmmmmmmmmmmmmmmmm]   32 bits
TF32  [s|eeeeeeee|mmmmmmmmmm         ---]   19 effective bits
BF16  [s|eeeeeeee|mmmmmmm]                  16 bits
FP16  [s|eeeee|mmmmmmmmmm]                  16 bits
FP8   [s|eeee|mmm]  (E4M3)                   8 bits
      [s|eeeee|mm] (E5M2)                    8 bits

The key visual: bf16 is the top 16 bits of fp32. They share the same 8-bit exponent, so the dynamic range is identical. You can truncate fp32 to bf16 by simply dropping the lower 16 bits of the mantissa — no reformatting, just a shift. This hardware-friendliness is why bf16 won.

Machine epsilon and ULP¶

Machine epsilon (\(\varepsilon_{\text{mach}}\)) is the smallest value such that \(1 + \varepsilon_{\text{mach}} \neq 1\) in the given format. It equals \(2^{-p}\) where \(p\) is the number of mantissa bits:

Unit in the last place (ULP) is the spacing between adjacent representable numbers at any given magnitude. For a normal fp32 number with magnitude around 1.0, 1 ULP ≈ \(1.2 \times 10^{-7}\). For a number around \(10^4\), 1 ULP ≈ \(1.2 \times 10^{-3}\) — precision scales with magnitude.

Dynamic Range vs. Precision: The Core Tradeoff¶

The two properties that determine a format’s fitness are independent:

• Dynamic range: the ratio of the largest to the smallest representable normal number, set by the exponent field.
• Relative precision: how finely you can resolve nearby values, set by the mantissa field.

fp16 has a 5-bit exponent, giving a maximum representable value of roughly 65504. That sounds large for activations, but gradients and weight updates during training can transiently exceed this — and once you hit overflow you get inf, which poisons the entire forward/backward pass.

bf16 keeps fp32’s 8-bit exponent (max value ~\(3.4 \times 10^{38}\)) but sacrifices mantissa bits. Its relative precision (\(\varepsilon_{\text{mach}} \approx 7.8 \times 10^{-3}\)) is roughly 10× worse than fp16, but that turns out to be acceptable for gradient noise in SGD. The dynamic range is what matters for training stability, and bf16 has it.

Overflow, underflow, and subnormals¶

Overflow occurs when a computed value exceeds the format’s maximum. In IEEE 754, the result is +inf or -inf. Arithmetic involving inf often produces NaN, which then propagates and kills training.

Underflow occurs when a value is too small to represent as a normal number. IEEE 754 defines subnormal (denormal) numbers that fill the gap between zero and the minimum normal value by using a leading 0. instead of 1.. Subnormals sacrifice range for a graceful flush to zero, but on hardware they are typically much slower to process — or flushed to zero (FTZ mode) entirely. On most GPU training setups, FTZ is enabled, so very small activations silently become zero.

Catastrophic Cancellation¶

Catastrophic cancellation is the most insidious floating-point pitfall. It occurs when you subtract two nearly-equal numbers, and the relative error in the result explodes.

Consider computing \(a - b\) where \(a = 1.0000001\) and \(b = 1.0000000\) in fp32. The true result is \(10^{-7}\), right at the edge of fp32’s precision. If both \(a\) and \(b\) carry even tiny rounding errors (say, each off by \(\pm 10^{-7}\)), the error in the difference is \(\pm 2 \times 10^{-7}\) — the same order of magnitude as the result itself. You have lost all significant digits.

Stable Softmax and the LogSumExp Identity¶

The softmax function is defined as:

The naive implementation overflows for large logits and underflows for very negative logits. The standard stabilization uses the identity:

This is valid because \(e^{z_i - m} / \sum_j e^{z_j - m} = (e^{-m} e^{z_i}) / (e^{-m} \sum_j e^{z_j})\) — the \(e^{-m}\) cancels.

The related operation logsumexp appears constantly in log-probabilities, cross-entropy, and attention:

The numerically stable version is:

Because \(z_j - m \leq 0\) for all \(j\), the exponentials are in \([0, 1]\) — no overflow possible, and underflow is graceful (a very negative \(z_j - m\) contributes essentially zero to the sum, which is correct behavior).

import torch
import numpy as np

# ------------------------------------------------------------------
# Numerically stable softmax — implemented from scratch in PyTorch.
# Matches torch.nn.functional.softmax to machine precision.
# ------------------------------------------------------------------

def naive_softmax(z: torch.Tensor) -> torch.Tensor:
    """Naive implementation — overflows for large inputs."""
    e = torch.exp(z)              # Can produce inf for large z
    return e / e.sum(dim=-1, keepdim=True)


def stable_softmax(z: torch.Tensor) -> torch.Tensor:
    """
    Numerically stable softmax via the max-subtraction trick.

    The key identity:
        softmax(z)_i = softmax(z - max(z))_i
    because the max subtraction cancels in numerator and denominator.
    """
    # Subtract maximum along last dimension; keepdim for broadcasting
    m = z.max(dim=-1, keepdim=True).values   # shape (..., 1)
    shifted = z - m                           # all values <= 0
    e = torch.exp(shifted)                    # all values in (0, 1]
    return e / e.sum(dim=-1, keepdim=True)


def stable_log_softmax(z: torch.Tensor) -> torch.Tensor:
    """
    Numerically stable log-softmax.

    log softmax(z)_i = z_i - logsumexp(z)
    logsumexp(z) = m + log(sum(exp(z - m)))

    Computing log after softmax loses precision; compute log-softmax directly.
    """
    m = z.max(dim=-1, keepdim=True).values
    shifted = z - m
    log_sum_exp = m.squeeze(-1) + torch.log(torch.exp(shifted).sum(dim=-1))
    # Broadcast log_sum_exp back for subtraction
    return z - log_sum_exp.unsqueeze(-1)


def stable_logsumexp(z: torch.Tensor) -> torch.Tensor:
    """Standalone logsumexp with the shift trick."""
    m = z.max(dim=-1).values              # shape (...)
    shifted = z - m.unsqueeze(-1)         # values in (-inf, 0]
    return m + torch.log(torch.exp(shifted).sum(dim=-1))


# ---------------------------------------------------------------
# Demonstrate fp16 overflow and how stable_softmax avoids it
# ---------------------------------------------------------------
if __name__ == "__main__":
    torch.set_printoptions(precision=6)

    # Large logits — pathological case
    z_large = torch.tensor([1000.0, 1001.0, 999.0])

    print("=== fp32 naive softmax (large logits) ===")
    z_fp32 = z_large.to(torch.float32)
    result_naive = naive_softmax(z_fp32)
    print(f"  Result: {result_naive}")          # [nan, nan, nan] — overflow!

    print("\n=== fp32 stable softmax (large logits) ===")
    result_stable = stable_softmax(z_fp32)
    print(f"  Result: {result_stable}")         # [0.2447, 0.6652, 0.0900] — correct

    # Show fp16 overflow with much smaller values
    print("\n=== fp16 naive softmax (overflow at ~88) ===")
    z_fp16 = torch.tensor([100.0, 101.0, 99.0], dtype=torch.float16)
    result_fp16_naive = naive_softmax(z_fp16)
    print(f"  Result: {result_fp16_naive}")     # [nan, nan, nan] — fp16 overflows at exp(88)

    print("\n=== fp16 stable softmax ===")
    result_fp16_stable = stable_softmax(z_fp16)
    print(f"  Result: {result_fp16_stable}")    # Correct softmax

    # Verify logsumexp
    print("\n=== logsumexp verification ===")
    z_test = torch.tensor([1.0, 2.0, 3.0])
    lse = stable_logsumexp(z_test)
    expected = torch.log(torch.exp(z_test).sum())
    print(f"  stable_logsumexp: {lse:.6f}")
    print(f"  torch reference:  {expected:.6f}")
    print(f"  match: {torch.allclose(lse, expected)}")

=== fp32 naive softmax (large logits) ===
  Result: tensor([nan, nan, nan])

=== fp32 stable softmax (large logits) ===
  Result: tensor([0.244728, 0.665241, 0.090031])

=== fp16 naive softmax (overflow at ~88) ===
  Result: tensor([nan, nan, nan], dtype=torch.float16)

=== fp16 stable softmax ===
  Result: tensor([0.244751, 0.665527, 0.090088], dtype=torch.float16)

=== logsumexp verification ===
  stable_logsumexp: 3.407606
  torch reference:  3.407606
  match: True

FlashAttention extends exactly this idea to the attention operation, fusing the online softmax into a single GPU pass without materializing the full \(N \times N\) attention matrix — see FlashAttention I: IO-Awareness & The Online Softmax for the full derivation.

Kahan Summation and Compensated Arithmetic¶

Summing \(N\) floating-point numbers naively accumulates \(O(N)\) rounding errors. For \(N = 10^9\) (a common size in LLM weight norms), this can become significant.

Kahan summation (Kahan, 1965) maintains a running “compensation” variable that tracks the error lost to rounding at each step, restoring it to the next addition:

Formally:

def kahan_sum(values):
    """
    Kahan compensated summation — reduces floating-point error from O(N*eps)
    to O(eps) regardless of N, at the cost of ~2x FLOPs per element.

    Reference: Kahan, W. (1965). "Further remarks on reducing truncation errors."
    Communications of the ACM, 8(1), 40.
    """
    total = 0.0
    compensation = 0.0   # tracks the "lost" low-order bits

    for v in values:
        # The compensation corrects for the rounding error from the previous step.
        # y is the value adjusted for previous lost bits.
        y = v - compensation
        # t is the running total after adding y; but it loses the low bits of y
        t = total + y
        # (t - total) recovers what was actually added — which differs from y
        # due to rounding. The difference is the new compensation.
        compensation = (t - total) - y
        total = t

    return total


# ---------------------------------------------------------------
# Demonstrate the error reduction on a large sum
# ---------------------------------------------------------------
import numpy as np

N = 1_000_000
# Create values that sum to exactly N (each value is 1.0)
# Then add a tiny perturbation to expose rounding errors
values_f32 = np.ones(N, dtype=np.float32)
values_f32[0] = 1.0 + 1e-7   # tiny perturbation

true_sum = float(N) + 1e-7    # exact answer

# Naive float32 accumulation
naive_result = float(np.sum(values_f32))

# Kahan in Python (pedagogical — slow)
kahan_result = kahan_sum(values_f32.tolist())

# float64 reference
ref_result = float(np.sum(values_f32.astype(np.float64)))

print(f"True sum:          {true_sum:.10f}")
print(f"Naive float32:     {naive_result:.10f}  error={abs(naive_result - true_sum):.2e}")
print(f"Kahan float32:     {kahan_result:.10f}  error={abs(kahan_result - true_sum):.2e}")
print(f"Float64 reference: {ref_result:.10f}  error={abs(ref_result - true_sum):.2e}")

Kahan summation matters most in gradient accumulation (summing many micro-batch gradients before an optimizer step) and in weight update computations. PyTorch’s optimizer implementations typically accumulate in fp32 even when parameters are stored in fp16/bf16; this is the “master weights” pattern described in the mixed-precision training chapter.

Stochastic rounding: the other way to beat accumulated bias¶

Kahan’s insight is to remember the discarded bits. The alternative is to randomize them. Default IEEE rounding is round-to-nearest-even, which is deterministic and therefore biased in a repeated accumulation: if every increment is smaller than half a ULP of the running total, every single one rounds to zero and the sum freezes forever. That is precisely the fp16/bf16 “stagnant update” failure from earlier in this chapter.

Stochastic rounding rounds \(x\) down to the representable value \(x_{\text{lo}}\) with probability \((x_{\text{hi}} - x)/(x_{\text{hi}} - x_{\text{lo}})\) and up otherwise, so that

The rounding error becomes zero-mean noise instead of a systematic drift, so errors cancel over many steps rather than compounding. Implementing it for fp32 → bf16 is a two-line bit trick: bf16 is the top 16 bits of fp32, so adding a uniform random 16-bit integer to the mantissa before truncating makes the carry into bit 16 happen with exactly the right probability.

import torch


def stochastic_round_to_bf16(x: torch.Tensor) -> torch.Tensor:
    """
    Round fp32 -> bf16 stochastically, so that E[result] == x.

    bf16 is literally the top 16 bits of fp32. Adding a uniform random
    16-bit integer to the low mantissa bits before truncating makes the
    carry into bit 16 fire with probability equal to the fractional
    distance to the next bf16 value -- i.e. exactly stochastic rounding.
    """
    assert x.dtype == torch.float32
    bits = x.view(torch.int32)
    noise = torch.randint(0, 1 << 16, x.shape, dtype=torch.int32, device=x.device)
    return ((bits + noise) & ~0xFFFF).view(torch.float32).bfloat16()


# Accumulate 0.01 one thousand times in a bf16 accumulator.
# Once the total passes ~8, one ULP is 0.0625 and 0.01 is below half
# a ULP -- round-to-nearest sends every further increment to zero.
torch.manual_seed(0)
N, step = 1000, 0.01

acc = torch.zeros((), dtype=torch.bfloat16)
for _ in range(N):
    acc = (acc.float() + step).bfloat16()          # round-to-nearest
nearest = acc.item()

acc = torch.zeros((), dtype=torch.bfloat16)
for _ in range(N):
    acc = stochastic_round_to_bf16(acc.float() + step)   # stochastic
stoch = acc.item()

print(f"exact sum:                 {N * step:.4f}")
print(f"bf16, round-to-nearest:    {nearest:.4f}")
print(f"bf16, stochastic rounding: {stoch:.4f}")

exact sum:                 10.0000
bf16, round-to-nearest:    4.0000
bf16, stochastic rounding: 9.1875

Round-to-nearest freezes at 4.0 — it never recovers, no matter how many more increments you add. Stochastic rounding lands near 10; because it is unbiased rather than exact, it fluctuates around the true value from run to run (roughly 9–11 across seeds here) instead of collapsing. That trade — a little variance in exchange for no drift — is why stochastic rounding is the standard tool for keeping weights or optimizer states in bf16 (and is a required ingredient in the FP4 training recipes discussed below). Some accelerators, notably Graphcore IPUs and several TPU generations, implement it in hardware; on NVIDIA GPUs it is done in software as above or inside a library like Transformer Engine.

Why bf16 Won Training¶

The story of how bf16 displaced fp32 (and then fp16) as the default training format is worth understanding in detail, because it illustrates every tradeoff discussed above.

The fp16 era and its problems¶

Around 2017–2018, the community began training in fp16 with the mixed-precision training recipe (Micikevicius et al., ICLR 2018). The idea: store weights and activations in fp16, accumulate gradient updates in fp32 (“master weights”), and use loss scaling to prevent gradient underflow. Hardware reasons: fp16 halves the bytes moved per tensor, and Volta’s tensor cores delivered several times the fp32 throughput on matmuls.

The problems: 1. Overflow: weight norms, gradient norms, and optimizer states can transiently exceed 65504. Loss scaling helps but requires dynamic tuning. 2. Precision (the “stagnant update” problem): with only 10 mantissa bits (\(\varepsilon_{\text{mach}} \approx 9.8 \times 10^{-4}\)), the update \(w \leftarrow w + \Delta w\) is a no-op whenever \(|\Delta w| < \tfrac{1}{2}\varepsilon_{\text{mach}}|w|\) — the increment falls below half a ULP at the weight’s magnitude and rounds away. With \(|w| \approx 0.02\) and \(\varepsilon_{\text{mach}} \approx 10^{-3}\), any update smaller than about \(10^{-5}\) vanishes. Late in training, when the learning rate has decayed, a large fraction of updates land in exactly this regime and the model quietly stops learning. This — not overflow — is why fp32 master weights are mandatory in fp16 training. 3. Engineering complexity: loss scaling, gradient unscaling, and inf/NaN checks add fragile infrastructure.

The bf16 solution¶

Google introduced bf16 for TPU training and NVIDIA added native bf16 support in Ampere (A100, 2020). The format keeps fp32’s 8-bit exponent, sacrificing mantissa bits:

• No overflow: the dynamic range matches fp32. A value that fits in fp32 fits in bf16. Loss scaling becomes unnecessary.
• Simpler hardware path: truncate the lower 16 bits of a fp32 value to get bf16. No reformatting, no bias adjustment.
• Acceptable precision loss: the 7-bit mantissa (\(\varepsilon_{\text{mach}} \approx 7.8 \times 10^{-3}\)) is coarser than fp16’s 10-bit mantissa, but empirically the optimization noise from SGD/Adam dominates this rounding error for large models. The signal-to-noise ratio of gradients at LLM scale is low enough that you cannot tell the difference.

The practical result: training LLMs in pure bf16 (no loss scaling, no master weights in many implementations) “just works.” This simplification dramatically reduced training infrastructure complexity.

The library layer: what torch.autocast actually does¶

“Training in bf16” almost never means every tensor is bf16. In practice you run a mixed-precision loop, and the component that decides what goes where is torch.autocast — PyTorch’s automatic mixed precision (AMP) context manager. It is worth knowing precisely what it does, because it is the single most-used numerics tool in the entire stack.

Autocast maintains an internal per-operator policy list. Roughly:

• Cast to bf16/fp16: the big, rounding-tolerant, tensor-core-bound ops — matmul, linear, conv*, bmm, einsum, scaled_dot_product_attention.
• Force fp32: ops that are reductions, involve exponentials/logs, or are otherwise precision-critical — softmax, log_softmax, cross_entropy, layer_norm, sum, norm, pow, exp, log.
• Promote to the widest input dtype: ops like cat, stack, addcdiv where mixing would be ambiguous.

This is exactly the taxonomy the rest of this chapter argues for: matmuls are error-averaging dot products where bf16’s coarse mantissa washes out, while softmax and normalization are the places where cancellation and small-variance division bite.

import torch

dev = "cuda" if torch.cuda.is_available() else "cpu"   # autocast supports both
model = torch.nn.Linear(512, 512).to(dev)           # params stay fp32
x = torch.randn(8, 512, device=dev)                 # input fp32
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)  # states fp32

with torch.autocast(device_type=dev, dtype=torch.bfloat16):
    h = model(x)                    # matmul -> runs in bf16
    print(h.dtype)                  # torch.bfloat16
    p = torch.softmax(h, dim=-1)    # softmax -> forced back to fp32
    print(p.dtype)                  # torch.float32
    loss = p.mean()                 # fp32 loss

# Backward and the optimizer step run OUTSIDE the autocast context.
# Gradients arrive in the dtype of each parameter (fp32 here), so the
# update happens at full precision -- the "master weights" pattern.
loss.backward()
opt.step()
opt.zero_grad(set_to_none=True)

Two details that trip people up. First, backward() belongs outside the autocast block: autograd replays each op in the dtype it was recorded with, so wrapping backward adds nothing and can mis-cast. Second, autocast does not change your parameter dtypes — the Linear weights are still fp32; autocast casts a bf16 copy on the fly for the matmul. That is what gives you fp32 master weights for free. With fp16 you additionally need torch.amp.GradScaler for loss scaling; with bf16 you do not, which is the whole point of the format.

For a ~100M model, this is the entire precision plan:

A useful sanity check on this table: fp32 master weights + fp32 Adam moments cost \(4 + 4 + 4 = 12\) bytes per parameter, so a 100M-parameter model needs ~1.2 GB of optimizer/parameter state before a single activation is stored. The full accounting, plus the FP8 recipe, is in Mixed Precision, bf16 & FP8 Training; the concrete loop that uses it is The Pretraining Run: A Complete Single-GPU Training Loop.

TF32: the tensor-core compromise¶

TF32 (TensorFloat-32) is an NVIDIA-specific format introduced with Ampere. It is not a storage format — it is a compute format used internally by tensor cores. When you do a matrix multiply in fp32, Ampere automatically: 1. Rounds each fp32 input to tf32 (8 exponent bits, 10 mantissa bits — same range as fp32, precision of fp16). 2. Performs the multiply-accumulate in the tensor core. 3. Accumulates the result in fp32.

The effect: on A100, TF32 tensor-core matmul is roughly 8× the non-tensor-core fp32 rate, with minimal accuracy loss because the accumulation is fp32.

There is a footgun here that catches people constantly: PyTorch does not enable TF32 for matmuls by default. Since PyTorch 1.12, torch.backends.cuda.matmul.allow_tf32 defaults to False (so that A @ B on fp32 tensors gives fp32-faithful results), while torch.backends.cudnn.allow_tf32 defaults to True (so convolutions do use it). If you benchmark “fp32 training” on an A100 and it is mysteriously slow, this is usually why. The modern, dtype-agnostic way to opt in is:

import torch

# Preferred modern API: "highest" = true fp32, "high" = TF32 (or bf16x3),
# "medium" = bf16 for the internal matmul. Affects fp32 matmuls only.
torch.set_float32_matmul_precision("high")   # turn TF32 matmuls on

# Equivalent low-level switches (still supported):
# torch.backends.cuda.matmul.allow_tf32 = True
# torch.backends.cudnn.allow_tf32 = True     # already True by default

Turning this on is essentially free accuracy-wise and is a standard first line in any training script — including the Stack-100M run in The Pretraining Run: A Complete Single-GPU Training Loop.

import torch

# Check TF32 settings
print(f"TF32 for matmul: {torch.backends.cuda.matmul.allow_tf32}")
print(f"TF32 for cudnn:  {torch.backends.cudnn.allow_tf32}")

# Demonstrate the precision difference between fp32, tf32, and bf16 matmul
# (requires CUDA GPU)
if torch.cuda.is_available():
    torch.manual_seed(42)
    A = torch.randn(1024, 1024, device='cuda', dtype=torch.float32)
    B = torch.randn(1024, 1024, device='cuda', dtype=torch.float32)

    # fp64 reference
    ref = (A.double() @ B.double()).float()

    # fp32 with TF32 disabled (pure fp32 matmul)
    torch.backends.cuda.matmul.allow_tf32 = False
    out_fp32 = A @ B
    err_fp32 = (out_fp32 - ref).abs().max().item()

    # fp32 with TF32 enabled (Ampere behavior)
    torch.backends.cuda.matmul.allow_tf32 = True
    out_tf32 = A @ B
    err_tf32 = (out_tf32 - ref).abs().max().item()

    # bf16 matmul (explicit cast)
    out_bf16 = (A.bfloat16() @ B.bfloat16()).float()
    err_bf16 = (out_bf16 - ref).abs().max().item()

    print(f"\nMax absolute error vs fp64 reference (1024x1024 matmul):")
    print(f"  fp32 (no TF32): {err_fp32:.2e}")
    print(f"  TF32:           {err_tf32:.2e}")
    print(f"  bf16:           {err_bf16:.2e}")

FP8, and the FP4 frontier¶

FP8 (introduced in the H100 Transformer Engine) pushes further. Two variants: - E4M3 (4 exponent, 3 mantissa): higher precision, for forward activations and weights. - E5M2 (5 exponent, 2 mantissa): larger range, for gradients.

Because FP8 has only 3 or 2 mantissa bits, it requires dynamic quantization — per-tensor or per-block scaling factors to keep values in the representable range. The scale is what supplies the dynamic range that the format itself lacks; the 8 bits only have to encode the shape of the local distribution.

At the library level you have two real choices, and you should know both. NVIDIA’s Transformer Engine (transformer_engine.pytorch) provides drop-in te.Linear / te.TransformerLayer modules and an fp8_autocast context that manages the scaling factors and their delayed-scaling history for you. The framework-native alternative is torchao (pytorch/ao), which converts an existing nn.Module tree in place and composes with torch.compile and FSDP:

# pip install torchao
from torchao.float8 import convert_to_float8_training

# Swaps eligible nn.Linear layers for float8 versions that pick scales
# dynamically and dispatch to torch._scaled_mm under the hood.
convert_to_float8_training(model)
model = torch.compile(model)   # fusion of the scaling/cast ops matters a lot

Both ultimately bottom out in the same primitive — a scaled FP8 matmul, exposed in PyTorch as torch._scaled_mm(a_fp8, b_fp8, scale_a=..., scale_b=..., out_dtype=torch.bfloat16) — which multiplies FP8 inputs on tensor cores while accumulating in fp32 and rescaling on the way out. See Mixed Precision, bf16 & FP8 Training for the full workflow, and Quantization II: INT4/INT8/FP8, GGUF, bitsandbytes & QAT for the inference-side story.

By 2026, NVIDIA’s Blackwell and Rubin generations push the same idea a step further with native 4-bit floating-point formats — NVFP4 and MXFP4 — applying finer-grained per-block scaling to keep a 4-bit mantissa usable. Transformer Engine now ships full NVFP4 training recipes (not just inference), typically pairing FP4 matmuls with higher-precision master weights and stochastic rounding to control the added quantization noise. FP8 remains the mainstream format for large-scale pretraining as of this writing; FP4 training is the current research and early-production frontier.

Demonstrating fp16 Overflow and Corrective Techniques¶

Let us make the overflow story concrete with runnable code:

import torch
import struct

def fp32_to_bits(x: float) -> str:
    """Return the IEEE 754 binary representation of a float32."""
    [bits] = struct.unpack('>I', struct.pack('>f', x))
    s = f"{bits:032b}"
    return f"[{s[0]}|{s[1:9]}|{s[9:]}]  (sign|exp|mantissa)"


def show_format_limits():
    """
    Print the representable range of fp32, fp16, and bf16.
    Useful for understanding where overflow occurs.
    """
    formats = [
        ("fp32",   torch.float32),
        ("fp16",   torch.float16),
        ("bf16",   torch.bfloat16),
    ]
    print(f"{'Format':<8} {'Max value':<16} {'Min normal':<16} {'Machine eps':<14}")
    print("-" * 56)
    for name, dtype in formats:
        info = torch.finfo(dtype)
        print(f"{name:<8} {info.max:<16.4e} {info.tiny:<16.4e} {info.eps:<14.4e}")


def demonstrate_fp16_overflow():
    """
    Show fp16 overflow and the bf16/loss-scaling remedies.
    fp16 max is ~65504; anything larger becomes inf.
    """
    # Values near and beyond fp16 max
    test_vals = [1e3, 1e4, 6.5e4, 6.55e4, 1e5, 1e38]

    print(f"\n{'Value':<12} {'fp16 result':<14} {'bf16 result':<14} {'fp32 result'}")
    print("-" * 56)
    for v in test_vals:
        fp16 = torch.tensor(v, dtype=torch.float16).item()
        bf16 = torch.tensor(v, dtype=torch.bfloat16).item()
        fp32 = torch.tensor(v, dtype=torch.float32).item()
        fp16_str = "inf" if fp16 == float('inf') else f"{fp16:.4e}"
        bf16_str = "inf" if bf16 == float('inf') else f"{bf16:.4e}"
        print(f"{v:<12.2e} {fp16_str:<14} {bf16_str:<14} {fp32:.4e}")


def demonstrate_loss_scaling():
    """
    Simulate fp16 loss scaling: scale up before backward,
    scale down gradient after, then check for inf/nan.
    This is the mechanism behind torch.cuda.amp.GradScaler.
    """
    # Simulate a tiny gradient that would underflow in fp16
    true_gradient = torch.tensor(1e-8)   # below fp16 min normal ~6e-5

    scale_factor = 2.0 ** 15            # = 32768, typical initial scale

    # Scale up: small gradient becomes representable
    scaled_grad_fp16 = (true_gradient * scale_factor).to(torch.float16)
    print(f"\nLoss scaling demo:")
    print(f"  True gradient:      {true_gradient.item():.2e}  (underflows fp16)")
    print(f"  Scaled (x{scale_factor:.0f}):    {(true_gradient * scale_factor).item():.2e}")
    print(f"  In fp16:            {scaled_grad_fp16.item():.2e}  (representable!)")

    # Unscale: divide out the scale factor in fp32
    unscaled = scaled_grad_fp16.float() / scale_factor
    print(f"  Unscaled fp32:      {unscaled.item():.2e}  (close to true)")

    # Check for inf (skip update if gradient overflowed)
    has_inf = not torch.isfinite(scaled_grad_fp16).all()
    print(f"  Has inf/nan:        {has_inf}  → {'skip update' if has_inf else 'apply update'}")


if __name__ == "__main__":
    show_format_limits()
    demonstrate_fp16_overflow()
    demonstrate_loss_scaling()
    print("\nIEEE 754 bit layout of 1.0 in fp32:")
    print(f"  1.0 = {fp32_to_bits(1.0)}")
    print("IEEE 754 bit layout of inf in fp32:")
    print(f"  inf = {fp32_to_bits(float('inf'))}")

Format   Max value        Min normal       Machine eps
--------------------------------------------------------
fp32     3.4028e+38       1.1755e-38       1.1921e-07
fp16     6.5504e+04       6.1035e-05       9.7656e-04
bf16     3.3895e+38       1.1755e-38       7.8125e-03

Value        fp16 result    bf16 result    fp32 result
--------------------------------------------------------
1.00e+03     1.0000e+03     1.0000e+03     1.0000e+03
1.00e+04     1.0000e+04     9.9840e+03     1.0000e+04
6.50e+04     6.4992e+04     6.5024e+04     6.5000e+04
6.55e+04     6.5504e+04     6.5536e+04     6.5500e+04
1.00e+05     inf            9.9840e+04     1.0000e+05
1.00e+38     inf            9.9692e+37     1.0000e+38

Loss scaling demo:
  True gradient:      1.00e-08  (underflows fp16)
  Scaled (x32768):    3.28e-04
  In fp16:            3.28e-04  (representable!)
  Unscaled fp32:      1.00e-08  (close to true)
  Has inf/nan:        False  → apply update

Numerical Stability in Practice: Patterns and Anti-Patterns¶

The section above covered two named algorithms (stable softmax, Kahan summation). Here we survey the broader landscape of numerical stability patterns that appear throughout the LLM stack.

Log-probability arithmetic¶

Language models assign probabilities to sequences. These probabilities are products of many per-token probabilities, each less than 1. For a 1000-token sequence with average per-token probability \(10^{-2}\), the sequence probability is \(10^{-2000}\) — far below any representable float. Always work in log space:

Cross-entropy loss is the negative log-probability, computed by torch.nn.CrossEntropyLoss, which internally uses log_softmax for numerical stability. Never compute log(softmax(logits)) — always call log_softmax directly.

Layer normalization and numerical stability¶

Layer normalization (Ba et al., 2016) computes:

The \(\varepsilon\) (typically \(10^{-5}\) or \(10^{-6}\)) prevents division by zero when the variance is tiny. In bf16, the computed variance can be zero even for non-constant inputs because the precision is \(7.8 \times 10^{-3}\). PyTorch’s LayerNorm implementation always computes the mean and variance in fp32 internally, regardless of the input dtype — a critical numerical safeguard. See The Transformer Block: Norms, Residuals, MLPs & Activations for the full normalization story.

Attention score overflow¶

Raw attention scores are \(q_i \cdot k_j / \sqrt{d_k}\). For large \(d_k\) (e.g., 128 in LLaMA-style heads), the dot product magnitudes grow as \(O(\sqrt{d_k})\) before scaling. After scaling they are \(O(1)\), which is stable. But in implementations that forget the \(1/\sqrt{d_k}\) denominator, or in attention variants with non-standard scaling, scores can be large enough to overflow fp16 or saturate softmax.

Gradient norms and clipping¶

The gradient of a large model is a concatenated vector of millions to billions of parameters. Its \(\ell_2\) norm,

is computed as a sum of squares across all layers, then square-rooted. Both operations are numerically sensitive: accumulating squares of large values can overflow, and taking the square root of a very large or very small value loses precision. Gradient clipping, torch.nn.utils.clip_grad_norm_, computes this norm in fp32 regardless of parameter dtype.

import torch

def stable_gradient_norm_clip(parameters, max_norm: float, norm_type: float = 2.0):
    """
    Numerically stable gradient norm clipping.
    Always computes in fp32, avoids overflow during squared sum accumulation.

    This is what torch.nn.utils.clip_grad_norm_ does internally.
    """
    params_with_grad = [p for p in parameters if p.grad is not None]
    if not params_with_grad:
        return torch.tensor(0.0)

    # Cast all gradient norms to fp32 before summing to avoid overflow
    # Each element's squared norm computed in fp32 even if grad is bf16
    total_norm_sq = sum(
        p.grad.detach().float().norm(norm_type) ** norm_type
        for p in params_with_grad
    )

    total_norm = total_norm_sq ** (1.0 / norm_type)   # fp32 result

    clip_coef = max_norm / (total_norm + 1e-6)
    if clip_coef < 1.0:
        for p in params_with_grad:
            p.grad.detach().mul_(clip_coef)

    return total_norm


# Illustrate: simulate a gradient spike that would overflow fp16
model = torch.nn.Linear(512, 512)
# Inject a large gradient to simulate a loss spike
with torch.no_grad():
    model.weight.grad = torch.randn_like(model.weight) * 1000.0
    model.bias.grad   = torch.randn_like(model.bias)   * 1000.0

norm_before = sum(p.grad.float().norm(2).item() ** 2
                  for p in model.parameters() if p.grad is not None) ** 0.5
print(f"Gradient norm before clipping: {norm_before:.2f}")

norm = stable_gradient_norm_clip(model.parameters(), max_norm=1.0)
print(f"Gradient norm reported:        {norm.item():.2f}")

norm_after = sum(p.grad.float().norm(2).item() ** 2
                 for p in model.parameters() if p.grad is not None) ** 0.5
print(f"Gradient norm after clipping:  {norm_after:.4f}  (target: ≤ 1.0)")

Interview Corner¶

Key Takeaways¶

Further Reading¶

• Goldberg, David. “What every computer scientist should know about floating-point arithmetic.” ACM Computing Surveys, 1991. The canonical reference — dense but complete.
• Kahan, W. “Further remarks on reducing truncation errors.” Communications of the ACM, 1965. The original compensated summation paper.
• Micikevicius et al. “Mixed Precision Training.” ICLR, 2018. The foundational paper that introduced the loss-scaling recipe for fp16 training.
• Kalamkar et al. “A Study of BFLOAT16 for Deep Learning Training.” arXiv:1905.12322, 2019. Intel and Google’s analysis of bf16 on large-scale training — the paper that anchored bf16 adoption.
• NVIDIA Transformer Engine documentation. Covers the FP8 training workflow with per-tensor scaling for H100 and later GPUs.
• Dao et al. “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.” NeurIPS, 2022. Applies stable logsumexp to fused attention in a single GPU kernel pass.
• Higham, Nicholas J. “Accuracy and Stability of Numerical Algorithms.” SIAM, 2002. The graduate-level reference for backward error analysis, condition numbers, and algorithm stability.

Exercises¶

1. bf16 and fp16 are both 16-bit formats, yet bf16 became the LLM training standard while fp16 did not. Explain, in terms of the exponent/mantissa split, (a) why bf16 has the same dynamic range as fp32, (b) why fp16 has a hard overflow wall at ~65504, and © why bf16 makes loss scaling unnecessary but pays for it somewhere else.

2. bf16 has 7 mantissa bits. Working by hand:

(a) What is its machine epsilon \(\varepsilon_{\text{mach}}\)? (b) What is 1 ULP (the spacing between adjacent representable values) at magnitude \(256\)? © What is the largest integer \(N\) such that every integer in \([0, N]\) is exactly representable in bf16? What does \(N+1\) round to?

3. You are given logits \(z = [10, 11, 12]\) (in fp32, so no overflow). Using the max-subtraction identity from the chapter, compute by hand: (a) the shifted exponentials, (b) the softmax probabilities, and © \(\operatorname{logsumexp}(z)\). Use \(e^{-1} \approx 0.3679\), \(e^{-2} \approx 0.1353\), \(\ln(1.5032) \approx 0.4076\).

4. The chapter’s stable softmax works by never exponentiating a positive argument. Quantify why this is essential.

(a) Above what input value \(x\) does \(e^{x}\) overflow fp16 (max ~\(6.55 \times 10^{4}\))? Above what value does it overflow fp32 (max ~\(3.40 \times 10^{38}\))? (b) Attention scores are \(q_i \cdot k_j / \sqrt{d_k}\). If a buggy implementation forgets the \(1/\sqrt{d_k}\) factor with \(d_k = 128\), and in the worst case where the per-dimension products do not cancel the raw dot product reaches ~\(O(d_k) = 128\), explain which formats’ softmax would break and how max-subtraction rescues all of them.

Use \(\ln(6.55 \times 10^{4}) \approx 11.09\) and \(\ln(3.40 \times 10^{38}) \approx 88.7\).

5. The chapter’s stable_logsumexp needs the whole logit vector in memory at once (it calls z.max() then torch.exp(z - m).sum()). FlashAttention instead processes logits in chunks, keeping only a small running state. Implement streaming_logsumexp(chunks) that consumes an iterable of 1-D tensors one chunk at a time and returns \(\log \sum_j e^{z_j}\) over the concatenation, without ever concatenating them. Maintain a running max \(m\) and running accumulator \(l = \sum e^{z - m}\), rescaling \(l\) whenever the max grows. Verify it against torch.logsumexp on a case where a late chunk contains the global maximum.

import torch


def streaming_logsumexp(chunks) -> torch.Tensor:
    """
    One-pass logsumexp over an iterable of 1-D tensors (chunks),
    never materializing the concatenation.

    Running state (m, l) satisfies  logsumexp(seen so far) = m + log(l),
    where m is the running max and l = sum(exp(z - m)). When a chunk
    pushes the max higher, l is rescaled by exp(m - m_new) before the
    new chunk's contribution is added -- the online-softmax update that
    FlashAttention applies across key blocks.
    """
    m = torch.tensor(float('-inf'))   # running max
    l = torch.tensor(0.0)             # running sum of exp(z - m)

    for chunk in chunks:
        m_chunk = chunk.max()
        m_new = torch.maximum(m, m_chunk)
        # exp(-inf - m_new) = 0 on the first iteration, so l stays 0 there
        l = l * torch.exp(m - m_new) + torch.exp(chunk - m_new).sum()
        m = m_new

    return m + torch.log(l)


if __name__ == "__main__":
    # Global maximum (50.0) lives in the LAST chunk -> exercises rescaling
    chunks = [
        torch.tensor([1.0, 2.0, 3.0]),
        torch.tensor([-5.0, 0.5]),
        torch.tensor([50.0, 49.0]),
    ]

    streamed = streaming_logsumexp(chunks)
    reference = torch.logsumexp(torch.cat(chunks), dim=-1)

    print(f"streaming_logsumexp: {streamed:.6f}")
    print(f"torch reference:     {reference:.6f}")
    print(f"match: {torch.allclose(streamed, reference)}")

streaming_logsumexp: 50.313263
torch reference:     50.313263
match: True