COSMOS: A Hybrid Adaptive Optimizer for Memory-Efficient Training of LLMs

• URL: http://arxiv.org/abs/2502.17410v2
• Date: Wed, 26 Feb 2025 01:23:19 GMT
• Title: COSMOS: A Hybrid Adaptive Optimizer for Memory-Efficient Training of LLMs
• Authors: Liming Liu, Zhenghao Xu, Zixuan Zhang, Hao Kang, Zichong Li, Chen Liang, Weizhu Chen, Tuo Zhao,
• Abstract summary: Large Language Models (LLMs) have demonstrated remarkable success across various domains.<n>Their optimization remains a significant challenge due to the complex and high-dimensional loss landscapes they inhabit.
• Score: 81.01082659623552
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Towards Efficient Optimizer Design for LLM via Structured Fisher Approximation with a Low-Rank Extension [16.037614012166063]
  This paper makes a step towards the systematic design of efficient approximations through the lens of Fisher information matrix (FIM)<n>We show that many state-of-the-art efficient approximations can be viewed as solutions to FIM (under the Frobenius norm) with specific structural assumptions.<n>We propose two design recommendations of practical efficients for LLMs, involving careful selection of structural assumptions to balance generality and efficiency.
  arXiv  Detail & Related papers  (2025-02-11T18:27:19Z)
• A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models [22.725326215887435]
  We introduce a Randomized Subspace Optimization framework for pre-training and fine-tuning Large Language Models.<n>Our approach decomposes the high-dimensional training problem into a series of lower-dimensional subproblems.<n>This structured reduction in dimensionality allows our method to simultaneously reduce memory usage for both activations and states.
  arXiv  Detail & Related papers  (2025-02-11T03:32:10Z)
• Sparse Gradient Compression for Fine-Tuning Large Language Models [58.44973963468691]
  Fine-tuning large language models (LLMs) for downstream tasks has become increasingly crucial due to their widespread use and the growing availability of open-source models.<n>High memory costs associated with fine-tuning remain a significant challenge, especially as models increase in size.<n>We propose sparse compression gradient (SGC) to address these limitations.
  arXiv  Detail & Related papers  (2025-02-01T04:18:28Z)
• Zeroth-Order Fine-Tuning of LLMs in Random Subspaces [66.27334633749734]
  As language models grow in size, memory demands for backpropagation increase. Zeroth-order (ZOZO) optimization methods offer a memory-efficient alternative. We show that SubZero enhances fine-tuning and achieves faster results compared to standard ZOZO approaches.
  arXiv  Detail & Related papers  (2024-10-11T17:01:43Z)
• Optima: Optimizing Effectiveness and Efficiency for LLM-Based Multi-Agent System [75.25394449773052]
  Large Language Model (LLM) based multi-agent systems (MAS) show remarkable potential in collaborative problem-solving. Yet they still face critical challenges: low communication efficiency, poor scalability, and a lack of effective parameter-updating optimization methods. We present Optima, a novel framework that addresses these issues by significantly enhancing both communication efficiency and task effectiveness.
  arXiv  Detail & Related papers  (2024-10-10T17:00:06Z)
• Memory-Efficient Optimization with Factorized Hamiltonian Descent [11.01832755213396]
  We introduce a novel adaptive, H-Fac, which incorporates a memory-efficient factorization approach to address this challenge. By employing a rank-1 parameterization for both momentum and scaling parameter estimators, H-Fac reduces memory costs to a sublinear level. We develop our algorithms based on principles derived from Hamiltonian dynamics, providing robust theoretical underpinnings in optimization dynamics and convergence guarantees.
  arXiv  Detail & Related papers  (2024-06-14T12:05:17Z)
• Revisiting Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning: A Benchmark [166.40879020706151]
  This paper proposes a shift towards BP-free, zeroth-order (ZO) optimization as a solution for reducing memory costs during fine-tuning. Unlike traditional ZO-SGD methods, our work expands the exploration to a wider array of ZO optimization techniques. Our study unveils previously overlooked optimization principles, highlighting the importance of task alignment, the role of the forward gradient method, and the balance between algorithm complexity and fine-tuning performance.
  arXiv  Detail & Related papers  (2024-02-18T14:08:48Z)
• AdaLomo: Low-memory Optimization with Adaptive Learning Rate [59.64965955386855]
  We introduce low-memory optimization with adaptive learning rate (AdaLomo) for large language models. AdaLomo results on par with AdamW, while significantly reducing memory requirements, thereby lowering the hardware barrier to training large language models.
  arXiv  Detail & Related papers  (2023-10-16T09:04:28Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.