Databricks、AI Runtime で高速・耐障害性の PyTorch 学習を実現

大規模な環境において、トレーニングの効率を決めるのは単一の指標です。それが「グッドプット（goodput）」で、GPU が待機や障害からの回復に費やす時間を除いた、生産的な計算に実際に使われる時間の割合を指します。

大規模運用では GPU の故障は想定内の出来事であり、高いグッドプットを維持し、総 GPU コストを管理するためには、障害から迅速かつ自動的に復旧できる能力が唯一の手段です。

この復旧プロセスを左右する 2 つのサブシステムがあります。しかし、これらは往々にして後回しにされがちです。1 つはアクセラレータへデータを供給するデータパイプライン、もう 1 つはジョブの再開のために状態のスナップショットを作成するチェックポイント機構です。どちらか一方でも設計を誤れば、障害発生時に想定以上のアイドル GPU 時間を無駄にしてしまいます。

故障シナリオ以外であっても、アクセラレータの処理速度に追いつけないデータパイプラインは、静かに GPU を飢餓状態に陥れ、クラッシュと同じようにグッドプットを損ないます。ここでは両者の仕組みとトレードオフについて解説し、それぞれがどのようにグッドプットや総 GPU コストに影響を与えるかを見ていきましょう。

コードのヒントや具体例については、補完資料である「トレーニングのパフォーマンスとレジリエンスガイド」Training performance and resiliency guide をご覧ください。

同様の問題におけるインフラストラクチャ側の課題、つまりジョブを停止させる前に不健康な GPU をどのように検出・隔離するかについては、関連記事「Databricks AI 全体で GPU の信頼性を維持する方法」をご覧ください。

スケール拡大では障害が当然の事態となる

ジョブに使用する GPU 数が増加すればするほど、中断なく完了する確率は急速に低下します。関連する Databricks の記事から引用される簡易な計算モデルでは、各 GPU の年間故障率が約 1% と仮定されています。この前提に基づくと、「30 日間実行される 256 GPU のジョブで障害が発生する確率は約 19% です。GPU 数が 1,024 に増えると、その確率は 57% に跳ね上がります」と指摘されています。これらはあくまでインフラレベルの事象に過ぎません。

この推計を現実と照らし合わせるため、608 基の H100 GPU を搭載したデルタ・スーパーコンピュータ（論文：https://arxiv.org/html/2503.11901v4）の事例を確認すると、障害は平均して 1.9 時間ごとに発生しています。つまり、32 GPU のジョブであれば、故障までの平均所要時間は約 36 時間となります。重要な点は、トレーニングジョブがいつか必ず失敗する可能性が高いという事実です。その際、適切な判断を下すことでモデルの耐障害性を高め、発生した際の損失時間を最小限に抑えることができます。

影響 1: チェックポイント形式が保存頻度を決定する

レジリエンス（回復力）の成否はチェックポイントにあり、採用するメカニズムが「どれくらいの頻度で保存できるか」に直結します。これはグッドプット（実効的な処理速度）を左右する最も大きな要因です。例えば 1 日に 1 回しかチェックポイントを保存しない場合、障害が発生した際に平均して 12 時間分の作業をやり直す必要があります。

torch.save の単一プロセスによるボトルネック

多くのチームが最初に採用するのは、ランク 0（主ノード）での単純な torch.save です。モデルの学習方法によっては、以下の 2 つの問題が発生する可能性があります。

• 分散学習の場合、すべての状態をランク 0 に集約して 1 つのファイルに書き出します。
• 単一のプロセスがチェックポイント全体を同期しながら書き出すため、Unity Catalog（UC）のようなリモートオブジェクトストアへの保存時にネットワーク転送などでブロックされることがあります。

このブロッキング動作により GPU がアイドル状態となり、グッドプットが低下してしまいます。しかし、GPU がチェックポイントに費やす時間を減らす方法は存在します。それが PyTorch の分散チェックポイント API です。

分散チェックポイント（DCP）：各ランクがシャードを個別に書き出す

PyTorch の分散チェックポイントは設計思想を逆転させています。すべてのランクが並列して独自のシャードを書き出し、それらがどのように組み合わさって完全なテンソルを構成するかを示す小さな .metadata ファイルも同時に作成します。

ラanks 数が増えるにつれて所要時間は概ね 1/N で短縮され、.metadata ファイルがグローバルなレイアウトを記録しているため、同じチェックポイントは異なる数の GPU に再ロード可能です。DCP は新しいランクに必要なバイトを再計画するため、ノードの喪失後に低容量クラスターへの回復もスムーズに機能します。

単純なデータ並列ジョブでも DCP の価値はある

「DCP はシャード化されたモデル専用で、すべてのランクが重みの同一レプリカを持つデータ並列（DDP）ジョブには恩恵がない」という考えが一般的ですが、それは誤りです。DCP は DDP 訓練タスクであってもモデル状態をシャードし、各ワーカーに並行して書き込みます。

また、将来 FSDP やテンソル並列に移行する際にも同じ API を使用することになります。早期に導入すれば、最も不適切なタイミングで耐障害性コードを書き換える必要はなくなります。

