PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient Quantization Training with Stochastic Rounding via Structured Subspace Randomness

Zeyu Wang, Yiheng Duan, Tao Chen, Donghyun Koo, Sihan Liu, Wenbo Zhang, Chuan He, Liang Liu

arXiv:2610.36587v1cs.LG

TL;DR

Independent stochastic rounding can make INT8 training slower than BF16 because per-element random-number generation bottlenecks accelerators. The paper proposes S3R, which shares thresholds and reuses a quantized gradient, and S3R-EC, which subtracts structured mean rounding error. In LLM pre-training, S3R-EC matches independent SR’s validation loss and speeds dense-model training steps on TPU v5p.

  • Problem

    Per-element random-number generation for unbiased stochastic rounding can bottleneck accelerator vector units and make INT8 training slower than BF16.

    Source: §I · p. 1

  • Method

    S3R reuses one quantized upstream gradient for both backward products and shares a threshold per token, while S3R-EC subtracts structured mean rounding error from the gradients.

    Source: §I · pp. 1–2, D. Structured Error Correction (S3R-EC) · p. 6, 1) Two-Way Correction · p. 6

  • Results

    S3R-EC matches independent SR’s validation loss, and accelerates dense LLM training steps on TPU v5p by 1.24x to 1.32x over BF16 and by 1.11x over independent SR.

    Source: §I · p. 2, C. Model Quality · p. 8

  • Takeaways & Limitations

    Per-token threshold sharing with structured error correction avoids the convergence degradation of uncorrected sharing while retaining the quality of independent SR.

    Source: C. Model Quality · p. 8, §V · p. 9

Abstract

from arXiv · show

Low-precision training often relies on stochastic rounding (SR) to keep quantized gradients unbiased and prevent small updates from rounding to zero. Standard SR draws an independent random threshold for every gradient element. However, accelerators generate random numbers on vector units whose throughput lags far behind their matrix units. Drawing one threshold per element saturates the vector unit and can make INT8 training with SR slower than BF16 training. Sharing random thresholds across tensor elements removes this cost but forces those elements to round in the same direction and shifts their rounding errors by a common offset. We find that this offset degrades training primarily through the activation gradient, which propagates the error backward into earlier layers. Because earlier layers cannot cancel an incoming offset using their own local rounding errors, the offset must be removed at the layer where it originates. We propose Structured Subspace Stochastic Rounding (S3R), which quantizes each upstream gradient once for both the activation and weight gradient matrix multiplications and shares a single threshold per token across the hidden dimension. Its error correction, S3R-EC, subtracts the structured mean rounding error from the activation and weight gradients while leaving the optimizer unchanged. In LLM pre-training experiments, uncorrected threshold sharing degrades convergence, whereas S3R-EC eliminates this degradation and matches the validation loss of independent SR. Furthermore, S3R-EC accelerates full INT8 training steps of dense LLMs on TPU v5p by 1.24x to 1.32x over BF16 and by 1.11x over independent SR.

I. INTRODUCTION

Low-precision integer training can preserve small gradient updates with stochastic rounding, but independent per-element randomness can bottleneck accelerators. The paper introduces S3R and S3R-EC to reduce that overhead while addressing the error correlations caused by shared thresholds.

I. INTRODUCTION

  • Heavy-tailed gradients can fall below half a quantization step, where deterministic round-to-nearest-even maps them to zero and can cause gradient stagnation.Stochastic rounding instead rounds toward neighboring grid points with probabilities proportional to distance, keeping gradients unbiased and small updates alive in expectation.Source: §I · p. 1
  • Per-element random-number generation can saturate the vector unit and make INT8 training with SR slower than BF16 despite INT8’s higher peak matrix-unit throughput.On TPU v5p, independent SR runs a small-dimension FFN block at 0.58x BF16 speed with JAX’s default Threefry generator.Source: §I · p. 1
  • S3R shares thresholds to reduce random-number generation, while S3R-EC subtracts the shared error offset where it originates to address convergence degradation.The offset primarily harms training through the activation gradient, which propagates it into earlier layers.Source: §I · p. 1

