はじめに

本記事はhuggingfaceブログ「Visualize and understand GPU memory in PyTorch」の紹介記事です。

RuntimeError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 7.93 GiB total capacity; 6.00 GiB already allocated; 14.88 MiB free; 6.00 GiB reserved in total by PyTorch)

GPU メモリがいっぱいであることは簡単にわかりますが、その理由と修正方法を理解することはより難しい場合があります。このチュートリアルでは、トレーニング中にPyTorchでGPUメモリの使用状況を視覚化して理解する方法を段階的に説明します。また、メモリ要件を推定し、GPUメモリ使用量を最適化する方法も確認します。

PyTorchビジュアライザー

PyTorchは、GPUメモリ使用量を視覚化するための便利なツールを提供します。

import torch
from torch import nn

# Start recording memory snapshot history
torch.cuda.memory._record_memory_history(max_entries=100000)

model = nn.Linear(10_000, 50_000, device ="cuda")
for _ in range(3):
    inputs = torch.randn(5_000, 10_000, device="cuda")
    outputs = model(inputs)

# Dump memory snapshot history to a file and stop recording
torch.cuda.memory._dump_snapshot("profile.pkl")
torch.cuda.memory._record_memory_history(enabled=None)

このコードを実行すると、実行中のGPUメモリ使用履歴を含むprofile.pklファイルが生成されます。この履歴は、https://pytorch.org/memory_viz で視覚化できます。

profile.pklファイルをドラッグ&ドロップすると、次のようなグラフが表示されます。

このグラフを重要な部分に分解しましょう。

モデル作成：メモリはモデルのサイズに応じて4GB増加します。
10,000×50,000の重み+float32（4バイト）の50,000バイアス⟹（5×10^8)×4バイト=4GB
このメモリ（青）は実行中ずっと持続します。
入力テンソル作成（第1ループ）：入力テンソルサイズに合わせてメモリが200MB増加：
5,000×float32の10,000要素（4バイト）⟹（5×10^7)×4バイト=0.2GB
フォワードパス（第1ループ）：出力テンソルのメモリが2GB増加します。
10,000×float32（4バイト）の50,000要素⟹（5×10^8)×4バイト=2GB
入力テンソル作成（2番目のループ）：新しい入力テンソルのメモリが200MB増加します。この時点で、ステップ2の入力テンソルが解放されることを期待できます。それでも、そうではありません。モデルはその活性化を保持しているため、テンソルが変数inputsに割り当てられなくなった場合でも、モデルのフォワードパス計算によって参照されたままです。これらのテンソルはニューラルネットワークの背伝播プロセスに必要であるため、モデルはその活性化を保持します。違いを確認するには、torch.no_grad()を試してください。
フォワードパス（2番目のループ）：ステップ3のように計算された新しい出力テンソルのメモリは2GB増加します。
1番目のループアクティベーションを解放する：2番目のループのフォワードパスの後、最初のループからの入力テンソル（ステップ2）を解放することができます。最初の入力テンソルを保持するモデルの活性化は、2番目のループの入力によって上書きされます。2番目のループが完了すると、最初のテンソルは参照されなくなり、そのメモリが解放されます。
outputの更新：ステップ3の出力テンソルは、変数outputに再割り当てされます。前のテンソルはもはや参照されず、削除され、メモリが解放されます。
入力テンソル作成（3番目のループ）：ステップ4と同じです。
フォワードパス（3rdループ）：ステップ5と同じです。
リリース2ndループアクティベーション：ステップ4の入力テンソルが解放されます。
outputを再度更新：ステップ5の出力テンソルが変数outputに再割り当てされ、前のテンソルが解放されます。
コード実行の終了：すべてのメモリが解放されます。

トレーニング中の記憶の視覚化

上で紹介した例は簡略化されたのもです。実際のシナリオでは、単一の線形層ではなく、複雑なモデルをトレーニングすることがよくあります。さらに、前の例にはトレーニングプロセスが含まれていませんでした。ここでは、実際の大型言語モデル (LLM) の完全なトレーニングループ中に GPU メモリがどのように動作するかを調べます。

import torch
from transformers import AutoModelForCausalLM

# Start recording memory snapshot history
torch.cuda.memory._record_memory_history(max_entries=100000)

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B").to("cuda")
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

