8月28日、Databricksが「Fast, fault-tolerant PyTorch training on AI Runtime」と題した記事を公開した。256基以上のGPUを使う大規模学習ではジョブ実行中に必ず障害が起きると前提を置き、チェックポイント設計・データパイプライン・自動リジュームを組み合わせてgoodputを91%まで高める実装手法を詳述している。
大規模学習では「障害は例外ではなく通常運転」
まずDatabricksが強調するのは、この前提認識の転換だ。
GPU1基あたりの年間障害率を約1%と仮定すると、256GPU×30日のジョブでは19%の確率で障害が発生し、1,024GPUでは57%に達する。実測値で言えば、608基のH100で構成されたDeltaスーパーコンピュータでは平均1.9時間に1回の割合で障害が発生した。32GPU構成に換算すると、平均的な無障害時間は36時間となる。
こうした状況下で訓練効率を測る指標として、記事は「GPU goodput」という概念を使う。MLコミュニティで広く参照される指標で、GPUが生産的な計算（＝順伝播・逆伝播）に費やしている時間の割合を指す。障害からの回復や待機に費やす時間が増えるほどこの値は下がる。
最大のレバー：チェックポイントの仕組みを根本から変える
障害が起きたとき、どれだけの計算をやり直さなければならないかはチェックポイント間隔に直結する。1日1回しか保存していなければ、障害発生時に平均12時間分の再計算が必要になる。
torch.saveの何が問題か
多くのチームが最初に書くのは、rank 0でtorch.saveを呼ぶシンプルな実装だ。しかしこれには2つの問題がある。
- 分散学習の場合、全ステートをrank 0に集約してから単一ファイルに書き出す
- 書き込み中はGPUがブロックされ、Unity Catalog（UC）などリモートストレージへの転送中はGPUがアイドル状態になる
DCP（Distributed Checkpoint）で並列保存へ
PyTorchのDistributed Checkpoint APIは設計を逆転させる。全rankが自分のシャードを並行して書き出す構成で、保存時間はrank数に反比例して短縮される。また.metadataファイルにグローバルレイアウトが記録されるため、異なるGPU数でのリストアも可能だ。
「DCPはFSDPや大規模モデル専用」と思われがちだが、DDP（Data Parallel）構成でも効果がある。全rankが同一のウェイトを持つDDPでも、DCPはモデルステートをシャード化して並列書き込みを行う。
非同期保存で「ほぼタダ」にする
さらに効果的なのがasync_saveだ。書き込みをステージングバッファへの高速コピーとバックグラウンドアップロードに分離し、アップロード中も学習を継続できる。
実測値は以下のとおりだ。
| ジョブ構成 | torch.save | DCP async_save | 削減率 |
|---|---|---|---|
| DDP LLM（2.8Bパラメータ、32×H100） | 66秒 | 36秒 | 1.8倍 |
| FSDP LLM（20Bパラメータ、32×H100） | 522秒 | 9秒 | 58倍 |
※上記はネットワークストレージへの転送時間を除いた値
FSDP構成での58倍という数字が示すように、モデルが大きくなるほど非同期保存の効果は顕著になる。
保存頻度がgoodputを直接決める
安価に保存できると、保存間隔を縮めることができる。Llama 3の事例（1日あたり約8.6回の中断）を例に取ると：
- 2時間ごとのチェックポイント：1日あたり再学習に8.6時間を浪費 → goodput 64%
- 30分ごとのチェックポイント：浪費は2.15時間 → goodput 91%
再起動時は最後に完全に書き終えたチェックポイントを自動検出してリジューム。DCPの.metadataファイルは「全シャードの書き込み完了後」にのみ作成されるため、これを「完全な保存の証明」として使える。
見落とされがちな罠：データパイプラインのチェックポイント忘れ
記事が「エラーも出ず、クラッシュもせず、ジョブも失敗しない」と表現するサイレント障害がある。
モデル・オプティマイザ・ステップ数はチェックポイントに含めても、データパイプラインの位置（どこまで読んだか）を含め忘れるパターンだ。
障害後にリジュームすると、データローダーはエポックの先頭から読み直す。すでに学習したサンプルを再度学習し、未学習のサンプルをスキップする。大規模学習では再起動が日常的に発生するため、これが積み重なるとデータ分布が静かに歪む。モデルは学習を完了するが、評価指標が想定より低い、という結果になる。
対処法はデータ位置をチェックポイントに含めることと、シャッフルや拡張処理に使うRNGの状態も含めることだ。シードと再現可能なデータ順序が揃って初めて、リジューム後のデータ順序が保証される。
データローディングがGPUを止める
記事は、データパイプラインがボトルネックになっているケースでもウォールクロック時間が20〜50%短縮されると報告している。
Databricks AI RuntimeではUCVolumeDatasetと専用DataLoaderを提供。初回アクセス時にローカルNVMeへキャッシュし、以降はキャッシュから読む構成で、プリフェッチとGPU計算を重複させる。
同一GPU・モデル・バッチサイズでの比較（Epoch 2はキャッシュ済み状態での数値であり、ストレージアクセスの差が最大限に現れる）：
| メトリクス（GPU1基あたり、定常状態） | 標準PyTorch DataLoader | UCVolumeDataset + Databricks DataLoader |
|---|---|---|
| Epoch 1スループット（画像/秒） | 57.2 | 417 |
| Epoch 2スループット（画像/秒）※キャッシュ済み | 371.6 | 6590 |
| GPU使用率 | 12.6% | 53.3% |
またDataLoaderのメトリクスはMLFlowに自動記録される。fetch_seconds（バッチ生成にかかった時間＝GPUがアイドルだった時間）を見れば、データパイプラインがボトルネックになっているかどうかをすぐ確認できる。
まとめ
記事が提示する原則はシンプルだ。「頻繁で、安価で、完全なチェックポイント」が、ハードウェア障害をジョブ終了イベントから軽微な中断に変える。
torch.saveを捨てDCPに移行する（DDP構成でも効果あり）async_saveで保存コストをほぼゼロにし、保存間隔を短くする- 再起動を自動化し、最新の完全なチェックポイントから自動リジュームする
- データ読み込みをコンピュートと重複させてGPUを止めない
- モデルだけでなくデータパイプラインとRNG状態もチェックポイントに含める
インフラ側（フリートレベルのGPU健全性監視）については、別記事「How we keep GPUs reliable across Databricks AI」で詳説されている。
詳細はFast, fault-tolerant PyTorch training on AI Runtimeを参照していただきたい。