A. Low-Precision Training and Stochastic Rounding

  • Prior low-precision training spans FP8, INT8, and 4-bit methods, while prior SR analyses commonly assume independent or mean-independent errors that shared thresholds violate.The cited 4-bit approaches include SR after a random Hadamard transform and unbiased rounding on a logarithmic grid.Source: A. Low-Precision Training and… · p. 2

B. The Cost of Randomness

The paper targets the vector-unit and memory costs of generating and materializing random thresholds, using broadcast sharing to reduce those costs. It then motivates structured correction because shared thresholds correlate errors, and S3R-EC removes their structured mean without changing the optimizer.

B. The Cost of Randomness

  • Counter-based random generation and materialized random tensors add vector-unit work and off-chip traffic, while S3R broadcasts thresholds in registers and writes the quantized gradient once for both backward products.The passage notes that independent-SR quantization issues more than a hundred vector operations per element on TPU v5p.Source: B. The Cost of Randomness · p. 2, Figure 1 · p. 3
  • S3R needs no special hardware and fewer random values, whereas hardware SR support is tied to particular chips and formats and still requires per-element random values.With few random bits per element, SR can become biased; S3R’s contribution is sharing one random value across many elements and addressing the resulting training effects.Source: B. The Cost of Randomness · p. 2

C. Shared Randomness and Error Correction

  • S3R-EC subtracts the token-constant and channel-constant parts of realized rounding error from the activation and weight gradients, leaving the AdamW update unchanged.It removes only mean error along the token and hidden dimensions, not each element’s exact error.Source: C. Shared Randomness and Error… · p. 3
  • Sharing a threshold within a token shifts its elements’ rounding errors by a common offset, while S3R shares within tokens rather than across them because token representations have persistent outlier channels.S3R-EC uses the realized backward-gradient error and subtracts its structured mean to eliminate the shared offset while keeping gradients unbiased.Source: C. Shared Randomness and Error… · pp. 2–3

A. Quantized Backpropagation with Stochastic Rounding

Quantized backpropagation uses the upstream gradient in both activation- and weight-gradient products. S3R reduces the associated random-number work by reusing its quantization and broadcasting compact thresholds, while structured correction addresses the resulting shared-error pattern.

A. Quantized Backpropagation with Stochastic Rounding

  • The backward pass applies stochastic rounding to G before the S8 products, using a step size and thresholds to decide whether each normalized residual rounds upward.The two backward products use the same quantized gradient when the scales are folded into it.Source: A. Quantized Backpropagation with… · p. 3
  • With shared thresholds, token 1 rounds all channels up and token 2 rounds all channels down; the structured error is separated from a residual with zero row and column sums.In the two-way correction, both variants subtract εWi1 from the activation gradient, while only two-way correction also adjusts the weight gradient.Source: Figure 2 · p. 4, A. Quantized Backpropagation with… · p. 4
  • Conventional QT draws independent full-size thresholds for both products, costing 2BSH 32-bit Threefry words per projection, one per element and product; its quantizer issues more than ten times as many vector operations as RTNE, while S8 GEMMs take the same time.The products are the activation-gradient product GW⊤ and weight-gradient product X⊤G.Source: A. Quantized Backpropagation with… · pp. 3–4

B. Structured Subspace Randomness

  • S3R uses one threshold tensor for both backward products and broadcasts a small seed, with 7-bit thresholds whose half-step shift bounds discretization bias by ∆/256.Partitionable Threefry gives bitwise-identical thresholds across pure-FSDP, mixed, and pure-TP factorizations without added collectives.Source: B. Structured Subspace Randomness · p. 4

1) Inter-Operand Sharing:

S3R first reuses one quantized gradient across both backward products, then broadcasts thresholds to reduce random draws. Marginally uniform thresholds preserve first-order unbiasedness, but the sharing topology determines the resulting error correlations.

1) Inter-Operand Sharing:

  • Reusing one threshold tensor makes the two backward quantizations coincide, computes Q once, halves backward quantization work and S8 writes, and leaves both products unbiased.XLA writes Q to HBM once for both GEMMs to read.Source: 1) Inter-Operand Sharing · p. 4

