4bit量化+分解优化,让大模型训练内存降一半
Memory-Efficient 4-bit Preconditioned Stochastic Optimization
- 用分解+4比特量化降低预条件矩阵内存占用
- 结合误差反馈,实现内存节省超50%且不降性能
- 适合大模型训练的高效优化器设计参考
预条件随机优化算法(如Shampoo)在大规模神经网络训练中表现优于一阶优化器,但其非对角预条件矩阵带来巨大内存开销。本文提出对Shampoo预条件器进行4比特量化:首先对Cholesky分解后的因子进行量化,利用其下三角结构减少存储;其次引入误差反馈机制,将因子与误差状态分别存于同一矩阵的上下三角部分。实验表明,该方法显著提升内存效率并保持算法性能。理论分析还给出了量化Shampoo在光滑与非光滑随机优化场景下的收敛性证明。
原文摘要 · Abstract (English)
Preconditioned stochastic optimization algorithms, exemplified by Shampoo, outperform first-order optimizers by offering theoretical convergence benefits and practical gains in large-scale neural network training. However, they incur substantial memory overhead due to the storage demands of non-diagonal preconditioning matrices. To address this, we introduce 4-bit quantization for Shampoo's preconditioners. We introduce two key methods: First, we apply Cholesky decomposition followed by quantization of the Cholesky factors, reducing memory usage by leveraging their lower triangular structure while better preserving spectral properties to minimize information loss. To our knowledge, this is the first quantization approach applied to Cholesky factors of preconditioners. Second, we incorporate error feedback in the quantization process, efficiently storing Cholesky factor and error state in the lower and upper triangular parts of the same matrix. Through extensive experiments, we demonstrate that combining Cholesky quantization with error feedback enhances memory efficiency and algorithm performance in large-scale deep-learning tasks. Theoretically, we also provide convergence proofs for quantized Shampoo under both smooth and non-smooth stochastic optimization settings.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。