PyTorch Monarch を AMD GPU に導入：ROCm 上での単一コントローラー分散トレーニング

Featured projects

• 数十億パラメータを持つ最先端の大規模言語モデル（LLM）を学習するには、数百から数千の GPU にわたる分散トレーニングが必要です。このスケールでは、ハードウェア障害は例外ではなく当然のこととして想定されます。GPU のメモリエラーが一つ発生しただけでも、ネットワークの分断やノードのクラッシュさえあれば、数日あるいは数週間かけて進んでいた学習全体が停止してしまいます。

私たちの 以前の成果 では、FP8 学習のスケーラビリティがほぼ線形に拡張されることを示しました（DeepSeekV3-671B モデルを用いた 1024 GPU の MI325 クラスターで 96.16% のスケーリング効率を達成）。しかし、依然として最大の課題は「スケールにおける信頼性」です。

これらの課題に対処するため、私たちは PyTorch Monarch を AMD Instinct GPU と ROCm に持ち込みました。これにより、単一コントローラーモデルが CUDA 環境に限定されることなく、この新興ランタイムをより広範なハードウェアエコシステムへと拡張しました。

本ブログでは、PyTorch Monarch のアーキテクチャを紹介し、Monarch の GPU ランタイムと分散通信スタックを ROCm に移植するために必要なエンジニアリングの取り組みについて解説します。さらに、学習ジョブ全体を停止させることなくノード障害から動的に回復する仕組みも実演します。読み終えた頃には、Monarch が AMD GPU 上でどのように弾力的で耐故障性の高い分散トレーニングを実現し、なぜこれが安定した大規模 AI インフラへの重要な一歩となるのかを理解いただけるでしょう。

The Challenge: Reliability at Scale

従来のフォールトトレランス戦略は、定期的なチェックポイントに大きく依存しています。これは、一定間隔でモデルの状態全体を永続ストレージに保存する手法です。障害が発生すると、ジョブは最後のチェックポイントから完全に再スタートします。概念的にはシンプルですが、このアプローチには重大な欠点があります。

大規模なトレーニングにおいて、単にスケーラビリティを高めるだけでは不十分です。失敗から回復できる仕組みも不可欠であり、健全なノードが計算を継続しながら、障害を起こしたノードが復旧して再参加できるような動的なアプローチが必要です。これにより、無駄な計算を最小限に抑え、GPU の利用率を最大化できます。ここで登場するのが PyTorch Monarch です。

PyTorch Monarch とは何か？

PyTorch Monarch は、開発者が単一の Python プログラムから GPU クラスタ全体をオーケストレーションできる新しい分散プログラミングのパラダイムを導入します。アクターベースのランタイム、プロセスメッシュ抽象化、非同期実行モデルを備えることで、大規模な分散トレーニングが簡素化され、トレーニング、評価、強化学習を組み合わせた複雑なワークフローも単一の統合スクリプト内で実現可能になります。

このアーキテクチャは、複数の異なるレベルで動作します：

• Python API: 開発者が単純な Python コードを書くだけで分散 GPU 実行が得られる、単一プログラムインターフェースです。
• Monarch Runtime: アクターやメッシュの管理、監督ツリー、テンソルシャードを担当します。
• Rust Runtime (Tokio): 高性能とメモリの安全性を保証します。
• インフラストラクチャ: RDMA、RCCL/NCCL、SLURM、Kubernetes、SkyPilot と統合されます。

図 1: PyTorch Monarch のアーキテクチャは、Python API を Rust ランタイムおよびインフラストラクチャから分離しています。

各トレーニングレプリカ内で使用する並列化戦略と、レプリカ間で使用するフォールトトレランス機構を分離することで、Monarch はよりクリーンなフォールトトレランスモデルを実現しています。障害は隔離されます（アクターはプライベート状態を持ち、クラッシュが伝播しない）、階層的です（可能な限り低いレベルで処理される）、そして回復も高速です（ローカル再起動なら数秒、エスカレーションが必要な場合でも数分）。

ROCm への移植：エコシステムとの統合

Monarch を AMD GPU で動かすには、GPU ランタイムと分散通信スタックを ROCm に移植するために多大なエンジニアリング effort が必要でした。

私たちは主に以下の 3 つの移植パスを成功させました。

• 集合通信: hipify_torch を使用して CUDA から HIP へ C++ ブリッジコードを変換し、NCCL の API と同等の RCCL にリンクしました。
• GPU メモリ管理: ビルドシステムを拡張してプラットフォームを自動検出し、CUDA ドライバー API の呼び出しを対応する HIP 経由にルーティングします。
• RDMA 統合: GPU_PLATFORM=rocm を設定することで、libibverbs ベースの RDMA パスを維持しつつ、GPU-Direct 転送用の GPU サイドバインディングを CUDA から HIP に切り替えます。

さらに、ポート化において二つの横断的な課題が影響しており、これらにはより深い検討が必要です。

