PyTorch MonarchをAMD GPUに移植：ROCmで実現する障害に強い分散学習

Bringing PyTorch Monarch to AMD GPUs

大規模なLLM学習では数百〜数千GPUにわたる分散学習が必要ですが、この規模ではハードウェア障害は例外ではなく日常茶飯事です。従来のチェックポイント方式はオーバーヘッドや計算の無駄が大きく、クラスタ全体が停止します。PyTorch Monarchは単一コントローラモデルを採用し、アクターベースのランタイムと監視ツリーで障害を局所化し、動的な復旧を可能にします。AMDはこのMonarchをROCmに移植し、CUDA環境以外でも動作するようにしました。本記事では、hipify_torchによるCUDAからHIPへの変換、RCCLとの統合、RDMA対応など移植の詳細と、SLURMとKubernetes上での障害注入テストの結果を紹介します。障害が発生してもトレーニングは中断されず、損失曲線は安定して収束しました。