AI Runtimeでの高速かつフォールトトレラントなPyTorchトレーニング

大規模な環境では、トレーニングの効率は「goodput」（GPUが待機や障害からの復旧ではなく、生産的な計算に費やす時間の割合）という単一の指標によって決定されます。大規模環境ではGPUの障害が発生することが前提となるため、障害から迅速かつ自動的に復旧できる能力が、高いgoodputを維持し、GPUの総支出を管理する唯一の方法です。

この復旧の成否を分ける2つのサブシステムがありますが、どちらも日常的に後回しにされがちです。それは、アクセラレータにデータを供給するデータパイプラインと、ジョブを再開できるように状態をスナップショットとして保存するチェックポイントメカニズムです。どちらか一方でも設定を誤ると、障害が発生するたびに、本来よりもはるかに多くのGPUのアイドル時間（無駄な時間）が発生することになります。障害が発生しない場合でも、アクセラレータの処理速度に追いつかないデータパイプラインは、クラッシュ時と同様に、GPUを密かにリソース不足（スターベーション）に陥らせ、goodputを低下させます。本稿では、これら両方のメカニズムとトレードオフ、そしてそれぞれがgoodputとGPUの総支出にどのように影響するかを説明します。コードの参照先や例については、関連するトレーニングのパフォーマンスとレジリエンスに関するガイドを参照してください。

同じ問題のインフラストラクチャ側、つまりフリートがジョブを停止させる前に異常なGPUを検出して隔離する方法については、関連ブログ記事「Databricks AI全体でGPUの信頼性を維持する方法」を参照してください。

大規模環境で障害が発生することが前提となる理由

ジョブで使用するGPUの数が増えるにつれて、中断することなく全期間を完了できる確率は急速に低下します。関連するDatabricksのブログ記事にある便利な簡易モデルでは、各GPUの年間故障率を約1%と仮定しています。この仮定に基づくと、同記事では「256個のGPUを使用するジョブを30日間実行した場合、障害が発生する確率は約19%です。1,024個のGPUでは、その確率は57%に上昇します」と指摘されています。しかも、これらはインフラストラクチャレベルの問題にすぎません。

この見積もりを現実的なものにするために、608台のH100 GPUを搭載したDeltaスーパーコンピューターでは1.9時間ごとに障害が発生しました。これは、32個のGPUを使用するジョブの場合、平均故障間隔が36時間になることを意味します。重要なポイントは、トレーニングジョブはいずれかの時点で失敗する可能性が高く、適切な決定を下すことでモデルのレジリエンス（回復力）を高め、障害発生時の総損失時間を削減できるということです。

影響1：チェックポイントのフォーマットが保存頻度を決定する

チェックポイントの作成はレジリエンスの成否を分けるポイントであり、選択するメカニズムは保存頻度に直接的な影響を与えます。これはgoodputを左右する最大の要因です。1日に1回しかチェックポイントを作成しない場合、障害が発生すると、障害発生時の状態に戻すために平均12時間の重複作業を再実行する必要があります。

モノリシックなtorch.saveのボトルネック

ほとんどのチームが最初に作成するチェックポイントは、ランク0でのシンプルなtorch.saveです。モデルのトレーニング方法によっては、主に2つの問題が発生する可能性があります。

1. 分散トレーニングの場合、すべての状態をランク0に集約し、単一のファイルを書き込みます。
2. 単一のプロセスがチェックポイント全体を同期的に書き込みます。Unity Catalog (UC) などのリモートオブジェクトストアに保存する場合、ネットワーク転送などによってこの処理がブロックされる可能性があります。

このブロッキング動作によりGPUがアイドル状態になり、goodputが低下します。しかし、GPUがチェックポイントの作成に費やす時間を短縮する方法があります。それがTorchの分散チェックポイント（DCP）APIです。

分散チェックポイント（DCP）：各ランクが独自のシャードを書き込む

PyTorchの分散チェックポイント（DCP）は、この設計を覆します。すべてのランクが独自の個別のシャードを並行して書き込み、それと同時に、シャードがどのように完全なテンソルを構成するかを記述した小さな.metadataファイルを書き込みます。

保存時間はランク数に応じてほぼ1/Nに短縮され、.metadataファイルがグローバルなレイアウトを記録しているため、同じチェックポイントを異なる数のGPUにリロードできます。DCPは新しい各ランクが必要とするバイトを再計画するため、ノードを失った後に容量が縮小したクラスターに復旧する場合でも問題なく動作します。

DCPは通常のデータ並列ジョブでも導入する価値がある

DCPはシャード化されたモデル専用であり、すべてのランクが重みの同一レプリカを保持するデータ並列（DDP）ジョブにはメリットがないと思われがちです。しかし、そうではありません。DCPはモデルの状態をシャード化し、DDPトレーニングタスクであっても各ワーカー間で並行して書き込みを行います。

