PyTorch MonarchをAMD GPUに導入する

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インフラストラクチャに向けた重要なステップをなぜ表すのかを理解できるようになるでしょう。 課題：大規模での信頼性 従来の耐障害性戦略は、定期的なチェックポインティングに大きく依存しています。これは、モデルの状態全体を定期的な間隔で永続ストレージに保存することです。障害が発生した場合、ジョブ全体が最後のチェックポイントから再開されます。概念的には単純ですが、このアプローチには重大な欠点があります。 課題の影響 チェックポイントのオーバーヘッド: 数百ギガバイトのモデル状態をストレージに書き込むことは、時間とI/O帯域幅を消費します。 無駄な計算: チェックポイント以降のすべての進捗は、障害時に失われます。 クラスタのアイドル時間: 障害が発生したノードが交換され、ジョブが再開される間、クラスタ全体がアイドル状態になります。 スケーラビリティの限界: クラスタサイズが大きくなるにつれて、チェックポイント間隔中に障害が発生する確率は増加します。真に大規模なトレーニングでは、スケーリングだけでは不十分です。トレーニングは障害から回復する必要もあります。 私たちは、より動的なアプローチを必要としています。これは、正常なノードがトレーニングを継続し、障害が発生したノードが回復して再参加できるようなアプローチであり、無駄な計算を最小限に抑え、GPU利用率を最大化します。ここでPyTorch Monarchが登場します。 PyTorch Monarchとは？ PyTorch Monarchは、開発者が単一のPythonプログラムからGPUクラスタ全体をオーケストレーションできる、新しい分散プログラミングパラダイムを導入します。そのアクターベースのランタイム、プロセスメッシュ抽象化、非同期実行モデルにより、Monarchは大規模分散トレーニングを簡素化し、トレーニング、評価、強化学習を1つの統合スクリプト内で組み合わせる複雑なワークフローを可能にします。 アーキテクチャは、複数の異なるレベルで動作します。 Python API: 開発者がシンプルなPythonコードを記述して分散GPU実行を取得する単一プログラムインターフェース。 Monarch Runtime: アクターとメッシュ、スーパービジョンツリー、テンソルシャーディングを管理します。 Rust Runtime (Tokio): 高性能とメモリ安全性を保証します。 インフラストラクチャ: RDMA、RCCL/NCCL、SLURM、Kubernetes、SkyPilotと統合します。 図1: Python APIをRustランタイムとインフラストラクチャから分離したPyTorch Monarchアーキテクチャ。 各トレーニングレプリカ内で使用される並列戦略を、レプリカ間の耐障害性メカニズムから分離することにより、Monarchはよりクリーンな耐障害性モデルを提供します。障害は分離され（アクターはプライベート状態を持ち、クラッシュは伝播しません）、階層的であり（可能な限り低いレベルで処理されます）、回復は高速です（ローカル再起動には数秒、エスカレーションした場合でも数分）。 図2: Monarchの階層的な障害処理モデルとスーパービジョンツリー。 ROCmへのMonarchの移植：エコシステム統合 AMD GPUにMonarchを導入するには、GPUランタイムと分散通信スタックをROCmに移植するために、かなりのエンジニアリング作業が必要でした。私たちは3つの主要な移植パスを成功裏に実装しました。 コレクティブ通信：hipify_torchを使用して、C++ブリッジコードをCUDAからHIPに変換し、RCCL（NCCLのAPIをミラーリング）にリンクしました。 GPUメモリ管理：ビルドシステムを拡張してプラットフォームを自動検出し、CUDAドライバーAPI呼び出しをHIPの同等物にルーティングしました。 RDMA統合：GPU_PLATFORM=rocmを設定することで、libibverbsベースのRDMAパスをそのまま維持しつつ、GPU側のバインディングをCUDAからHIPに切り替えてGPUダイレクト転送を行いました。 図3: hipify_torchと自動検出によるCUDAからROCmへのMonarchの移植。 さらに、2つの横断的な問題が移植に影響を与え、詳しく検討する価値があります。 HIPランタイムの静的リンクなし：NVIDIAはlibcudart_static.aを配布しているため、CUDAパスはcudart_staticを直接リンクします。ROCmはlibamdhip64の静的同等物を配布していないため、ROCmビルドはamdhip64を動的にリンクします。両方のプラットフォームは、hipMemCreate、cuMemCreate、および関連する呼び出しを含むGPUドライバーAPI関数をさらにdlopenし、両側でランタイム契約を同一に保ちます。 フォークする代わりにRust互換性シムを使用：hipify_torchがC/C++ヘッダーを書き換えると、bindgenはhipError_t、hipDeviceptr_t、hipStream_tなどのHIP名前付き型を出力します。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は、アクターランタイム、RDMA、スーパービジョン、テンソルシャーディングを含む完全なエコシステムサポートを提供します。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ステップを実行しながら、チェックポイントとメトリクスを管理します。 図4: Monarch、TorchFT、TorchTitanを統合したAMD GPU上の回復力のあるトレーニングスタック。 このセットアップでは、Monarchは細かい障害検出と分離のためのスーパービジョンツリーを提供します。トレーニングアクターに障害が注入されると、Lighthouseによって検出され、TorchFTによって処理されます。正常なレプリカは、グローバルな中断を必要とせずに、ピアの障害にもかかわらず独立してトレーニングを継続します。 動的な障害回復ワークフロー 4つのレプリカグループでの回復ワークフローを理解するために、具体的なシナリオを説明します。 通常のトレーニング: OrchestrationManagerは4つのReplicaActor（Monarch Supervisor）と1つのLighthouseを起動します。各ReplicaActorは、TorchTitanトレーナーを実行する8 GPUプロセスのReplicaを起動します。すべての4つのレプリカが準備完了（quorum_id=1）であり、DiLoCo勾配同期は20ステップごとに発生します。 障害検出: Replica 0のGPUプロセスがクラッシュします。Monarchスーパーバイザーは、プロセスが終了する前にreport_training_error（完全なトレースバック付き）をキャプチャします。レプリカ1、2、3は影響を受けないとマークされ、トレーニングを継続します。 ローカル再起動: ReplicaActor 0は、インプレース再起動（_s