PaperlayerView exact revision on arXiv

Source-linked AI summary

ECO: Quantized Training without Full-Precision Master Weights

Mahdi Nikdan, Amir Zandieh, Dan Alistarh, Vahab Mirrokni

arXiv:2601.22101v1cs.CLcs.AIcs.LG

TL;DR

Quantized LLM training still relies on high-precision master weights, creating substantial memory overhead, particularly for SMoE models. ECO removes these weights by injecting post-quantization errors into optimizer momentum, forming an error-feedback loop. Across diverse training and fine-tuning settings, ECO nearly matches master-weight baselines while reducing static memory usage by up to 25%.

  • Problem

    Quantized training commonly retains high-precision master weights for gradient accumulation, leaving substantial memory overhead, especially in SMoE models.

  • Method

    ECO applies updates directly to quantized parameters and injects each step’s quantization error into the optimizer’s momentum buffer.

  • Results

    ECO nearly matches master-weight baselines across evaluated models and reduces static memory usage by up to 25%.

  • Takeaways & Limitations

    ECO provides a scalable approach to quantized LLM training without master weights or additional memory overhead.

  • Takeaways & Limitations

    ECO performs best with stochastic rounding; with round-to-nearest, it can have a higher noise floor and a slight accuracy ceiling relative to master-weight baselines.

Abstract

from arXiv · show

Quantization has significantly improved the compute and memory efficiency of Large Language Model (LLM) training. However, existing approaches still rely on accumulating their updates in high-precision: concretely, gradient updates must be applied to a high-precision weight buffer, known as $\textit{master weights}$. This buffer introduces substantial memory overhead, particularly for Sparse Mixture of Experts (SMoE) models, where model parameters and optimizer states dominate memory usage. To address this, we introduce the Error-Compensating Optimizer (ECO), which eliminates master weights by applying updates directly to quantized parameters. ECO quantizes weights after each step and carefully injects the resulting quantization error into the optimizer momentum, forming an error-feedback loop with no additional memory. We prove that, under standard assumptions and a decaying learning rate, ECO converges to a constant-radius neighborhood of the optimum, while naive master-weight removal can incur an error that is inversely proportional to the learning rate. We show empirical results for pretraining small Transformers (30-800M), a Gemma-3 1B model, and a 2.1B parameter Sparse MoE model with FP8 quantization, and fine-tuning DeepSeek-MoE-16B in INT4 precision. Throughout, ECO matches baselines with master weights up to near-lossless accuracy, significantly shifting the static memory vs validation loss Pareto frontier.

1. Introduction

Quantized training reduces compute and activation memory, but master weights preserve a major memory bottleneck, especially for SMoE models. ECO removes this buffer by feeding quantization error into momentum, retaining convergence while improving memory–loss trade-offs.

  • Motivation: Master weights remain a key overhead because quantized pipelines typically retain high-precision parameter copies for gradient accumulation.Many updates fall below low-precision discretization gaps, causing vanished updates or quantization noise when applied directly.
  • Empirical results: Figure 1 shows ECO with stochastic rounding establishing a significantly better static-memory versus validation-loss Pareto frontier.Gradient accumulation is disabled in all compared cases.
  • Motivation: The memory burden is especially pronounced in SMoE models, where all master weights remain resident despite only subsets of parameters being active per token.
  • ECO: ECO quantizes updated parameters after each step and injects the resulting quantization error into optimizer momentum, creating an error-feedback loop without extra memory.This carries lost updates forward so subsequent updates can be applied directly to quantized parameters.
  • Theory: Under standard non-convex assumptions and a decaying learning rate, ECO converges to a constant-radius neighborhood of the optimum.The bound is within a factor 1/(1−β)^2 of the best master-weight bound, while naive removal produces error inversely proportional to the learning rate.
  • Empirical results: Across models from 30M to 16B parameters, ECO nearly matches master-weight baselines and can reduce static memory usage by up to 25%.Evaluations include small Transformers, Gemma-3 1B, a 2.1B SMoE model, and DeepSeek-MoE-16B.

2. Related Work

Prior work has advanced quantized training, optimizer-state compression, and error feedback, but master-weight removal remains difficult at LLM scale. ECO complements these approaches by targeting the remaining dependence on high-precision accumulation.

  • Quantized training: Quantization-aware training and large-scale FP8 systems generally retain high-precision accumulators to maintain training stability.
  • Quantized training: ECO complements existing quantized-training methods by removing their remaining dependence on master weights.
  • Low-precision accumulation: Prior attempts at low-precision accumulation have been restricted by scale or setting, including small convolutional networks and edge-device training.FP8 accumulation has been reported to fail at large LLM scales.
  • Optimizer-state quantization: Optimizer-state quantization compresses first and second moments, with reported settings ranging from FP8 to 4-bit precision.
  • Error feedback: Error-feedback methods accumulate quantization residuals and add them back later to preserve the sum of updates over time.

