PyTorch Monarch brings fault-tolerant training to AMD GPUs

Training large language models across hundreds or thousands of GPUs makes hardware failures inevitable. A single GPU memory error, network partition, or node crash can otherwise erase days or weeks of work.

PyTorch says it has brought PyTorch Monarch to AMD Instinct GPUs running ROCm, extending the runtime’s single-controller model beyond CUDA. The project is designed to let healthy workers continue training while failed nodes recover and rejoin, rather than forcing a full job restart from a checkpoint.

How Monarch handles distributed training

Monarch lets developers orchestrate an entire GPU cluster from a single Python program. Its actor-based runtime, process-mesh abstraction, and asynchronous execution model support workflows that combine training, evaluation, and reinforcement learning in one script.

Recommended reading

The AI jobs apocalypse may be farther away

The architecture separates the parallelism strategy inside each training replica from fault tolerance across replicas. Actors hold private state, failures are isolated, and supervision trees handle errors at the lowest possible level. Local restarts can complete in seconds, while escalated recovery takes minutes.

The stack includes a Python API, a Monarch runtime for actors, meshes, supervision trees, and tensor sharding, and a memory-safe Rust runtime built on Tokio. It integrates with RDMA, RCCL/NCCL, SLURM, Kubernetes, and SkyPilot.

Porting Monarch from CUDA to ROCm

The ROCm port required changes to GPU execution and distributed communications:

• Collective communications: PyTorch used hipify_torch to convert the C++ bridge from CUDA to HIP, then linked it against RCCL, which mirrors NCCL’s API.
• GPU memory management: The build system now detects the platform and routes CUDA driver calls through their HIP equivalents.
• RDMA: Setting GPU_PLATFORM=rocm preserves the libibverbs-based RDMA path while changing the GPU bindings from CUDA to HIP for GPU-direct transfers.

×

ROCm does not provide a static equivalent to NVIDIA’s libcudart_static.a, so the ROCm build links amdhip64 dynamically. Both platforms still load GPU driver functions with dlopen, keeping the runtime contract consistent.

The project also added a Rust compatibility layer instead of forking bindings. HIP types such as hipError_t and hipStream_t are re-exported under CUDA names, allowing the rest of the Rust code to remain platform-agnostic. PyTorch reports that all 1,171 tests passed, supporting ROCm 7.0+. The changes were upstreamed through PR #2393 and PR #2891.

Monarch on ROCm now supports the Actor runtime, RDMA, Supervision, and Tensor sharding across SLURM, Kubernetes, and SkyPilot, with downstream support for TorchTitan and TorchFT.

Fault recovery with TorchTitan and TorchFT

For its demonstration, PyTorch combined Monarch, TorchFT, and TorchTitan into a checkpoint-less distributed training setup:

• Monarch orchestrates the cluster, spawning ReplicaActors and a Lighthouse service.
• TorchFT manages step-level fault tolerance, quorum coordination, Quorum AllReduce, and failed-node skipping.
• TorchTitan runs the forward, backward, and optimizer steps while handling checkpoints and metrics.

×

In a test with four replica groups, each ReplicaActor launched a replica with eight GPU processes. When a GPU process in Replica 0 failed, Monarch isolated the error while Replicas 1, 2, and 3 continued training. Replica 0 restarted in place, and the Lighthouse selected Replica 1 to transfer the model, optimizer, scheduler, and trainer state.

After a brief pause at the quorum boundary, all four replicas rejoined and DiLoCo synchronization resumed. No manual intervention or global checkpoint reload was required.

×

Results on MI300 and MI355 clusters

On a 16-node SLURM cluster with 128 MI300 GPUs, PyTorch trained a Llama 3 8B model while injecting RCCL failures every 180 seconds and synchronizing the quorum every 20 steps. Active workers varied between 8 and 16, but training continued without a full restart and the loss curve closely matched a failure-free baseline.

A second test used a 32-node Kubernetes cluster with 256 MI355 GPUs. Participation stayed between 30 and 32 nodes, while global average loss declined smoothly from 12 to approximately 4.

PyTorch’s next steps include extending NIC support, improving runtime performance, adding more pre-training and reinforcement-learning frameworks on ROCm, reducing rejoin reload latency, overlapping recovery with computation, and continuing open-source work with the PyTorch community.