Why Long bf16 Training Runs Stall: The Residual Stream Has No Room for Small Updates

2026-10-22 · VIDRAFT AI Media

aimachinelearningpytorchllm

Why Long bf16 Training Runs Stall: The Residual Stream Has No Room for Small Updates

TL;DR

bf16 keeps 8 significant bits. Once a residual-stream value grows to around 600, the gap between neighbouring bf16 numbers is 4, so any single update smaller than 2 is rounded away completely. Matrix multiplies tolerate this. The running sum that carries the residual stream does not, because it adds many small updates to one large number. In a simulation (illustrative, code below) a bf16 residual received between 0.3% and 43% of the drift an fp32 residual received over 4,000 steps, depending on the noise level, and in one setting it stopped moving entirely. The fix used by mainstream training code is cheap: keep the residual path in fp32, keep the matmuls in bf16. If you cannot, stochastic rounding recovers the average in expectation. Whatever you choose, watch residual and gate maxima over time, not just loss.

Why does a bf16 residual stream lose small updates?

A bf16 number has 1 sign bit, 8 exponent bits and 7 stored mantissa bits, which is 8 significant bits including the implicit leading one. The relative resolution is therefore about 2^-8 to 2^-7, roughly 0.4% to 0.8%. That is fine for a weight or an activation that gets consumed once.

A transformer's residual stream is different. Each block computes x = x + f(x). The same tensor is read-modified-written once per block, for dozens of blocks, and the quantity you care about is the sum. Floating-point addition rounds the result to the nearest representable number. If the update is smaller than half the gap at the current magnitude, the sum rounds back to the old value and the update has no effect.

I measured the gaps with torch.nextafter on bf16 tensors:

bf16 spacing by magnitude

Magnitude of the value Gap to the next bf16 number Smallest update that registers
1 0.0078 above 0.0039
64 0.5 above 0.25
128 1 above 0.5
256 2 above 1
600 4 above 2
1,024 8 above 4
4,096 32 above 16

Round-to-nearest-even resolves an exact tie toward the even mantissa, so an update of exactly half the gap is also lost. At 600 an update of 2.0 does nothing, while 2.1 jumps the value by 4.

Are residual values really that large?

Yes, in specific dimensions. Sun et al. (Massive Activations in Large Language Models) report a handful of activations that are up to 100,000 times larger than typical ones, stable across inputs, and acting as implicit biases that shape attention. These live in the hidden states that flow between layers, which is the residual stream. Any model whose stream norm grows with depth can push some coordinates into the hundreds. A median coordinate near 1 and a few coordinates near 600 can coexist in the same tensor, and the large ones are exactly the ones that gate behaviour.

The consequence is a precision profile that is invisible in loss curves. A coordinate near 600 only moves in steps of 4. If the true gradient-driven drift on that coordinate is 0.3 per step, it takes a long run of same-direction updates before anything registers, and what registers arrives in lumps.

What does the failure look like in a simulation?

The numbers below are an illustrative simulation, not a model run. Each of 4,096 coordinates starts at 600. Every step adds mu + sigma * noise, with mu the small coherent drift and sigma the per-step noise. I compare three residual paths: a full fp32 residual, an fp32 residual that receives bf16-rounded updates (what you get when only the matmul output is bf16), and a fully bf16 residual.

import torch

def run(mu, sigma, steps=4000, d=4096, start=600.0):
    g = torch.Generator().manual_seed(1)
    x32 = torch.full((d,), start)                      # fp32 residual
    xmix = torch.full((d,), start)                     # fp32 residual, bf16 updates
    xb = torch.full((d,), start).to(torch.bfloat16)    # bf16 residual
    for _ in range(steps):
        u = mu + sigma * torch.randn(d, generator=g)
        x32 = x32 + u
        xmix = xmix + u.to(torch.bfloat16).float()
        xb = xb + u.to(torch.bfloat16)
    return [float((x - start).mean()) for x in (x32, xmix, xb.float())]

print(run(0.5, 1.0))   # [2000.5, 2000.4, 427.5]

Drift after 4,000 steps for three residual paths

Per-step update fp32 residual fp32 residual, bf16 updates bf16 residual
mean 0.2, sd 0.5 800.2 800.2 2.4
mean 0.2, sd 1.0 800.5 800.5 348.5
mean 0.5, sd 1.0 2000.5 2000.4 427.5

Three things stand out.

First, the bf16 residual is not just noisy, it is biased. With a mean of 0.2 and a standard deviation of 0.5, almost no single update exceeds the rounding threshold, so the coherent drift of 800 collapses to 2.4. The signal is the small part, and it is the part that gets deleted.

Second, noise can hide the problem. With a standard deviation of 1.0 the bf16 residual still drifts, because occasional large noise draws cross the threshold and get rounded up to a full step of 4. It recovers 43% of the fp32 drift for mean 0.2 and 21% for mean 0.5. The result looks like progress, but it is progress that depends on how noisy the updates happen to be, which is not a property you want in an optimizer.

Third, in the third row the bf16 residual flattens at about 427. That is where the running value crosses 1,024 and the gap doubles to 8. At that point the updates of roughly 1 are far below half the gap and nothing registers. The residual freezes.