2) Threshold Broadcasting:

  • Per-token broadcasting draws one threshold per token and reduces PRNG work by a factor of H, while per-example broadcasting needs only B draws but correlates errors across all tokens of an example.The paper uses per-example sharing only as an ablation because the common error reaches every token’s activation gradient.Source: 2) Threshold Broadcasting · p. 4

C. Why Sharing Hurts and How S3R-EC Repairs It

  • Sharing thresholds ties together rounding errors, and subtracting their known mean removes the induced correlation; the analysis fixes G and requires only uniform thresholds.Under per-token sharing, errors remain positively correlated across channels within a token, while errors are independent across tokens.Source: C. Why Sharing Hurts and How S3R-EC… · p. 4, 2) Threshold Broadcasting · p. 4

1) First-Order Unbiasedness:

  • Marginally uniform thresholds keep each gradient unbiased even when thresholds are correlated, and the guarantee composes across layers when each layer’s thresholds are independent of noise from above.This permits choosing the threshold topology for its cost and second moments.Source: 1) First-Order Unbiasedness · pp. 4–5

2) Broadcasting Makes the Errors Move Together:

Broadcast thresholds correlate rounding errors within each token, adding covariance when backward products sum those errors. This extra cost falls on the activation gradient, while the weight-gradient error retains the variance of independent SR.

2) Broadcasting Makes the Errors Move Together:

  • Shared-threshold errors are nonnegatively correlated, with covariance largest when both residuals are near 1/2 and zero when either is nearly representable.For evenly spread residuals, pairwise covariance is about ∆2/12, half the individual error variance of about ∆2/6.Source: 2) Broadcasting Makes the Errors Move… · p. 5

3) Where the Correlation Adds Error:

  • A backward product sums weighted errors, so correlated errors add pairwise covariance instead of canceling as independent SR errors can.Under Φ2D, each token’s D=Hout channels share a threshold, while different tokens have independent thresholds.Source: 3) Where the Correlation Adds Error · p. 5, 2) Broadcasting Makes the Errors Move… · p. 5
  • Per-token threshold sharing preserves each weight-gradient error’s variance but adds pairwise-covariance terms to activation-gradient error variance.The added term is largest when weight rows share a common component.Source: 3) Where the Correlation Adds Error · p. 5
  • The added activation-gradient error travels into the lower layer, where a correction based on that layer’s independent rounding errors cannot remove it.In a Transformer, RMSNorm rescales each token and attention mixes tokens, so the proposition describes a mechanism rather than an exact model.Source: 3) Where the Correlation Adds Error · p. 5

4) Subtracting the Mean Removes the Correlation:

Subtracting the group mean removes shared error components and reduces average covariance and variance. The average improvement does not guarantee that every individual output becomes less noisy.

4) Subtracting the Mean Removes the Correlation:

  • For any covariances, subtracting the mean lowers average pairwise covariance and average variance by the variance of the mean error.With shared thresholds, the remaining errors are slightly negatively correlated; under equal covariances, centering deletes the shared component and retains private components.Source: 4) Subtracting the Mean Removes the… · p. 5
  • For evenly spread residuals, centered errors have variance about ∆2/12, compared with ∆2/6 for independent SR.This comparison follows in the equal-covariance case.Source: 4) Subtracting the Mean Removes the… · p. 5
  • With unequal covariances, centering can increase an individual element’s variance even as average variance falls.An element becomes less noisy exactly when 2 Cov(Xi, X̄) > Var(X̄).Source: 4) Subtracting the Mean Removes the… · p. 5
  • 17–27% of activation-gradient outputs become noisier after centering, while summed variance falls by 40–63% and stays below independent SR.The exact computation used a synthetic layer with unequal residuals.Source: 4) Subtracting the Mean Removes the… · pp. 5–6

5) Which Axis to Center:

With per-token thresholds, centering over channels removes the covariance introduced by threshold sharing, whereas centering over tokens has little to remove because tokens are independent.

