Tian Yuandong らによる新作: メモリのボトルネックを突破し、4090 で事前トレーニングされた 7B 大型モデルを可能にする

Meta FAIR Tian Yuandong が参加した研究プロジェクトは、先月広く称賛されました。彼らの論文「MobileLLM: Optimizing Subbillion Parameter Language Models for On-Device Use Cases」では、モバイル デバイス上で大規模な言語モデルを実行するという目標を達成することを目的として、10 億未満のパラメータを持つ小規模なモデルを最適化する方法の検討を開始しました。 。

3 月 6 日、Tian Yuandong のチームは最新の研究結果を発表しました。今回は LLM メモリの効率向上に焦点を当てました。研究チームには、Tian Yuandong氏自身に加えて、カリフォルニア工科大学、テキサス大学オースティン校、CMUの研究者も含まれています。この研究は、LLM メモリのパフォーマンスをさらに最適化し、将来の技術開発へのサポートと指針を提供することを目的としています。

彼らは共同で、完全なパラメータ学習を可能にする GaLore (Gradient Low-Rank Projection) と呼ばれるトレーニング戦略を提案しました。LoRA、適応法、GaLore などの一般的な低ランク自動手法と比較メモリ効率が高くなります。

この調査では、モデルの並列処理を使用せずに、NVIDIA RTX 4090 などの 24 GB のメモリを備えたコンシューマ GPU で 7B モデルを正常に事前トレーニングできることが初めて示されました (Checkpoint)またはオフロード戦略。

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

論文タイトル: GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection

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

現在、大規模言語モデル (LLM) は多くの分野で優れた可能性を示していますが、現実的な問題にも直面する必要があります。それは、LLM の事前トレーニングと微調整だけではありません。大量のコンピューティング リソースが必要です。また、大量のメモリ サポートも必要です。

LLM のメモリ要件には、数十億のパラメータだけでなく、勾配やオプティマイザ状態 (Adam の勾配の勢いや分散など) も含まれており、ストレージ自体よりも大きくなる可能性があります。たとえば、単一のバッチ サイズを使用して最初から事前トレーニングされた LLaMA 7B には、少なくとも 58 GB のメモリが必要です (トレーニング可能なパラメータに 14 GB、Adam Optimizer の状態と重み勾配に 42 GB、アクティベーションに 2 GB)。そのため、24 GB のメモリを搭載した 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 のメモリ効率のおかげで、LLaMA 7B は初めて、高価なメモリ オフロード技術を使用せずに、24 GB のメモリを備えた単一の GPU (NVIDIA RTX 4090 など) で最初からトレーニングできるようになりました (図 1)。

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

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

実験と結果

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

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

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

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

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

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

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

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

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

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

以上がTian Yuandong らによる新作: メモリのボトルネックを突破し、4090 で事前トレーニングされた 7B 大型モデルを可能にするの詳細内容です。詳細については、PHP 中国語 Web サイトの他の関連記事を参照してください。