Efficient Quantization Training with Stochastic Rounding via Structured Subspace Randomness

Introduction

Training large language models (LLMs) demands massive compute and memory bandwidth [1, 2, 3]. To reduce both costs during pre-training and fine-tuning, full quantization training (QT) accelerates the training step itself by executing both the forward- and backward-pass matrix multiplications in 8-bit or 4-bit precision [4, 5, 6, 7, 8, 9]. On GPUs, fine-grained FP8 recipes rely on an exponential quantization grid whose step size shrinks near zero, preserving small values under standard deterministic rounding [8]. In contrast, integer formats such as INT8 (the native 8-bit format on accelerators like Google Cloud TPU v5p [9]) use a uniform grid with a constant step size, and emerging 4-bit floating-point formats (NVFP4 and MXFP4) [10], [11] have an extremely coarse grid, making the backward pass in both settings prone to severe underflow.

Specifically, because backward gradients are heavy-tailed and span a wide dynamic range, many of their entries fall below half a quantization step. Deterministic round-to-nearest-even (RTNE) maps these entries to zero every time, causing the affected parameters to stall in a failure mode known as gradient stagnation [12]. Stochastic rounding (SR) prevents this stagnation by rounding each value up or down with probabilities proportional to its distance from the two neighboring grid points. The quantized gradient remains unbiased, and small updates survive in expectation [13, 14, 15].

Standard SR, however, requires an independent uniform random threshold for every gradient element. On modern accelerators, general matrix multiplications (GEMMs) execute on dedicated matrix units (MXUs, such as systolic arrays or Tensor Cores), whereas counter-based pseudo-random number generators (PRNGs, such as Threefry and Philox [16]) execute on the vector processing unit (VPU). Because matrix throughput has scaled much faster than vector throughput across hardware generations, generating a 32-bit random word per gradient element saturates the VPU and stalls the MXU. Consequently, INT8 training with SR can run slower than BF16 training despite INT8 offering twice the peak MXU throughput of BF16. For example, using JAX’s default Threefry generator [17] on TPU v5p, independent SR executes a single feed-forward network (FFN) block at only 0.58x the speed of BF16 for smaller model dimensions.

To eliminate this random number generation bottleneck, we propose Structured Subspace Stochastic Rounding (S3R), which shares each random threshold across multiple tensor elements. Because each threshold remains uniformly distributed, every element is still rounded without bias. Unlike standard SR with independent draws [14], sharing a threshold forces elements in the same group to round in the same direction, shifting their rounding errors by a common offset. We show that this offset degrades convergence primarily through the activation gradient, which propagates the error backward into earlier layers. Because earlier layers cannot cancel an incoming offset using only their own local rounding errors, we pair S3R with a structured error correction method (S3R-EC) that subtracts the offset directly at the layer where it is generated. Our contributions are as follows.

1. We design S3R to eliminate the random number generation overhead of SR. Standard quantized training quantizes the upstream gradient twice in the backward pass, once for the activation gradient and once for the weight gradient, drawing a separate random threshold

for every tensor element. S3R instead quantizes each upstream gradient once for both matrix multiplications and draws a single random threshold per token, sharing it across the hidden dimension. Implemented in standard JAX without custom kernels, S3R allows XLA to broadcast these thresholds directly inside vector registers without materializing the full random tensor in memory.

2. We analyze the error covariance induced by shared thresholds and derive S3R-EC (Sections III-C and III-D). Because sharing a threshold across a sequence correlates errors across tokens, S3R shares one threshold per token across the hidden dimension. We prove that per-token sharing preserves the weight gradient variance of independent SR and confines the positive error covariance to each token’s activation gradient, which propagates into earlier layers (Lemma 2 and Proposition 1). We then show that subtracting the mean rounding error removes this covariance (Theorem 1 and Corollary 1). Building on this result, S3R-EC subtracts the mean error along the token and hidden dimensions without extra random draws or optimizer changes (Proposition 2), while its dX-only variant corrects only the activation gradient to capture nearly the full accuracy gain at lower overhead.

3. We evaluate convergence in Llama-style language model pre-training and show that while uncorrected threshold sharing degrades validation loss, S3R-EC eliminates this degradation and matches the quality of independent SR (Section IV-C). In end-to-end benchmarks on TPU v5p, S3R-EC accelerates full training steps of three dense LLMs [3], [18], [19] by 1.24x to 1.32x over BF16 and by 1.11x over independent SR (Section IV-D). On two mixture-of-experts (MoE) models [18], [20], S3R and S3R-EC achieve speedups of 1.09x to 1.11x and 1.04x to 1.06x over independent SR, respectively.

Related Work

Low-Precision Training and Stochastic Rounding

Full quantization training runs the forward product and both backward products in low precision [4]–[6]. FP8 recipes train LLMs at scale with fine-grained scaling and stability fixes [7], [8], [21] and also compress optimizer states and activations [22]. INT8 recipes quantize the whole data flow [9] or keep the weight-gradient product in 16 bits [23]. At 4 bits, NVFP4 and MXFP4 recipes apply SR to gradients after a random Hadamard transform [10], [11], and LUQ rounds them without bias on a logarithmic grid [24]. Other 4-bit recipes sample gradients by leverage score [25] or, like Quartet, prefer round-to-nearest [26].

