When Small Updates Disappear

Why low-precision optimizer states ignore small updates, and a Triton kernel that fixes it

Kristi Topollai

September 2026 · Paper (EMNLP 2026) · Code on GitHub

This post is a companion to our EMNLP 2026 paper, Understanding Quantization of Optimizer States in LLM Pre-training^[1].

1Small updates can vanish

Add a small positive number to a stored value, over and over. With exact arithmetic the additions pile up. In low precision, each new value can be so close to the old one that rounding puts it right back. The next step starts from the same stored value, so the lost pieces never add up.

This is state-update stalling: the optimizer computes an update, but the stored state does not change. Figure 1 shows it with repeated BF16 additions. The purple paths use stochastic rounding, which Section 3 explains.

1.1 The optimizer has a memory

Adam keeps two running averages (EMAs) for every parameter^[2]: the first moment $m_t$, an average of the gradients, and the second moment $v_t$, an average of the squared gradients:

Both move by small steps, and both can stall in the same way. This post focuses on the second moment, because it allows a simple closed-form analysis. Its update is driven by $g_t^2$, and the stored value has the same scale as $g_t^2$, so the scale cancels and the chance of stalling depends mostly on the number format and $\beta_2$ (Section 2).

Write $v^q_{t-1}$ for one coordinate of the second moment, stored in low precision, and $\widetilde v_t$ for the new value computed in full precision:

Reading the stored value into FP32 is exact, so new error only enters when the result is written back. Compared with the stored value, the new value moves by

With $\beta_2=0.999$, only a thousandth of the difference between the new squared gradient and the stored value reaches the update. In low precision that is often smaller than the gap to the next representable number. A long memory and a coarse grid are a bad combination.

This model only looks at the rounding when the state is written back; the arithmetic itself is treated as exact. "Stalled" means that one coordinate of the state did not change, not that training stopped.

2Nearest rounding has a threshold

Let $u$ be the distance between neighbouring representable values near the stored value, also called one unit in the last place (ULP). Nearest rounding returns to the stored value whenever

Half a ULP is the threshold. An update below it is lost, and the next one starts from the same place.

2.1 How big is one ULP?

For a positive value $v=m2^e$ with mantissa $m\in[1,2)$ and $p$ stored fraction bits, the spacing is

At the same scale, an FP8 grid is 16 times coarser than BF16, and this FP4 grid is 32 times coarser. The FP4 format here is unsigned E2M2, for a state that is never negative; it is not the signed E2M1 format. The spacing changes at powers of two, and exact ties depend on the last kept bit; the model ignores these edge cases.

2.2 A second moment that can only go up

Put the two together for Adam in BF16. Near a value $v$ the grid step is $u=v/(128m)$ with $m\in[1,2)$, so half a step is more than $v/512$. With $\beta_2=0.999$, the largest possible downward update, when $g_t=0$, is only $0.001\,v$. That is always below the threshold. With nearest rounding, the stored second moment can never decrease.

If the gradients get smaller at some point in training, the state stays at its old value, and Adam keeps dividing by a number that is too large, which shrinks its steps. In BF16 the state can still move up when a large squared gradient arrives, so it slowly drifts upward. FP8 is more extreme: half a step is more than $v/32$, and even an upward move needs a squared gradient more than about 30 times the stored value, so the state barely moves at all. Figure 2 simulates both.

This depends on $\beta_2$. The argument needs $1-\beta_2<1/512$, so in BF16 it holds for $\beta_2$ above about 0.998. With a shorter memory, such as 0.99, large downward updates do get through. In FP8 the condition is $1-\beta_2<1/32$, so it holds for $\beta_2$ above about 0.97.

2.3 From the threshold to a probability

To turn the threshold into a probability, model one gradient coordinate over a short window as Gaussian, $g_t/\sigma\approx\mathcal N(0,1)$, with the state near its typical value, $v^q_{t-1}\approx\sigma^2$. Then

The mantissa $m$ changes from coordinate to coordinate. Replace it by a typical value, $\bar m=1/\ln 2$. The ratio of update size to grid spacing then depends on a single number,

