Bringing PyTorch Monarch to AMD GPUs: Single-Controller Distributed Training on ROCm
A single-controller distributed training solution on ROCm.
1Challenges of Distributed Training
Training large language models requires distributed systems across thousands of GPUs. At this scale, hardware failures are inevitable and can severely impact the training process.
2The Role of PyTorch Monarch
PyTorch Monarch simplifies large-scale distributed training by enabling GPU cluster management through a single Python program, offering fault tolerance and flexibility.
3Porting to ROCm
Bringing Monarch to AMD GPUs required adapting the GPU runtime and distributed communication stack to ROCm, a significant engineering effort.
4Collective Communications
hipify_torch was used to transition from CUDA to HIP, and integration with RCCL was achieved, facilitating the adaptation of NCCL API to ROCm.
5GPU Memory Management
CUDA driver API calls were routed through their HIP equivalents, enhancing cross-platform compatibility.
6RDMA Integration
The RDMA path was preserved with GPU_PLATFORM=rocm, and GPU-side bindings were switched from CUDA to HIP, supporting GPU-direct transfers.
7Static Linking Issue
ROCm lacks static linking for HIP runtime. This was resolved using dynamic linking, ensuring consistency between CUDA and ROCm.
8Rust Compatibility
HIP symbols were re-exported under CUDA names, allowing Rust code to remain platform-agnostic.
9Tests and Community Contribution
1,171 tests passed successfully, ensuring full support for ROCm 7.0+. These contributions were shared with the open-source community.
10The Future of Monarch
Monarch enables flexible, fault-tolerant distributed training on AMD GPUs, marking a significant step toward large-scale AI infrastructure.
Original article: PyTorch Blog — PyTorch Monarch ile AMD GPU'larda Dağıtılmış Eğitim