By preventing gradient stagnation [12], unbiased SR enabled early 8-bit floating-point training [13]. Under independent draws, its errors are mean-independent, which gives probabilistic error bounds that grow with the square root of the problem size [14]. SR also implicitly regularizes tall-and-thin matrices [27], and Croci et al. survey its uses [15]. In training, SR yields unbiased gradient estimators whose variance depends on the quantizer [28], and its interaction with Adam has been analyzed for LLMs [29]. These analyses assume independent or mean-independent errors, which shared thresholds violate (Lemma 2, Proposition 1).

The Cost of Randomness

Counter-based generators such as Philox and Threefry compute each random word as a keyed function of a counter [16], and JAX uses Threefry by default [17]. Each word takes many rounds of integer operations on the vector unit, which has scaled far more slowly than the matrix unit. In inference, a few extra vector instructions per multiply-accumulate can already make a kernel vector-bound [30]. On TPU v5p, the independent-SR quantizer issues more than a hundred vector operations per element (Section IV-D). Materializing random tensors also adds off-chip traffic, which costs far more energy than arithmetic [31], [32]. S3R broadcasts its thresholds in vector registers without materializing them, and it writes the quantized INT8 gradient to HBM once for both backward products.

Hardware support lowers the cost of SR. Graphcore’s IPU supports SR in binary16 and binary32 arithmetic [15], proposed SR units accelerate SR in fixed-point arithmetic [33], and Blackwell GPUs provide conversion instructions (cvt.rs) that round to FP16, BF16, FP8, FP6 or FP4 with random bits supplied by software [34]. Such support, however, is tied to particular chips and number formats, and it still needs a random value for every element, which software must generate for cvt.rs. S3R needs no special hardware and far fewer random values.

With few random bits per element, SR becomes biased unless the implementation compensates. El Arar et al. analyze limited-precision SR [35], and Fitzgibbon and Felix show that natural few-bit implementations can be significantly biased [36]. S3R uses 7-bit thresholds shifted by half a grid step (Section III) and thus falls in this regime. Slicing one random word into several narrow values, as S3R also does, is existing practice in JAX’s legacy Threefry path [17] and in cvt.rs [34]. Generator-free SR instead takes random bits from values already present in training [37]. Our contribution concerns the far larger saving from sharing one random value across many elements, and what this sharing does to training.

Shared Randomness and Error Correction

SR is equivalent to round-to-nearest with non-subtractive dither, which adds a uniform offset over one quantization step before rounding [38]. Prior work uses non-subtractive dither to sparsify backward gradients [39] and couples random draws across distributed clients so their quantization errors partially cancel upon aggregation [40]. Rather than coupling draws across clients to cancel errors, S3R shares random thresholds within a tensor to reduce generation cost. Because LLM token representations are anisotropic with persistent outlier channels across tokens [41]–[43], S3R shares each threshold across the hidden dimension within a token rather than across tokens.

Within each token, this shared threshold shifts element rounding errors by a common offset (Proposition 1).

To reduce quantization error after rounding, classical subtractive dither subtracts the exact dither signal from every element [38], while post-training bias correction subtracts the expected output error of quantized weights [44]. S3R-EC instead operates on the realized backward gradient error without per-element subtraction, removing only the mean rounding error along the token and hidden dimensions to eliminate the shared offset (Section III-D). Because these mean errors have zero expectation, subtracting their linear effect acts as a control variate that keeps the gradients unbiased while complementing few-bit SR techniques.

Methodology and Theoretical Framework

In this section, we first introduce SR in quantized backpropagation, then describe how S3R works and shares thresholds, and analyze the effect of sharing, which motivates the structured error correction S3R-EC for both the dX-only and two-way variants.

Quantized Backpropagation with Stochastic Rounding

Consider a linear projection $\mathbf{Y} = \mathbf{X}\mathbf{W}$ of a Transformer with activations $\mathbf{X} \in\mathbb{R}^{B \times S \times H_{\mathrm{in}}}$, weights $\mathbf{W} \in\mathbb{R}^{H_{\mathrm{in}} \times H_{\mathrm{out}}}$, micro-batch size $B$, and sequence length $S$. Writing $H = H_{\mathrm{out}}$ for brevity, let $\mathbf{G} = \partial\ell/\partial\mathbf{Y} \in\mathbb{R}^{B \times S \times H}$ be the upstream gradient of the summed per-token loss $\ell= \sum_{b,s} \ell_{b,s}$, where the training loss is the token mean $\mathcal{L} = \ell/(BS)$. The backward pass computes two matrix products: the activation-gradient (dX) product $\mathbf{G}\mathbf{W}^{\top}$ and the weight-gradient (dW) product $\mathbf{X}^{\top}\mathbf{G}$, where $\mathbf{G}$ serves as the left-hand (LHS) and right-hand (RHS) operand, respectively.

