はじめに
RTX 3090（VRAM 24GB）1枚、WSL2上のUbuntuという環境で、7Bパラメータ規模のTransformer系モデルを事前学習させるプロジェクトに取り組んでいます。
クラウドGPUではなく手元のハードウェアだけで完結させたい。その動機は単純で、研究の全工程を自分の手の届く範囲に置いておきたかったからです。
結論から言えば、最終的にはVRAM 22.6GBで安定動作し、約1770 tok/sのスループットに到達しました。しかし、そこに至るまでに踏んだ落とし穴は8つ。optimizer選定だけで6方式を試し、すべて失敗しています。
この記事では、各段階で何が起き、なぜ失敗し、どう判断したかを記録します。7Bモデルを「作った話」ではなく、7Bモデルを個人GPU環境で「安定して学習させるまでの工学的失敗記録」です。
同じ環境で大規模モデルの学習を考えている方の参考になれば幸いです。
環境
- GPU: NVIDIA RTX 3090 (VRAM 24GB)
- OS: WSL2 Ubuntu 24 on Windows
- フレームワーク: PyTorch 2.6 + CUDA 12.4
- 精度: BF16
- モデル: 独自設計の7B Transformer系モデル（詳細非公開）
- データ: 事前処理済み学習データを使用
落とし穴 1: 8bit Optimizerは「ギリギリ収まる」計算が裏切る
最初に試したのは8bit AdamWです。
通常のAdamW (fp32) では、モデルのパラメータ数 × 2（一次・二次モーメント）× 4バイトのoptimizer stateが必要になります。7Bパラメータだと、これだけで数十GB。当然24GBには収まりません。
8bit量子化すればstateは1/4になるので、理論上は収まるはず——と思って実行したところ、学習開始直後に CUDA driver error: device not ready で停止しました。
原因は、optimizer stepの瞬間に一時的なfp32バッファが生成され、それがモデル本体＋activationsと合わせてVRAMの物理限界を超えたことです。
「理論上の見積もり」と「実行時のピークVRAM」は別物です。PyTorchのmemory allocatorは断片化を起こしますし、CUDA内部のワークスペースも見えないメモリを消費します。
教訓: 「理論上収まる」は信用しない。実行時のピークVRAMで判断する。
落とし穴 2: CPU Offloadは転送速度が壁になる
次に試したのがCPU Offload方式です。optimizer stateをCPU RAM上に置き、毎step GPU↔CPU間で転送する方法です。
動きました。しかし、スループットは71 tok/s。PCIe 4.0 x16の帯域（理論32GB/s）があっても、毎stepで数GBのパラメータを往復させるオーバーヘッドは無視できません。
しかも、WSL2環境ではCPU側のメモリ管理にも癖があり、長時間実行するとRAM側でOOMが発生しました。
教訓: CPU Offloadは「動く」と「使える」の間に大きな溝がある。RTX 3090クラスの計算速度に対して、PCIeの帯域は圧倒的に不足する。
落とし穴 3: DeepSpeed ZeROはWSL2と相性が悪い
分散学習フレームワークDeepSpeedのZeRO-2/3を単一GPU環境で試しました。
ZeRO-3はVM（仮想メモリ）使用量が296GBまで膨張してOOM。WSL2はWindowsのページファイルを仮想メモリとして使うため、この膨張が致命的です。
ZeRO-2でもoptimizer stateの生成時にRAMが95%に達し、やはりOOM。
単一GPU + WSL2という構成では、ZeRO系は「分散しないのに分散フレームワークのオーバーヘッドだけ払う」結果になりました。
教訓: DeepSpeedはマルチGPU前提の設計。単一GPU + WSL2では、VMとRAMの管理が想定外の挙動を示す。
落とし穴 4: NVMe直接読み書きは桁違いに遅い
optimizer stateをmmapでNVMeに直接配置するStreaming方式も試しました。
スループットは14 tok/s。RTX 3090の計算能力を1/100も使えていない状態です。NVMe SSDはシーケンシャルリードで数GB/s出ますが、optimizerの更新パターンはランダムアクセスに近く、実効帯域は劇的に低下します。
教訓: NVMe I/Oはシーケンシャルでは速いが、optimizer stateのようなランダムアクセスパターンでは実用的な速度が出ない。
落とし穴 5: 既存のAdafactorもfp32コピーで詰まる
ここで発想を転換しました。「巨大なoptimizer stateをどう運ぶか」ではなく、「そもそもstateが小さいoptimizerを使えばよい」。
Adafactorは二次モーメントを行ベクトルと列ベクトルに因子分解して保持します。大規模な重み行列では、この因子分解によるstate削減効果は非常に大きく、理論上はほとんどメモリを消費しません。
これは正しい方向でした。しかし、PyTorchのAdafactor実装をそのまま使ったところ、内部で p.data.float() によるfp32コピーが発生し、内部処理で大きな一時メモリ確保が発生し、VRAM制約により device not ready になりました。
教訓: 「optimizer stateを減らす」という方向性は正しい。ただし、既存実装の内部挙動まで確認しないと、別の場所でメモリが爆発する。
落とし穴 6: 解決策はoptimizer stateを「運ぶ」のではなく「無くす」こと
最終的に、fp32一時コピーを回避するようoptimizer実装を環境に合わせて調整しました。BF16のまま直接更新し、メモリのピーク使用量を最小限に抑える構成です。
結果、optimizer stateをGB単位からMB単位まで削減できました。モデル本体と合わせてVRAM 22.6GBで安定動作しています。
6方式の失敗を振り返ると、問題の構造が見えます。
| 方式 | 失敗の本質 |
|---|---|
| 8bit AdamW | stateは小さくなったが、まだ大きすぎた |
| CPU Offload | stateを移動 → 転送がボトルネック |
| DeepSpeed ZeRO | stateを分散 → 分散先がない |
| NVMe Streaming | stateを退避 → I/Oが遅すぎる |
| 標準Adafactor | stateは小さいが、実装が別のメモリを消費 |
| 環境適合版 | stateそのものを削減 + 実装レベルでメモリ制御 |
すべての失敗に共通するのは、「巨大なstateが存在する前提で、それをどう扱うか」を最適化していたことです。正解は前提そのものを変えることでした。
教訓: 制約が厳しい環境では、「既存の仕組みを工夫して使う」より「制約に合った仕組みを選ぶ（または作る）」方が効果的。
落とし穴 7: DataLoaderの「見えないボトルネック」
optimizer問題を解決してGPU計算が回り始めると、次のボトルネックが現れました。
スループットは約80 tok/s。GPUベンチマークではforward単体で1800 tok/s以上出ているのに、20倍以上遅い。
時間内訳を計測したところ、学習ループの88%がDataLoader待ちでした。
原因は torch.tensor() の呼び出しです。学習データをNumPy配列から読み込む際に、サンプルごとに torch.tensor() を呼んでいました。この関数はPython→C++のブリッジコストが高く、小さな配列に対して毎回呼ぶと積み重なります。
修正は、shard（数千サンプルの塊）単位で torch.from_numpy() を1回だけ呼び、あとはインデックスで取り出す方式に変更。DataLoaderの占有率は88% → 2%に低下し、スループットは約1770 tok/sに到達しました。
教訓: PyTorchのtensor生成は意外に高コスト。大量のサンプルを扱う場合は、まとめて変換してからインデックスアクセスする。GPUがボトルネックだと思い込まず、必ず計測する。
落とし穴 8: torch.compileと動的モデルの構造的非互換
最後にtorch.compileを試しました。mode="reduce-overhead" はCUDA Graphを活用してカーネル起動のオーバーヘッドを削減するモードで、5〜15%の高速化が期待できます。
起動直後は正常に動作。しかし、約20分後に CUDA driver error: device not ready で停止しました。
原因は、CUDA Graphが「固定された実行グラフ」を前提とするのに対し、動的な実行経路を持つモデルでは相性問題が発生したことです。実行時の形状や処理経路が変化するモデルでは、CUDA Graphのキャッシュ効率が低下し、最終的にGPUドライバがリセットされる場合があります。
さらに、gradient checkpointingとの組み合わせも問題を悪化させます。checkpointingはbackward時にforwardを再実行しますが、この再実行パターンもCUDA Graphとは相性が悪い構成です。
torch.compileを外した状態で約1770 tok/sが出ているため、compileによる追加の高速化は不要と判断しました。
なお、torch.compileを有効にした状態で保存されたチェックポイントには、state_dictのキーに _orig_mod. というプレフィックスが付加されます。compile無効版でロードする際はこのプレフィックスを除去する処理が必要です。
教訓: torch.compileの reduce-overhead モードは、実行グラフが動的に変わるモデルには使えない。これはバグではなく、CUDA Graphの設計上の制約との構造的な非互換。
プロファイルに関する補足: sync無しの計測は嘘をつく
落とし穴ではありませんが、デバッグ中に遭遇した重要な知見を記録しておきます。
学習ループの時間内訳を time.perf_counter() で計測したところ、以下のような結果が出ました。
Forward=24% Loss=55% Backward=9% Optimizer=12%
Loss計算が55%？Backward（逆伝播）が9%？これは直感に反します。
原因は torch.cuda.synchronize() を挟んでいなかったことです。PyTorchのGPU演算は非同期で実行されるため、time.perf_counter() で区間を測っても、前の区間のGPUカーネルがまだ実行中の場合があります。結果として、Backwardの計算時間がLossの計測区間に「漏れ出す」現象が起きます。
synchronize() を挟んで正確に計測すると、実際の内訳は以下でした。
Forward=25% Loss=1% Backward=75%
Backward=75%でForward=25%、つまりBWD/FWD比=3.0x。これはgradient checkpointingの理論値と完全に一致します（backward中にforwardが再実行されるため、約3倍）。
ただし、synchronize() 自体にオーバーヘッドがあるため、常時有効にするとスループットが低下します。最初の数十iterationだけ有効にして計測し、以降は無効にする実装が実用的です。
教訓: GPUプロファイルでは torch.cuda.synchronize() を挟まないと、非同期実行により計測区間がずれる。ただし常時syncは性能を落とすので、計測時のみ有効にする。
最終構成
| 項目 | 選択 |
|---|---|
| Optimizer | Adafactor系（環境に合わせた低メモリ実装） |
| Gradient Checkpointing | ON（VRAM節約のため必須） |
| torch.compile | OFF（動的モデルとの非互換） |
| DataLoader | shard単位tensor変換 + streaming |
| バッチサイズ | 固定値（adaptive制御はCUDAコンテキスト破壊の原因） |
| プロファイル | 最初の数十iterのみsync付き計測 |
VRAM 22.6GB安定、約1770 tok/s。
まとめ: 8つの教訓
- 理論上のVRAM見積もりは信用しない。 ピークVRAMは実行してみないと分からない。
- 「stateを運ぶ」より「stateを無くす」。 CPU offload、分散、退避はすべて帯域制約に衝突する。
- 既存実装の内部挙動まで確認する。 Adafactorでもfp32コピーが隠れていた。
- DeepSpeedはマルチGPU前提。 単一GPU + WSL2では期待通りに動かない。
- DataLoaderのtensor生成は意外に重い。 まとめて変換する。
- GPUがボトルネックだと思い込まない。 計測すれば本当の原因が見える。
- torch.compileは動的実行グラフに使えない。 CUDA Graphの構造的制約。
- adaptive制御（バッチサイズ、メモリ管理等）は避ける。 CUDAコンテキストの安定性を優先する。
ひとつひとつは「言われてみれば当然」かもしれません。しかし、実際にコードを動かしてエラーログと向き合うまでは、どれも「自分には関係ない」と思っていたものばかりです。
同じ道を歩む方が、同じ穴に落ちずに済むことを願っています。