33. LLM訓練に必要なGPUメモリはどれくらい?計算方法を徹底解説!

スポンサーリンク

33. LLM訓練に必要なGPUメモリはどれくらい?計算方法を徹底解説!

こんにちは!今回は、LLM(大規模言語モデル)のトレーニングに必要なGPUメモリについて、詳しく解説していきます。

前回までLoRAという効率的な手法を学びましたが、「もし完全なファインチューニングをしたら?」「もっと大きなモデルを使ったら?」という疑問が湧いてきますよね。

  • 24GBのメモリを持つRTX 3090で十分なのか?
  • 40GBのA100が必要なのか?
  • それとも80GBのH100まで必要なのか?

この記事では、これらの疑問に答えるために、トレーニングに必要なGPUメモリを正確に計算する方法を学んでいきます!

なぜメモリが最初のスケーリング問題なのか

LLMのトレーニングは、計算コストだけでなく、メモリコストも非常に高いんです。

ここで興味深い事実があります:

推論(モデルを使うだけ)の場合:
– 13億パラメータのモデルを、一般的なGPUで実行できる

ファインチューニング(モデルを学習させる)の場合:
– 同じモデルが、最初のバッチを処理する前にメモリ不足でクラッシュする可能性がある

この非対称性が、LLMエンジニアリングの根本的な課題なんですね。ファインチューニングしたいモデルが、利用可能なハードウェアに収まらないことがよくあるんです。

ハードリミットとしてのGPUメモリ

計算(つまり処理時間)は、お金を払って時間を買えば解決できます。でも、GPUメモリはハードリミットです。限界に達したら、トレーニングは容赦なく停止します。

そのため、メモリ要件を理解することは:
– オプションではなく、必須
– 真剣なLLMトレーニングプロジェクトの最初のステップ

モデルを選ぶ前に、ハイパーパラメータを設定する前に、クラウドGPUをレンタルする前に、1つの質問に答える必要があります:

「これは収まるのか?」

この記事は、その質問に自信を持って答えるためのツールを提供します!

メモリ消費の4つのコンポーネント

GPUメモリは、トレーニング中に4つの要素によって消費されます。それぞれを理解することが、メモリ計画の鍵になるんですね。

1. モデルパラメータ(重み)

これはモデルの「脳」にあたる部分、つまり学習された数値パラメータです。

パラメータのメモリ使用量は、データ型(精度)によって変わります:

  • FP32(32ビット浮動小数点): パラメータあたり4バイト
  • FP16 / BF16(16ビット): パラメータあたり2バイト
  • INT8(8ビット整数): パラメータあたり1バイト

:FP32で10億パラメータのモデルの場合
– 1,000,000,000パラメータ × 4バイト = 4GB

シンプルですよね!でも、これはまだ始まりに過ぎません…

2. 勾配(バックプロパゲーションの作業メモリ)

トレーニング中、モデルは「勾配」を計算します。これは、各パラメータをどのように調整すべきかを示す値です。

勾配は、モデルパラメータと同じ精度で保存されるので、同じだけのメモリが必要になります。

重要なポイント
LoRAなどの手法でモデルを凍結する場合、ほとんどのレイヤーの勾配は不要になる
完全ファインチューニングの場合、パラメータのメモリコストが倍になる

:FP32で10億パラメータモデルの完全ファインチューニング
– 勾配 = さらに4GB

3. オプティマイザー状態(Adamの隠れたコスト)

最新のトレーニングでは、Adamオプティマイザーが広く使われています。Adamは、各パラメータに対して2つの追加値を追跡します:

  • モーメンタム(1次モーメント推定):過去の勾配の移動平均
  • 分散(2次モーメント推定):勾配のばらつき情報

これらは通常、数値安定性のためにFP32で保存されます(たとえトレーニング自体がFP16やBF16で行われていても)。

:10億パラメータモデルの場合
– 1,000,000,000パラメータ × 2状態 × 4バイト = 8GB

これが、トレーニング中の最大のメモリ消費者であることが多いんです!

補足:SGDのようなシンプルなオプティマイザーはモーメンタムのみを保存するので、これを半分に削減できます。でも、Adamの優れた収束性能は、通常このコストに見合う価値があるんですよね。

4. 活性化(順伝播フットプリント)

順伝播(Forward Pass)中、モデルは中間テンソル(活性化)を保存します。これらは、後のバックプロパゲーション(逆伝播)で勾配を計算するために必要なんです。

活性化メモリは、以下に依存します:

  • バッチサイズ:より多くのサンプルを一度に処理するほど、メモリ増加
  • シーケンス長:長いテキストほど、メモリ増加(しかも2乗で増える!)
  • モデルアーキテクチャ:レイヤー数、アテンションヘッド数、隠れ次元など

重要な特徴
– 活性化はバッチサイズに対して線形に成長
– シーケンス長に対して2次的に成長(アテンション機構のため)

これが、予想外の「CUDA out of memory」エラーの原因になることがよくあるんです!

まとめ:メモリ式

