提出混合优化器COSMOS,高效训练大模型同时节省内存。
COSMOS: A Hybrid Adaptive Optimizer for Memory-Efficient Training of LLMs
- 结合SOAP与MUON,分层处理梯度特征空间。
- 在多种架构上实现更高每词效率,内存消耗显著降低。
- 适合超大规模语言模型的高效训练,开源可用。
大型语言模型(LLMs)在多个领域表现卓越,但其优化面临复杂高维损失曲面的挑战。尽管自适应优化器如AdamW广泛应用,却存在难以捕捉坐标间依赖性及内存占用高的问题。后续工作如SOAP虽改进了依赖建模,但内存开销更大,限制了对超大规模模型的可扩展性。另一类方法通过低维投影降内存,却引入显著近似误差,影响优化效果(如每词效率下降)。本文提出COSMOS,一种新型混合优化器,利用梯度矩阵中不同特征子空间的重要性差异,在不牺牲优化性能的前提下实现内存高效。具体地,将SOAP应用于主导特征子空间以捕捉主要优化动态,而对剩余次要子空间采用更轻量的MUON策略。该设计显著降低内存需求并保持强优化能力,特别适用于大规模语言模型。实验在多个数据集和Transformer架构上验证了COSMOS的有效性。代码已开源:https://github.com/lliu606/COSMOS。
原文摘要 · Abstract (English)
Large Language Models (LLMs) have demonstrated remarkable success across various domains, yet their optimization remains a significant challenge due to the complex and high-dimensional loss landscapes they inhabit. While adaptive optimizers such as AdamW are widely used, they suffer from critical limitations, including an inability to capture interdependencies between coordinates and high memory consumption. Subsequent research, exemplified by SOAP, attempts to better capture coordinate interdependence but incurs greater memory overhead, limiting scalability for massive LLMs. An alternative approach aims to reduce memory consumption through low-dimensional projection, but this leads to substantial approximation errors, resulting in less effective optimization (e.g., in terms of per-token efficiency). In this paper, we propose COSMOS, a novel hybrid optimizer that leverages the varying importance of eigensubspaces in the gradient matrix to achieve memory efficiency without compromising optimization performance. The design of COSMOS is motivated by our empirical insights and practical considerations. Specifically, COSMOS applies SOAP to the leading eigensubspace, which captures the primary optimization dynamics, and MUON to the remaining eigensubspace, which is less critical but computationally expensive to handle with SOAP. This hybrid strategy significantly reduces memory consumption while maintaining robust optimization performance, making it particularly suitable for massive LLMs. Numerical experiments on various datasets and transformer architectures are provided to demonstrate the effectiveness of COSMOS. Our code is available at https://github.com/lliu606/COSMOS.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。