Claude CodeでCUDA out of memoryを切り分け、バッチサイズから順に直す

学習ジョブが CUDA out of memory で落ちたら、Claude Codeには「直して」と頼まず、調べる順番ごと渡すと早く収まります。順番は、バッチサイズの半減、勾配累積、混合精度、勾配チェックポイントです。1回に1つだけ変えて、変えるたびに短い試走で結果を見ます。

この記事は、GPUのメモリ不足が原因のOOMを扱います。コンテナがメモリ超過で exit code 137 になる場合は、原因がGPUではなくホストのRAMなので、exit code 137のOOMを切り分ける手順の領分です。

OOMはまず「どの時点で落ちたか」を渡す

CUDAのOOMは、モデルが大きすぎる場合だけでなく、バッチが大きい、検証時に勾配を切っていない、長い系列が混ざった、といった理由でも起きます。Claude Codeに最初に渡すのは次の3点です。

• 落ちたときのトレースバック全文(どの行の forward か backward か optimizer.step か)
• 落ちる直前のnvidia-smiの出力(搭載メモリと、他のプロセスが使っている量)
• 学習コマンドと、バッチサイズ・系列長・精度設定が書かれた設定ファイル

トレースバックの形は、例えば次のようなものです(説明用の例で、数値は環境によって変わります)。

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB.
  File "train.py", line 88, in training_step
    loss.backward()

落ちた行が forward なら、そのバッチの活性化がすでに載り切っていません。backward なら、勾配と活性化の保持が効いています。optimizer.step なら、オプティマイザの状態を含めた保持量が効いています。この3つで、次に効く手が変わります。

nvidia-smiの数字は「確保済み」を含む

PyTorchはキャッシュ付きのメモリアロケータを使うため、使われていない確保済みメモリもnvidia-smiでは使用中に見えます。テンソルが実際に占めている量は memory_allocated() と max_memory_allocated()、アロケータが確保している総量は memory_reserved() と max_memory_reserved() で見ます。

empty_cache() は未使用のキャッシュを他のGPUアプリに返すだけで、テンソルが占める分は解放しません。Claude Codeが「torch.cuda.empty_cache() を毎ステップ呼ぶ」という案を出したら、OOMの対策にはならないと返してください。

学習ループの末尾に、ステップごとのピーク値を1行出させておくと、以降の比較が数字でできます。

peak = torch.cuda.max_memory_allocated() / 2**30
print(f"step={step} peak_allocated={peak:.2f} GiB")
torch.cuda.reset_peak_memory_stats()

1変更ずつ試す4段階

順番の理由は、副作用の小さい手から使うことです。バッチを半分にするのは、精度にも収束にもほとんど手を付けずに済みます。チェックポイントは計算時間が増えるので最後に回します。

勾配累積は損失を割るのを忘れやすい

PyTorchのAMPの例では、累積の実効バッチは batch_per_iter * iters_to_accumulate(分散学習なら、さらにプロセス数を掛けた値)です。損失を iters_to_accumulate で割ってから backward を呼びます。

loss = loss_fn(output, target)
loss = loss / iters_to_accumulate
loss.backward()
if (i + 1) % iters_to_accumulate == 0:
    optimizer.step()
    optimizer.zero_grad()

割り忘れると、累積した勾配が iters_to_accumulate 倍の大きさのまま更新に使われます。OOMは消えたのに損失が発散した場合は、まずここを疑ってください。AMPと併用するときは、勾配のスケールを累積の途中で変えない決まりがあります。unscale_ は step の直前、update は実際に step を呼んだ回だけです。

混合精度とチェックポイントは副作用を確かめる

autocastはfloat16を既定とし、float16ではGradScalerによる勾配スケーリングが収束を助けます。精度が落ちるモデルでは、損失曲線が変わっていないかを数ステップ分比べます。

勾配チェックポイントは、順伝播の中間テンソルを保持せず、逆伝播のときに関数を呼び直して再計算します。メモリは減りますが、計算時間は増えます。PyTorchは use_reentrant を明示するよう求めており、False を使うことが推奨されています。Claude Codeに入れさせるときは、この引数が省略されていないかをレビューで見てください。

Claude Codeに渡す指示の書き方

指示には、順番と止める条件を書いておきます。CLAUDE.mdに次のように置くと、毎回説明せずに済みます。

## 学習ジョブがOOMで落ちたとき
- 原因の切り分けが先。トレースバックの落ちた行とnvidia-smiを読んでから案を出す
- 変更は1回に1つ。順番は「バッチ半減 → 勾配累積 → AMP → 勾配チェックポイント」
- 各変更の後に `--max-steps 20` の試走を1回だけ行い、peak_allocatedを表で報告する
- GPUを使っているプロセスをkillしない。止めたいときは理由を書いて確認を求める

最後の1行は、記事の後半で権限としても固めます。CLAUDE.mdの記述は、Claudeが試す内容を方向づけるだけで、実行を止める仕組みではないためです。

落ちた行から最初の一手を決める