for _ in range(3):
    inputs = torch.randint(0, 100, (16, 256), device="cuda")  # Dummy input
    loss = torch.mean(model(inputs).logits)  # Dummy loss
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

# Dump memory snapshot history to a file and stop recording
torch.cuda.memory._dump_snapshot("profile.pkl")
torch.cuda.memory._record_memory_history(enabled=None)

ヒント：プロファイリングするときは、ステップの数を制限してください。すべてのGPUメモリイベントが記録され、ファイルが非常に大きくなる可能性があります。たとえば、上記のコードは8 MBのファイルを生成します。

この例のメモリプロファイルは次のとおりです。

このグラフは前の例よりも複雑ですが、それでも段階的に分解できます。トレーニングループの反復に対応する3つのスパイクに注目してください。グラフを単純化して、解釈しやすくしましょう。

モデルの初期化（model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B").to("cuda")
最初のステップは、モデルをGPUにロードすることです。モデルパラメータ（青）はメモリを占有し、トレーニングが終了するまでそこに残ります。
フォワードパス（model(inputs)
フォワードパス中に、活性化（各レイヤーの中間出力）が計算され、バックプロパグのためにメモリに保存されます。オレンジ色で表されるこれらの活性化は、最後の層まで層ごとに成長します。損失はオレンジゾーンのピーク時に計算されます。
バックワードパス（loss.backward()
グラデーション（黄色）は、このフェーズ中に計算され、保存されます。同時に、活性化はもはや必要とされなくなったため破棄され、オレンジゾーンが縮小します。黄色のゾーンは、勾配計算のメモリ使用量を表します。
オプティマイザーステップ（optimizer.step()
グラデーションは、モデルのパラメータを更新するために使用されます。最初に、オプティマイザ自体が初期化されます（グリーンゾーン）。この初期化は一度だけ行われます。その後、オプティマイザはグラデーションを使用してモデルのパラメータを更新します。パラメータを更新するために、オプティマイザは中間値（レッドゾーン）を一時的に保存します。更新後、グラデーション（黄色）と中間オプティマイザー値（赤）の両方が破棄され、メモリが解放されます。

この時点で、1つのトレーニング反復が完了します。プロセスは残りの反復を繰り返し、グラフに3つのメモリスパイクが表示されます。

このようなトレーニングプロファイルは通常、一貫したパターンに従うため、特定のモデルとトレーニングループのGPUメモリ要件を推定するのに役立ちます。

メモリ要件の推定

上記のセクションから、GPUメモリ要件の見積もりは簡単に見えます。必要な総メモリは、フォワードパス中に発生するメモリプロファイルの最高ピークに対応する必要があります。その場合、メモリ要件は（青+緑+オレンジ）です：
モデルパラメータ+オプティマイザー状態+アクティベーション

しかしそんなに単純ではなく、実際には罠があります。プロファイルは、トレーニングのセットアップによって異なる場合があります。たとえば、バッチサイズを16から2に減らすと、様子が変わります。

- inputs = torch.randint(0, 100, (16, 256), device="cuda")  # Dummy input
+ inputs = torch.randint(0, 100, (2, 256), device="cuda")  # Dummy input

現在、最高のピークは、フォワードパスではなく、オプティマイザーステップ中に発生します。この場合、メモリ要件は（青+緑+黄色+赤）になります。

モデルパラメータ+オプティマイザー状態+グラデーション+オプティマイザー中間体

メモリ推定を一般化するには、フォワードパスまたはオプティマイザーステップ中に発生するかどうかに関係なく、すべての可能なピークを考慮する必要があります。

モデルパラメータ+オプティマイザー状態+最大（グラデーション+オプティマイザー中間体、アクティベーション）

方程式がわかったので、各コンポーネントを推定する方法を見てみましょう。

モデルパラメータ

モデルパラメータは推定するのが最も簡単です。
モデルメモリ = N × P

Nはパラメータの数です。
Pは精度（precision）です（バイト単位、例えば、float32の場合は4）。

たとえば、15億のパラメータと4バイトの精度を持つモデルには、次のようになります：

モデルメモリ = 1.5 × 10^9 × 4バイト = 6GB

オプティマイザー状態

