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

要約

大規模な言語モデルのパラメーター数の増加に伴い、事前トレーニングと微調整のプロセスでは、ますます大量の GPU メモリが必要になります。
通常、このメモリの大部分はオプティマイザ状態によって消費されます。
この課題を克服するために、低ランク適応 (LoRA (Hu et al., 2021))、低ランク勾配投影 (GaLore (Zhao et al., 2024))、ブロックごとの最適化 (BAdam (Luo et al., 2024)) などの最近のアプローチが使用されています。
al., 2024)) が提案されています。
ただし、これらすべてのアルゴリズムでは、$\textit{重み更新の有効ランク}$ が低ランクのままであり、勾配からの情報が大幅に失われる可能性があります。
この損失は、特にトレーニング前の段階では非常に重要になる可能性があります。
この論文では、$\texttt{FRUGAL}$ ($\textbf{F}$ull-$\textbf{R}$ank $\textbf{U}$pdates with $\textbf{G}$r$\ を紹介します
textbf{A}$dient sp$\textbf{L}$itting)、新しいメモリ効率の高い最適化フレームワーク。
$\texttt{FRUGAL}$ は、勾配分割を利用して、高度なアルゴリズム (Adam など) を使用して低次元の更新を実行しますが、残りの方向に沿った更新は、SGD や SignSGD などのステートフリーのメソッドを介して実行されます (Bernstein et al., 2018)
。
私たちのフレームワークは、GaLore や BAdam などのさまざまな低ランク更新選択手法と統合できます。
低次元更新に SGDM を使用し、ステートフリー更新に SGD を使用する場合、フレームワークに対して理論的な収束保証を提供します。
さらに、私たちの手法は、さまざまな固定メモリ バジェットにわたって一貫して同時アプローチよりも優れたパフォーマンスを発揮し、メモリ効率とパフォーマンス メトリクスのバランスをとりながら、事前トレーニングとタスクの微調整で最先端の結果を達成します。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google