トレースバックの落ちた位置で、Claude Codeに最初に試させる手が変わります。PyTorchのアロケータの説明に沿って整理すると、次の対応になります。

試走で比べる条件は揃えます。乱数シードと入力データの先頭バッチを固定し、1回の変更で変わったのがメモリと時間だけか、損失まで動いたかを見分けられるようにします。Claude Codeに「試走は毎回同じシードと同じ20ステップで行う」と指示しておくと、変更の効果が混ざりません。

表の最後の行は、同じ設定で最初は通り、途中のステップから落ちるときの典型です。アロケータの説明にも、バッチサイズが頻繁に変わるジョブで効くと書かれています。

許可ルールの効く範囲を決めておく

学習は長く走るので、nohup や timeout を付けて起動することがあります。Claude Codeは、timeout / time / nice / nohup / stdbuf を照合の前に取り除きます。Bash(python train.py --max-steps 20 *) は timeout 600 python train.py --max-steps 20 --bs 8 にも当たります。

一方、docker exec や npx のような実行環境のランナーは取り除かれません。コンテナの中で学習する構成では、Bash(docker exec trainer python train.py --max-steps 20 *) のように、ランナーと中のコマンドを並べて書きます。

承認画面で「Yes, and don't ask again」を選ぶと、Bashコマンドのルールがリポジトリルートの .claude/settings.local.json に保存されます。試走用の許可を個人の設定として残したいときはこの経路が早く、チームで共有するなら .claude/settings.json に書きます。どの階層でもdenyが優先されるので、プロジェクト側のdenyをユーザー設定のallowが上書きすることはありません。

学習コマンドだけを許可し、killは許可しない

nvidia-smiは、読み取り専用として無確認で通るコマンドの一覧(ls / cat / grep など)に挙がっていません。そのままだと毎回承認を求められるので、参照用の呼び出しだけ許可します。学習コマンドは、試走の引数まで含めて固定します。

{
  "permissions": {
    "allow": [
      "Bash(nvidia-smi)",
      "Bash(python train.py --max-steps 20 *)"
    ],
    "deny": [
      "Bash(kill *)",
      "Bash(pkill *)",
      "Bash(killall *)"
    ]
  }
}

ワイルドカードの * は、サブコマンドの後ろに置きます。Bash(python train.py --max-steps 20 *) なら、試走の指定がついたコマンドだけが対象になります。引数なしの本番学習は、この許可に当たらないため、承認画面が出ます。

denyは「いつもの書き方」しか止めない

注意が1点あります。Bashのdenyルールは、Claudeが書いたコマンド文字列に対して照合されます。/bin/kill 1234 や bash -c 'kill 1234' のように別の形で呼ばれると、Bash(kill *) では止まりません。権限のルールは、その呼び出しの形を止める安全弁であって、プログラムそのものを囲む境界ではない、という位置づけです。

確実に止めたいなら、ルールの上にフックかサンドボックスを重ねます。PreToolUseフックでコマンド全文を検査する方法と、コマンド文字列に依存しないファイル・ネットワークの隔離にはサンドボックスがあります。

また、複合コマンドは区切りごとに分解して判定されます。python train.py --max-steps 20 && kill 1234 のような連結は、後半がdenyに当たるので止まります。deny・askのルールは、複合コマンドの中のどのサブコマンドに当たっても効きます。

共有GPUでは「自分のジョブ」の見分けがつかない

共有サーバーでは、GPUを使っている他人のプロセスが同じnvidia-smiに出てきます。Claude Codeは、メモリを食っているプロセスを見つけると、それを落とす案を出しがちです。killのdenyを入れておけば、案が出ても実行はされず、その場で拒否されます。空きが足りない場合の判断は、人が持つようにしてください。

OOMが消えた後に確認すること

OOMが出なくなっても、直せたとは限りません。試走の結果は次の表でまとめさせると見比べやすくなります。

損失が元の設定と大きく食い違うときは、損失を割り忘れていないか、AMPで精度が落ちていないかを順に疑います。ピークが下がらないのに通ったときは、アロケータのフラグメンテーションかもしれません。PyTorchのアロケータは PYTORCH_ALLOC_CONF で調整できます。max_split_size_mb は、ほかの手を尽くしても「inactive split blocks」が大きいまま落ちる場合の最後の手段とされています。expandable_segments(実験的)は、バッチサイズが頻繁に変わるジョブで効く設定です。

ホスト側のメモリ不足なら原因が違うため、Pythonのメモリリークをtracemallocで特定する手順や、Claude Code自体のメモリ使用量が増える問題のほうが近い話題です。学習が通った後にモデルを書き出して検証する流れは、PyTorchモデルのONNX書き出しの一致検証に続きます。

まとめ

OOMの修正は、設定を1つ変えて試走して記録する作業の繰り返しです。Claude Codeには、手順と記録の形を固定して渡し、実行権限は学習の試走と参照に絞ります。killのdenyは入口の1枚で、止めたい対象が本当に危ないなら、フックかサンドボックスまで重ねるのが次の選択肢になります。