関係マップ：学習時の GPU メモリ管理（再計算とスワップ）

この地図は、DNN / LLM の学習中に GPU メモリを節約する技術（テンソルの再計算とスワップ、およびその組み合わせ）の系譜を整理する。 出発点はノート化した4本（Capuchin、HOME、XEngine、ProTrain）。 情報源は各論文のアブストラクトと公開論文（ノート化済みの4本は各ノートの # AI解説 に準じる）に基づく。本文未取得の論文は憶測で補わず、アブストラクトの範囲で記す。

「チェックポイント」の二つの意味

この地図の「チェックポイント」は、深層学習のチェックポイントの「チェックポイント」とは別物である。

同じ語だが、前者は時間軸方向の永続化、後者は時間とメモリのトレードオフである。 本ドキュメント群では前者を扱う地図が既にあるため、この地図は後者だけを扱う。

前提：学習時のメモリはどこで詰まるか

学習のメモリ消費は大きく2種類に分かれる。

これを GPU に収める基本手段が2つある。

どちらを・どのテンソルに・いつ使うかの決め方が各研究の個性であり、この地図は次の2軸で整理する。

軸 選択肢
手段 再計算のみ ／ スワップのみ ／ ハイブリッド（両方）
決め方 静的計画（実行前に最適化。ILP 等の厳密解法か、GA・PSO 等のヒューリスティックか）／ 動的ランタイム（実行時に判断）

学習時 GPU メモリ管理の系譜

補足図（AI生成）：2016年の2つの源流（再計算＝Chen ら、スワップ＝vDNN）から、ハイブリッド化、最適計画、動的ランタイム、LLM 時代の統合へ至る系譜。太枠はノート化した4本。

1. 源流（2016）：再計算とスワップが別々に生まれる

Chen らの sublinear memory cost2 が再計算の系譜の起点である。 N 層のネットワークを √N 個の区間に分け、区間境界の活性だけを保持して残りを backward 時に再計算することで、O(√N) のメモリで学習できることを示した。 自動微分の分野ではさらに古く、Griewank の Revolve（逆モード自動微分のチェックポイント法）に遡る3。

同じ2016年、vDNN4 がスワップの系譜を開いた。 畳み込みネットワークの層ごとの活性を CPU メモリへ退避し、GPU メモリを仮想化する。

以後の研究はすべて、この2つの手段の組み合わせ方と決め方の探求として読める。

2. ハイブリッド化とランタイム方式（2018〜2020）

SuperNeurons（PPoPP 2018）5 が、生存期間解析、統合テンソルプール（スワップ）、コストを考慮した再計算の3技法を組み合わせた最初期の動的ランタイムである。 ピークメモリを「個々の層の最大メモリ使用量」（層単位実行の理論下限）まで下げると主張し、12GB の K40c で 10^4 層の ResNet を学習した。

SwapAdvisor（ASPLOS 2020）6 はスワップのみを突き詰めた静的計画で、データフローグラフから演算子スケジュール、メモリ割り当て、スワップ決定を遺伝的アルゴリズムで同時最適化し、GPU メモリの12倍のモデルを学習できる。

同じ ASPLOS 2020 の Capuchin が、この地図の中心にあるハイブリッドの代表である。 テンソル粒度でスワップ（eviction/prefetch）と再計算を組み合わせ、判断の根拠を「学習イテレーション中のテンソルアクセスパターンは規則的」という観察に置く。 1バッチ走らせてアクセスパターンを観測し、以後のイテレーションでいつ・どちらの手段を使うかを決める動的ランタイム方式で、スワップのみ（vDNN）比で最大 286%、再計算のみ（gradient checkpointing）比で最大 55% 速い。 単一手段に対するハイブリッドの優位を数字で示した点が、後続（HOME、DELTA、STR など）の共通の出発点になっている。

3. 最適計画の系譜（2020〜2023）：Checkmate とその子孫

ハイブリッド化と並行して、「再計算の計画を厳密に最適化する」系譜が現れた。

Checkmate（MLSys 2020）7 は、時間とメモリのトレードオフを「テンソル再材料化最適化問題」として定式化し、MILP（混合整数線形計画）ソルバで最適解を得る。 動機は、Chen らの分割法が「層ごとのコストが一様」「線形チェーン構造」という仮定に依存し、実在のネットワークで使いにくいことだった。

