无需高精度权重缓冲,实现高效低内存大模型量化训练
ECO: Quantized Training without Full-Precision Master Weights
- 直接在量化参数上更新,用误差反馈机制替代高精度主权重
- 30-800M小模型、1B/Gemma-3、2.1B稀疏专家模型均达近无损精度
- 适合追求低内存占用的模型训练与部署,尤其适用于稀疏MoE架构
量化显著提升了大语言模型训练的计算与内存效率。然而,现有方法仍需在高精度中累积更新:具体而言,梯度更新必须作用于称为‘主权重’的高精度权重缓冲区。该缓冲区带来显著内存开销,尤其在稀疏专家混合(SMoE)模型中,模型参数与优化器状态主导内存使用。为此,我们提出误差补偿优化器(ECO),通过直接对量化参数应用更新,消除主权重。ECO在每步后量化权重,并将量化误差注入优化器动量,形成无需额外内存的误差反馈循环。我们证明,在标准假设和衰减学习率下,ECO收敛至最优解的常数半径邻域,而简单移除主权重会导致与学习率成反比的误差。我们在预训练30-800M小Transformer、Gemma-3 1B模型及2.1B参数稀疏MoE模型(均采用FP8量化),以及深度求索-16B MoE模型微调(INT4精度)上进行实证。结果表明,ECO在几乎所有场景中均达到与带主权重基线近乎无损的精度,显著推动静态内存与验证损失之间的帕累托前沿。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。