arXiv:2502.07222cs.LG2025-02ICML被引 18

提出一种低内存的随机子空间优化方法,高效训练大模型。

A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models

  • 将高维训练分解为随机低维子问题逐次优化
  • 内存使用显著降低,性能媲美Adam和GaLore
  • 适用于长序列与大批次场景,适合资源受限训练

训练大语言模型(LLMs)的内存挑战日益突出,尤其在使用Adam优化器时。尽管已有多种内存高效技术,如GaLore,能减少优化器状态内存占用,但无法缓解激活内存压力,因此不适用于长上下文或大批次场景。此外,其收敛性尚未充分研究。本文提出一种针对预训练与微调的随机子空间优化框架,将高维训练问题分解为一系列低维子问题。每轮迭代中随机选取一个子空间,并优化其中参数。该结构化降维方法同时降低了激活与优化器状态的内存开销。我们建立了完整的收敛性保证,并推导了不同策略下的收敛速率。大量实验验证了该方法在内存与通信效率上的优越性,性能可媲美GaLore和Adam。

原文摘要 · Abstract (English)

The memory challenges associated with training Large Language Models (LLMs) have become a critical concern, particularly when using the Adam optimizer. To address this issue, numerous memory-efficient techniques have been proposed, with GaLore standing out as a notable example designed to reduce the memory footprint of optimizer states. However, these approaches do not alleviate the memory burden imposed by activations, rendering them unsuitable for scenarios involving long context sequences or large mini-batches. Moreover, their convergence properties are still not well-understood in the literature. In this work, we introduce a Randomized Subspace Optimization framework for pre-training and fine-tuning LLMs. Our approach decomposes the high-dimensional training problem into a series of lower-dimensional subproblems. At each iteration, a random subspace is selected, and the parameters within that subspace are optimized. This structured reduction in dimensionality allows our method to simultaneously reduce memory usage for both activations and optimizer states. We establish comprehensive convergence guarantees and derive rates for various scenarios, accommodating different optimization strategies to solve the subproblems. Extensive experiments validate the superior memory and communication efficiency of our method, achieving performance comparable to GaLore and Adam.

大模型训练内存优化随机优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。