Extending PyTorch Monarch to AMD GPUs

To address the hardware failure issues that occur during large-scale LLM training, PyTorch Monarch has been ported to AMD Instinct GPUs and the ROCm environment. The existing checkpoint-based approach has drawbacks such as heavy storage overhead and wasted computation.

PyTorch Monarch offers the following features:

• Actor-based runtime: Orchestrates a GPU cluster from a single Python program, and when a failure occurs, it quickly recovers only the affected node without halting the entire training run.
• Hierarchical fault-tolerance: Isolates and handles failures at a low level, enabling local restarts within seconds and recovery within minutes even in cases of widespread failure.
• Architecture: Consists of the Python API, the Monarch runtime, a Rust runtime (Tokio) that ensures performance and memory safety, and an infrastructure layer including RDMA/RCCL.

With this extension, developers can build resilient and stable large-scale distributed training infrastructure across a broader hardware ecosystem, beyond the CUDA environment.