In the INT8 pipeline (Figure 1), each operand of the forward and backward products is scaled, rounded from BF16 to S8, multiplied as an S8 general matrix multiplication (GEMM) with S32 accumulation, and rescaled. In the backward pass, SR is applied to $\mathbf{G}$. For a step size $\Delta> 0$ and thresholds $\mathbf{R} \in[0,1]^{B \times S \times H}$, the stochastic rounding operator $\widetilde{\mathbf{G}} = \operatorname{SR}(\mathbf{G};\mathbf{R})$ rounds each element up when its normalized residual $P = G/\Delta- \lfloor G/\Delta\rfloor$ reaches its threshold:

where $\mathbb{I}(\cdot)$ is the indicator function. With rounding error $E(\mathbf{G},\mathbf{R}) = \widetilde{\mathbf{G}} - \mathbf{G}$ and uniform thresholds, each entry $E_{b,s,h} = \Delta(\mathbb{I}(P_{b,s,h} \ge R_{b,s,h}) - P_{b,s,h})$ is a centered Bernoulli variable scaled by $\Delta$. For example, with $\Delta= 1$, the value $0.3$ is rounded to 1 with probability 0.3 and to 0 otherwise, so its mean is exactly 0.3. Our quantizer computes $\Delta\lfloor G/\Delta+ R\rfloor$, which is (1) with the threshold $1-R$; this is uniform and shared wherever $R$ is, so all results apply.

The forward product applies RTNE with per-token scales $\mathbf{s}_{x}$ and per-output-channel scales $\mathbf{s}_{w}$, so that $\mathbf{X} \approx\operatorname{diag}(\mathbf{s}_{x})\widehat{\mathbf{X}}$ and $\mathbf{W} \approx\widehat{\mathbf{W}}\operatorname{diag}(\mathbf{s}_{w})$ with S8 operands $\widehat{\mathbf{X}}$ and $\widehat{\mathbf{W}}$. When both backward products share a quantized gradient (as in RTNE and S3R), the backward pass folds both scales into the gradient and quantizes the result once:

where $\Delta= \max|\mathbf{G}'|/127$ and $\mathbf{Q} = \operatorname{SR}(\mathbf{G}';\mathbf{R})/\Delta$ holds S8 codes. With $\mathbf{G}'$ in place of $\Delta\mathbf{Q}$, the two products equal $\mathbf{G}\mathbf{W}^{\top}$ and $\mathbf{X}^{\top}\mathbf{G}$ for the dequantized operands. One S8 tensor feeds both GEMMs, which contract its channels (dX) and its tokens (dW), so only a per-tensor step can be applied after S32 accumulation; fine-grained recipes such as DeepSeek-V3 [8] instead requantize an operand for its transposed use. Below, $\mathbf{G}$, $\mathbf{W}$ and $\mathbf{X}$ denote $\mathbf{G}'$, $\mathbf{W}^{c}$ and $\mathbf{X}^{b}$, and the fixed rescalings by $1/s_x$ and $1/s_w$, which preserve unbiasedness, are left implicit.

Conventional QT implementations draw independent full-size thresholds $\mathbf{R}^{\mathrm{LHS}}, \mathbf{R}^{\mathrm{RHS}} \sim U([0,1]^{B \times S \times H})$ for the two products: $\widehat{\nabla}_{X} = \operatorname{SR}(\mathbf{G};\mathbf{R}^{\mathrm{LHS}})\mathbf{W}^{\mathsf{T}}$ estimates $\mathbf{G}\mathbf{W}^{\mathsf{T}}$, and $\widehat{\nabla}_{W} = \mathbf{X}^{\mathsf{T}}\operatorname{SR}(\mathbf{G};\mathbf{R}^{\mathrm{RHS}})$ estimates $\mathbf{X}^{\mathsf{T}}\mathbf{G}$, so $\bar{\mathbf{g}}_W = \widehat{\nabla}_W/(BS)$ estimates $\nabla_W\mathcal{L}$. This costs $2BSH$ 32-bit Threefry words per projection, one per element and product. Each word is hashed on the VPU, so the quantizer issues more than ten times as many vector operations as an RTNE quantizer, while the S8 GEMMs on the MXU take the same time.

Structured Subspace Randomness

S3R reduces this cost in two steps (Figure 1): it uses one threshold tensor for both backward products, and it broadcasts a small seed to the full tensor. Only the second step changes the error distribution of each product. Each threshold is a 7-bit bf16 value; as in existing implementations [17], [34], S3R cuts four thresholds from each 32-bit Threefry word [16], and a half-step shift of the threshold grid bounds the resulting rounding bias by $\Delta/256$. Lemma 1 thus holds up to this discretization. S3R builds on partitionable Threefry, which gives bitwise-identical thresholds under pure-FSDP, mixed and pure-TP factorizations without added collectives.

Inter-Operand Sharing

