650億のパラメータ、すべてのパラメータを8つのGPUで微調整可能：Qiu Xipengのチームは大規模モデルの閾値を下げました

大規模モデルに向けて、テクノロジー大手はより大規模なモデルをトレーニングしており、学界はそれらを最適化する方法を模索しています。最近、コンピューティング能力を最適化する方法は新たな高みに達しました。

大規模言語モデル (LLM) は、自然言語処理 (NLP) の分野に革命をもたらし、出現や洞察などの優れた機能を発揮しました。しかし、ある程度の汎用性を備えたモデルを構築しようとすると、数十億のパラメータが必要となり、NLP研究の敷居は大きく上がります。 LLM モデルのチューニング プロセスには通常、8×80GB GPU デバイスなどの高価な GPU リソースが必要になるため、小規模な研究室や企業がこの分野の研究に参加するのは困難です。

最近では、限られたリソースで LLM を調整するためのソリューションを提供するために、LoRA やプレフィックスチューニングなどのパラメータ効率の高い微調整技術 (PEFT) が研究されています。しかし、これらの方法は、パラメータ効率の良い微調整よりも強力なアプローチとして認識されている完全なパラメータの微調整に対する実用的なソリューションを提供しません。

先週、復旦大学のQiu Xipeng氏のチームが提出した論文「限られたリソースでの大規模言語モデルの完全なパラメータの微調整」の中で、研究者らは新しい最適化ツールである低メモリ最適化（LOMO）を提案した。

LOMO を既存のメモリ節約技術と統合することで、新しい方法では標準方法 (DeepSpeed ソリューション) と比較してメモリ使用量が 10.8% 削減されます。その結果、新しいアプローチにより、それぞれ 24 GB のメモリを搭載した 8 基の RTX 3090 を搭載したマシンで 65B モデルの完全なパラメータ微調整が可能になります。

論文リンク: https://arxiv.org/abs/2306.09782

この研究では、著者らは LLM におけるメモリ使用量の 4 つの側面 (アクティベーション、オプティマイザーの状態、勾配テンソル、パラメーター) を分析し、次の 3 つの側面でトレーニング プロセスを最適化しました。

1. 私たちはアルゴリズムの観点からオプテ​​ィマイザーの機能を再考し、SGD が LLM の完全なパラメータを微調整するための優れた代替手段であることを発見しました。 SGD は中間状態を保存しないため、これにより著者はオプティマイザー状態のセクション全体を削除することができました。
2. 新しく提案された最適化装置LOMOは、勾配テンソルのメモリ使用量をO(1)に削減します。これは、最大の勾配テンソルのメモリ使用量に相当します。
3. LOMO による混合精度トレーニングを安定させるために、著者らは勾配正規化、損失スケーリングを統合し、トレーニング中に特定の計算を完全精度に変換しました。

新しい技術により、メモリ使用量は、パラメータ使用量とアクティベーションおよび最大勾配テンソルの合計に等しくなります。完全なパラメータの微調整のメモリ使用量は極限まで押し上げられ、推論の使用量と同等になります。これは、forward+backward プロセスのメモリ使用量が、forward プロセスのみのメモリ使用量よりも少なくならないためです。 LOMO を使用してメモリを節約しながらも、新しい方法では、パラメータ更新プロセスが SGD と同等であるため、微調整プロセスが影響を受けないことが保証される点に注目すべきです。

この研究では、LOMO のメモリとスループットのパフォーマンスを評価し、LOMO を使用すると、研究者は 8 つの RTX 3090 GPU で 65B のパラメータを持つモデルをトレーニングできることを示しました。さらに、下流タスクにおける LOMO のパフォーマンスを検証するために、LOMO を適用して SuperGLUE データセット セット上の LLM のすべてのパラメータを調整しました。結果は、数十億のパラメータを持つ LLM を最適化するための LOMO の有効性を示しています。

方法の紹介

この論文の方法セクションでは、LOMO (LOW-MEMORY OPTIMIZATION) について詳しく説明します。一般的に、勾配テンソルはパラメーター テンソルの勾配を表し、パラメーターと同じサイズを持つため、メモリのオーバーヘッドが高くなります。 PyTorch などの既存のディープラーニング フレームワークは、すべてのパラメータの勾配テンソルを保存します。現在、勾配テンソルは、オプティマイザーの状態の計算と勾配の正規化という 2 つの理由で保存されます。

この研究ではオプティマイザとして SGD を採用したため、勾配に依存するオプティマイザの状態はなく、勾配正規化のためのいくつかの代替手段があります。

彼らは、アルゴリズム 1 に示すように、勾配計算とパラメータ更新を 1 つのステップで組み合わせ、勾配テンソルの保存を回避する LOMO を提案しました。

下の図は、バックプロパゲーションとパラメータ更新フェーズにおける SGD と LOMO の比較を示しています。 Pi はモデルパラメータであり、Gi は Pi に対応する勾配です。 LOMO は、勾配計算とパラメータ更新を 1 つのステップに統合して、勾配テンソルを最小化します。

LOMO の対応するアルゴリズム疑似コードは次のとおりです。

これは 2 段階のプロセスで、最初に勾配を計算し、次にパラメータを更新します。フュージョンバージョンは