これらをまとめると、以下のメンタルモデルが得られます:

合計GPUメモリ ≈ モデルパラメータ + 勾配 + オプティマイザー状態 + 活性化

実例:LLaMA 3.1 1Bの完全ファインチューニング

前回QLoRAを使ってLLaMA 3.1 1Bをファインチューニングしましたよね。もし代わりに、完全精度(FP32)でモデル全体を完全ファインチューニングしたらどうなるか、計算してみましょう!

メモリの内訳

  • モデルパラメータ: 1B × 4バイト = 4GB
  • 勾配: 1B × 4バイト = 4GB
  • オプティマイザー状態: 1B × 8バイト = 8GB
  • 活性化: 約1GB(勾配チェックポイントなどの最適化を想定)

小計: 約17〜18GB

追加オーバーヘッド

さらに、以下のようなオーバーヘッドがあります:
– CUDAカーネル
– フレームワークバッファー
– メモリフラグメンテーション

これらで約1〜2GBが追加されます。

合計: 約18〜20GB

驚くべき事実

4GBで読み込める10億パラメータモデルが、トレーニングにはほぼ5倍のメモリが必要になるんです!

パラメータ自体は始まりに過ぎません。トレーニング中は、勾配、オプティマイザー状態、活性化のための個別のストレージが必要で、これらが一緒になってメモリ予算を支配するんですね。

これが、前回完全ファインチューニングの代わりにQLoRAを使った正確な理由です!

ビデオウォークスルー:LLaMA 3.1 8Bトレーニングメモリ

元のレッスンでは、80億パラメータモデルの4つの異なるファインチューニングシナリオを比較しています:

  1. 完全精度セットアップ:最大メモリ使用
  2. 混合精度トレーニング:メモリ削減
  3. LoRA:さらなるメモリ削減
  4. QLoRA:最小メモリ使用

精度、PEFT(パラメータ効率的ファインチューニング)メソッド、メモリ最適化が、要件を100GB以上から10GB未満に押し下げる様子がわかるんです!

次に来るもの:メモリボトルネックを軽減するツール

この時点で、パターンは明確ですよね:

メモリが真の制約であり、計算ではない

トレーニングフットプリントは急速に成長し、GPUは追いついていません。

そこで、コミュニティはGPUメモリを物理的限界をはるかに超えて引き伸ばすツールキットを構築しました。今週は、モデルが単一GPUの容量を超えたときにトレーニングを可能にするテクニックを学んでいきます!

メモリを直接削減する技術

活性化チェックポイント

計算とメモリをトレードオフします。順伝播中に保存する活性化を減らし、バックプロパゲーション中に再計算するんです。

「必要になったら計算し直せばいいじゃん!」というアプローチですね。時間はかかるけど、メモリは節約できます。

勾配累積

多くの小さなマイクロバッチを使用して大きなバッチサイズをシミュレートします。メモリを一定に保ちながら、有効バッチサイズを増やせるんです。

「一度に少しずつ処理して、最後にまとめて更新する」というイメージですね!

FSDPとZeRO

各GPUがモデルの一部(スライス)のみを保持するように、GPU全体でパラメータ、勾配、オプティマイザー状態を分割(シャード)します。

これにより、どこでも80GBカードを必要とせずに、70億パラメータ以上のモデルをファインチューニングできるんです!

速度に焦点を当てた技術

分散データ並列化(DDP)

トレーニングを高速化するために、GPU全体でバッチを分割します。

GPUあたりのメモリを削減するわけではありませんが、複数GPUで協力して作業する分散エコシステムの一部です。

極端なスケールのための技術

テンソルおよびパイプライン並列化

個々のレイヤーまたは実行ステージをGPU全体に分割します。

これは、単一のレイヤーでさえGPUメモリを超える場合に必要になる、最先端の技術です!

まとめ

今回の記事では、LLMトレーニングに必要なGPUメモリの計算方法について詳しく学びました。

重要なポイント

  1. メモリは4つの要素で構成される
    • モデルパラメータ(重み)
    • 勾配
    • オプティマイザー状態
    • 活性化
  2. トレーニングは推論の約5倍のメモリが必要
    • 4GBで読み込めるモデルが、トレーニングには18〜20GB必要になることも
  3. オプティマイザー状態が最大の消費者
    • Adamは各パラメータに2つの追加状態を保持
  4. 活性化はバッチサイズとシーケンス長に依存
    • シーケンス長に対しては2次的に増加
  5. メモリ制約を克服する多様なツールが存在
    • 活性化チェックポイント
    • 勾配累積
    • FSDP/ZeRO
    • 分散データ並列化

週の終わりまでに、研究所が数千億のパラメータを持つモデルをどのようにトレーニングするかを正確に理解し、自分でそれを行うツールを持つことになります。

次回は、すぐに適用できるテクニック(活性化チェックポイントと勾配累積)から始めていきましょう!

「これは収まるのか?」という質問に、自信を持って答えられるようになりましたね!

コメント

タイトルとURLをコピーしました