大規模言語モデルでのメモリ最適化と推論およびファインチューニング

大規模言語モデル（LLM）であるGPT-4、Bloom、LLaMAは、数十億のパラメータにスケールアップすることで、驚くべき能力を達成しました。ただし、これらの巨大なモデルを推論またはファインチューニングすることは、巨大なメモリ要件のため、課題です。この技術ブログでは、LLMの推論とファインチューニング中にメモリ消費を推定および最適化するためのテクニックを探究します。

メモリ要件の理解

LLMをロードするために必要なメモリは、主にパラメータの数とパラメータを保存するために使用される数値精度によって決定されます。簡単なルールは次のとおりです。

• X十億パラメータを持つモデルをロードするには、32ビット浮動小数点精度で約4X GBのVRAMが必要です。
• X十億パラメータを持つモデルをロードするには、16ビットbfloat16/float16精度で約2X GBのVRAMが必要です。

たとえば、175BパラメータのGPT-3モデルをロードするには、bfloat16精度で約350GBのVRAMが必要です。現在、最大の商用GPUであるNVIDIA A100とH100は、80GBのVRAMしか提供していないため、テンソル並列性とモデル並列性のテクニックが必要になります。

推論中、メモリフットプリントは、モデルパラメータと一時的なアクティベーションテンソルによって支配されます。推論中のピークメモリ使用量の高レベル推定は、モデルパラメータのメモリ要件とアクティベーションのメモリの合計です。

推論メモリの量化

OctoCodeモデルを使用して、推論のメモリ要件を量化してみましょう。OctoCodeモデルは、bfloat16形式で約15億パラメータ (~ 31GB) を持っています。Transformersライブラリを使用してモデルをロードし、テキストを生成します。

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

<p>model = AutoModelForCausalLM.from_pretrained(&quot;bigcode/octocoder&quot;,
torch_dtype=torch.bfloat16,
device_map=&quot;auto&quot;,
pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained(&quot;bigcode/octocoder&quot;)
pipe = pipeline(&quot;text-generation&quot;, model=model, tokenizer=tokenizer)</p>

<p>prompt = &quot;質問：Pythonでバイトをギガバイトに変換する関数を書いてください。\n\n答え：&quot;
result = pipe(prompt, max_new_tokens=60)[0][&quot;generated_text&quot;][len(prompt):]</p>

<p>def bytes_to_gigabytes(bytes):
return bytes / 1024 / 1024 / 1024</p>

<p>bytes_to_gigabytes(torch.cuda.max_memory_allocated())

出力：

ピークGPUメモリ使用量は約29GBで、bfloat16形式でモデルパラメータをロードするための推定31GBと一致しています。

量化による推論メモリの最適化

bfloat16は、LLMをトレーニングするための一般的な精度ですが、研究者は、モデル重みを8ビット整数（int8）または4ビット整数に量化することで、推論タスクの精度の低下と引き換えにメモリ使用量を大幅に削減できることを発見しました。

OctoCodeモデルの8ビットと4ビット量化からのメモリ節約を見てみましょう。

</div>
# 8ビット量化
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_8bit=True,
pad_token_id=0)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
bytes_to_gigabytes(torch.cuda.max_memory_allocated())</pre>

出力：

# 4ビット量化
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_4bit=True,
low_cpu_mem_usage=True, pad_token_id=0)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
bytes_to_gigabytes(torch.cuda.max_memory_allocated())

出力：

8ビット量化では、メモリ要件が31GBから15GBに低下し、4ビット量化ではさらに9.5GBに低下します。これにより、15BパラメータのOctoCodeモデルを、24GBのVRAMを持つRTX 3090などの消費者向けGPUで実行できます。

ただし、4ビットなどのより激しい量化は、8ビットまたはbfloat16精度と比較して、精度の低下につながる可能性があることに注意してください。メモリ節約と精度のトレードオフを評価する必要があります。

量化は、メモリフットプリントを大幅に削減できる強力なテクニックであり、クラウドインスタンス、エッジデバイス、またはモバイルフォンなどのリソース制約のある環境でLLMを展開できるようにします。

ファインチューニングのメモリ推定

量化は主に推論に使用される一方で、テンソル並列性やモデル並列性などのテクニックは、LLMのトレーニングまたはファインチューニング中にメモリを管理するために不可欠です。

ファインチューニング中のピークメモリ消費量は、推論よりも3〜4倍高くなります。これは、次の追加メモリ要件のためです。

• 勾配
• オプティマイザ状態
• 逆伝播のために保存されたフォワードパスのアクティベーション

保守的な推定では、X十億パラメータを持つLLMをファインチューニングするには、bfloat16精度で約4 * (2X) = 8X GBのVRAMが必要です。

たとえば、7BパラメータのLLaMAモデルをファインチューニングするには、bfloat16精度で約7 * 8 = 56GBのVRAMが必要です。これは、現在のGPUのメモリ容量を超え、分散ファインチューニングテクニックが必要になります。

分散ファインチューニングテクニック

GPUメモリ制約を克服するために、LLM用のいくつかの分散ファインチューニング方法が提案されています。

1. データ並列性：従来のデータ並列アプローチでは、モデルを複数のGPUに複製し、トレーニングデータバッチを分割して分配します。これにより、GPUの数に応じてトレーニング時間が線形に短縮されますが、各GPUのピークメモリ要件は削減されません。
2. ZeROステージ3：データ並列性の高度な形式で、モデルパラメータ、勾配、オプティマイザ状態をGPUに分割します。各GPUは、トレーニングの異なる段階で必要なデータのみを保持するため、クラシックデータ並列性よりもメモリを削減します。
3. テンソル並列性：モデルを複製するのではなく、テンソル並列性ではモデルパラメータを行または列に分割し、GPUに分配します。各GPUは、パラメータ、勾配、オプティマイザ状態のパーティション化されたセットを操作し、メモリを大幅に節約します。
4. パイプライン並列性：このテクニックでは、モデルレイヤーを異なるGPU/ワーカーに分割します。各デバイスは、レイヤーのサブセットを実行し、アクティベーションをワーカー間で渡します。これにより、ピークメモリが削減されますが、通信オーバーヘッドが増加します。

これらの分散方法のメモリ使用量を推定することは簡単ではありません。パラメータ、勾配、アクティベーション、オプティマイザ状態がテクニックによって異なるように分配されるためです。また、トランスフォーマーボディと言語モデリングヘッドなどの異なるコンポーネントは、メモリアロケーションの異なる動作を示す場合があります。

LLMemソリューション

研究者は、LLMemと呼ばれるソリューションを提案しました。これは、分散ファインチューニング方法をLLMに適用する際のGPUメモリ消費を正確に推定します。

事前トレーニング済みLLMのファインチューニングのためのGPUメモリ使用量の推定

LLMemは、計算の前にパラメータを再結合すること（ZeROステージ3）、バックワードパスの出力収集（テンソル並列性）、トランスフォーマーボディと言語モデリングヘッドのメモリアロケーション戦略の違いなどの要素を考慮します。

実験結果は、LLMemが、1つのGPUでLLMのファインチューニングのピークGPUメモリ使用量を、最大1.6%のエラーレートで推定できることを示しています。これは、DNNMemの平均エラーレート42.6%を上回ります。LLMに分散ファインチューニング方法を適用する場合、LLMemは、1億を超えるパラメータを持つLLMに対して、平均3.0%のエラーレートを達成します。