Blockwise 8-Bit Optimizer State Without Losing the Second Moment

Adam stores two numbers per parameter: an exponentially decaying average of the gradient, and an exponentially decaying average of the squared gradient. Both are conventionally fp32, so the optimiser state costs 8 bytes per parameter, twice the fp32 weight footprint and the largest single term in the memory budget of a full-precision run.

Storing those moments in eight bits would divide that term by four. The obstacle is that the second moment is not uniformly distributed: its values span several orders of magnitude across parameters and over the course of training, and a fixed-point eight-bit representation of a value that ranges over $10^{-12}$ to $10^{-2}$ either clips the large end or underflows the small end. Because the second moment appears in the denominator of the update, underflow there is not a small numerical error; it produces an enormous effective step. The same asymmetry between exponent range and mantissa precision governs the choice of compute format discussed in the fp16 and bf16 comparison, but the optimiser has an extra constraint that format choice does not: the stored value participates in a division.

The update and why the denominator is fragile

The Adam update for parameter $\theta$ at step $t$ is

$$\theta_{t+1} = \theta_t - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$$

where $\hat{m}$ and $\hat{v}$ are bias-corrected estimates of the first and second moments. The ratio $|\hat{m}|/\sqrt{\hat{v}}$ is bounded in magnitude by roughly one, because the first moment is at most the square root of the second by Cauchy-Schwarz. That bound is what makes the step size predictable: the learning rate $\eta$ sets the scale, and the ratio modulates it within a factor of a few.

If $\hat{v}$ is quantised to zero for some parameter, the denominator collapses to $\epsilon$, and the step for that parameter becomes $\eta \hat{m} / \epsilon$. With $\epsilon$ typically $10^{-8}$ and $\hat{m}$ of order $10^{-4}$, the step is four orders of magnitude larger than intended. The parameter jumps, the gradient at the new point is large, and the run is corrupted within a few steps. A quantisation scheme that can produce a zero second moment is therefore unusable, regardless of how well it preserves the average error.

The first moment has a different problem: its values are signed and centred near zero, with a distribution that is roughly symmetric and whose magnitude decays as training progresses. Quantising a signed quantity near zero with a fixed scale loses the small values, but the consequence is a slightly wrong step direction rather than an explosion.

Blockwise scaling

The mechanism that makes eight-bit storage viable is to choose the quantisation scale per block rather than globally. A block of consecutive parameters, typically a few hundred to a few thousand elements, shares one scale factor, and within the block the values are represented as integers against that scale.

Block size is the trade. A large block amortises the scale factor’s storage and reduces the overhead, but a block spanning parameters with very different magnitudes forces the small ones to quantise coarsely. A small block adapts better but stores more scales. With 256 elements per block and a 32-bit scale, the overhead is 4 bytes per 256 elements, or 0.015 bytes per element, which is negligible against the 1 byte per element the quantised value costs.

The block layout must match how the tensors are laid out in memory, because the scale is applied during the optimiser update and that update reads the tensor contiguously. Blocks that straddle tensor boundaries in a concatenated buffer are fine; blocks that straddle non-contiguous strides are not.

Sharing the exponent

Within a block, the scale handles the overall magnitude. What remains is the distribution of values inside the block, and here a second trick applies: rather than storing eight bits of mantissa per value, store the values as integers and rely on the block scale to place them. This is effectively a shared-exponent representation, where the block’s scale plays the role of the exponent and the eight bits play the role of a signed mantissa.

A signed eight-bit integer spans $[-127, 127]$, giving roughly 7 bits of precision across the block’s dynamic range. That is two orders of magnitude less precision than fp32, which sounds fatal for a quantity that participates in a division. It is tolerable because the second moment is a long-running average: the bias correction and the $\epsilon$ floor mean the ratio is dominated by the recent gradient history rather than by the exact value, and the noise introduced by quantisation is averaged over the block rather than concentrated on a single parameter.

For the second moment, the quantisation is applied after taking the square root. Storing $\sqrt{v}$ rather than $v$ compresses the dynamic range, because a square root halves the exponent. A second moment spanning $10^{-12}$ to $10^{-2}$ becomes a root spanning $10^{-6}$ to $10^{-1}$, a range of five orders of magnitude, which is within what a shared block scale can cover. The stored value is then squared again at read time, and the rounding error in the root becomes a proportional error in the denominator rather than an absolute one.

The update path

def quantize_block(x, block_size=256):
    """Symmetric int8 quantisation with one fp32 scale per block."""
    flat = x.reshape(-1, block_size)
    scale = flat.abs().amax(dim=1, keepdim=True) / 127.0
    scale = scale.clamp_min(1e-12)          # avoid a zero scale for all-zero blocks
    q = (flat / scale).round().clamp(-127, 127).to(torch.int8)
    return q, scale

def dequantize_block(q, scale):
    return (q.to(torch.float32) * scale).reshape(-1)

Two details in that code carry the correctness. The clamp on the scale prevents a block of zeros from producing a division by zero, which would otherwise yield NaN and poison the whole tensor. And the rounding is to nearest rather than truncation, because truncation toward zero introduces a systematic bias in a signed distribution: the first moment would shrink in magnitude every step, which acts as an unintended decay.

The dequantisation happens inside the optimiser step, so the moments are held as bytes between steps and expanded to fp32 only for the arithmetic. The expansion is transient, which means the peak memory during the update is the block size times 4 bytes, not the whole tensor.

What is lost and when it matters

The error introduced is a relative error of about $1/127$ per element, applied to the moments. Because the update is a ratio, errors in numerator and denominator partially cancel when they are correlated, which they are: both moments are driven by the same gradients and a block of parameters with consistently small gradients has small values in both.

The failure mode appears when a block contains parameters with very different gradient scales. A mixture of a large embedding row and a small projection weight in one block forces the small one to quantise against the large scale, and its second moment may round to a small integer that carries little information. The remedy is to align blocks with tensor boundaries so that parameters of similar scale share a block, which is why the implementation detail of block layout matters more than the block size itself.

A second consideration is stability of the scale over time. Recomputing the scale every step means the quantisation grid shifts, which introduces a small amount of noise into the moment estimates even when the underlying values are unchanged. Recomputing less often, or storing the scale alongside the block, trades freshness for stability. Implementations that store the scale per block per step pay 4 bytes per block, which is the overhead already counted.

The place of 8-bit state in a budget

The saving is 6 bytes per parameter against fp32 Adam: from 8 to 2 for the two moments. For a model with $N$ parameters that is $6N$ bytes, which for a 7-billion-parameter model is 42 GB. That is the difference between needing optimiser sharding and not, in many configurations.

It is not a substitute for sharding, and the two compose. Sharding divides the state across devices; quantisation divides it by four within each device. Applying both gives a factor of $4P$, and the communication cost of the sharding is unchanged because the reduce-scatter operates on gradients, not on the moments.

The technique also interacts with mixed precision in a way worth noting: the master weights in an fp32-master scheme are 4 bytes per parameter and are not part of the optimiser moments, so an 8-bit moment scheme leaves the master copy as the dominant per-parameter term. Removing that copy requires a different approach to the update, such as keeping the weights in a higher-precision format only where the update needs it, which is a separate design problem. The composition of all these terms into a single plan is where the parallelism layout begins.