An update stalls when $|z_t-1|<\widehat\rho$, and the probability of that is

The gradient scale cancels. Both the update and the grid spacing grow with $\sigma^2$, so their ratio does not depend on it. The prediction depends only on the precision $\varepsilon$ and the decay $\beta_2$, not on the model size or the size of the gradients. That makes it easy to test on very different models (Section 4).

Push $\beta_2$ toward one, or remove fraction bits, and the stalling interval grows. Figure 3 shows the effect. The purple curve is stochastic rounding, covered in the next section.

At $\beta_2=0.999$ the prediction is about 94.6% for BF16 and almost 100% for FP8 and FP4. The model leaves out clipping, very small numbers (subnormals), the start of training and changing gradient scales.

3Stochastic rounding keeps small updates on average

Take two neighbouring representable values $q_-<x<q_+$. Nearest rounding always picks the closer one. Stochastic rounding picks the upper one with probability equal to how far $x$ is across the gap^[3]:

A value 30% of the way across rounds up 30% of the time. A value that is already representable stays as it is. Figure 4 shows this with repeated draws.

3.1 Unbiased in every step

For each rounding, the expected result is the input itself:

The price is noise: the variance of one rounding is $(x-q_-)(q_+-x)$^[3]. This holds for each rounding on its own. It does not make a whole training run through a nonlinear optimizer unbiased.

3.2 Stochastic rounding still stalls

If an update is smaller than one grid step, stochastic rounding leaves the state unchanged with probability $1-|\Delta|/u$. Only updates of a full step or more always move it. Averaging over the same gradient model gives

At $\beta_2=0.999$ this still predicts about 82.5% stalled updates in BF16, 98.9% in FP8 and 99.4% in FP4. A small chance of moving helps, but in very low precision that chance is small.

3.3 Stalled is not the same as biased

Take an update of $0.2u$. Nearest rounding always delivers zero. Stochastic rounding also delivers zero 80% of the time, but the other 20% of the time it moves a full step $u$, so on average it delivers exactly $0.2u$.

A high stalled fraction alone therefore does not mean stochastic rounding behaves like nearest rounding. The updates that get through are larger, and on average they point the right way.

Above half a step the picture flips, as Figure 5 shows: nearest rounding moves every time, while stochastic rounding sometimes stays put. Stochastic rounding does not always stall less; what it guarantees is the correct average. This is also why the two curves in Figure 3 can cross.

4The prediction matches training

The model throws away almost everything about the network: architecture, data and size. It keeps only the format, the rounding mode and $\beta_2$. If it is right, one prediction per format and rounding mode should match many different models.

Figure 6 compares the prediction with the stalled fraction measured during training, for seven models at $\beta_2=0.999$^[1]. Each dot is a measurement and each dashed line is the prediction. The prediction stays within 2.8 percentage points of every measurement.

5What stalling does to training

Stalling is measured one state coordinate at a time. Does it matter for the model? In the paper we pretrain LLaMA models on C4 with the optimizer states stored in BF16, FP8 or FP4^[1]; the model weights and activations are in BF16. We compare plain AdamW with AdamW plus two fixes: stochastic rounding, and periodic resets of the states. For the 100M-parameter model, the two fixes lower the final validation loss from 3.27 to 3.21 with BF16 states, from 3.33 to 3.22 with FP8 states, and from 3.32 to 3.26 with FP4 states.

Why resets help. Stalling gets worse as a state ages. Right after it starts from zero, the second moment is small, so new squared gradients are large compared with its grid spacing and the updates go through. As the state grows, more and more updates fall under the threshold. A reset sets the moments back to zero, which restarts this cycle and lets a state that got stuck high, as in Figure 2, come back down. The stalling model gives the reset period: about every 1000 steps for BF16, 300 for FP8 and 200 for FP4. The cost is that a reset throws away what the state had averaged.

Table 2 separates the two fixes. Here the model itself trains in FP32 and only the optimizer states are stored in low precision, so any difference comes from the states.