• HIP ランタイムの静的リンクは不可：NVIDIA は libcudart_static.a を同梱しており、CUDA パスでは cudart_static が直接リンクされる。

ROCm には libamdhip64 の静的版が用意されていないため、ROCm ビルドでは amdhip64 を動的にリンクします。

両プラットフォームとも、GPU ドライバの API 関数（hipMemCreate、cuMemCreate、および関連する呼び出しなど）を dlopen で読み込み、いずれの側でもランタイム契約が同一に保たれています。

バインディングのフォークではなく、Rust 互換性の shim を用意する：hipify_torch が C/C++ ヘッダーを書き換えた後、bindgen は HIP に準拠した型を出力します。

hipError_t、hipDeviceptr_t、そして hipStream_t。

Rust の各呼び出し箇所に #ifdef ブランチを追加するのではなく、nccl-sys と rdmaxcel-sys に rocm_compat モジュールを追加し、HIP シンボルを CUDA 名義で再エクスポートしました。

例として、pub type cudaError_t = hipError_t や pub use hipSetDevice as cudaSetDevice があります。

残りのRustコードはプラットフォーム非依存のままです。

これらの取り組みは、Rust における HIP タイプエイリアスの導入へと結実し、1,171 件のテストをすべてパスさせることで、ROCm 7.0 以降の完全なサポートを実現しました。また、これらの貢献はオープンソースコミュニティへアップストリームされ（PR #2393 および PR #2891 を参照）

本日、ROCm 上の Monarch は Actor ランタイム、RDMA、Supervision、Tensor シャーディングを含むフルエコシステムサポートを提供します。これにより、SLURM（HPC）、Kubernetes（クラウドネイティブ）、SkyPilot（マルチクラウド）上でシームレスに動作し、TorchTitan（トレーニングエンジン）や TorchFT（フォールトトレランス）などの下流エンジンが本番環境のワークロードで利用可能になります。

ケーススタディ：大規模な耐障害性トレーニング

AMD GPU 上の Monarch の能力を示すため、TorchTitan と TorchFT を統合し、チェックポイントレスの堅牢な分散トレーニングアーキテクチャを構築しました。

アーキテクチャ概要

このアーキテクチャは以下の 3 つのレイヤーで構成されています。

• Monarch: オーケストレーターとして機能し、プロセスとクラスターの調整を担当します。ReplicaActor と Lighthouse サービスを起動し、GPU を Process Mesh に整理します。

• TorchFT: ステップレベルでの耐障害性を処理します。Lighthouse へ連絡してクォーラム調整を行い、Quorum AllReduce を実行して失敗したノードをスキップします。

• TorchTitan: トレーニングエンジンとして機能し、Forward（FSDP）、Backward、Optimizer の各ステップを実行するとともに、チェックポイントとメトリクスの管理を行います。

この構成では、Monarch が微細な障害検出と隔離のための監視ツリーを提供します。トレーニングアクターに障害を注入すると、Lighthouse によって検知され、TorchFT によって処理されます。他の正常なレプリカは、ピアの障害が発生してもグローバルな中断を必要とせず、独立してトレーニングを継続します。

Dynamic Fault Recovery Workflow

4 つのレプリカグループを持つ具体的なシナリオを通じて、回復ワークフローを理解してみましょう。

• 通常トレーニング: OrchestrationManager は 4 つの ReplicaActor（Monarch スーパーバイザー）と Lighthouse を起動します。各 ReplicaActor は、8 つの GPU プロセスで TorchTitan トレーナーを実行する Replica を生成します。4 つのレプリカすべてが準備完了し（quorum_id=1）、DiLoCo による勾配同期は 20 ステップごとに行われます。

• 障害検出: レプリカ 0 の GPU プロセスがクラッシュしました。Monarch スーパーバイザーは、プロセスが終了する前に report_training_error（完全なスタックトレース付き）をキャプチャします。レプリカ 1、2、3 は影響を受けないと判断され、トレーニングを継続します。

• ローカル再起動: ReplicaActor 0 がインプレイス再起動（_stop_and_restart()）を開始し、古いプロセスメッシュを停止して新しいものを生成します。一方、他の 3 つのレプリカは同期を継続します（quorum_id=2）。

ピアチェックポイント転送：ライトハウスは Replica 1 をドナーとして選択します。Replica 1 から回復中の Replica 0 へ、モデル、オプティマイザ、スケジューラ、トレーナの状態を含むピアチェックポイント転送が開始されます。新しいクォーラムが形成される間、すべてのレプリカはクォーラム境界で一時的に停止します。

トレーニングの再開：Replica 0 の同期が完了すると、全 4 つのレプリカを備えた新しいクォーラム（quorum_id=3）が確立され、DiLoCo の同期が再開されます。

図 5：グローバルチェックポイントの再読み込みを伴わないピアチェックポイント転送を示す動的障害回復ワークフロー。

