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

📄 arXiv: 2411.07837v3 📥 PDF

作者: Philip Zmushko, Aleksandr Beznosikov, Martin Takáč, Samuel Horváth

分类: cs.LG

发布日期: 2024-11-12 (更新: 2025-08-14)

💡 一句话要点

FRUGAL：通过减少状态开销实现内存高效优化，用于可扩展训练

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 内存高效优化 梯度分裂 低秩更新 大规模语言模型 预训练 微调 无状态优化 深度学习

📋 核心要点

  1. 现有低秩优化方法在大型模型训练中损失大量梯度信息，尤其在预训练阶段表现明显。
  2. FRUGAL通过梯度分裂，对重要维度进行高级优化，其余维度采用无状态方法，兼顾效率与性能。
  3. 实验表明，FRUGAL在固定内存预算下优于现有方法，并在预训练和微调任务中取得领先成果。

📝 摘要（中文）

随着大型语言模型参数数量的增加，预训练和微调过程对GPU内存的需求越来越大。优化器状态通常会消耗大部分内存。为了克服这一挑战，最近提出了低秩适应（LoRA）、低秩梯度投影（GaLore）和分块优化（BAdam）等方法。然而，在所有这些算法中，权重更新的有效秩仍然是低秩的，这可能导致梯度信息的大量损失，尤其是在预训练阶段。本文介绍了一种新的内存高效优化框架FRUGAL（具有梯度分裂的全秩更新）。FRUGAL利用梯度分裂，使用高级算法（如Adam）执行低维更新，而沿剩余方向的更新则通过无状态方法（如SGD或signSGD）执行。我们的框架可以与各种低秩更新选择技术（包括GaLore和BAdam）集成。我们为使用SGDM进行低维更新和使用SGD进行无状态更新时，框架的收敛性提供了理论保证。此外，我们的方法在各种固定内存预算下始终优于现有方法，在预训练和微调任务中实现了最先进的结果，同时平衡了内存效率和性能指标。

🔬 方法详解

问题定义：大型语言模型训练过程中，优化器状态占用大量GPU内存，限制了模型规模和训练效率。现有低秩优化方法（如LoRA、GaLore、BAdam）虽然降低了内存占用，但由于权重更新的有效秩较低，导致梯度信息损失，影响模型性能。尤其是在预训练阶段，这种信息损失可能更为严重。

核心思路：FRUGAL的核心思路是将梯度分解为两个部分：一个低维部分和一个剩余部分。对低维部分，使用高级优化算法（如Adam）进行更新，以保留重要的梯度信息；对剩余部分，使用无状态优化算法（如SGD或signSGD）进行更新，以减少内存占用。通过这种梯度分裂的方式，FRUGAL能够在内存效率和模型性能之间取得平衡。

技术框架：FRUGAL框架主要包含以下几个步骤：1. 梯度计算：计算模型参数的梯度。2. 梯度分裂：将梯度分解为低维部分和剩余部分。可以使用各种低秩更新选择技术（如GaLore或BAdam）来选择低维部分。3. 低维更新：使用高级优化算法（如Adam）更新低维部分的参数。4. 剩余更新：使用无状态优化算法（如SGD或signSGD）更新剩余部分的参数。5. 参数更新：将低维更新和剩余更新的结果合并，更新模型参数。

关键创新：FRUGAL的关键创新在于梯度分裂策略，它允许对不同的梯度分量采用不同的优化算法。这种策略能够有效地利用有限的内存资源，同时保留重要的梯度信息。与现有低秩优化方法相比，FRUGAL能够实现全秩更新，从而避免了梯度信息损失。

关键设计：FRUGAL的关键设计包括：1. 低维部分的维度选择：可以使用各种低秩更新选择技术（如GaLore或BAdam）来选择低维部分的维度。2. 低维优化算法的选择：可以使用各种高级优化算法（如Adam、SGDM）来更新低维部分的参数。论文给出了使用SGDM进行低维更新和使用SGD进行无状态更新时的收敛性保证。3. 剩余优化算法的选择：可以使用各种无状态优化算法（如SGD、signSGD）来更新剩余部分的参数。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

FRUGAL在各种固定内存预算下，在预训练和微调任务中均优于现有方法。实验结果表明，FRUGAL能够在平衡内存效率和模型性能的同时，实现最先进的结果。具体性能数据和对比基线在论文中有详细展示，表明FRUGAL在内存受限场景下具有显著优势。

🎯 应用场景

FRUGAL适用于大规模语言模型的预训练和微调，尤其是在GPU内存资源有限的情况下。该方法可以帮助研究人员和工程师训练更大规模的模型，提高模型性能，并降低训练成本。此外，FRUGAL还可以应用于其他需要内存高效优化的机器学习任务，例如图像识别、语音识别等。

📄 摘要（原文）

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.