Training with Fewer Bits: Unlocking Edge LLM s Training with Stochastic Rounding (2025)

Abstract: LLM training is resource-intensive. Quantized training improves computational and memory efficiency but introduces quantization noise, which can hinder convergence and degrade model accuracy. Stochastic Rounding (SR) has emerged as a theoretically attractive alternative to deterministic rounding, offering unbiased gradient estimates. However, its interaction with other training factors—especially batch size—remains underexplored. In this paper, we present a theoretical and empirical study of mini-batch stochastic gradient descent (SGD) with SR, showing that increased batch sizes can compensate for reduced precision during backpropagation. Furthermore, we show that quantizing weights and activations impacts gradient variance in distinct ways. Our experiments validate these theoretical insights. Our experiments validate these theoretical insights.

## introduction training large language models (llms) demands significant computational and memory resources. mixed-precision training, using lower-precision formats, offers a crucial path to efficiency @xcite @xcite , enabling training on edge @xcite . however, reduced precision introduces quantization noise, potentially hindering convergence. stochastic rounding (sr) @xcite provides an attractive, unbiased alternative to deterministic rounding methods for managing this noise. as shown in figure 1 , sr can maintain training stability and performance where round-to-nearest (rtn) fails, especially under aggressive quantization. despite sr's advantages, its practical interaction with mini-batch stochastic gradient descent (sgd), particularly the role of batch size in mitigating srinduced noise, remains underexplored. key questions persist: can larger batch sizes-a common variance reduction tool-effectively counteract sr noise? how does this apply distinctly to quantizing shared model weights versus per-sample activations and gradients? and crucially, how does sr's variance impact sgd convergence guarantees? addressing these is vital for establishing theoretically grounded guidelines for optimal precision, rounding, and batch size selection, especially to unlock aggressive quantization for training llms on edge devices. although prior work has explored mixedprecision training empirically @xcite @xcite and theoretically analyzed aspects like gradient quantization for efficient communication @xcite @xcite or weight quantization @xcite , a rigorous understanding of low-precision arithmetic within the backpropagation process itself, combined with the specific variance characteristics introduced by sr in a mini-batch context, remains limited. @xcite analyzed the impact of quantization error but did not explicitly incorporate batch size as a variable interacting with different sources of sr noise. this paper addresses this gap through a theoretical and empirical analysis of sr within mini-batch sgd. we explicitly model and differentiate the statistical properties of noise from quantizing shared weights versus per-sample activations and gradients during backpropagation. our core theoretical finding, supported by empirical validation, is that increased batch sizes can effectively compensate for reduced precision during backpropagation by mitigating sr-induced variance from persample operations. specifically, we demonstrate that reducing precision in activation and gradient quantization can be offset by a quantifiable increase in batch size-for instance, a 1-bit reduction may be balanced by at most a fourfold batch increase to maintain convergence quality, with practical increases often being milder. our main contributions are: from these results, we derive practical guidelines quantifying the trade-off, e.g., showing that a 1-bit precision reduction in activation/gradient quantization can be compensated by approximately doubling the batch size in practice to maintain similar convergence behavior. ## a general framework for mixed-precision sgd in mixed-precision sgd, matrix multiplications are performed using lower precision. this involves quantizing the layer's input activations, its weight matrix, and the gradients backpropagated from the subsequent layer before they are used in the forward and backward pass matrix multiplications. the choice of quantization parameters (precision levels ∆ and rounding thresholds ϵ) for each of these affects training dynamics. our analysis focuses on the quantization within linear layers (including convolutional layers viewed as matrix multiplications), as they typically dominate the computational cost @xcite . consider a network with n linear layers, indexed @xmath0, . . . , n. the weights for linear layers are @xmath1, . . . , w n . the input activation to layer i is a i-1 , the weight is w i , and the output activation is a i . x and y are the network input and label, respectively. for simplicity, we exclude the bias. we denote ∇a i as the gradient of the loss l with respect to a i . we also denote f (w) = e x,y [l(w, x, y)]. for each layer, 5 quantization operations are involved. the exact quantization operation behavior is exactly defined through quantization thresholds ϵ and quantization precision ∆. we define and ϵ = {ϵ i } n @xmath2. definition 4 (gradient approximation). the gradient approximation g(w, x, y, ∆, ϵ) for a sample (x, y) is computed via the forward and backward passes described in algorithms 1,2 and figure 2 . ## convergence analysis this section presents a theoretical analysis of the convergence properties of mini-batch mixedprecision sgd in a non-convex setting. we aim to understand how weight quantization and activation/gradient quantization interact with the minibatch size and hence affect convergence. the sgd update rule is given by: another crucial aspect of our analysis involves quantifying the total variance of this estimator g(w t ) and understanding how its components behave, particularly in relation to the mini-batch size b. this variance arises from two primary sources: sampling variance as well as variance introduced by stochastic rounding. our analysis assumes that quantization errors are i.i.d. per-sample. this is a reasonable assumption in our framework, as stochastic rounding is performed independently on each number before multiplication, and accumulation is performed in high precision. a key insight is that sr variance, similar to sampling variance, diminishes as the mini-batch size b increases. intuitively, because sr is applied independently to each of the b samples' activations/gradients, the errors introduced by these quantization steps tend to average out across the mini-batch. lemma 3 provides a concrete derivation of this 1/b scaling. the interplay between the bias and the two sources of variance (sampling and per-sample sr, both influenced by batch size) is the core of our convergence analysis. we operate on assumptions: assumption 1 (smoothness and boundedness). the true loss function l(w, x, y) is l-smooth, meaning its gradient is l-lipschitz continuous: ∥∇l(w, x, y) -∇l(v, x, y)∥ ≤ l∥w -v∥ for all w, v. furthermore, l(w, x, y) is bounded below by l min , i.e., l(w, x, y) ≥ l min . ## experiment our experiments evaluate both gradient norms and downstream task performance, allowing us to verify the theoretical findings and examine their practical implications. figure 5 reports results on cifar-10 and figure 6 reports results on lmsys-chat @xcite , showing that the degradation caused by sr in activation and gradient quantization can be substantially mitigated by increasing the batch size, whereas rtn does not benefit from such scaling. for downstream evaluation, we finetuned llama-3.2-3b (dubey et al., 2024) on gsm8k @xcite using both adam and sgd, and bert-base-uncased @xcite on the glue benchmark. the gsm8k results (tables 1 and 2 ) are consistent with our theoretical predictions, showing clear accuracy improvements as the batch size increases across quantization formats. the glue @xcite ) benchmark results ( ## conclusion this paper studied sr in mixed-precision training and provided a theoretical analysis of minibatch sgd under quantization. we showed that the variance introduced by sr in activation and gradient quantization decays inversely with batch size, while bias from weight quantization remains unaffected. empirical results confirmed that sr consistently outperforms deterministic rtn, particularly under aggressive quantization. these results offer practical guidance for mixedprecision training: larger batch sizes can offset sr-induced noise, enabling more aggressive quantization without loss of convergence. this insight has direct implications for fine-tuning and training llms under resource constraints. future work may examine adaptive batch sizing and broader quantization schemes.