この研究の重要なアイデアは、勾配を計算するときにパラメータを即座に更新し、勾配テンソルがメモリに保存されないようにすることです。このステップは、バックプロパゲーションにフック関数を挿入することで実現できます。 PyTorch はフック関数を挿入するための API を提供していますが、現在の API では正確な即時更新を実現することはできません。代わりに、この研究では最大 1 つのパラメータの勾配をメモリに保存し、バックプロパゲーションが進むにつれて各パラメータを 1 つずつ更新します。この方法では、すべてのパラメータの勾配を保存することから、1 つのパラメータのみの勾配を保存することに勾配のメモリ使用量が削減されます。

LOMO のメモリ使用量の大部分は、パラメータ効率の高い微調整方法のメモリ使用量と一致しており、LOMO とこれらの方法を組み合わせても、勾配によって占有されるメモリがわずかに増加するだけであることがわかります。これにより、PEFT メソッドに対してさらに多くのパラメータを調整できるようになります。

実験結果

実験部分では、研究者は提案された方法をメモリ使用量、スループット、ダウンストリームパフォーマンスという 3 つの側面から評価しました。さらなる説明なしに、すべての実験は 7B から 65B までの LLaMA モデルを使用して実行されました。

メモリ使用量

研究者らはまず、さまざまな設定でのトレーニング中のモデルの状態とアクティベーションのメモリ使用量をプロファイリングしました。表 1 に示すように、LOMO オプティマイザーを使用すると、AdamW オプティマイザーと比較してメモリ使用量が 102.20 GB から 14.58 GB に大幅に削減されます。また、SGD と比較すると、LLaMA-7B モデルをトレーニングする際のメモリ使用量は 51.99 GB から 14.58 GB に削減されます。メモリ使用量が大幅に削減されたのは、主に勾配とオプティマイザーの状態に対するメモリ要件が削減されたためです。したがって、トレーニング中はメモリの大部分がパラメータによって占有され、これは推論中のメモリ使用量に匹敵します。

図 2 に示すように、LLaMA-7B トレーニングに AdamW オプティマイザーを使用すると、メモリのかなりの割合 (73.7%) がオプティマイザー状態に割り当てられます。 AdamW オプティマイザーを SGD オプティマイザーに置き換えると、オプティマイザー状態によって占有されるメモリの割合を効果的に削減でき、GPU メモリの使用量を軽減できます (102.20 GB から 51.99 GB)。 LOMO を使用すると、パラメータの更新と後方処理が 1 つのステップに統合され、オプティマイザの状態に対するメモリ要件がさらに排除されます。

スループット

研究者らは、LOMO、AdamW、SGDのスループット性能を比較した。実験は、8 個の RTX 3090 GPU を搭載したサーバー上で実施されました。

7B モデルの場合、LOMO のスループットは AdamW および SGD を約 11 倍上回り、大きな優位性を示しています。この大幅な改善は、LOMO が単一の GPU で 7B モデルをトレーニングし、GPU 間の通信オーバーヘッドを削減する機能のおかげです。 SGD は AdamW と比較してスループットがわずかに高くなっています。これは、SGD が運動量と分散の計算を除外していることに起因します。

13B モデルに関しては、メモリ制限のため、既存の 8 つの RTX 3090 GPU 上で AdamW を使用してトレーニングすることはできません。この場合、LOMO にはモデルの並列処理が必要ですが、スループットの点では SGD よりも優れています。この利点は、LOMO のメモリ効率の良さと、同じ設定でモデルをトレーニングするのに 2 つの GPU しか必要ないため、通信コストが削減され、スループットが向上するという事実に起因しています。さらに、SGD は 30B モデルをトレーニングする際に 8 個の RTX 3090 GPU でメモリ不足 (OOM) の問題が発生しましたが、LOMO は 4 個の GPU のみで良好なパフォーマンスを発揮しました。

最終的に、研究者らは 8 つの RTX 3090 GPU を使用して 65B モデルのトレーニングに成功し、4.93 TGS のスループットを達成しました。このようなサーバー構成と LOMO を使用すると、1000 個のサンプル (各サンプルには 512 個のトークンが含まれます) のモデル トレーニング プロセスには約 3.6 時間かかります。

下流のパフォーマンス

大規模言語モデルの微調整における LOMO の有効性を評価するために、研究者らは一連の大規模な実験を実施しました。彼らは、LOMO を他の 2 つの方法と比較しました。1 つは微調整を必要としないゼロショット、もう 1 つは現在人気のあるパラメータ効率の高い微調整技術である LoRA です。

• LOMO は Zero-shot よりも大幅に優れたパフォーマンスを発揮します。
• ほとんどの実験では、LOMO は一般に LoRA よりも優れています。
• LOMO は、650 億のパラメータを持つモデルまで効果的に拡張できます。

LOMO と LoRA は本質的に互いに独立しています。この主張を検証するために、研究者らは LLaMA-13B を使用して BoolQ および MultiRC データセットで実験を実施しました。結果は図3に示されています。

彼らは、LoRA によって達成された結果がどれほど高いかに関係なく、LOMO が一貫して LoRA のパフォーマンスを向上させることを発見しました。これは、LOMO と LoRA で採用されている異なる微調整方法が補完的であることを示しています。具体的には、LOMO は事前トレーニング済みモデルの重みを微調整することに重点を置き、LoRA は他のモジュールを調整します。したがって、LOMO は LoRA のパフォーマンスには影響しません。むしろ、下流のタスクに合わせてモデルをより適切に調整するのに役立ちます。

詳細については、元の論文を参照してください。