1. Inter-Operand Sharing: With separate PRNG keys, the gradient is quantized twice. S3R uses one threshold tensor for both products, $\mathbf{R}^{\mathrm{LHS}} = \mathbf{R}^{\mathrm{RHS}} = \mathbf{R}$, so the quantizations coincide, as under RTNE, and $\mathbf{Q}$ in (2) is computed once. XLA compiles the quantizer into a fusion that writes $\mathbf{Q}$ to HBM once, at one byte per element, for both GEMMs to read. This halves the backward quantization work and the S8 writes ($\alpha_{\mathrm{share}} = 1/2$), and both products remain unbiased (Lemma 1, Part 2). Their errors become correlated, but the correlation enters none of the second moments analyzed below.

Threshold Broadcasting

2. Threshold Broadcasting: The second step draws a compact seed $\mathbf{u} \sim U([0,1]^{d_0})$ with $d_0 \ll BSH$ and expands it to $\mathbf{R} = \Phi(\mathbf{u}) \in[0,1]^{B \times S \times H}$ with a broadcast $\Phi$. XLA fuses the broadcast into the quantizer, so thresholds are broadcast in vector registers and $\mathbf{R}$ is never materialized. Every $\Phi$ keeps thresholds marginally uniform; sharing more saves draws but correlates more errors, because elements that share a threshold tend to round in the same direction. Independent SR shares nothing; we study two choices:

• Per token, $\Phi_{2\mathrm{D}}$ (default). Here $\mathbf{u} \in[0,1]^{B \times S \times1}$ and $R_{b,s,h} = u_{b,s}$. This cuts PRNG work by a factor of $H$ and keeps rounding errors independent across tokens. Errors within a token remain positively correlated across channels (Corollary 1); S3R-EC removes this covariance (Proposition 2).

• Per example, $\Phi_{\mathrm{bcast}}$ (ablation). Here $\mathbf{u} \in[0,1]^{B \times1 \times1}$ and $R_{b,s,h} = u_b$. This needs only $B$ draws per layer, but it correlates the errors of all tokens of an example, and their common part reaches the activation gradient of every token (Proposition 1), so we use it only as an ablation.

Why Sharing Hurts and How S3R-EC Repairs It

Sharing a threshold biases no gradient, but it ties together the rounding errors of the elements that share it. We show in turn that these errors are positively correlated, where the correlation adds error, and that subtracting their known mean, which is what S3R-EC does, removes the correlation. Throughout, $\mathbf{G}$ is fixed and only the thresholds are random; we assume nothing about the residuals and require only that every threshold be uniform on $[0,1]$.

First-Order Unbiasedness

1. First-Order Unbiasedness: SR is unbiased whenever each threshold is marginally uniform. Expectation is linear, so correlations among the thresholds do not matter.

The lemma composes across layers: backpropagation is linear in the upstream gradient, and each layer draws its thresholds independently of the noise from above, so by the tower property every layer’s gradient is unbiased with respect to the backward pass of the executed forward graph (with $\hat{X},\hat{W}$ if it is quantized). The topology $\Phi$ can therefore be chosen for its cost and its second moments alone.

Broadcasting Makes the Errors Move Together

With independent thresholds, every element chooses its rounding direction on its own, and the errors of different elements are independent. With a broadcast threshold, one draw chooses for the whole group: a large draw rounds most of its elements up, and a small draw rounds most of them down.

Every pair of errors in a group that shares a threshold is thus nonnegatively correlated. The covariance is largest when both residuals are near $\frac{1}{2}$ and vanishes when either element is almost exactly representable. If the residuals of a group are spread evenly over $[0,1)$, each error has variance about $\Delta^2/6$ and each pair has covariance about $\Delta^2/12$, so half of the variance of every element is shared with the rest of its group. Under $\Phi_{2D}$, a group is the $D=H_{\mathrm{out}}$ channels of one token, and different tokens have independent thresholds.

Where the Correlation Adds Error

A backward product adds up errors with weights, and the variance of such a sum is the sum of the variances plus the sum of all pairwise covariances. Independent SR pays only the first part. Correlated errors also pay the second, because they no longer cancel.

Per-token sharing thus keeps the weight gradient exactly as accurate as independent SR and puts the whole cost into the activation gradient, which sums over exactly the channels that share a threshold. The extra term is large when the rows of $\mathbf{W}$ have a common component, because $(\mathbf{W}_i1_D)^2$ then grows with the square of the width while $\lVert\mathbf{W}_i\rVert^2$ grows linearly. The shared part of a token’s error is its mean $\varepsilon_{b,s}=\frac{1}{D}\mathbf{1}_D^{\mathsf{T}}\mathbf{E}_{b,s}$, which enters the activation gradient as $\varepsilon_{b,s}\mathbf{W}1_D$, along the same direction for every token. By the last statement, this error has to be removed in the activation gradient, where it arises. In a Transformer it does not arrive unchanged: the RMSNorm backward rescales each token and removes its component along the normalized input, and the attention backward mixes tokens. The proposition is therefore a mechanism, not an exact model. If all tokens of an example shared one threshold, the weight gradient would pay the same price over the tokens.

Subtracting the Mean Removes the Correlation

The quantizer knows every rounding error exactly, so it can compute the mean error of a group and subtract it. The theorem below shows what this does to correlated errors; it holds for any covariances.