この回復プロセス全体は、手動介入なしで、フルチェックポイントからの再起動なしに完了し、全体のトレーニングスループットへの影響も最小限に抑えられます。

パフォーマンス特性

AMD Instinct MI300 クラスターを使用した SLURM および Kubernetes 環境の両方で、このアプローチを検証しました。

SLURM 16 ノード MI300 クラスタ（128 GPU）

16 ノード構成の SLURM クラスタ（MI300 GPU 合計 128 基）で Llama 3 8B モデルをトレーニングしました。RCCL の障害を 180 秒ごとに注入し、20 ステップごとにクォーラム同期を行いました。その結果は目覚ましいものでした。

• 注入された障害により、アクティブなワーカー数は動的に変動（8〜16 基の間）しました。
• 頻繁な障害が発生してもトレーニングはシームレスに継続し、フルリスタートは行われませんでした。
• ロス曲線は着実に収束し、障害注入なしのベースライン実行とほぼ同じ結果を示しました。

Figure 6: SLURM MI300 クラスターで頻繁な障害が発生してもトレーニングは継続します。1 つのレプリカも 30 分以上停止することはなく、回復は迅速でした。また、異なるレプリカが動的に終了・回復されています。

Kubernetes 32 ノード MI355 クラスター（GPU 256 基）

実験を 32 ノードの Kubernetes クラスター（MI355 GPU を合計 256 基使用）に拡張しました。参加ノード数は非常に安定しており、障害回復イベント中も 30〜32 の間でわずかに変動するのみでした。また、グローバル平均損失は 12 から約 4 まで滑らかに減少しています。これは、Monarch のフォールトトレラントモデルが、大規模な SLURM および Kubernetes の両方で確実に機能することを示しています。

Figure 7: 32 ノードの Kubernetes MI355 クラスターにおける安定した回復と、滑らかな損失の収束。

まとめと今後の方向性

大規模な AI モデルをスケールしてトレーニングするには、単に計算能力が優れているだけでは不十分です。避けられないハードウェア障害を優雅に処理できる、堅牢なインフラストラクチャが必要です。PyTorch Monarch を ROCm 上の AMD Instinct GPU に導入したことで、無駄な計算を最小限に抑え、GPU の利用率を最大化する実用的なフォールトトレラント分散トレーニングのアプローチを実証しました。

この統合は、AMD GPU 上での大規模トレーニングにおいて以下の重要な成果をもたらします：

AMD ハードウェアにおける大規模な検証を初めて実施しました。Monarch を TorchTitan および TorchFT と共に AMD GPU で正常にデプロイし、ROCm ソフトウェアスタックが高度な障害耐性メカニズムを完全にサポートしていることを実証しました。

Monarch は、堅牢な監視ツリーとプロセスメッシュの抽象化を提供する、よりクリーンな障害耐性モデルです。これにより障害が隔離され、迅速なローカル回復が可能になります。

このアプローチは SLURM と Kubernetes の両方でシームレスに動作するため、本番環境でのワークロードにも即座に対応可能です。

アーキテクチャの鍵となる洞察は、Monarch のアクターベースランタイムと監視ツリーを活用して障害を隔離し、TorchFT のクォーラムに基づく同期によって健全なノードがトレーニングを継続できるようにすることです。AMD GPU 上で大規模なトレーニングワークロードを実行するチームにとって、この統合により、より安定で効率的かつコスト効果の高いモデル開発への道が開かれます。

今後の予定は以下の通りです。

• NIC サポートの拡張とランタイムパフォーマンスの向上
• ROCm 上での事前学習や強化学習（RL）フレームワークへの対応拡大
• 障害耐性パフォーマンスのさらなる最適化、特に再参加時のリロードレイテンシ短縮と計算との回復処理の重なり
• PyTorch コミュニティとのオープンソース協力の継続

追加リソース

• PyTorch Monarch GitHub リポジトリ

• TorchTitan ドキュメント

TorchFT の GitHub リポジトリ

https://github.com/pytorch/torchft

AMD GPU 上で TorchTitan と TorchFT を統合した、大規模なトレーニングの耐障害性向上

https://rocm.blogs.amd.com/artificial-intelligence/primus-torchft/README.html

免責事項

サードパーティ製のコンテンツは、その所有権を持つ第三者から直接ライセンス供与されるものであり、AMD からのライセンスではありません。リンク先のすべてのサードパーティ製コンテンツは「現状有姿」で提供され、いかなる保証もありません。これらのコンテンツの利用はご自身の判断で行うものとし、いかなる場合においても AMD はサードパーティ製コンテンツについてあなたに対して一切の責任を負いません。サードパーティ製コンテンツの使用に伴って生じる可能性のあるすべてのリスクをあなたが負い、その利用による損害についてもあなたが単独で責任を負います。

© 2026 Advanced Micro Devices, Inc.