Three things stand out. Stochastic rounding alone closes most of the gap in BF16 and FP8. With stochastic rounding and resets together, BF16 and FP8 states reach the same loss as FP32 states, 3.155, with half or a quarter of the memory. And resets slightly hurt with FP32 states (3.161 against 3.155), so they help because low-precision states go stale, not because resets are good in general.

6Stochastic rounding in PyTorch

How do we implement it, and is it worth the time? Section 5 shows what stochastic rounding gains in loss. Nearest rounding needs no extra work: the cast to BF16 is built into PyTorch and the GPU. Stochastic rounding is not, so we have to write it ourselves. Written directly in PyTorch, it takes about 70 times longer than the built-in nearest cast on an H200 (Section 8). At that cost, the better loss comes with a much slower optimizer step.

The rounding itself is only a few integer operations per element. The time goes into moving the tensor through memory, once for every PyTorch operation. If all the steps run in a single kernel, the tensor is read and written once, and stochastic rounding costs almost the same as nearest rounding in the same kernel. Inside a fused AdamW step the difference is about 0.3%.

The rest of the post shows how to get there. We start with a standalone cast from FP32 to BF16 in PyTorch, to see what has to be computed, and then write it as a single kernel in Triton. Section 8.1 puts the same rounding inside an optimizer step. The code follows our implementation^[4], simplified in a few places.

Nearest rounding is one line, x.to(torch.bfloat16). Stochastic rounding has to choose between two BF16 neighbours, so it takes a few more steps:

The idea is to read the FP32 bits as an integer, add a random 16-bit number, and clear the low 16 bits. The lines with special keep infinities and NaNs intact. This works, but on a GPU each PyTorch operation here is a separate kernel that reads and writes the whole tensor, so the data makes many trips through memory instead of one. This is the gap that a single Triton kernel closes.

6.1 Why the low 16 bits are enough

FP32 has 1 sign bit, 8 exponent bits and 23 fraction bits. BF16 keeps the sign, the exponent and the first 7 fraction bits. The 16 bits it drops tell you where the FP32 value sits between its two BF16 neighbours.

Let $\ell$ be the integer formed by the 16 dropped bits. Add a uniform random integer $r\in\{0,\ldots,2^{16}-1\}$ to the FP32 word and clear the low 16 bits. A carry into the kept bits happens exactly when $\ell+r\ge 2^{16}$, which is true for $\ell$ of the $2^{16}$ possible values of $r$:

A carry moves away from zero, so this is the stochastic rounding rule for both positive and negative numbers. Exact BF16 values have $\ell=0$ and never carry.

7A short introduction to Triton

Triton lets you write GPU kernels in Python^[5]. Instead of writing code for a single GPU thread, you write a program that works on a whole block of elements at once, and Triton decides how to spread that work over the threads.

For $N$ elements and block size $B$, you launch $\lceil N/B\rceil$ copies of the program, and copy $j$ handles elements $jB$ to $jB+B-1$. Each element slot of a block is called a lane. The last block can run past the end of the tensor, so a mask switches off the lanes past the end.

7.1 The kernel

A Triton kernel is a Python function marked with @triton.jit. Ours has five steps: find this program's elements, load them, give each one a counter for its random number, round, and store.

To run it, allocate the output and launch one program per block. The grid says how many programs to start:

kernel[grid](...) is how Triton launches a kernel. The grid in square brackets sets how many programs run: here a one-element tuple, one program per block, and tl.program_id(0) tells each program its position. The call then passes the arguments. Tensors arrive in the kernel as pointers to their first element, which is why the kernel names them input_ptr and output_ptr, and numbers are passed as they are. BLOCK_SIZE is declared tl.constexpr, so it is fixed when the kernel is compiled, and Triton compiles one version per block size.

7.2 The rounding helpers

The rounding is the bit trick from Figure 7. It lives in two small helpers, which the kernel calls like ordinary functions. _stochastic_bf16_bits returns the chosen 32-bit word, and _stochastic_bf16_value reads it back as an FP32 number.

