arXiv:2411.07837cs.LG2024-11ICML被引 10

FRUGAL通过梯度拆分实现高效训练,大幅降低显存占用且不损失性能。

FRUGAL: Memory-Efficient Optimization by Reducing State Overhead for Scalable Training

  • 将梯度拆分为低维和高维两部分,分别用Adam和无状态SGD更新
  • 在固定显存下比现有方法更优,预训练与微调均达顶尖效果
  • 可兼容LoRA、GaLore等技术,适合大规模模型训练场景

随着大语言模型参数量增长,预训练和微调对GPU显存的需求日益增加,其中优化器状态占用了大量内存。尽管低秩适配(LoRA)、低秩梯度投影(GaLore)和分块优化(BAdam)等方法被提出,但这些方法的权重更新始终为低秩,导致梯度信息显著损失,尤其在预训练阶段影响严重。本文提出新框架FRUGAL(Full-Rank Updates with Gradient Splitting),通过梯度拆分,利用先进算法(如Adam)在低维方向进行更新,同时在剩余方向采用无状态方法(如SGD、signSGD)完成更新。该框架可与多种低秩更新选择技术结合,并提供理论收敛性保证:当低维更新使用SGDM、高维更新使用SGD时成立。实验表明,在不同固定显存预算下,FRUGAL持续优于当前方法,在预训练与微调任务中均达到最优表现,兼顾内存效率与性能。

原文摘要 · Abstract (English)

With the increase in the number of parameters in large language models, the process of pre-training and fine-tuning increasingly demands larger volumes of GPU memory. A significant portion of this memory is typically consumed by the optimizer state. To overcome this challenge, recent approaches such as low-rank adaptation (LoRA (Hu et al., 2021)), low-rank gradient projection (GaLore (Zhao et al., 2024)), and blockwise optimization (BAdam (Luo et al., 2024)) have been proposed. However, in all these algorithms, the $\textit{effective rank of the weight updates remains low-rank}$, which can lead to a substantial loss of information from the gradient. This loss can be critically important, especially during the pre-training stage. In this paper, we introduce $\texttt{FRUGAL}$ ($\textbf{F}$ull-$\textbf{R}$ank $\textbf{U}$pdates with $\textbf{G}$r$\textbf{A}$dient sp$\textbf{L}$itting), a new memory-efficient optimization framework. $\texttt{FRUGAL}$ leverages gradient splitting to perform low-dimensional updates using advanced algorithms (such as Adam), while updates along the remaining directions are executed via state-free methods like SGD or signSGD (Bernstein et al., 2018). Our framework can be integrated with various low-rank update selection techniques, including GaLore and BAdam. We provide theoretical convergence guarantees for our framework when using SGDM for low-dimensional updates and SGD for state-free updates. Additionally, our method consistently outperforms concurrent approaches across various fixed memory budgets, achieving state-of-the-art results in pre-training and fine-tuning tasks while balancing memory efficiency and performance metrics.

优化器显存优化大模型训练

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