Applied update versus intended update at a residual of 600

The middle column is the important control. Rounding the update itself to bf16, which is what a bf16 matmul output does, costs almost nothing as long as the accumulator is fp32. The loss comes from rounding the sum.

How do production training stacks handle this?

They keep the residual path wide. Megatron-LM exposes an fp32 residual connection option that is only valid together with fp16 or bf16 (its argument validation asserts exactly that), see NVIDIA/Megatron-LM. The Mamba reference implementation has a residual_in_fp32 flag that casts the residual to float32, see state-spaces/mamba. Both choices say the same thing: matmuls can be narrow, the running sum cannot.

The pattern in PyTorch is short. This block was run on CPU (PyTorch 2.4 or newer for nn.RMSNorm) to confirm it returns an fp32 residual:

import torch, torch.nn as nn

class Block(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.norm = nn.RMSNorm(d)
        self.ff = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x):                       # x is fp32 and stays fp32
        h = self.norm(x).to(torch.bfloat16)
        with torch.autocast(x.device.type, dtype=torch.bfloat16):
            u = self.ff(h)
        return x + u.float()                    # the sum happens in fp32

The cost is memory and bandwidth on one tensor per block boundary, not on the weights or the matmul inputs. For most training setups that is a small price compared with the weights, optimizer state and attention activations, though you should measure it on your own batch shape. The same idea applies on any low-precision backend: FP8 matmuls, fp16 on older GPUs, and accelerators with narrow accumulators.

What if you cannot keep the residual in fp32?

Use stochastic rounding for the residual add. Instead of rounding to the nearest bf16 value, round up with probability proportional to the fractional distance. The expected value of the result equals the true sum, so small updates survive on average, at the price of extra variance. A reference implementation in PyTorch adds random bits below the bf16 cutoff and truncates:

def sr_bf16(x32, g):
    b = x32.contiguous().view(torch.int32)
    r = torch.randint(0, 1 << 16, b.shape, generator=g, dtype=torch.int32)
    return ((b + r) & ~0xFFFF).view(torch.float32).to(torch.bfloat16)

In the same simulation, stochastic rounding of the bf16 residual recovered a drift of 797.5 against 800.2 for fp32 (mean 0.2, sd 0.5) and 1,998.8 against 2,000.5 (mean 0.5, sd 1.0). That is the average over 4,096 coordinates. Individual coordinates are noisier, so check that the variance is acceptable for your task. Hardware support for stochastic rounding varies, and a software version like the one above is slower than a plain add.

How do you catch this before the run is lost?

The trap is that nothing breaks early. The loss falls normally for thousands of steps, then a gate or norm starts creeping, and by the time the loss reacts the run is hard to recover. Three signals are worth logging:

  1. Per-layer residual maxima. Log the max absolute value of the stream at each layer every N steps. A steady climb in a few coordinates is the warning that your resolution is shrinking.
  2. Gate and router maxima. Anything that multiplies the residual (learned gates, per-channel scales, MoE router logits) amplifies a quantization bias. Track the max, not only the mean.
  3. Gradient norm trend. A slow, monotone rise in the global gradient norm without a matching loss change is the typical signature of a stalled or biased state.

One more caution from experience: a single-step check will not find this. Compare the fp32 and bf16 gradient for one step and the cosine similarity looks excellent, because one rounding error per step is tiny. The failure is an accumulation effect over thousands of steps, so it needs a long-horizon test: run a short reference in fp32 and compare residual and gate maxima along the trajectory.

At VIDRAFT we saw a long low-precision run drift with exactly this signature, and it disappeared when the residual path was held in fp32 while the matrix multiplies stayed in bf16. This post stays on the general mechanism and the public evidence rather than our run details.

FAQ

Does this affect inference? Much less. At inference each residual add happens once per layer on a fixed-length path, and errors do not accumulate across thousands of optimizer steps. Very deep models or long generations can still show small drift, so some inference stacks keep the residual in higher precision too.

Is fp16 better than bf16 here? fp16 has 11 significant bits, so its gap at 600 is 0.5 instead of 4, which is better on resolution but has a much smaller range (maximum 65,504) and needs loss scaling. bf16 trades resolution for range. Neither is safe for an unbounded running sum, which is why the sum belongs in fp32.

Do I need fp32 for the weights too? Separate question. Master weights in fp32 for the optimizer step is the standard defense against the same effect in parameter updates (a learning-rate-sized update added to a large weight). The residual stream is the activation-side version of the same arithmetic.

Will this show up as NaN? Usually not. The typical outcome is silent: stuck coordinates, creeping gates, a loss curve that flattens or diverges slowly. NaN comes later or not at all.

How large is the memory cost of an fp32 residual? One tensor of shape batch x sequence x hidden at each block boundary doubles from 2 to 4 bytes per element. If you use activation checkpointing the stored boundaries are exactly these tensors, so measure with your configuration rather than trusting a general estimate.

Can I reproduce your numbers? Yes. Both scripts are about 20 lines, need only PyTorch on CPU, and use fixed seeds. The values above are from torch 2.12 on CPU. The simulation is illustrative: it shows the arithmetic, not the training dynamics of any particular model.

Further reading