これは、将来FSDPやテンソル並列処理に移行する際にも必要となるのと同じAPIです。そのため、早期に導入しておけば、最悪のタイミングでレジリエンスコードを書き直す必要がなくなります。

非同期保存により、保存頻度のコストがほぼゼロに

並行書き込みを使用しても、同期保存ではデータがストレージに永続化されるまでトレーニングがブロックされます。リモートボリュームへの大規模なチェックポイントの場合、アクセラレータのアイドル時間は数十秒に及びます。async_saveはこの操作を分割します。ステージングバッファへの高速なコピーを実行した後、バックグラウンドでアップロードを行い、継続するトレーニングと重ね合わせます。

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

AI Runtimeでは、UCVolumeWriterとUCVolumeReaderがUCボリュームに対してDCPを実装し、ローカルNVMeを介してI/Oをステージングし、データが完全に保存された後にのみチェックポイントを完了としてマークします。詳細とコード例については、パフォーマンスとレジリエンスに関するガイドを参照してください。

上記には、torch.saveのネットワークストレージへの保存時間は含まれていません。

影響2：チェックポイントの頻度が復旧コストを決定する

ここで各要素が相乗効果を発揮します。ジョブが失敗すると、最後の有効なチェックポイント以降のすべてのデータが失われ、再計算が必要になります。そのため、障害1回あたりの予想される無駄な作業時間はチェックポイント間隔の約半分になり、低コストの非同期保存（async save）を使用することで、その間隔を短くすることができます。

間隔を10分の1に短縮すると、予想される復旧時間も10分の1に短縮されます。1日あたり約8.6回の中断が発生するというLlama 3の数値を思い出してください。この障害率では、2時間ごとにチェックポイントを作成すると、1日あたり再トレーニングに平均8.6時間を無駄にすることになり、goodputは64%になります。30分ごとにチェックポイントを作成すれば、無駄にする時間はわずか2.15時間になり、goodputは91%に向上します。

また、復旧は自動的に行われなければなりません。再起動時、ジョブは書き込みが完了した最新のチェックポイントを検出し、クラッシュによって書き込みが途中で終了したものをスキップして、人の手を介さずにそこから再開する必要があります。DCPはこれを確実なものにします。.metadataファイルはすべてのシャードが保存された後にのみ書き込まれるため、その存在は「この保存は完了している」という信頼できるマーカーとして選択の基準になります。

影響3：データロードがGPUのアイドル状態の有無を決定する

トレーニングジョブは、最も遅い入力の速度で進行します。アクセラレータが次のバッチを待機している間、GPUは単にアイドル状態になるため、goodputが低下します。この問題を解決する唯一の方法は、以下の図に示すように、入力パイプラインが次のステップのデータ準備と現在のステップの計算を確実にオーバーラップ（並行処理）させることです。

データロードと計算のオーバーラップに移行したお客様では、実時間（ウォールクロック時間）が20〜50%短縮されるケースがよく見られます。

リモートストレージから直接読み取る場合のコスト

ガバナンスが効いたプラットフォームでは、トレーニングデータはリモートオブジェクトストレージに保存されます。AI Runtimeでは、Unity Catalog (UC) ボリュームはネットワークマウントとして表示されます。アクセスのたびにそのマウントから直接ファイルを読み取ると、ステップ時間がネットワークレイテンシに縛られ、エポックごとに同じファイルが再ダウンロードされることになります。解決策は、最初のアクセス時に各ファイルを高速なローカルストレージにコピーし、それ以降の読み取りはそのローカルキャッシュから提供し、GPUが計算している間に次のファイルを並行して取得するデータローダーを使用することです。

AI Runtimeを使用すると、UCVolumeDatasetとDataLoaderがまさにこれを行います（コード例については、ガイドを参照してください）。UCVolumeDatasetはUCボリュームからファイルをストリーミングし、最初のアクセス時に各ファイルをローカルのNVMeにキャッシュします。そして、すべてのアクセラレータが重複のない独立したスライスを取得できるように、ランクとワーカー間でファイルをパーティショニングします。DatabricksのDataLoaderは、PyTorchのDataLoaderのドロップインサブクラスであり、デフォルト設定がこのパス向けにチューニングされています。そのため、トレーニングスレッド上で1つずつ処理するのではなく、GPUが計算を行っている間にファイルが並行して取得およびキャッシュされます。

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

UCボリュームからJPEGをデコードし、拡張（オーグメンテーション）して、ビジョンモデルをトレーニングするという、シンプルな画像分類ワークロードを考えてみましょう。同じGPU、モデル、バッチサイズでこれを行う2つの方法を比較してみます。1つはUCボリュームから読み取る標準のPyTorch Datasetを使用する方法、もう1つはUCVolumeDatasetとDatabricksのDataLoaderのデフォルト設定を組み合わせて使用する方法です。

時間がどこに費やされているか推測する必要はありません

DataLoaderの開発の一環として、メトリックがMLflowに記録されるようにしました。これにより、データパイプラインがトレーニングを妨げているかどうかを一目で簡単に確認できます。

