8位量化让穆子优化器内存减62%,性能不降反而更省显存。
Effective Quantization of Muon Optimizer States
- 用分块量化将穆子优化器状态压缩至8位
- 27亿参数模型上验证损失与下游任务表现持平
- 兼容简单线性量化,无需复杂动态缩放,适合大模型部署
基于矩阵正交化的穆子优化器在大语言模型预训练中已展现比AdamW更快的收敛速度和更高的计算效率。然而,维持高精度优化器状态带来的内存开销仍是大规模部署的挑战。本文提出8位穆子优化器,采用分块量化方法。在最大规模达27亿参数的模型上进行大量Chinchilla最优预训练与指令微调实验,结果表明,8位穆子优化器在验证损失和下游基准测试上与原版穆子优化器相当,同时优化器状态存储量最多减少62%。关键的是,我们证明穆子更新机制天然适配简单的线性量化方案,无需像量化AdamW那样依赖复杂的动态缩放。我们还通过理论分析说明了穆子对量化噪声具有鲁棒性。
原文摘要 · Abstract (English)
The Muon optimizer, based on matrix orthogonalization, has recently shown faster convergence and better computational efficiency over AdamW in LLM pre-training. However, the memory overhead of maintaining high-precision optimizer states remains a challenge for large-scale deployment. In this paper, we introduce the 8-bit Muon optimizer using blockwise quantization. In extensive Chinchilla-optimal experiments on pre-training models of up to 2.7B in size and fine-tuning them for instruction following, we demonstrate that 8-bit Muon achieves parity with Muon in terms of validation loss and downstream benchmarks, while achieving up to a 62\% reduction in optimizer state footprint. Crucially, we show that Muon's update mechanism is uniquely compatible with a simple linear quantization scheme, bypassing the complex dynamic scaling required for quantized AdamW. We supplement our empirical findings with a theoretical analysis of Muon's robustness to quantization noise.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。