The average covariance and the average variance fall by the same amount, the variance of the mean. For errors that share a threshold, $\bar{c}\ge0$ by Lemma 2, and for large $n$ this amount is about $\bar{c}$. Subtracting the mean therefore removes the positive correlation that broadcasting created, and the centered errors are even slightly negatively correlated, like draws without replacement. The case of equal covariances shows the mechanism. Each error is then a shared part with variance $c$ plus a private part with variance $\sigma^2-c$. Subtracting the mean deletes the shared part and keeps the private parts, which are uncorrelated, so any weighted sum of centered errors has at most the variance of the same sum of independent errors with variance $\sigma^2-c$. For evenly spread residuals, this is about $\Delta^2/12$, half of the $\Delta^2/6$ of independent SR.

The first two statements are averages and hold for any covariances. The third needs equal covariances, which fixed residuals do not give: By Lemma 2, an element that is almost exactly representable is nearly uncorrelated with the rest of its group, and subtracting the group mean adds the noise of the group to it. In general, element $i$ becomes less noisy exactly when $2\operatorname{Cov}(X_i,\bar{X})>\operatorname{Var}(\bar{X})$. For $X_1=-Z$ and $X_2=X_3=Z$, the variance of the first element grows to $\frac{16}{9}\operatorname{Var}(Z)$, while the average variance falls to $\frac{8}{9}\operatorname{Var}(Z)$. Single outputs of a GEMM can behave in the same way. In an exact computation on a synthetic layer with unequal residuals, 17–27% of the activation-gradient outputs become noisier after centering, while their summed variance falls by 40–63% and stays below that of independent SR.

Which Axis to Center

Centering removes covariance only along an axis whose errors are correlated.

Under per-token thresholds, the activation-gradient term of S3R-EC, which centers each token, therefore does nearly all of the work, and the weight-gradient term adds little. If the tokens of an example also shared a threshold, centering over tokens would remove their covariance as well, and both terms would matter. Section IV-C confirms both predictions.

Structured Error Correction (S3R-EC)

By Proposition 1, the cost of sharing is the covariance of the errors within each token, which enters the activation gradient and travels to every lower layer. S3R-EC removes this covariance by subtracting the known mean error, the operation analyzed in Theorem 1, along both axes of the error of each example.

Two-Way Correction

The idea is simple (Figure 2). The quantizer knows $\mathbf{G}$ and $\tilde{\mathbf{G}}$, and hence the rounding error $\mathbf{E}=\tilde{\mathbf{G}}-\mathbf{G}$, exactly. S3R-EC takes two averages of $\mathbf{E}$: the mean error of each token over its channels, which carries the covariance created by the shared threshold, and the mean error of each channel over the tokens of an example. After the unchanged INT8 GEMMs, it subtracts their effect: from the activation gradient, the averages times $\mathbf{W}$, which for the token mean is the mean times the precomputed row sums $\mathbf{W}1_D$; from the weight gradient, the inputs times the same averages. What remains sums to zero along every row and every column. No extra random numbers are needed.

Formally, with $\mathbf{G}$, $\mathbf{X}$ and $\mathbf{W}$ as in (2), the averages of example $b$ form the orthogonal projection of $\mathbf{E}_b \in\mathbb{R}^{S\times D}$ ($D=H_{\mathrm{out}}$) onto the subspace $\mathcal{S}=\{1_S\mathbf{a}^{\mathsf{T}}+\mathbf{c}1_D^{\mathsf{T}}\}$ of sums of row-constant and column-constant matrices:

Passing $\mathbf{E}^{\mathrm{str}}_b$ through both backward products and subtracting the result gives

where $\bar{\mathbf{x}}_b=\frac{1}{S}\mathbf{X}_b^{\mathsf{T}}1_S$ and $\mathbf{e}_b=S\boldsymbol{\epsilon}_b$. In (6), the first term removes the mean error of each channel over the tokens of an example, and the second removes the mean error that each token shares across its channels. The corrected gradients go to the unchanged optimizer, AdamW [45] in our experiments. The correction keeps both gradients unbiased and leaves the remaining errors of every token and every channel with a nonpositive average covariance:

The covariance statement is the first statement of Theorem 1 applied to every row and every column of the error, and under per-token thresholds the rows carry nearly all of the removed covariance, by Corollary 1.

Relation to dithering. The correction is a control variate, akin to bias correction in post-training quantization [44]. SR rounds to nearest after adding the dither $\Delta(r-\frac{1}{2})$; because the dither is shared, the error it causes is known and can be subtracted, as in subtractive dither [38]. Correlated quantization [40] likewise gains from shared randomness. The diagnostics agree in direction: with the two-way correction, $\Phi_{2D}$ has a bottom-layer activation-gradient error of 0.476 vs. 0.572 for independent SR (Table II).

Cost and the dX-Only Variant

