GPUメモリ不足サバイバル — OOMの原因と対処の全パターン

`CUDA out of memory` は、最後に失敗した確保だけを報告する不親切なエラーです。犯人はたいてい、すでに載っている側にいます。何がVRAMを食うのかを5つの箱に分けて数え、パラメータ1個あたり16バイトという学習の固定費を出し、そこから勾配チェックポイント・オプティマイザ圧縮・オフロード・KVキャッシュ管理・断片化まで、打ち手を副作用の小さい順に。

Training Deep Nets with Sublinear Memory Cost

荷台の大きさは、走り出す前に決まっている

引っ越しトラックの荷台は、契約した時点で広さが決まっていて、走り出してから伸びることはありません。入り切らなければ、積み方を変えるか、荷物を減らすか、車をもう1台出すか。打てる手はこの3つだけです。

GPUのメモリ（VRAM）はこの荷台です。そして学習や推論を回していれば、いつか必ずこれを見ます。

torch.cuda.OutOfMemoryError: CUDA out of memory.
Tried to allocate 2.00 GiB (GPU 0; 79.15 GiB total capacity;
74.23 GiB already allocated; 1.02 GiB free;
77.31 GiB reserved in total by PyTorch)

親切なようで不親切なメッセージです。報告しているのは最後に失敗した2GBの確保だけで、犯人はまず間違いなく、すでに載っている74GBのほうにいます。2GBを削る努力は、たいてい徒労に終わる。

なので手順は1つです。何が何GB食っているかの内訳を先に出す。 内訳さえ出れば、打ち手はほとんど自動的に決まります。

VRAMを食っているのは、5つの箱

学習中にGPUに載っているものは、5つに分けられます。

1. 重み — パラメータ数 × 1個あたりのバイト数。最初から最後まで居座る
2. 勾配 — 重みとまったく同じ形の配列。逆伝播で埋まる
3. オプティマイザ状態 — Adamなら1次・2次モーメントの2本。やはり重みと同じ形が2つ
4. 活性値 — 順伝播の途中結果。逆伝播で使うので、それが終わるまで捨てられない
5. 一時バッファ・断片化・CUDAコンテキスト — 行列積のワークスペース、通信バッファ、CUDAランタイム自体が持っていく数百MB

決定的なのは、1〜3がモデルとオプティマイザを決めた瞬間に固定されることです。バッチサイズを1に落としても1バイトも減りません。動くのは4だけ。OOMの話がいつも「バッチを下げろ」で始まり、「バッチ1にしてもまだ落ちる」で行き詰まるのは、この非対称性のせいです。

学習の固定費は、パラメータ1個あたり16バイト

$\Psi$ はパラメータの個数、$b_w$ は重み1個のバイト数、$b_g$ は勾配1個ぶん、$b_{\text{opt}}$ はオプティマイザ状態1個ぶんです。言い換えれば「パラメータの数に、1個を養うのに要るバイト数を掛けただけ」。

fp32でAdamなら 4+4+8 = 16バイト/パラメータ。では混合精度にすれば半分かというと、なりません。bf16の重み2 + bf16の勾配2 + fp32のマスター重み4 + モーメント2本で8。やはり16バイトです（ZeRO論文と同じ勘定）。

ここは踏みやすい穴です。半精度にすればメモリも半分と思って切り替え、固定費が1バイトも動かずに驚く。混合精度が減らすのは4番の活性値のほうで、そもそもの主目的はTensor Coreに載せて計算を速くすることでした（混合精度学習）。

70億パラメータなら $7\times10^9 \times 16 = 112$ GB。80GBのGPU 1枚には、活性値を1バイトも積む前に乗りません。「7Bだから小さい」という直感が最初に裏切られる場所です。

推論では、重みが固定費でKVキャッシュが変動費

推論に切り替えると勾配もオプティマイザ状態も消えます。7Bをbf16で載せれば重みは14GB、80GBのうち66GBが空く。その66GBに入るのは、ほぼ全部がKVキャッシュです。

先頭の2はKとVの2本ぶん、$L$ は層数、$H_{kv}$ はKVヘッドの数、$d_{\text{head}}$ は1ヘッドの次元、$b$ は1要素のバイト数、$T$ はサーバ上に同時に存在するトークンの総数です。要するに「層ごと・ヘッドごとに、トークン1個につきKとVを1組ずつ取っておく」。

32層・KVヘッド32・ヘッド次元128・bf16なら、トークン1個あたり $2\times32\times32\times128\times2 = 524{,}288$ バイト、約512KBです。4096トークンの会話1本で2GB。66GBの余裕は33本ぶんで、そこから先の同時リクエストは受けられません。GQA（KVヘッドだけを減らす設計）が標準になったのは、この式の $H_{kv}$ を直接叩けるからです（KVキャッシュを1から理解する）。

動かせるのは、活性値だけ

活性値は「バッチ × 文長 × 隠れ次元」に比例する項と、attention内部の「バッチ × ヘッド数 × 文長²」に比例する項の和として溜まります。後者をメモリから消したのがFlashAttentionで、長文脈の学習が現実的になった理由の半分はここです。