オプティマイザの状態に必要なメモリは、オプティマイザの種類とモデルパラメータによって異なります。たとえば、AdamWオプティマイザは、パラメータごとに2つのモーメント（最初と2番目）を保存します。これにより、オプティマイザの状態サイズが次のようになります。

オプティマイザーの状態サイズ = 2 × N × P

アクティベーション

活性化に必要なメモリは、フォワードパス中に計算されたすべての中間値が含まれているため、推定が困難です。起動メモリを計算するには、フォワードフックを使用して出力のサイズを測定できます。

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B").to("cuda")

activation_sizes = []

def forward_hook(module, input, output):
    """
    Hook to calculate activation size for each module.
    """
    if isinstance(output, torch.Tensor):
        activation_sizes.append(output.numel() * output.element_size())
    elif isinstance(output, (tuple, list)):
        for tensor in output:
            if isinstance(tensor, torch.Tensor):
                activation_sizes.append(tensor.numel() * tensor.element_size())

# Register hooks for each submodule
hooks = []
for submodule in model.modules():
    hooks.append(submodule.register_forward_hook(forward_hook))

# Perform a forward pass with a dummy input
dummy_input = torch.zeros((1, 1), dtype=torch.int64, device="cuda")
model.eval()  # No gradients needed for memory measurement
with torch.no_grad():
    model(dummy_input)

# Clean up hooks
for hook in hooks:
    hook.remove()

print(sum(activation_sizes))  # Output: 5065216

Qwen2.5-1.5Bモデルの場合、入力トークンごとに5,065,216回のアクティベーションが得られます。入力テンソルの合計起動メモリを推定するには、次のようにします。

アクティベーションメモリ = A × B × L × P

Aはトークンあたりのアクティベーション数です。
Bはバッチサイズです。
Lはシーケンスの長さです。

しかし、この方法を直接使用することは必ずしも実用的ではありません。理想的には、モデルを実行せずにアクティベーションメモリを推定するヒューリスティックを望みます。さらに、より大きなモデルにはより多くのアクティベーションがあることが直感的にわかります。これは質問につながります：モデルパラメータの数とアクティベーションの数の間に関連性はありますか？

トークンあたりのアクティベーション数はモデルアーキテクチャに依存するため、直接的ではありません。しかし、LLMは同様の構造を持つ傾向があります。さまざまなモデルを分析することで、パラメータの数と活性化の数の間の大まかな線形関係を観察します。

この線形関係により、ヒューリスティックを使用して活性化を推定できます。

A = 4.6894 × 10^4 × N + 1.8494 × 10^6

これは近似値ですが、各モデルに対して複雑な計算を実行することなく、アクティベーションメモリを推定する実用的な方法を提供します。

Gradients

Gradientsは推定しやすく、必要なメモリはモデルパラメータと同じです。

Gradientsメモリ= N × P

オプティマイザー中間体

モデルパラメータを更新するとき、オプティマイザは中間値を保存します。これらの値に必要なメモリは、モデルパラメータと同じです。

オプティマイザー中間メモリ=N×P

要約すると、モデルをトレーニングするために必要な総メモリは次のとおりです。

トータルメモリ=モデルメモリ+オプティマイザー状態+max（グラデーション、オプティマイザー中間体、アクティベーション）

以下のコンポーネントで：

モデルメモリ： N×P
オプティマイザーの状態：2×N×P
Gradients：N×P
オプティマイザー中間体：N×P
アクティベーション：A×B×L×P
ヒューリスティックを使用して推定

A=4.6894×10^4×N+1.8494×10^6

この計算を簡単にするために、私はあなたのために小さなツールを作成しました：

次のステップ

メモリの使用を理解するためのあなたの最初の動機は、おそらくある日、あなたがメモリを使い果たしたという事実によって推進されました。このブログは、それを修正するための直接的な解決策を提供しましたか？たぶんそうではない。しかし、メモリ使用量がどのように機能し、それをプロファイリングする方法をよりよく理解したので、それを減らす方法を見つける準備が整います。

まとめ

GPUメモリがどのようなパラメータによって支配されているかを確認し、可視化することでメモリオーバーに対する理解を深めることができます。また、Gradioで提供されているシンプルな計算ツールを使うことでメモリの適切な使用（過剰でも過小でもない）を行うことができます。
生成AIを使う際は、どのようなマシンがあれば適切なのかを判断するための材料として考え方を身につけておきましょう。