5) Which Axis to Center:

  • Centering each token’s D channels lowers their average pairwise covariance from a nonnegative value to a nonpositive one.The stated result concerns the average covariance within a token.Source: 5) Which Axis to Center · p. 6
  • Because tokens have independent thresholds, centering each channel over the S tokens finds no covariance to remove and lowers average variance only slightly.The reduction is by 1/S of the channel’s average variance.Source: 5) Which Axis to Center · p. 6
  • Per-token thresholds make activation-gradient centering over each token’s channels do nearly all the work; weight-gradient centering adds little.If tokens also shared a threshold, centering over tokens would remove their covariance and both correction terms would matter.Source: 5) Which Axis to Center · p. 6

D. Structured Error Correction (S3R-EC)

S3R-EC subtracts structured mean rounding errors from activation and weight gradients to remove covariance created by threshold sharing. Its two-way correction preserves unbiased gradients and sends the corrected gradients to the unchanged optimizer.

D. Structured Error Correction (S3R-EC)

1) Two-Way Correction:

  • The correction computes each token’s channel-mean error and each channel’s across-token mean error, then subtracts their effects from the activation and weight gradients after the unchanged INT8 GEMMs.The activation-gradient correction uses the averages multiplied by W; the weight-gradient correction uses the inputs multiplied by those averages.Source: 1) Two-Way Correction · p. 6
  • The corrected gradients remain unbiased and go to the unchanged optimizer, AdamW in the experiments.No additional random numbers are needed.Source: 1) Two-Way Correction · p. 6
  • After correction, residual errors sum to zero along each token and channel, with nonpositive average pairwise covariance on both axes.The removed structured component is an orthogonal projection of the original error.Source: 1) Two-Way Correction · p. 6

4) Variance Reduction:

S3R-EC removes structured covariance introduced by shared thresholds, with per-token sharing leaving most removable covariance in within-token errors. The dX-only correction retains nearly all Φ2D benefit at lower communication and step-time cost.

4) Variance Reduction:

  • Under Φ2D, rows carry nearly all of the covariance removed by S3R-EC, matching the analysis of shared-threshold error covariance.The covariance statement applies across every row and column of the error.Source: 4) Variance Reduction · p. 6
  • 0.476 vs. 0.572 in bottom-layer activation-gradient error: two-way correction lowers Φ2D error below independent SR, consistent with the shared-dither control-variate interpretation.SR is described as rounding after adding dither; the shared dither’s error is known and subtracted.Source: 4) Variance Reduction · p. 6

2) Cost and the dX-Only Variant:

  • Under Φ2D, dX-only correction matches two-way correction within seed noise and adds 1.1–1.8% to dense-model S3R step time.Different tokens’ errors are uncorrelated, so centering over tokens removes almost no covariance; the activation-gradient term retains nearly all the benefit.Source: 2) Cost and the dX-Only Variant · p. 6
  • The dX-only correction also reduces a compiled two-layer model’s f32 all-reduce volume from 422.6 MiB to 3.1 MiB per step, with about 8 MiB of small-factor all-gathers.For Φbcast, where tokens share thresholds, both correction terms are needed.Source: 2) Cost and the dX-Only Variant · pp. 6–7

IV. EXPERIMENTS

Experiments assess convergence, FFN-block speed, and full training-step speed on TPU v5p using JAX/XLA. They compare BF16 and INT8 baselines with independent SR, threshold-sharing variants, and corrected S3R.

A. Setup

  • All experiments run on Google Cloud TPU v5p with JAX/XLA; full-step timing uses medians over 20 steps after 12 warm-up steps.The full-step setup uses FSDP and tensor parallelism, with four synthetic 4,096-token sequences per chip.Source: A. Setup · p. 7
  • Three workloads comprise 12-layer Llama-style pre-training on C4, single-block SwiGLU FFN benchmarks, and full training steps for five models.The convergence run uses 2,000 steps, 32.8M tokens, and three seeds; the full-step workloads include dense and MoE models.Source: A. Setup · p. 7