Checkmate の直系が XEngine（ACM TACO 2022）である。 Checkmate が単一デバイス前提であるのに対し、XEngine は全変数・全制約にデバイス次元を加えた MIQP として定式化し、チェックポイントと再計算に加えて演算子ごとのデバイス配置（CPU で計算するか GPU で計算するか）まで同時に決める。 CPU と GPU の性能が拮抗する環境では、単一デバイスの最速 Checkmate スケジュールより最大 22.5% 速く、GPU 単体でメモリが足りないネットワークにも実行可能スケジュールを与える。 詳細はノートを参照。

同じ最適計画の系譜に、エッジデバイス向けに再計算とページングを MILP で組み合わせてエネルギー効率を狙う POET（Patil ら 2022）8、Checkmate の O(n²) 個のブール変数を「保持区間」に基づく O(n) 個の整数変数に置き換え、制約プログラミング（CP-SAT）で n=1000 規模へスケールさせた Moccasin（ICML 2023）9 がある（Checkmate は n=500 で3時間制限でもタイムアウトする）。 再計算とオフロードの最適な組み合わせを動的計画法で解く POFO（Beaumont ら、NeurIPS 2021）10 も、厳密解法でハイブリッドを扱う点でこの系譜とハイブリッド系譜の交点にある。

4. 動的ランタイムの系譜（2021〜2024）：DTR とその子孫

静的計画と対照的に、実行時にその場で判断する系譜がある。

DTR（Dynamic Tensor Rematerialization）（ICLR 2021）11 は、メモリが枯渇したら貪欲にテンソルを追い出し、必要になったら親演算子を再帰的に再計算するオンラインアルゴリズムである。 静的計画が前提にする静的グラフが不要になるため、動的計算グラフを扱える。 N 層の線形ネットワークを Ω(√N) メモリで O(N) 演算で学習できるという保証を持ち（Chen らの結果のオンライン版）、シミュレーションでは最適な静的計画に肉薄する。 PyTorch へはテンソル確保と演算子呼び出しへの割り込みだけで組み込める。

DTR の後続として、テンソル分割と再材料化を協調させる MegTaiChi（ICS 2022）12、既存手法がフレームワークのメモリアロケータを無視している（別アドレスの空きブロックを同一視して断片化を招く）ことを突き、スライディングウィンドウ内での追い出しでテンソル確保と再材料化を共最適化する Coop（NeurIPS 2023）13 がある。 テンソルを分割して部分単位でスワップする TSPLIT（ICDE 2022）14 もこの近傍にある。

ハイブリッドを動的に行う集大成が DELTA（ACM TACO 2024）15 で、スワップと再計算の両方を実行時に細粒度で判断する初の動的スケジューラを名乗る。 問題を 0/1 ナップサックとして定式化し、多項式時間ヒューリスティックと双方向プリフェッチで、最新手法比 40〜72% のメモリ削減、GPT2-XL のバッチ 6 倍を達成する。 「スワップを常に優先すべき」という一部先行研究の主張を実験で反証し、テンソルごとの動的な使い分けを支持している点は、Capuchin 以来のハイブリッド路線の追認である。

5. ハイブリッドの静的計画（2023）：シードの残り2本の位置

Capuchin（動的ランタイム）と Checkmate 系（静的・再計算のみ）の間を埋めるのが、ハイブリッドを静的に計画する系譜である。

HOME（IEEE TC 2023）は、既存手法が「モデルの部分的な情報」しか使わないことを問題とし、モデル全体の情報を入力に粒子群最適化（PSO）で各テンソルの配置（保持／スワップ／再計算）を大域的に決める。 vDNN 比 5.7 倍、Capuchin 比 1.3 倍のスループットで、シード4本の中では「vDNN → Capuchin → HOME」という改善の連鎖を明示的に作っている。 ILP でなくメタヒューリスティックを選ぶことで、厳密性を諦めて探索空間の大きさに対処する構図である（選択理由の詳細は本文未取得）。

STR（IEEE TPDS 2023）16 は、スワップと再計算の組み合わせを制約付き最適化（ILP ソルバ）で解き、さらに「スワップ済みテンソルを再計算の起点（host checkpoint）として使い回す」機構で再計算コストを削る。 参考文献に Capuchin と XEngine の両方を引いており、シード間をつなぐ結節点になっている。ResNet 系で従来のハイブリッド最適化比 28.1% のスループット改善を報告する。

6. LLM 時代（2021〜）：対象と粒度の移動

LLM の登場で、この分野の重心が3つの方向に動いた。

