Tian Yuandong らの新しい研究: メモリのボトルネックを突破し、4090 で 7B の大規模モデルを事前トレーニング可能に

先月、Meta FAIR の Tian Yuandong が参加した研究が大きな称賛を受けました。彼らの論文「MobileLLM: オンデバイス ユースケース向けの 10 億未満のパラメータ言語モデルの最適化」では、モバイル デバイス上での LLM の実行に焦点を当て、10 億未満のパラメータを持つ小規模なモデルの展開を開始しました。

3月6日、Tian Yuandong氏は、今度はLLMメモリ効率に焦点を当てた別の研究を発表しました。田元東氏本人に加え、カリフォルニア工科大学、テキサス大学オースティン校、CMU の研究者も参加しています。

彼らは協力して、完全なパラメータ学習を可能にしながら、LoRA などの一般的な低ランク適応方法よりもメモリ効率の高いトレーニング戦略である GaLore (Gradient Low-Rank Projection) を提案しました。

この研究では、モデルの並列処理、チェックポイント、オフロード戦略を必要とせずに、24GB のメモリを搭載したコンシューマーグレードの GPU (NVIDIA RTX 4090 など) で 7B モデルを事前トレーニングできることを初めて実証しました。

論文アドレス: https://arxiv.org/abs/2403.03507

論文タイトル: GaLore: 勾配低ランク射影によるメモリ効率の良い LLM トレーニング

次に、記事の主な内容を見てみましょう。

現在、大規模言語モデル (LLM) は多くの分野で素晴らしい成果を上げています。しかし、LLM の事前トレーニングと微調整には大量の計算だけでなく大量のメモリも必要になるという問題に直面しています。

LLM のメモリ要件には、数十億単位で計算されるパラメータだけでなく、勾配やオプティマイザ状態 (Adam の勾配運動量や分散など) も含まれ、これらはストレージ自体よりも大きくなる場合があります。たとえば、単一のバッチ サイズを使用して最初から事前トレーニングされた LLaMA 7B には、少なくとも 58 GB のメモリが必要です (トレーニング可能なパラメーター用に 14 GB、Adam Optimizer 状態と重み勾配用に 42 GB、アクティベーション用に 2 GB)。このため、24GB のメモリを搭載した NVIDIA RTX 4090 などのコンシューマーグレードの GPU では LLM のトレーニングは実行不可能になります。

上記の問題に対処するために、研究者は事前トレーニングと微調整中のメモリ使用量を削減するためのさまざまな最適化手法の開発を続けています。

このアプローチにより、オプティマイザー状態でのメモリ使用量が 65.5% 削減され、LLaMA 1B および 7B アーキテクチャで最大 197 億トークンの C4 データセットの事前トレーニングと、GLUE タスクでの RoBERTa の微調整の効率とパフォーマンスが維持されます。 BF16 ベースラインと比較すると、8 ビット GaLore ではオプティマイザー メモリがさらに 82.5% 削減され、トレーニング メモリの合計が 63.3% 削減されます。

この研究を見たネットユーザーはこう言った。「クラウドと HPC のことは忘れるべきだ。GaLore があれば、AI4Science はすべて 2,000 ドルの消費者向け GPU で実行されるだろう。」

Tian Yuandong 氏は次のように語っています。「GaLore を使用すると、24G メモリを搭載した NVidia RTX 4090 で 7B モデルを事前トレーニングできるようになりました。

LoRA のように低ランクの重み構造を想定する代わりに、重み勾配は自然に低ランクであり、したがって (変化する) 低次元空間に投影できることを示します。したがって、勾配、アダム運動量、分散のメモリを同時に節約できます。

したがって、LoRA とは異なり、GaLore はトレーニング ダイナミクスを変更せず、メモリを消費するウォームアップなしで 7B モデルを最初から事前トレーニングするために使用できます。 GaLore は微調整にも使用でき、LoRA に匹敵する結果を生成します。」

方法の紹介

前述したように、GaLore は完全なパラメータ学習を可能にするトレーニング戦略ですが、LoRA などの一般的な低ランク適応型方法よりもメモリ効率が優れています。 GaLore の重要なアイデアは、重み行列を低ランク形式に直接近似するのではなく、勾配がゆっくり変化する重み行列 W の低ランク構造を活用することです。