The random number comes from a third helper, _u32_hash. It is a hash, not a random number generator with a state: it mixes the bits of the counter with shifts, XORs and two multiplications by large constants, so that neighbouring counters give numbers that look unrelated. The same counter always gives the same number. Nothing has to be stored between calls, and every element computes its own random number on its own. The rounding keeps the low 16 bits of the hash as the random number $r$ from Section 6.1.

A bitcast is not a numeric cast. tl.cast(value, tl.uint32, bitcast=True) reads the same 32 bits as an integer. Without bitcast=True, Triton converts the number itself, so 1.25 would become 1^[6].

7.3 Two things to get right

Tie the random number to the element, not the program. The counter is built from the element's index plus a seed and an offset. Because it does not depend on the program, changing the block size does not change the result, and the same seed and offset always give the same rounding. Use a new offset for each new cast so the random numbers change.

Mask the loads and stores themselves. An out-of-range load has already happened by the time a later tl.where throws its value away. Pass mask=mask to tl.load and tl.store.

8How fast is it?

Figure 9 compares three ways to run the same stochastic rounding on an NVIDIA H200: eager PyTorch (the code from Section 6), the same code compiled with torch.compile, and the Triton kernel. It also shows the same Triton kernel with nearest rounding, to isolate the cost of the randomness.

For the cast, the Triton kernel is about 36 times faster than eager PyTorch and about twice as fast as torch.compile. It is also less than 1% slower than the same kernel with nearest rounding, so the randomness is almost free; the kernel spends its time moving data, not on the extra integer operations. PyTorch's built-in nearest cast is about twice as fast again, so a more carefully tuned kernel still has room to improve.

8.1 Inside an optimizer

In an optimizer, the rounding belongs inside the update step, so that each tensor is read and written only once. The kernel below is a complete AdamW step for one block of a parameter tensor. It keeps both moments in BF16 and rounds them with the same helper as the cast.

It is launched like the cast, with one program per block of the parameter tensor. Figure 9 also times this step. The fused step with stochastic rounding is about 2.6 times faster than the same step compiled from PyTorch, and stochastic rounding adds about 0.3% over nearest rounding.

9Conclusion

Storing optimizer states in low precision does more than add noise: it can stop a state from changing at all. An update smaller than half a grid step is rounded away, and with Adam's long memory that is true for most updates in BF16 and nearly all in FP8 and FP4. In BF16 the second moment cannot even go down. A model that only knows the format and $\beta_2$ predicts how often this happens, and it matches measurements from seven different models.

The fixes are simple. Stochastic rounding keeps every update correct on average, and periodic resets stop the states from going stale. Together they bring LLaMA-100M with BF16 or FP8 states to the same loss as FP32 states. Stochastic rounding is also cheap: in a Triton kernel it is one addition of a random number and one mask per element, and it costs less than 1% more than the same kernel with nearest rounding.

The paper has the full analysis, including how to choose the reset period and results at larger scales^[1]. The code is in our implementation^[4].

References

1. K. Topollai, A. Choromanska. Understanding Quantization of Optimizer States in LLM Pre-training: Dynamics of State Staleness and Effectiveness of State Resets. EMNLP 2026. arXiv:2603.16731.
2. D. P. Kingma, J. Ba. Adam: A Method for Stochastic Optimization. ICLR 2015. arXiv:1412.6980.
3. M. Croci, M. Fasi, N. J. Higham, T. Mary, M. Mikaitis. Stochastic rounding: implementation, error analysis and applications. Royal Society Open Science, 2022. doi:10.1098/rsos.211631.
4. K. Topollai. Stochastic Rounding for PyTorch (sr-states 0.1.0). github.com/KristiTopollai/stochastic-rounding-pytorch. The code in this post follows commit 69d7ccb.
5. Triton documentation. Vector Addition. triton-lang.org.
6. Triton documentation. triton.language.cast. triton-lang.org.

Citation

To cite the results, please cite the paper^[1]:

To cite this post:

@misc{topollai2026smallupdates,
  title        = {When Small Updates Disappear},
  author       = {Topollai, Kristi},
  year         = {2026},
  month        = {September},
  howpublished = {Blog post},
  url          = {https://kristitopollai.github.io/blog/quantization/}
}