関係マップ:深層学習のチェックポイント
この地図は、深層学習の学習チェックポイントの全体像を描く。 「何を保存するのか」と「単一ノードと分散の違い」を整理したうえで、分散深層学習の詳細は distributed-dl-checkpoint.md に委ねる。 上位の地図(GPU チェックポイント全般)は gpu-checkpoint.md を参照。 情報源は各論文のアブストラクト、公式ドキュメント、公開情報に基づく。
深層学習のチェックポイントとは
深層学習の学習チェックポイントとは、学習途中のモデルの状態をストレージに保存し、障害や中断が起きたときにその地点から学習を再開できるようにする仕組みである。 汎用 GPU C/R(gpu-checkpoint.md §1)がプロセス全体を透過的に保存するのに対し、学習チェックポイントはアプリケーション層で必要な状態だけを選んで保存する。
保存する状態
学習チェックポイントで保存する状態は、おもに次の4つである。
- モデルパラメータ(重み):ニューラルネットワークの全レイヤーの重みとバイアス。モデルの学習成果そのもの。
- オプティマイザ状態:Adam であればパラメータごとの一次モーメント(m)と二次モーメント(v)。SGD with momentum であればモメンタム項。モデルパラメータと同等かそれ以上のサイズになることが多い(Adam の場合は重みの2倍)。
- 学習率スケジューラの状態:現在のステップ数、ウォームアップの進捗など。
- その他のメタデータ:現在のエポック番号、乱数シードの状態、データローダの位置など。
これらをすべて保存すれば、同じ地点から学習を再開してビット単位で同一の結果を得られる(決定論的な設定の場合)。
チェックポイントの基本的な流れ
- 学習ループ中の所定のタイミング(N イテレーションごと、またはエポックごと)で学習を一時停止する
- GPU メモリ上のモデルパラメータとオプティマイザ状態を CPU メモリにコピーする
- CPU メモリからストレージ(ローカルディスク、NFS、S3 など)に書き出す
- 学習を再開する
この流れ自体は単純だが、モデルの大規模化に伴い、保存すべきデータ量が数十 GB から数 TB に達し、チェックポイントのオーバーヘッドが無視できなくなった。
単一ノードと分散の違い
深層学習のチェックポイントは、学習の規模によって難しさが大きく変わる。
単一ノード(GPU 1枚、または1ノード内の複数 GPU)
単一ノードでは、保存すべき状態が1つのプロセス(または同一マシン上の少数プロセス)に閉じている。
チェックポイントは「1プロセスがモデル重みとオプティマイザ状態をローカルディスクに書く」だけで完了する。
PyTorch の torch.save(model.state_dict()) のような1行のAPI呼び出しで済む。
単一ノードでも、モデルサイズが大きければ書き出し時間が問題になる。 CheckFreq1 はチェックポイント頻度を自動チューニングし、復旧時間を秒単位に短縮しつつオーバーヘッドを 3.5% 以内に抑える。 DeepFreeze2 は非同期でのシリアライズと書き出しにより、学習と I/O を重畳させる。
分散(複数ノード × 複数 GPU)
分散学習では、チェックポイントの難しさが質的に変わる。 その理由は3つある。
第一に、各 GPU が持つ状態の断片が異なる。 データ並列ではモデルパラメータは各 GPU で同一だが、オプティマイザ状態は分割されている(ZeRO Stage 2/3 の場合)。 モデル並列やパイプライン並列では、モデルパラメータ自体が GPU ごとに異なる断片に分かれている。 すべての断片を矛盾なく同時に保存しなければ、復元時に一貫性が崩れる。
第二に、規模が大きいほど障害確率が上がる。 Meta の OPT-175B の学習では100回以上のハードウェア障害が発生した3。 障害頻度が高いほどチェックポイント間隔を短くする必要があるが、短くするほどオーバーヘッドが増える。 このトレードオフが単一ノードより格段にきつくなる。
第三に、並列戦略の変更への対応が求められる。 学習途中で GPU 数を増減させたり、並列戦略を切り替えたりする場合、チェックポイントのレイアウトを変換する必要がある。 Universal Checkpointing4 が解こうとしているのがこの問題である。
分散深層学習のチェックポイントに関する論文の詳細な調査は distributed-dl-checkpoint.md を参照。
この地図の位置づけ
GPU チェックポイント全般(gpu-checkpoint.md)
├── 汎用 CUDA C/R(CRIUgpu, PhoenixOS)
├── 深層学習のチェックポイント(この地図)
│ ├── 単一ノード(CheckFreq, DeepFreeze)
│ └── 分散深層学習(distributed-dl-checkpoint.md)← 論文調査
└── ノートブックの状態管理(checkpoint.md)
-
Jayashree Mohan et al., “CheckFreq: Frequent, Fine-Grained DNN Checkpointing,” FAST 2021. ↩
-
Bogdan Nicolae et al., “DeepFreeze: Towards Scalable Asynchronous Checkpointing of Deep Learning Models,” CCGrid 2020. ↩
-
Susan Zhang et al., “OPT: Open Pre-trained Transformer Language Models,” 2022. 学習ログに障害の詳細が記録されている。 ↩
-
Xinyu Lian et al., “Universal Checkpointing: Efficient and Flexible Checkpointing for Large Scale Distributed Training,” 2024. ↩