The correction reuses $\mathbf{E}$ and adds only per-example reductions and small products. With the batch shared, its weight-gradient term all-gathers small per-example factors such as $\bar{\mathbf{x}}_b$ and $\bar{\mathbf{e}}_b$ instead of reducing a weight-sized $\mathrm{f32}$ partial, which lowers the $\mathrm{f32}$ all-reduce volume of a compiled two-layer model from 422.6 MiB to 3.1 MiB per step, at the cost of about 8 MiB of small factor all-gathers. Under $\Phi_{2D}$, the errors of different tokens are uncorrelated, so centering over tokens finds almost no covariance to remove (Corollary 1), and the activation-gradient term (5) alone keeps nearly all of the benefit, including the removal of the shared term of every token’s activation-gradient error. Intuitively, the token offsets are independent under $\Phi_{2D}$, so they largely average out in the weight gradient, a sum over tokens, but not in each token’s activation gradient (Figure 2). This dX-only correction matches the two-way correction within seed noise (Section IV-C) and adds 1.1–1.8% to the S3R step time of the three dense models. $\Phi_{\mathrm{bcast}}$ whose tokens share thresholds, needs both terms; Section IV-D reports the remaining costs.

Experiments

We report model quality first (Section IV-C) and speed second (Section IV-D), because a faster configuration is useful only if it trains as well as its baseline.

Setup

All experiments run on Google Cloud TPU v5p with JAX/XLA [17]. We use three workloads.

1. Convergence runs. A 12-layer Llama-style decoder is trained on English C4 [46] for 2000 steps ($S = 2048$, $B = 8$, 32.8M tokens). We report the mean and sample standard deviation of the final validation loss over three seeds, with each configuration at its best learning rate on seed 0.

2. Single-block FFN benchmarks. Forward and backward passes of one SwiGLU FFN block at nine shapes ($W_1$–$W_9$) taken from models of 1B to 70B parameters (Table III).

3. Full training steps. MaxText training steps of five models from four families [3, 18, 19, 20]: the dense Llama 3.1-8B, Qwen3-8B, and Gemma 3-12B, and the mixture-of-experts (MoE) models Qwen3-30B-A3B and DeepSeek-V2-Lite (16B). The devices form a two-dimensional mesh with Fully Sharded Data Parallelism (FSDP) [47] and Tensor Parallelism (TP) [48], orchestrated via Pathways [49]. Each chip processes 4 synthetic sequences of 4096 tokens per step, with full rematerialization and flash attention [50]. Step times are medians over 20 steps after 12 warm-up steps and differ by at most 0.6% between repeated runs.

Compared Configurations

All INT8 configurations quantize as in Section III: RTNE in the forward pass, one per-tensor step $\Delta$ for the backward gradient, and int32 accumulation. In the full training steps, each GEMM rounds its output to bf16 per shard, so the TP partial sums are reduced in bf16, as in the BF16 baseline. The flash-attention products, the embedding, and the logits projection stay in bf16. They differ only in how the upstream gradient is rounded:

• BF16 and INT8 RTNE: the 16-bit baseline and deterministic rounding.

• Independent SR (Threefry): one 32-bit random threshold per element from JAX’s default Threefry generator, which S3R also uses. In full training steps, it quantizes the upstream gradient once per backward product, as in AQT.

• S3R ($\Phi_{2\mathrm{D}}$), the default with one threshold per token, and S3R ($\Phi_{\mathrm{bcast}}$), the ablation with one threshold per example.

• S3R-EC, two-way and S3R-EC, dX-only: S3R with the correction of Section III-D, applied to both backward products or to the activation gradient only.

Model Quality

Table I lists the final validation loss of every configuration, and Figure 3 shows the loss curves of all of them.

Configuration	Best LR	Val. loss	Gap to BF16
BF16	$6 \times10^{-4}$	$4.892 \pm0.015$	—
INT8 RTNE	$6 \times10^{-4}$	$5.106 \pm0.013$	$+0.214$
Independent SR	$6 \times10^{-4}$	$4.929 \pm0.019$	$+0.037$
S3R ($\Phi_{2\mathrm{D}}$)	$3 \times10^{-4}$	$5.055 \pm0.017$	$+0.163$
S3R-EC, dX-only ($\Phi_{2\mathrm{D}}$)	$6 \times10^{-4}$	$4.924 \pm0.021$	$+0.032$
S3R-EC, two-way ($\Phi_{2\mathrm{D}}$)	$6 \times10^{-4}$	$4.919 \pm0.022$	$+0.027$
S3R ($\Phi_{\mathrm{bcast}}$)	$3.75 \times10^{-5}$†	$6.239 \pm0.026$	$+1.347$
S3R-EC, dX-only ($\Phi_{\mathrm{bcast}}$)	$6 \times10^{-4}$	$5.094 \pm0.018$	$+0.202$
S3R-EC, two-way ($\Phi_{\mathrm{bcast}}$)	$6 \times10^{-4}$	$5.002 \pm0.024$	$+0.109$

SR is still needed at 8 bits. With deterministic rounding, INT8 RTNE ends 0.176 above independent SR (Welch $t \approx13.2$) and 0.214 above BF16 (Table I).

Sharing hurts, mainly through the activation gradient.