B. Compared Configurations

  • Comparisons include BF16 and INT8 RTNE, independent per-element SR, per-token Φ2D and per-example Φbcast threshold sharing, and two S3R-EC correction variants.S3R-EC is tested with correction on both backward products or on the activation gradient only.Source: B. Compared Configurations · p. 7
  • The convergence visual tracks validation loss over training for configurations at their best learning rates; Φbcast ends off-scale at 6.24.The accompanying table reports final loss after 2,000 steps on 32.8M C4 tokens, as mean ± sample standard deviation over three seeds.Source: Figure 3 · p. 7, Table I · p. 7

LEARNING RATE, CHOSEN ON SEED 0. ALL CONFIGURATIONS USE

Per-token S3R-EC restores validation quality to within seed noise of independent SR, while uncorrected threshold sharing degrades convergence. The results also distinguish activation-gradient correction needs and bound what the three-seed comparisons can resolve.

C. Model Quality

  • Without correction, Φ2D ends 0.126 above independent SR, while Φbcast does not train usefully; an ablation locates the damage in the activation-gradient product.Per-example thresholds there reproduce the Φbcast gradient error, while per-token thresholds remove almost all of it regardless of weight-gradient thresholds.Source: C. Model Quality · p. 8
  • 4.919 ± 0.022 for two-way S3R-EC (Φ2D), against 4.929 ± 0.019 for independent SR; the seed bands overlap over the final quarter of training.The corrected configuration’s best learning rate returns to that of BF16.Source: C. Model Quality · p. 8
  • The dX-only correction matches two-way correction within seed noise for Φ2D but remains 0.092 above it for Φbcast.Per-token errors across tokens are uncorrelated, whereas per-example sharing correlates them, leaving covariance for the weight-gradient correction to remove.Source: C. Model Quality · p. 8
  • Under Φ2D, two-way correction lowers every listed gradient error below independent SR, including bottom-layer activation-gradient error of 0.48 against 0.57.Under Φbcast, bottom-layer activation-gradient error falls from 3.16 to 0.52, while weight-gradient error remains higher; rounding variance contributes at most 0.094 of AdamW’s late-training second-moment estimate.Source: C. Model Quality · p. 8, Table II · p. 8
  • Three seeds resolve differences of about 0.05 or more; smaller differences, including independent SR versus two-way Φ2D S3R-EC, remain undetermined.Each configuration’s learning rate was selected on seed 0, which is also included in the reported mean.Source: C. Model Quality · p. 8

D. Throughput and Memory

Full-step tests show S3R-EC retains independent-SR quality while improving dense-model throughput over BF16 and independent SR. The section also reports block-level and MoE timing trade-offs.

D. Throughput and Memory

  • S3R (Φ2D) costs −2 to +6 operations per element, runs within 1% of INT8 RTNE at W2–W9, and is 1.08–1.92× faster than BF16.It hashes one word per four tokens and broadcasts thresholds in vector registers inside the quantizer fusion.Source: D. Throughput and Memory · p. 8
  • S3R-EC adds at most 15–23 operations per element and 8–14% temporary memory; only W1 becomes slower than BF16, at 0.94–0.96×.The correction adds 13–15% of S3R device time at W1, 7–10% at W2, and at most 5% from W3 onward.Source: D. Throughput and Memory · p. 9
  • At independent-SR quality, S3R-EC is 1.24–1.32× faster than BF16 and 1.11× faster than independent SR on dense models.The full-step comparison covers forward, backward, and AdamW update on TPU v5p.Source: D. Throughput and Memory · p. 9, §V · p. 9, Table IV · p. 9
  • On the MoE models, S3R-EC is 1.04× and 1.06× faster than independent SR with two-way correction, or 1.05× and 1.07× with dX-only correction.The two-way expert-path correction adds 4.6% and 4.9% to the S3R step.Source: D. Throughput and Memory · p. 9

V. CONCLUSIONS

  • S3R-EC removes the quality degradation from uncorrected threshold sharing by subtracting structured rounding error from activation and weight gradients while leaving the optimizer unchanged.In pre-training, its quality matches independent SR.Source: §V · p. 9
Loading 2610.36587v1…