この論文ではまず、勾配行列 G がトレーニングプロセス中に低ランクになることを理論的に証明します。この理論に基づいて、この論文では GaLore を使用して 2 つの射影行列を計算し、勾配行列 G を低ランク形式 P^⊤GQ に射影します。この場合、コンポーネントの勾配統計に依存するオプティマイザー状態のメモリ コストを大幅に削減できます。表 1 に示すように、GaLore は LoRA よりもメモリ効率に優れています。実際には、これにより、LoRA と比較して事前トレーニング中のメモリ使用量が最大 30% 削減されます。

この論文では、GaLore が事前トレーニングと微調整の両方で優れたパフォーマンスを発揮することを示しています。 C4 データセットで LLaMA 7B を事前トレーニングする場合、8 ビット GaLore は 8 ビット オプティマイザーとレイヤーごとの重み更新手法を組み合わせて、オプティマイザー状態メモリ コストの 10% 未満でフルランクに匹敵するパフォーマンスを実現します。

事前トレーニングでは、GaLore は ReLoRA のようなフルランク トレーニングを必要とせず、トレーニング プロセス全体を通じて低いメモリを維持することに注目すべきです。 GaLore のメモリ効率のおかげで、高価なメモリオフロード技術を使用せずに、24GB のメモリを搭載した単一の GPU (NVIDIA RTX 4090 など) で LLaMA 7B をゼロからトレーニングできるのは今回が初めてです (図 1)。

勾配投影法である GaLore は、オプティマイザーの選択に依存せず、アルゴリズム 1 に示すように、わずか 2 行のコードで既存のオプティマイザーに簡単に組み込むことができます。

次の図は、GaLore を Adam に適用するアルゴリズムを示しています。

実験と結果

研究者らは、GaLore の事前トレーニングと LLM の微調整を評価しました。すべての実験は NVIDIA A100 GPU で実行されます。

研究者らは、そのパフォーマンスを評価するために、GaLore を適用して、C4 データセット上で大規模な LLaMA ベースの言語モデルをトレーニングしました。 C4 データセットは、Common Crawl Web クローリング コーパスの大規模なクリーンアップ バージョンであり、主に言語モデルと単語表現の事前トレーニングに使用されます。実際の事前トレーニング シナリオを最適にシミュレートするために、研究者は、最大 70 億のパラメータを持つモデル サイズで、データを重複させずに十分な量のデータでトレーニングを行いました。

この論文は、Lialin らの実験設定に従い、RMSNorm と SwiGLU アクティベーションを備えた LLaMA3 ベースのアーキテクチャを採用しています。各モデル サイズに対して、学習率を除く同じハイパーパラメータ セットを使用し、すべての実験を BF16 形式で実行してメモリ使用量を削減しながら、同じ計算予算で各手法の学習率を調整し、最高のパフォーマンスを報告しました。

さらに、研究者らは、GaLore と LoRA のメモリ効率の高い微調整のベンチマークとして GLUE タスクを使用しました。 GLUE は、感情分析、質問回答、テキスト関連性など、さまざまなタスクにおける NLP モデルのパフォーマンスを評価するためのベンチマークです。

この論文では、まず Adam オプティマイザーを使用して、GaLore と既存の低ランク手法を比較します。結果は表 2 に示されています。

研究者らは、GaLore をさまざまな学習アルゴリズム、特にメモリ効率の高いオプティマイザーに適用して、メモリ使用量をさらに削減できることを実証しました。研究者らは、GaLore を AdamW、8 ビット Adam、および Adafactor オプティマイザーに適用しました。パフォーマンスの低下を避けるために、一次統計を備えた Adafactor を採用しています。

LLaMA 1B アーキテクチャで 10K のトレーニング ステップを使用して評価し、各設定の学習率を調整して、最高のパフォーマンスを報告します。図 3 に示すように、下の図は、GaLore が AdamW、8 ビット Adam、Adafactor などの一般的な最適化ツールに適用できることを示しています。さらに、ごく少数のハイパーパラメータを導入しても、GaLore のパフォーマンスには影響しません。

表 4 に示すように、ほとんどのタスクでは、GaLore はメモリ使用量を抑えながら LoRA よりも高いパフォーマンスを実現できます。これは、GaLore が LLM 事前トレーニングおよび微調整のためのフルスタックのメモリ効率の高いトレーニング戦略として使用できることを示しています。

図 4 に示すように、8 ビット GaLore は BF16 ベースラインや 8 ビット Adam よりもはるかに少ないメモリしか必要とせず、GPU あたりのトークン バッチ サイズが小さい (最大 500 トークン) LLaMA 7B を事前トレーニングする場合、必要なメモリは 22.0G のみです。

より技術的な詳細については、元の論文をお読みください。