3. Method

ECO removes full-precision master weights by injecting quantization error into optimizer momentum, creating a memory-free error-feedback loop. Its analysis establishes convergence to a bounded neighborhood, while experiments and lower-bound results characterize when quantization and naive removal limit accuracy.

  • ECO method: ECO quantizes updated parameters and injects the resulting error into momentum, carrying lost updates forward without master weights or an error buffer.The method applies updates directly to quantized parameters and uses the current quantization error as the memory-free injection signal.
  • Optimizer variants: ECO supports both SGDM and Adam by applying the same error-injection principle, with Adam using an element-wise effective step size.The algorithmic overview separates optimizer-specific implementations while preserving the shared quantization-plus-momentum-injection design.
  • ECO method: ECO’s memory-free rule approximates the exact residual correction by substituting the previous residual with the current one because successive quantization errors are typically close.The exact rule requires storing the previous residual, whereas ECO removes that persistent buffer through the e_t≈e_t+1 heuristic.
  • Convergence analysis: As the learning rate vanishes, stochastic quantization retains a noise floor, while deterministic rounding can produce a larger O(L^2δ^2/(1−β)^2) floor because systematic bias is harder to average out.The stochastic floor is O(L^2σ^2/(1−β^2)); the deterministic dependence follows when σ≈δ.
  • Lower-bound behavior: Without error compensation, annealing the learning rate cannot achieve high accuracy because naive master-weight removal exhibits an error that grows inversely with the learning rate.ECO prevents the corresponding 1/η explosion, and its noise-floor bound is tight up to a factor of 4.

4. Experiments

Experiments show that ECO preserves quantized-training quality while removing master-weight storage across scaling studies, dense and SMoE pretraining, and DeepSeek-MoE fine-tuning. It improves memory–loss trade-offs and avoids the divergence seen with naive master-weight removal.

  • Scaling Law Experiments: ECO nearly recovers master-weight performance with stochastic rounding, while providing smaller gains with round-to-nearest.Round-to-nearest introduces bias into the momentum buffer.
  • Scaling Law Experiments: ECO adds negligible runtime overhead because error injection is a simple element-wise operation.
  • Error Similarity Study: Consecutive quantization errors remain similar during training, with relative norms near 1 and consistently high cosine similarity.Differences between consecutive errors diminish as the learning rate decays.
  • Gemma 3 1B Pre-training: ECO confirms its effectiveness on Gemma 3 1B validation loss, particularly with stochastic rounding.
  • Mixture of Experts Pre-training: ECO outperforms naive master-weight removal on the 2.1B SMoE model while incurring only minimal loss relative to master-weight methods.
  • Mixture of Experts Pre-training: Reducing master-weight precision from FP32 to FP8 lowers peak memory from 12 bytes per parameter to 9, approximately 25%.In this SMoE setting, peak memory is dominated by master weights and optimizer states.
  • DeepSeek-MoE-16B Fine-tuning: ECO matches master-weight baselines in DeepSeek-MoE-16B fine-tuning, whereas naive removal diverges under both RTN and SR.ECO also recovers master-weight zero-shot benchmark accuracy.

5. Conclusion

ECO removes master weights through momentum-based error feedback without additional memory overhead. Across dense Transformers and SMoE models, it nearly matches high-precision baselines while improving the static-memory versus loss trade-off, with strongest performance under stochastic rounding.

  • ECO forms an error-feedback loop through optimizer momentum to remove high-precision accumulation without additional memory overhead.
  • Across dense Transformers and SMoE models, ECO nearly matches high-precision baselines while improving the static-memory versus loss trade-off.
  • ECO performs best with stochastic rounding, while round-to-nearest can produce a higher noise floor.
  • With master weights available, round-to-nearest generally slightly outperforms stochastic rounding, creating an accuracy ceiling for ECO relative to the best RTN master-weight baselines.

A. Exact Error Injection

The paper shows that ideal momentum injection can reproduce master-weight SGDM exactly while storing only quantized weights. The injected method therefore preserves the same quantized iterates and gradients as the master-weight procedure.

  • Exact equivalence: Ideal momentum correction depending on current and previous quantization residuals enables exact reproduction of SGDM with high-precision master weights.The correction is applied to the momentum buffer while the method stores only quantized weights.
  • Exact equivalence: The injected method is initialized to match the master-weight state and follows the same SGDM momentum recurrence.The proof introduces implicit master weights and momentum corresponding to the injected method.
  • Exact equivalence: For all t ≥ 0, the injected method produces quantized iterates identical to those of SGDM with master weights.This follows from matching initial conditions and the shared update recurrence.
  • Exact equivalence: Identical quantized iterates yield identical gradients at every step.Because both procedures use the same quantized parameters for forward and backward computation, their gradients coincide.