第一に、スワップの対象が活性からモデル状態へ移った。 Adam のオプティマイザ状態がパラメータの何倍ものメモリを占めるためである。 ZeRO-Offload（USENIX ATC 2021）17 はオプティマイザ状態と更新計算を CPU へ移し、V100 1枚で 13B モデルを学習可能にした（CPU 側の Adam を PyTorch 標準の5〜7倍に最適化した点が、オフロードでは CPU 側の計算最適化も必要という教訓になっている）。 ZeRO-Infinity（SC 2021）18 はこれを NVMe まで拡張した。 この線は、オプティマイザ状態の多層ストレージオフロードで 280B モデルの学習を 2.5 倍速くする MLP-Offload（SC 2025）19 や、ZeRO-Infinity の CPU 側メモリの浪費（断片化、pinned buffer、ピークスパイク）を削る MemAscend（2025）20 まで続いている。

第二に、再計算そのものを減らす方向が生まれた。 Megatron-LM の selective activation recomputation（MLSys 2023）21 は、sequence parallelism と選択的再計算（メモリを食うが計算が軽い部分だけ再計算する）で、全層再計算の冗長計算の 90% 以上を削り、530B モデルの学習を 29% 速くした。 FlashAttention（NeurIPS 2022）22 は self-attention のメモリを O(n²) から O(n) にし、「そもそも保持すべき活性」を減らした。 再計算を最適化する系譜（Checkmate ら）に対し、再計算の必要自体を減らす系譜と言える。

第三に、複数手法の設定の自動化が問題になった。 ZeRO 分割、スワップ、checkpointing が全部使える時代には、どれをどこに効かせるかの設定空間が人手に負えなくなる。 LoHan（ICDE 2024）23 は民生 GPU 1枚（RTX 4090＋256GB DRAM）で 175B モデルをファインチューニングするために、活性のスワップと再計算をモデル状態のオフロードと一体で計画する。 ProTrain（MLSys 2026）は、この統合をブロック粒度の構造化とコストモデルによる自動探索で解く。 テンソル粒度（Capuchin 流）は LLM では探索空間が爆発し、全ブロック一律（DeepSpeed 流）はメモリを使い残す、という粒度の議論を明示的にしており、transformer ブロック単位でスワップ／再計算／無最適化を混載する。 詳細はノートを参照。

シード4本の相互関係

4本は同じ分野の別世代・別象限として関係づけられる。

論文 手段 決め方 一言
Capuchin（ASPLOS 2020） ハイブリッド 動的ランタイム アクセスパターン観測でスワップと再計算を使い分け
HOME（IEEE TC 2023） ハイブリッド 静的計画（PSO） モデル全体の情報で配置を大域最適化。Capuchin を改善対象に
XEngine（TACO 2022） 再計算＋デバイス配置 静的計画（MIQP） Checkmate のヘテロ多デバイス拡張。CPU を計算資源として使う
ProTrain（MLSys 2026） 全部（ZeRO＋スワップ＋再計算） 静的計画（コストモデル探索） LLM 向けにブロック粒度で統合し設定を自動化

引用関係で確認できたのは次のとおり。

ノートブック状態管理との橋（ElasticNotebook との独立再発見）

このリポジトリの中心分野であるノートブックの状態管理と、この地図の分野は、同じ「保存 vs 再計算」のトレードオフを別世界で独立に定式化している。

ElasticNotebook（VLDB 2023）は、セッション内の変数集合 X を「移行する部分集合 S」と「再計算する部分集合 X−S」に分け、コスト w(S) = wM(S) + wR(X−S) を最小化する複製プランを求める。 これは Capuchin が GPU テンソルに対して行う「スワップ（保存）か再計算か」の判断と構造的に同型である。

ただし両分野に引用関係はない。 ElasticNotebook の参考文献（Semantic Scholar で全件確認）に Capuchin、vDNN、Checkmate、DTR、Chen らの sublinear memory など DNN メモリ管理系の論文は一切含まれず、後続の Kishu も同様である。 ElasticNotebook が「読み込みと再計算のバランス」の先行研究として挙げるのはデータサイエンスワークフロー分野の Helix であり、引用の世界が完全に分かれている。

解法も対照的で、比較すると面白い。

  ElasticNotebook Checkmate / XEngine Capuchin / DTR / DELTA
対象 セッション変数 活性テンソル 活性テンソル
定式化 s-t 最小カット MILP / MIQP ランタイムヒューリスティック
解の性質 多項式時間で厳密 NP 困難だが厳密（ソルバ次第で数時間） 近似だがオーバーヘッド小

ElasticNotebook が多項式時間の厳密解を持てるのは、変数間依存（AHG）の構造が min-cut への帰着を許すからで、演算子のスケジューリング順序まで決める必要がある DNN 側は整数計画になる。 同じトレードオフでも、問題構造の違いが解法の複雑さを分けている。