Without a correction, S3R ($\Phi_{2\mathrm{D}}$) ends 0.126 above independent SR. S3R ($\Phi_{\mathrm{bcast}}$) does not train usefully, and its loss keeps improving as the learning rate decreases, down to the lowest rate we tried (Table I). A causal ablation in the last two rows of Table II locates the damage. These rows quantize $G$ separately for the two backward products. Per-example thresholds on the activation-gradient product alone reproduce the gradient error of S3R ($\Phi_{\mathrm{bcast}}$), and per-token thresholds there remove almost all of it, whichever thresholds the weight-gradient product uses.

With per-token thresholds, S3R-EC restores quality within seed noise. S3R-EC, two-way ($\Phi_{2\mathrm{D}}$) reaches 4.919 $\pm$ 0.022, against 4.929 $\pm$ 0.019 for independent SR. Its seed band overlaps that of independent SR over the last 25% of training, and its best learning rate returns to that of BF16. With per-example thresholds, the correction closes most of the gap but not all of it. It stays 0.073 above independent SR (Welch $t \approx4.1$) and 0.109 above BF16 (Table I, Fig. 3).

The dX-only correction suffices for $\Phi_{2\mathrm{D}}$ but not for $\Phi_{\mathrm{bcast}}$. With per-token thresholds, the dX-only and two-way variants end within seed noise of each other. Because it skips the weight-gradient term, the dX-only variant is also cheaper in full training steps (Section IV-D). With per-example thresholds, correcting the activation gradient alone recovers most of the loss, but the dX-only variant stays 0.092 above the two-way correction (Welch $t = 5.4$; Table I). This pattern agrees with the analysis. Per-token thresholds leave the errors of different tokens uncorrelated, so the weight-gradient half of the correction, which centers over tokens, has almost no covariance to remove; per-example thresholds correlate the tokens, and this half then removes their covariance too (Corollary 1).

The correction lowers the gradient error itself, whereas the rounding variance barely affects the optimizer. Table 2 measures every configuration on the same BF16 checkpoint. Under $\Phi_{2\mathrm{D}}$, the two-way correction brings every listed error below that of independent SR, for example to 0.48 against 0.57 in the bottom-layer activation gradient. Under $\Phi_{\mathrm{bcast}}$, it lowers the bottom-layer activation-gradient error from 3.16 to 0.52, below that of independent SR, although the weight-gradient error stays higher. The rounding variance makes up at most 0.094 of the AdamW second-moment estimate $v_t$ late in training, so it barely changes the per-coordinate step size of AdamW.

Configuration	dX error layer 0	dX error layer 11	dW error layer 0	Emb. error
Independent SR	0.572	0.111	0.456	0.457
S3R ($\Phi_{2\mathrm{D}}$)	0.877	0.228	0.691	0.703
S3R ($\Phi_{2\mathrm{D}}$), dW-only corr.	0.882	0.228	0.694	0.706
S3R-EC, two-way ($\Phi_{2\mathrm{D}}$)	0.476	0.103	0.377	0.378
S3R ($\Phi_{\mathrm{bcast}}$)	3.157	0.313	5.141	5.285
S3R ($\Phi_{\mathrm{bcast}}$), dW-only corr.	3.150	0.313	4.529	5.262
S3R-EC, two-way ($\Phi_{\mathrm{bcast}}$)	0.521	0.105	0.551	0.558
dX $\Phi_{2\mathrm{D}}$, dW $\Phi_{\mathrm{bcast}}$	0.893	0.227	0.746	0.714
dX $\Phi_{\mathrm{bcast}}$, dW $\Phi_{2\mathrm{D}}$	3.494	0.261	5.979	6.849

Limits. Three seeds per configuration resolve differences of about 0.05 or more. Smaller differences remain undetermined, including the one between independent SR and the two-way S3R-EC with $\Phi_{2\mathrm{D}}$. The learning rate of each configuration was selected on seed 0, which is also part of the mean.

Throughput and Memory

Shape	$(H, D_{\mathrm{ffn}})$	INT8 RTNE	Indep. SR (Threefry)	S3R ($\Phi_{2\mathrm{D}}$)	S3R-EC, two-way	S3R-EC, dX-only
$W_1$	(512, 2048)	1.12	0.58	1.08	+13.0	+14.8
$W_2$	(1024, 4096)	1.45	0.91	1.44	+7.3	+9.6
$W_3$	(2048, 8192)	1.73	1.29	1.73	−0.6	+1.7
$W_4$	(3072, 12288)	1.92	1.54	1.92	+4.8	+4.1
$W_5$	(3840, 15360)	1.84	1.53	1.83	+2.8	+3.3
$W_6$	(4096, 12288)	1.89	1.60	1.90	+2.7	+3.1
$W_7$	(4096, 14336)	1.87	1.56	1.87	+3.4	+3.1
$W_8$	(5120, 17920)	1.82	1.58	1.81	+0.7	+2.7
$W_9$	(8192, 28672)	1.81	1.66	1.81	+3.8	+2.2

