Bringing PyTorch Monarch to AMD GPUs

Featured projects

Training state-of-the-art large language models (LLMs) with billions of parameters requires distributed training across hundreds or thousands of GPUs. At this scale, hardware failures are not exceptional events—they are expected. A single GPU memory error, network partition, or node crash can bring down an entire training run that has been progressing for days or weeks. While our previous work demonstrated near-linear scaling of FP8 training at scale (achieving 96.16% scaling efficiency on a 1024-GPU MI325 cluster with DeepSeekV3-671B), the key challenge remains: reliability at scale.

To address these challenges, we have brought PyTorch Monarch to AMD Instinct GPUs with ROCm, expanding the single-controller model beyond CUDA environments and bringing this emerging runtime to a broader hardware ecosystem.

In this blog, we will explore the architecture of PyTorch Monarch, walk through the engineering effort required to port Monarch’s GPU runtime and distributed communication stack to ROCm, and demonstrate how the system dynamically recovers from node failures without halting the entire training job. By the end, you will understand how Monarch enables elastic, fault-tolerant distributed training on AMD GPUs and why this represents a significant step toward stable, large-scale AI infrastructure.

The Challenge: Reliability at Scale

Traditional fault-tolerance strategies rely heavily on periodic checkpointing: saving the full model state to persistent storage at regular intervals. When a failure occurs, the entire job restarts from the last checkpoint. While conceptually simple, this approach has significant drawbacks:

• Checkpoint overhead: Writing hundreds of gigabytes of model state to storage consumes time and I/O bandwidth.
• Wasted computation: All progress since the last checkpoint is lost upon failure.
• Cluster idle time: The entire cluster sits idle while the failed node is replaced and the job restarts.
• Scalability limits: As cluster size grows, the probability of failure during any checkpoint interval increases.

For truly large-scale training, scaling is not enough—training must also recover from failures. We need a more dynamic approach, one that allows healthy nodes to continue training while failed nodes recover and rejoin, minimizing wasted computation and maximizing GPU utilization. This is where PyTorch Monarch comes in.

What is PyTorch Monarch?

PyTorch Monarch introduces a new distributed programming paradigm that enables developers to orchestrate entire GPU clusters from a single Python program. With its actor-based runtime, process mesh abstraction, and asynchronous execution model, Monarch simplifies large-scale distributed training and enables complex workflows that combine training, evaluation, and reinforcement learning within one unified script.

The architecture operates at multiple distinct levels:

• Python API: A single-program interface where developers write simple Python code to get distributed GPU execution.
• Monarch Runtime: Manages actors and meshes, supervision trees, and tensor sharding.
• Rust Runtime (Tokio): Ensures high performance and memory safety.
• Infrastructure: Integrates with RDMA, RCCL/NCCL, SLURM, Kubernetes, and SkyPilot.

By decoupling the parallelism strategy used within each training replica from the fault-tolerance mechanism used across replicas, Monarch provides a cleaner fault-tolerance model. Failures are isolated, hierarchical, and recovery is fast.

Porting Monarch to ROCm: Ecosystem Integration

Bringing Monarch to AMD GPUs required significant engineering effort to port the GPU runtime and distributed communication stack to ROCm.

We successfully implemented three main porting paths:

• Collective Communications: We used hipify_torch to convert the C++ bridge code from CUDA to HIP and linked against RCCL, which mirrors NCCL’s API.
• GPU Memory Management: We extended the build system to auto-detect the platform and route CUDA driver API calls through their HIP equivalents.
• RDMA Integration: Configuring GPU_PLATFORM=rocm keeps the libibverbs-based RDMA path intact while swapping the GPU-side bindings from CUDA to HIP for GPU-direct transfers.

Moreover, two cross-cutting issues shaped the port and deserve a closer look:

• No static link for the HIP runtime: NVIDIA ships libcudart_static.a, so the CUDA path links cudart_static directly. ROCm ships no static equivalent for libamdhip64, so the ROCm build links amdhip64 dynamically.
• Rust compatibility shim instead of forking the bindings: We added a rocm_compat module in nccl-sys and rdmaxcel-sys that re-exports HIP symbols under their CUDA names (e.g. pub type cudaError_t = hipError_t).

These efforts culminated in full support for ROCm 7.0+ with all 1,171 tests passing.

Case Study: Fault-Tolerant Training at Scale

To demonstrate the power of Monarch on AMD GPUs, we integrated it with TorchTitan and TorchFT to build a resilient, checkpoint-less distributed training architecture.

Architecture Overview

• Monarch: Acts as the orchestrator, managing process and cluster orchestration.
• TorchFT: Handles fault tolerance at the step level, performing Quorum AllReduce and skipping failed nodes.
• TorchTitan: Serves as the training engine, executing Forward, Backward, and Optimizer steps.