体系 › 推論・高速化
·無料·論文·15分で読めます

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

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

要点

  • OOMメッセージは最後に失敗した確保だけを報告し、犯人はすでに載っているほうにいます。
  • VRAMは重み・勾配・オプティマイザ状態・活性値・一時バッファ等の5つの箱に分けられます。
  • AdamのAdamでの学習の固定費は16バイト/パラメータで、混合精度にしても減りません。
  • 5つのうち自由に動かせるのは活性値だけで、打ち手は活性値を削るか固定費を分割・圧縮・退避するかの2系統です。
  • 推論のOOMはほぼKVキャッシュで、空きがあるのに乗らないときは断片化で、reservedとallocatedの差で見分けます。
対象textタスクinference

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バイト

Mstatic=Ψ⋅(bw+bg+bopt)M_{\text{static}} = \Psi \cdot (b_w + b_g + b_{\text{opt}})
(1)

Ψ\Psi はパラメータの個数、bwb_w は重み1個のバイト数、bgb_g は勾配1個ぶん、boptb_{\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×109×16=1127\times10^9 \times 16 = 112 GB。80GBのGPU 1枚には、活性値を1バイトも積む前に乗りません。「7Bだから小さい」という直感が最初に裏切られる場所です。

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

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

MKV=2⋅L⋅Hkv⋅dhead⋅b⋅TM_{\text{KV}} = 2 \cdot L \cdot H_{kv} \cdot d_{\text{head}} \cdot b \cdot T
(2)

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

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

FIG 1文脈長 n を伸ばすと、KVキャッシュは n に比例して伸び、素朴なattentionのスコア行列は n² で伸びる。線形と対数を切り替えると、長文脈でどちらが先に天井を打つかが見える

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

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

そして活性値は、5つのうち唯一こちらの都合で動かせる箱です。だからOOMの打ち手は、突き詰めると2系統しかありません。(A) 活性値を削るか、(B) 固定費そのものを分割・圧縮・退避するか。以下は効き目順ではなく、副作用の小さい順です。最初の2つで足りるなら、その先へは行かないほうがいい。

対処1: バッチを削り、勾配累積で埋め戻す

活性値はバッチサイズにほぼ比例するので、半分にすれば半分になります。小さいバッチは勾配が暴れますが、勾配累積で買い戻せます。バッチ4を4回ぶん流して勾配を足し込み、4回目にだけ更新する。数学的にはバッチ16の1ステップと同じ勾配です。

optimizer.zero_grad()
for i, batch in enumerate(loader):
    loss = model(**batch).loss / ACCUM        # 平均が欲しいので割る
    loss.backward()                            # 勾配は .grad に足し込まれる
    if (i + 1) % ACCUM == 0:
        optimizer.step(); optimizer.zero_grad()

/ ACCUM を忘れると実効学習率が累積回数ぶん跳ね上がり、しかも静かに壊れます。もう1つ、BatchNormの統計や、バッチ内の他サンプルを負例に使う対照学習では、小バッチ×累積は大バッチと別物になります。

対処2: 勾配チェックポイント — 覚える代わりに、作り直す

活性値を持つ理由は逆伝播で使うからでした。ならば捨ててしまい、必要になった時点で順伝播をやり直して作り直せばいい。nn 層をセグメントに区切って境界だけを保存し、セグメント数を n\sqrt{n} に取ると、活性値メモリはこう落ちます。

O(n)  ⟶  O(n)O(n) \;\longrightarrow\; O(\sqrt{n})
(3)

層数に比例して増えていたものが、層数の平方根までしか増えなくなる、という意味です。100層なら10ぶんの規模になる勘定。代償は追加の順伝播1回ぶんの計算で、順伝播と逆伝播の比はおおまか1:2なので3が4になる感覚ですが、実際の遅さは実装とモデル形状で変わります。必ず自分の環境で測ってください。

model.gradient_checkpointing_enable()   # Hugging Face
model.config.use_cache = False          # ← これを忘れると警告付きで効かない

2行目が地味な事故ポイントです。use_cache=True(生成用のKVキャッシュ)はチェックポイントと両立しません。警告は出ますが処理は続くので、「有効にしたのにメモリが減らない」という形で気づくことになります。PyTorchを直に使うなら torch.utils.checkpoint.checkpoint(fn, x, use_reentrant=False) を。旧来の再入方式はRNGの再現性やフックまわりに罠があります。

対処3: オプティマイザ状態を殴る

16バイトのうち8バイト、つまり半分がAdamのモーメント2本です。ここは3方向から削れます。

ここで、非常に多い誤解を潰しておきます。素のDDP(データ並列)はメモリを1バイトも節約しません。 各GPUがモデル・勾配・オプティマイザ状態の完全なコピーを持ち、勾配だけを平均するからです。8枚に増やしても1枚あたりの固定費は同じ。「枚数を増やしたのに乗らない」の答えがこれで、ZeROやFSDPが必要になる理由もこれです。

対処4: オフロード — CPUへ、そしてNVMeへ

それでも乗らないなら、いま使っていないものをGPUの外に出します。オプティマイザ状態は更新の瞬間しか触らないので、残りの時間はCPUのメインメモリに置いておける。これがZeRO-Offloadで、さらにNVMe SSDまで退避させるのがZeRO-Infinityです。

# DeepSpeed
zero_optimization:
  stage: 2
  offload_optimizer: { device: cpu, pin_memory: true }

効き目は絶大ですが、ボトルネックがPCIeバスの帯域に移ります。GPU内部のメモリ帯域とバスの帯域には桁の開きがあるので、ステップごとに何十GBも往復させれば当然遅くなる。「速度を犠牲にしてでも1枚で回したい」ときの手段であって、常用する設定ではありません。

対処5: そもそも、学習する重みを減らす

ここまでは全パラメータを学習する前提でした。外せば、固定費の構造ごと変わります。LoRAは元の重みを凍結し、低ランクの差分行列だけを学習します。凍結された重みには勾配もオプティマイザ状態も要らないので、16バイトのうち14バイトが消える。7Bの全パラメータ学習が112GBだったのに対し、LoRAなら重み14GBとわずかな追加分で済み、80GB 1枚に活性値ごと収まります。

FIG 2784個の重みを r×56 個で近似する。rを下げるとパラメータ数は桁で落ちるのに、誤差はゆるやかにしか増えない — LoRAが学習対象を1%未満まで絞れる理由が、この比の中にある

低ランク近似がなぜ成立するのかはLoRA論文解説で扱っています。注意点は、LoRAは活性値を減らさないことです。順伝播は全層を通るので、長い文脈で落ちているならLoRAにしても落ちます。そこはチェックポイントの仕事です。

推論のOOMは、ほぼKVキャッシュの管理問題

推論サーバが落ちるとき、犯人は9割方KVキャッシュです。触るパラメータは決まっています(vLLMの名前で挙げます)。

構造的な改善がPagedAttentionです。従来はリクエストごとに最大長ぶんの連続領域を先に確保していたため、実際には使われない領域が大量に死んでいました。これをOSの仮想メモリと同じくページ単位で管理し、書かれたぶんだけ割り当てる。同じVRAMで捌ける同時リクエスト数が大きく変わります。

断片化 — 空いているのに、乗らない

最後に、いちばん気味の悪いOOMです。「10GB空いているのに2GBの確保に失敗する」。

PyTorchのアロケータは、確保したブロックをOSに返さずキャッシュします。確保と解放を繰り返すと空きが細切れになり、合計は足りているのに連続した2GBが取れない状態になる。可変長のバッチで起きやすく、症状は「最初の数百ステップは通るのに、途中で突然落ちる」です。見分け方は2つの数の差を見るだけ。

torch.cuda.memory_allocated()   # テンソルが実際に使っている量
torch.cuda.memory_reserved()    # PyTorchがドライバから借りている量

開いていれば断片化で、対策は環境変数1行 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True です。ここで torch.cuda.empty_cache() に手を出したくなりますが、キャッシュを手放すだけなので次の確保が遅くなり、本当に足りていない場合は何も解決しません。ループの中に置くのは、まず間違いです。

なお nvidia-smi の数字を鵜呑みにしないこと。あれが表示するのは借りている量(≒ reserved)で、テンソルが使っている量ではありません。

現場ではこう使う

誰が、いつ触るか。 モデル基盤/MLOpsの担当者が、新しいモデルを手持ちのGPUに載る形へ見積もるとき。研究エンジニアが、走らせている実験が3時間目に落ちたとき。推論サーバの担当者が同時接続数の上限を決めるとき——3つ目はそのままインフラ費用の見積もりになります。

手順として固定しておくとよいもの。 ①始める前に Ψ×16\Psi \times 16 バイトを暗算して、そもそも乗るかを確かめる。②落ちたら torch.cuda.memory_summary() と max_memory_allocated() で内訳を出す。③固定費が原因なら 8bitオプティマイザ→ZeRO→LoRA、活性値が原因なら 勾配チェックポイント→バッチ削減+累積。④reserved と allocated が開いていたら expandable_segments:True。当てずっぽうにバッチを下げるのは時間の無駄です。

知らないと事故になる落とし穴。

設計レビューで問われる形。 「7Bをフルファインチューニングしたい、80GBは何枚要るか」。固定費112GB、DDPでは分割されない、ZeRO stage 3なら枚数で割れる——この3点を順に出せば筋が通ります。数字そのものより、内訳に分けて考えられるかが見られています。

まとめ

メモリは、速度と違って足りるか足りないかの二値です。だからこそ、走り出す前に暗算できます。

参考文献

  1. Training Deep Nets with Sublinear Memory Cost. arXiv:1604.06174論文ページ·PDF
  2. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. arXiv:1910.02054論文ページ·PDF
  3. 8-bit Optimizers via Block-wise Quantization. arXiv:2110.02861論文ページ·PDF
  4. Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180論文ページ·PDF

本記事は上記論文の本文にもとづいて執筆しています。数値・主張は原典を優先してください。

運営: TechAthletes · 編集方針 · 出典の付け方 · 訂正の記録

コメント

コメントにはログインが必要です