非同期保存により頻度をほぼ無料に

並行書き込みを行っても、同期保存ではストレージへの永続化完了まで訓練がブロックされ、大規模チェックポイントをリモートボリュームへ転送する際、数十秒ものアイドル状態のアクセラレーター時間が発生します。async_save はこの操作を分割し、まず高速なコピーをステージングバッファへ行い、その後に背景処理でアップロードを実行して訓練継続と重畳させます。

訓練ループが負担するのはステージングへのコピーのみであり、アップロードのコストは発生しません。かつて数十秒のアイドル時間を要していたチェックポイントがほぼゼロコストとなり、これが次のセクションで説明する頻繁なチェックポイントを可能にする理由です。

AI Runtime 上では、UCVolumeWriter と UCVolumeReader が UC ボリュームに対して DCP（データ整合性プロトコル）を実装しており、ローカルの NVMe を介して入出力をステージングします。チェックポイントの完了は、そのデータが完全に書き込まれてからのみマークされます。

詳細な仕様やコード例については、パフォーマンスと耐障害性のガイド をご覧ください。

上記の計算には、torch.save のネットワークストレージへの書き込み時間は含まれていません。

インパクト 2：チェックポイント頻度が復旧コストを決める

ここが重要なポイントです。ジョブが失敗すると、直前の有効なチェックポイント以降のすべての進捗が失われ、そこから再計算する必要があります。つまり、1 回の失敗で失われる作業量の期待値は、およそチェックポイント間隔の半分になります。非同期保存を安価に利用できれば、この間隔を短く設定できるのです。

間隔を 10 分の 1 に短縮すれば、復旧にかかる時間の期待値も 10 分の 1 に減ります。Llama 3 の事例で 1 日あたり約 8.6 回の中断が発生すると仮定しましょう。この失敗頻度の場合、2 時間ごとにチェックポイントを取得すると、1 日あたり 8.6 時間を再学習に費やすことになり、実効性能（goodput）は 64% に留まります。一方、30 分ごとのチェックポイントであれば、必要な時間は 2.15 時間に減り、実効性能は 91% に向上します。

復旧プロセスも自動化されている必要があります。ジョブを再起動した際、最も直近で書き込みが完了したチェックポイントを自動的に見つけ、クラッシュによって半分の状態で残されたファイルは無視し、人間の介入なしにそこから再開できる仕組みが必要です。DCP（Distributed Checkpoint）はこの信頼性を担保します。.metadata ファイルはすべてのシャードの保存が完了してからのみ作成されるため、このファイルが存在することは「保存が完了した」という確実な合図として機能し、復旧時に選択する対象として信頼できます。

インパクト 3: データ読み込みが GPU のアイドル時間を決める

トレーニングの速度は、最も遅い入力部分によって決定されます。アクセラレーターが次のバッチを待っている間、GPU が単にアイドル状態にあるため、実効的な処理能力（グッドプット）は低下します。この問題を解決する唯一の方法は、図に示すように、次のステップのデータ準備と現在のステップでの計算を重畳させることです。

多くの顧客が、データ読み込みと計算の重畳化に移行した結果、実測時間（ウォールクロックタイム）で 20〜50% の短縮を実現しています。

リモートストレージからの直接読み出しのコスト

ガバナンスされたプラットフォームでは、トレーニングデータはリモートのオブジェクトストレージに保存されています。AI Runtime では、Unity Catalog (UC) ボリュームがネットワークマウントとして公開されます。

毎回このマウントからファイルを直接読み出すと、ステップの所要時間がネットワークレイテンシに縛られ、各エポックで同じファイルが再ダウンロードされてしまいます。解決策は、最初のアクセス時にファイルを高速なローカルストレージへコピーし、その後の読み出しをローカルキャッシュから提供し、GPU が計算を行っている間に次のファイルを並列で取得するデータローダーを採用することです。

AI Runtime では、UCVolumeDataset と DataLoader が以下のような動作を行います（コード例についてはこちらのガイドをご覧ください）。

UCVolumeDataset は UC ボリュームからファイルをストリーミングし、最初にアクセスした際に各ファイルをローカルの NVMe にキャッシュします。また、ファイルはランクとワーカー間で分割されるため、すべてのアクセラレーターが重複しない独立したデータスライスを受け取ることができます。

当社の DataLoader は PyTorch の標準 DataLoader を継承したサブクラスであり、このパスに最適化されたデフォルト設定を備えています。これにより、GPU が計算を行っている間にファイルの取得とキャッシュが並行して実行され、トレーニングスレッドで一つずつ処理される従来の方式よりも高速化されます。

例：画像モデルのトレーニング（UC ファイルから）

単純な画像分類ワークロードを想定してみましょう。具体的には、UC ボリュームから JPEG をデコードし、データ拡張を行った上でビジョンモデルを訓練するケースです。

同じ GPU、モデル、バッチサイズでこの処理を行う場合、2 つのアプローチを比較してみます。1 つは標準的な PyTorch の Dataset を用いて UC ボリュームから直接読み込む方法、もう 1 つは UCVolumeDataset と Databricks の DataLoader デフォルト設定を組み合わせた方法です。