メトリックfetch_secondsは、データローダーがバッチを生成するのにかかる時間を明示的に測定します。この間、GPUはアイドル状態（待機状態）になっています。

影響4：データパイプラインへの配慮を怠ると、気づかないうちにモデルが破損する

エラーメッセージもクラッシュもジョブの失敗も発生せず、ただモデルの精度が本来よりもわずかに悪くなるという、レジリエンス（回復力）に関する最後のバグがあります。これは、モデル、オプティマイザ、ステップのチェックポイントは作成するものの、データセット内におけるデータパイプラインの位置のチェックポイントを保存しない場合に発生します。

エポックの途中で中断されたジョブを考えてみましょう。モデルは正しく復元され、トレーニングループは再開されますが、データローダーはデータセットの最初からやり直してしまいます。

再開されたジョブは、このエポックですでに学習したサンプルを再トレーニングし、まだ到達していなかったサンプルをスキップしてしまう可能性があります。大規模な環境では日常茶飯事となる何度も行われる再起動の中で、これによりデータ分布に気づかないうちに偏り（バイアス）が生じます。モデルのトレーニング自体は継続されますが、データの誤ったサンプリングに基づいてトレーニングされてしまいます。これこそが、最もコストのかかるサイレントな失敗の典型例です。なぜなら、ジョブは完了するため、メトリックを見て落胆するまで誰も問題に気づかないからです。

解決策は、データの位置をチェックポイントの一部として扱うことです。パイプラインに応じて、サンプルまたはシャードのオフセットを追跡して再開時にスキップする、カスタムデータセットに独自の位置をシリアル化させる、またはエポックの境界でチェックポイントを作成することを意味します。これらすべては、1つの前提条件に基づいています。それは、決定論（再現性）です。シャッフルや拡張（オーグメンテーション）は乱数ジェネレータを使用するため、これらのシードとRNGの状態もチェックポイントの一部にする必要があります。そうしないと、再起動後のデータ順序が再起動前の順序と一致しなくなり、保存された位置が誤ったサンプルを指すことになります。

シード、再現可能な順序、そして再開可能なデータパイプラインは、1つのアイデアを3つの側面から表したものです。ガイドでは、コードを交えてそれぞれの戦略を説明しています。

まとめ

高速でフォールトトレラント（耐障害性のある）なトレーニングは、いくつかの決定が相乗効果を生むことで実現します。

1. DDPの場合でも、torch.saveの代わりにDistributed Checkpointを使用することで、保存処理がシリアルなボトルネックにならず、並列かつ低コストで行えるようになります。
2. 非同期で保存し、チェックポイントのコストをほぼゼロにすることで、頻繁な保存が可能になります。
3. 最新の有効なチェックポイントに自動的に回復することで、障害発生時の再計算コストを数時間ではなく数分に抑えられます。
4. リモートストレージからのキャッシュとプリフェッチにより、データロードと計算をオーバーラップ（並行処理）させ、アクセラレータが入力待ちでアイドル状態にならないようにします。これにより、すべてのステップでGPU時間を継続的に節約できます。
5. データパイプラインとRNGの状態のチェックポイントを作成することで、再開されたジョブが、気づかないうちにモデルを破損させることなく、正しいデータで継続できるようにします。

統一された原則：頻繁で低コスト、かつ完全なチェックポイントを作成することで、ハードウェア障害はジョブを終了させる致命的なイベントから、単なる「端数誤差」程度の影響に変わります。また、オーバーラップされた入力パイプラインにより、その間もアクセラレータを稼働させ続けることができます。低コストな（非同期）保存によって頻繁な保存が可能になり、完全な保存（モデル、データ、RNG）によって正しい回復が可能になります。これら両方が整い、さらに障害が発生しているハードウェアを検出して隔離するフリートがあれば、基盤となるクラスターがどれほど不安定であっても、実質的なトレーニング時間はハードウェアが許容する限界値に近づきます。

参考文献

• Meta, The Llama 3 Herd of Models (2024): §3.3.1、最大16,384台のH100 GPUにおけるトレーニングの信頼性と中断の内訳。
• Characterizing the Resilience of Hopper H100 and Ampere A100 GPUs (2025): Deltaシステムの2.5年間にわたるフィールド研究。ジョブサイズ別のMTBEと故障確率。
• Databricks, How we keep GPUs reliable across Databricks AI：フリートレベルのGPUヘルスチェックと確率的故障モデル。

試してみる準備はできましたか？コードの全文については、Databricks AI Runtimeドキュメントのトレーニングのパフォーマンスとレジリエンスに関するガイドを参照してください。また、インフラストラクチャ側の詳細については、How we keep GPUs reliable across Databricks AI（Databricks AI全体でGPUの信頼性を維持する方法）をお読みください。

(このブログ記事はAI翻訳ツールを使用して翻訳されています) 原文記事