B.1. Proof of Lemma 3.5

The proof establishes that momentum injection cancels the quantization error in the virtual parameter sequence. It then bounds the momentum through separate gradient and error accumulations.

  • Virtual-sequence cancellation: Substituting the injected momentum update and quantization residual relation into the parameter update isolates the virtual sequence.The algebra rewrites the quantized update in terms of the virtual parameters and injected momentum.
  • Virtual-sequence cancellation: The quantization-error term cancels exactly, leaving the virtual sequence governed by the corresponding master-weight-style update.The cancellation uses the coefficient identities for the injected momentum and residual term.
  • Descent analysis: The proof applies the standard descent lemma to the virtual sequence under L-smoothness and a step-size condition η ≤ 1/(2L).The condition ensures 1 − Lη ≥ 1/2 in the descent argument.
  • Momentum bound: The momentum expansion separates into gradient accumulation S1 and quantization-error accumulation S2.The proof bounds the two components separately before combining them.
  • Momentum bound: The error-accumulation bound uses conditional unbiasedness of residuals and a uniform second-moment bound E[||e_k||^2] ≤ σ^2.These assumptions control the contribution of quantization errors in the geometric momentum sum.

B.4. Proof of Theorem 3.8

The proof derives a bound on the momentum magnitude by combining earlier descent and momentum lemmas. Under bounded quantization error and bounded gradients, the resulting recurrence remains bounded.

  • Assumptions: The analysis assumes an absolute quantization-error bound ||e_t|| ≤ δ and a bounded-gradient condition ||∇f(θ)|| ≤ G.These assumptions provide the control terms used in the theorem’s bound.
  • Recurrence bound: The momentum bound reduces to a linear recurrence x_{t+1} ≤ βx_t + K, whose solution is bounded by a geometric series.The bound assumes zero initial injected momentum.

B.6. Proof of Theorem 3.10

The proof of Theorem 3.10 uses the previously established lemmas and rearranges the resulting inequality to complete the argument.

  • Theorem proof: The proof invokes Lemmas 3.6 and 3.9 as the ingredients for Theorem 3.10.These lemmas supply the descent and momentum bounds used in the theorem.
  • Theorem proof: Rearranging the inequality and using f* ≤ f(θ_T) completes the proof.The final step applies the objective lower-bound relation before closing the argument.

C. Formal Analysis of the Worst-Case Lower-Bounds

The appendix derives stationary second-moment formulas for master-weight SGDM, naive quantized updates, and ECO. The analysis shows that naive removal produces a 1/η error blow-up, whereas ECO yields a finite noise floor.

  • C. Formal Analysis of the Worst-Case Lower-Bounds: The appendix casts all three regimes in a common linear dynamical form and derives second-moment update equations.The stationary analysis uses constants that depend on the regime and applies a dedicated second-moment lemma.
  • C. Formal Analysis of the Worst-Case Lower-Bounds: A unique stationary second moment requires ρ(A) < 1, and the SGDM calculations assume a condition ensuring positive closed-form denominators.Here ρ(A) is the largest absolute eigenvalue of A.
  • C.1. Fundamental limits on f(x) = L: The analysis measures the stationary squared gradient of the quantized parameter seen by the forward and backward passes.The model uses quantized weights ˆx_t during computation, while the master-weight baseline stores a full-precision x_t.
  • C.1.1. SGDM with master weights: For SGDM with master weights, stationary second moments are obtained by solving the linear system induced by the regime parameters (B1, B2) = (−ηc, c).The resulting expression gives the stationary parameter second moment u = E[x^2].
  • C.1.2. Naive master-weight removal: Naive master-weight removal stores only quantized weights and updates them directly with SGDM before requantization.Its dynamics correspond to (B1, B2) = (1, 0), with the state identified as x_t ≡ ˆx_t.
  • C.1.3. ECO: momentum injection eliminates the 1/η blow-up: ECO performs the same quantized SGDM step but injects the post-quantization error into momentum, using ˆm_t+1 = ˜m_t+1 + γξ_t+1.The error is defined as e_t+1 = ˜x_t+1 − ˆx_t+1 = −ξ_t+1, and ECO corresponds to (B1, B2) = (1, γ).
  • C.1.3. ECO: momentum injection eliminates the 1/η blow-up: Comparing the stationary solutions, ECO stabilizes the dynamics with a finite noise floor controlled by 1/(1 − β^2), unlike naive removal’s 1/η blow-up.The ECO stationary squared-gradient expression follows from solving the corresponding stationary second-moment system.
Loading 2601.22101v1…