未整理・今後の課題

  1. Ren らの ZeRO-Offload 論文の見積もりとして ProTrain（arXiv:2406.08334）が引用。fp16/fp32 パラメータ、fp16 勾配、fp32 モーメント・分散の合計。 ↩

  2. Tianqi Chen et al., “Training Deep Nets with Sublinear Memory Cost,” arXiv:1604.06174, 2016. ↩

  3. XEngine（arXiv:2212.09290）の related work による。Andreas Griewank, “Achieving Logarithmic Growth of Temporal and Spatial Complexity in Reverse Automatic Differentiation,” 1992 ほか。 ↩

  4. Minsoo Rhu et al., “vDNN: Virtualized Deep Neural Networks for Scalable, Memory-Efficient Neural Network Design,” MICRO 2016, doi:10.1109/MICRO.2016.7783721. ↩

  5. Linnan Wang et al., “SuperNeurons: Dynamic GPU Memory Management for Training Deep Neural Networks,” PPoPP 2018. arXiv:1801.04380. ↩

  6. Chien-Chin Huang, Gu Jin, and Jinyang Li, “SwapAdvisor: Pushing Deep Learning Beyond the GPU Memory Limit via Smart Swapping,” ASPLOS 2020, doi:10.1145/3373376.3378530（オープンアクセス）. ↩

  7. Paras Jain et al., “Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization,” MLSys 2020. arXiv:1910.02653. ↩

  8. Shishir G. Patil et al., “POET: Training Neural Networks on Tiny Devices with Integrated Rematerialization and Paging,” 2022. Moccasin と XEngine が Checkmate の後続として言及。 ↩

  9. Burak Bartan et al., “Moccasin: Efficient Tensor Rematerialization for Neural Networks,” ICML 2023 (PMLR 202:1826-1837). arXiv:2304.14463. ↩

  10. Olivier Beaumont, Lionel Eyraud-Dubois, and Alena Shilova, “Efficient Combination of Rematerialization and Offloading for Training DNNs,” NeurIPS 2021, pp. 23844-23857. ↩

  11. Marisa Kirisame et al., “Dynamic Tensor Rematerialization,” ICLR 2021. arXiv:2006.09616. ↩

  12. Zhongzhe Hu et al., “MegTaiChi: Dynamic Tensor-Based Memory Management Optimization for DNN Training,” ICS 2022, doi:10.1145/3524059.3532394（CC-BY オープンアクセス）. ↩

  13. Jianhao Zhang et al., “Coop: Memory is not a Commodity,” NeurIPS 2023 (spotlight). arXiv:2311.00591. ↩

  14. Xiaonan Nie et al., “TSPLIT: Fine-grained GPU Memory Management for Efficient DNN Training via Tensor Splitting,” ICDE 2022. ↩

  15. Yu Tang et al., “DELTA: Memory-Efficient Training via Dynamic Fine-Grained Recomputation and Swapping,” ACM TACO 21(4), 2024, doi:10.1145/3689338. 旧題の arXiv 版（”DELTA: Dynamically Optimizing GPU Memory beyond Tensor Recomputation,” arXiv:2203.15980, 2022）と同一の系統。 ↩

  16. Zan Zong et al., “STR: Hybrid Tensor Re-Generation to Break Memory Wall for DNN Training,” IEEE TPDS 34(8):2403-2418, 2023, doi:10.1109/TPDS.2023.3266110（オープンアクセス版なし）. ↩

  17. Jie Ren et al., “ZeRO-Offload: Democratizing Billion-Scale Model Training,” USENIX ATC 2021. ↩

  18. Samyam Rajbhandari et al., “ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning,” SC 2021. arXiv:2104.07857. ↩

  19. Avinash Maurya et al., “MLP-Offload: Multi-Level, Multi-Path Offloading for LLM Pre-training to Break the GPU Memory Wall,” SC 2025. arXiv:2509.02480. ↩

  20. “MemAscend: System Memory Optimization for SSD-Offloaded LLM Fine-Tuning,” arXiv:2505.23254, 2025. ↩

  21. Vijay Korthikanti et al., “Reducing Activation Recomputation in Large Transformer Models,” MLSys 2023. arXiv:2205.05198. ↩

  22. Tri Dao et al., “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness,” NeurIPS 2022. ↩

  23. Changyue Liao et al., “LoHan: Low-Cost High-Performance Framework to Fine-Tune 100B Model on a Consumer GPU,” ICDE 2024. arXiv:2403.06504. ↩