PyTorchにおけるGPUメモリの可視化と理解 – Hugging Face Blog サマリー

TL;DR

この投稿では、PyTorchでGPUメモリスナップショットを記録・可視化する方法を示し、メモリプロファイルの各部分が何を表すか（モデルパラメータ、オプティマイザ状態、活性化、勾配、オプティマイザ中間結果）を説明し、トレーニング中のピークメモリ使用量を見積もるための一般的な公式を提示します。

🔎 PyTorchビジュアライザー

PyTorchには、GPUメモリ使用量を記録・可視化する組み込みツールがあります。コードを実行する前に torch.cuda.memory._record_memory_history を呼び出し、実行後に torch.cuda.memory._dump_snapshot を呼び出すと、profile.pkl ファイルが得られ、これは https://pytorch.org/memory_viz で閲覧できます。ビジュアライザーは、メモリの割り当てと解放のタイムラインを表示します。

単純な nn.Linear 層の例では、グラフは次のように示されます：

  • モデルの作成では、重みとバイアス（float32）に約2 GBが割り当てられます。
  • 各入力テンソルは約200 MBを追加します。
  • 各フォワードパスの出力は約1 GBを追加します。
  • 前のステップの活性化は、バックプロパゲーションで必要なくなるまで保持され、その後メモリが解放されます。
  • 変数を再代入すると、以前参照していたテンソルが解放されます。

📊 トレーニング中のメモリ可視化

大規模言語モデル（Qwen/Qwen2.5-1.5B）とAdamWオプティマイザーを使用した現実的なトレーニングループでは、メモリプロファイルはトレーニングイテレーションごとに1つずつ、合計3つのスパイクを示します。

プロファイルは次のように分解できます：

  1. モデル初期化 – モデルパラメータ（青）がメモリを占有し、トレーニングが終了するまで割り当てられた状態が続きます。
  2. フォワードパス – 活性化（オレンジ）が層ごとに計算され、保存されます；これらは損失計算時にピークに達します。
  3. バックワードパス – 勾配（黄色）が計算されます；活性化が破棄され、オレンジゾーンが縮小します。
  4. オプティマイザーステップ – オプティマイザ状態（緑）は一度だけ初期化されます；オプティマイザーは勾配を使ってパラメータを更新し、一時的にオプティマイザ中間結果（赤）を保存します。更新後、勾配と中間結果は解放されます。

このパターンは各イテレーションで繰り返され、観測されたスパイクが生じます。

📐 メモリ要件の見積もり

ピークメモリ使用量はプロファイルにおける最高点であり、バッチサイズに応じてフォワードパスまたはオプティマイザーステップで発生する可能性があります。

両方のケースをカバーする一般的な表現は次のとおりです：

Total Memory = Model Memory + Optimizer State + max(Gradients + Optimizer Intermediates, Activations)

モデルパラメータ

モデルメモリ = N × P

  • N = パラメータ数
  • P = バイト単位の精度（例：float32の場合は4）

Qwen2.5-1.5Bの例（1.5 Bパラメータ、float32）では： モデルメモリ = 1.5 × 10⁹ × 4 バイト = 6 GB。

オプティマイザ状態

AdamWはパラメータごとに2つのモーメントを保存します： オプティマイザ状態サイズ = 2 × N × P

勾配

勾配メモリ = N × P（モデルパラメータと同じサイズ）。

オプティマイザ中間結果

オプティマイザ中間結果メモリ = N × P（モデルパラメータと同じサイズ）。

活性化

活性化メモリはバッチサイズ (B)、シーケンス長 (L)、およびトークンあたりの活性化数 (A) に依存します。 活性化メモリ = A × B × L × P

Aはフォワードフックを使って直接測定できますが、この投稿ではモデル全体に対する線形フィットから導き出されたヒューリスティックを提供しています：

A = 4.6894 × 10⁻⁴ × N + 1.8494 × 10⁶

このヒューリスティックを使えば、フルフォワードパスを実行せずに活性化メモリを推定できます。

総メモリ公式（組み合わせ）

コンポーネントの公式を代入すると：

Total Memory = N×P + 2×N×P + max(N×P, N×P, A×B×L×P)
             = 3×N×P + max(N×P, A×B×L×P)

max項は常にN×P以上となるため、式は次のように簡略化されます：

Total Memory = 3×N×P + A×B×L×P

活性化が優占する場合；それ以外の場合、オプティマイザーステップが優占し、総メモリは4×N×Pになります。

この投稿には、N、P、B、Lを入力して見積もりを得る小さなツールが含まれています。

🚀 次のステップ

メモリプロファイルを理解することで、バッチサイズを小さくする、グラディエントチェックポイントを使用する、または混合精度トレーニングに切り替えるなど、使用量を削減する方法を見つけることができます。TRLドキュメントの「メモリ使用量の削減」セクションの提案は、PyTorchベースのトレーニング全般に適用されます。

🤝 謝辞

ブログ記事へのフィードバックと提案に対して、Kashif Rasul氏に感謝します。

Sources