時間をどこに費やしているか、推測する必要はありません

DataLoader のエンジニアリングの一環として、MLFlow にメトリクスをログ出力する機能を追加しました。これにより、データパイプラインがトレーニングのボトルネックになっているかどうかを一目で確認できるようになりました。

fetch_seconds というメトリクスは、Dataloader がバッチを生成するまでに実際にどれだけの時間を要したかを計測します。この間、GPU はアイドル状態になっています。

影響 4：データパイプラインの忘却がモデルを静かに劣化させる

エラーメッセージもクラッシュもなく、ジョブ失敗として表示されることもない、最後の耐障害性のバグがあります。それが「モデルが微妙に期待値より劣っている」という現象です。これは、モデル、オプティマイザ、ステップ数はチェックポイントするものの、データセット内におけるデータパイプラインの位置情報を保存しない場合に発生します。

例えば、エポック途中まで実行中のジョブが中断された場合を考えましょう。システムはモデルを正しく復元してトレーニングループを再開しますが、Dataloader はデータセットの先頭から再度読み始めます。

再開したジョブは、すでにこのエポックで学習済みのサンプルを再学習し、まだ到達していなかったサンプルをスキップする可能性があります。スケーラビリティのために頻繁にリスタートが行われる環境では、これがデータの分布に静かにバイアスを生じさせます。モデル自体はトレーニングを継続しますが、実際には不適切なサンプリングされたデータで学習していることになります。これは最もコストがかかる「静かな失敗」の典型です。ジョブは完了し、メトリクスが期待外れになるまで誰も問題に気づかないからです。

解決策は、データ位置をチェックポイントの一部として扱うことです。パイプラインの構成によっては、サンプルまたはシャードのオフセットを追跡して再開時にスキップする、カスタムデータセットが自身の位置をシリアライズする、あるいはエポック境界でチェックポイントを保存するなどの方法があります。これらすべてに共通する前提条件は決定性（determinism）です。

シャッフルやデータ拡張は乱数生成器に依存するため、その乱数シードと状態もチェックポイントに含まれる必要があります。そうでなければ、再開後のデータ順序が再開前と一致しなくなり、保存された位置情報が誤ったサンプルを指すことになります。

「シード」「再現可能な順序」「再開可能なデータパイプライン」は、すべて同じ概念の異なる表現です。ガイドでは、各戦略についてコード付きで解説しています。

まとめ

高速かつ障害耐性のあるトレーニングを実現するには、いくつかの意思決定を積み重ねることが重要です。

DDP（Distributed Data Parallel）であっても、torch.save を使うのではなく分散チェックポイントを活用しましょう。これにより保存処理が並列化され、シリアルのボトルネックではなく高速かつ低コストな運用が可能になります。

非同期で保存を行うことで、チェックポイントの取得はほぼ無負荷となり、頻繁に保存することが可能になります。また、最も直近の有効なチェックポイントから自動的に復元できるため、障害発生時の再計算にかかる時間は数分に抑えられ、数時間かかるような事態を回避できます。

リモートストレージからのキャッシュとプリフェッチによりデータ読み込みと計算処理を重畳させることで、アクセラレータが入力待ちでアイドル状態になることを防ぎます。この手法は、すべてのステップで GPU 使用時間を節約する効果をもたらします。

さらに、データパイプラインの状態と RNG（乱数生成）のステートもチェックポイントに含める必要があります。これにより、ジョブを再開した際に正しいデータから継続でき、モデルが静かに破損するリスクを防げます。

これらの手法を貫く基本原則は、「頻繁で安価かつ完全なチェックポイント」です。ハードウェア障害を「ジョブ終了の重大事象」から「計算上の丸め誤差レベルの問題」へと変えることができます。また、重畳された入力パイプラインにより、アクセラレータ間の待ち時間を最小限に抑えつつ稼働し続けます。

安価な非同期保存が頻繁なチェックポイントを可能にし、モデル・データ・RNG のすべてを含む完全な保存が正確な復元を実現します。これら二つの要素を整備し、Databricks のブログで紹介されているような「故障したハードウェアの検出と隔離」機能を備えたファームウェアを運用すれば、クラスタの信頼性が低くても、理論上の最大性能に近いトレーニング時間を確保できます。

参考文献

• Meta, The Llama 3 Herd of Models (2024): 最大 16,384 台の H100 GPU を用いたトレーニングにおける信頼性と中断時の詳細分析（セクション 3.3.1）。
• Characterizing the Resilience of Hopper H100 and Ampere A100 GPUs (2025): Delta システムにおける 2 年半にわたるフィールドスタディ。ジョブサイズ別の MTBE（平均故障間隔）と故障確率の分析。
• Databricks, How we keep GPUs reliable across Databricks AI: フリートレベルでの GPU 健全性チェックおよび確率的故障モデルについて。

実際に試してみたい方は、Databricks AI Runtime のドキュメントにある Training performance and resiliency guide で完全なコードを確認できます。また、インフラストラクチャ側の仕組みについては、How we keep GPUs reliable across Databricks AI をご覧ください。