INT8 nearly halves the GEMM time, and vector work decides how much of this saving remains. The S8 GEMMs run on the MXU and take the same time in every uncorrected INT8 configuration, 47.55–47.63 ms for one FFN block at $W_9$ against 90.29 ms for BF16. Scaling, rounding, threshold generation and the correction run on the VPU and do not get cheaper with a narrower format. We convert extra device time to int32 VPU operations per upstream-gradient element at the measured $14.3 \times10^{12}$ operations/s.

Per-element Threefry is the expensive step, and S3R removes it. Independent SR hashes one 32-bit word per gradient element, which costs 120–133 operations per element over INT8 RTNE at every shape and takes 56.5% of the busy time at $W_1$ and 9.9% at $W_9$. It is therefore slower than BF16 at $W_1$ and $W_2$ (Table III). S3R ($\Phi_{2\mathrm{D}}$) hashes one word per four tokens and broadcasts the thresholds in vector registers inside the quantizer fusion, so they are never written to memory. It costs −2 to +6 operations per element, runs within 1% of INT8 RTNE at $W_2$–$W_9$ and 4.3% slower at $W_1$, and is 1.08–1.92$\times$ faster than BF16.

S3R-EC adds small reductions and products. It reuses the rounding error inside the quantizer fusion, and its per-token and per-channel reductions, matrix-vector products and rank-1 updates leave the GEMM time nearly unchanged. On a single block, both variants add 15–23 operations per element at most shapes and 8–14% temporary memory. This is 13–15% of the device time of S3R at $W_1$, 7–10% at $W_2$ and at most 5% from $W_3$ on, so only $W_1$ becomes slower than BF16 (0.94–0.96×; Table III). The two-way variant does more arithmetic but is not always slower, because both variants change how XLA fuses the GEMMs.

On the dense models, the GEMM saving reaches the full step. In a Llama 3.1-8B step, INT8 RTNE spends 1182 ms in the GEMMs against 2236 ms for BF16, and 249 against 174 ms in vector fusions. S3R matches INT8 RTNE to within 1% on all five models and is 1.29–1.35× faster than BF16 on the dense models (Table IV), but without a correction it trails independent SR in quality (Section IV-C). Independent SR also quantizes the upstream gradient twice in full steps, doubling the quantization work and the S8 writes, and is only 1.11–1.19× faster than BF16. S3R-EC, two-way, is 1.24–1.32× faster than BF16 and 1.11× faster than independent SR. Its weight-gradient term needs extra reductions and small factor all-gathers (Section III-D) and adds 2.2–4.5% to the S3R step; the dX-only correction skips it and adds 1.1–1.8%. By XLA’s compile-time estimate, S3R needs the same full-step device memory as RTNE on the dense models, and S3R-EC, two-way, at most 0.4 GB more.

Configuration	Dense models			MoE models
	Llama 3.1-8B	Qwen3-8B	Gemma 3-12B	Qwen3-30B-A3B	DeepSeek-V2-Lite (16B)
BF16	3862	4018	6550	6454	3659
INT8 RTNE	2926 (1.320×)	3110 (1.292×)	4822 (1.358×)	7305	3842
Independent SR (Threefry)	3350 (1.153×)	3607 (1.114×)	5526 (1.185×)	7914	4269
S3R ($\Phi_{2D}$)	2903 (1.331×)	3104 (1.294×)	4853 (1.350×)	7286	3840
S3R-EC, two-way	3008 (1.284×)	3243 (1.239×)	4961 (1.320×)	7621	4029
S3R-EC, dX-only	2935 (1.316×)	3160 (1.272×)	4911 (1.334×)	7548	3984
Speedup over independent SR (Threefry)
S3R	1.154×	1.162×	1.139×	1.086×	1.112×
S3R-EC, two-way	1.114×	1.112×	1.114×	1.038×	1.060×
S3R-EC, dX-only	1.141×	1.141×	1.125×	1.049×	1.072×

On the MoE models, S3R and S3R-EC are also faster than independent SR. S3R is 1.09–1.11× faster, runs as fast as INT8 RTNE, and needs 0.4 and 0.7 GB more device memory than RTNE. With the expert-path correction fused into Pallas kernels over the expert-sorted rows, the two-way correction adds 4.6% and 4.9% to the S3R step, so S3R-EC is 1.04× and 1.06× faster than independent SR; with the dX-only correction, it is 1.05× and 1.07× faster.

Conclusions

Stochastic rounding keeps low-precision gradients unbiased, but standard SR draws one random threshold per element on the vector units, which can make INT8 training slower than BF16 training. Sharing thresholds removes this cost but shifts the rounding errors of all elements that share a threshold by a common offset. We showed that this offset harms training mainly through the activation gradient, which carries it to every lower layer, so it must be removed where it is produced. S3R quantizes each upstream gradient once for both backward products and uses one threshold per token. S3R-EC subtracts the structured part of the known rounding error from the activation and weight gradients, keeps them unbiased, and leaves the optimizer unchanged. In pre-training, uncorrected sharing degrades quality, whereas S3R-EC removes this loss and reaches the quality of independent SR. At this quality, S3R-EC makes full training steps of three dense LLMs 1.24–1.32× faster than BF16 and about 1.11× faster than independent SR.

References