Expanding PyTorch Monarch to the AMD ROCm Ecosystem

By AMD: Chaojun Hou, Liz Li, Zachary Streeter, Xinyu Kang, Lei Zhang, Yuankai Chen, Yao Fu, Wen Chen, Zhenyu Gu, Andy Luo │ Meta: Matthias Reso, Hamid Shojanazeri, Monarch team Date: July 6, 2026

Training state-of-the-art Large Language Models (LLMs) with billions of parameters necessitates distributed training across thousands of GPUs. At this immense scale, hardware failures are ~~rare accidents~~ expected certainties.

A single node crash, a network partition, or a GPU memory error can instantly terminate a training run that has been active for weeks. While previous benchmarks showed impressive scaling—specifically achieving a scaling efficiency of $\text{Efficiency} = 96.16\%$ on a 1024-GPU MI325 cluster using DeepSeekV3-671B—the primary hurdle remains reliability at scale.

To solve this, PyTorch Monarch has been ported to AMD Instinct GPUs via ROCm. This move extends the single-controller model beyond the CUDA-only environment, bringing a sophisticated runtime to a more diverse hardware landscape.

Goal: Enable elastic, fault-tolerant distributed training on AMD GPUs to create stable, large-scale AI infrastructure.

The Challenge: Reliability at Scale

Historically, fault tolerance has relied on periodic checkpointing, where the entire model state is saved to persistent storage. If a failure occurs, the system restarts from the most recent save.

The Drawbacks of Traditional Checkpointing

Challenge	Impact
Checkpoint Overhead	Massive I/O bandwidth and time consumed writing hundreds of GBs.
Wasted Computation	All progress made between the last checkpoint and the crash is lost.
Cluster Idle Time	All GPUs sit idle while the failed node is replaced and the job re-initializes.
Scalability Limits	As the cluster grows, the statistical likelihood of a failure during a checkpoint interval increases.

For massive clusters, simply scaling hardware is insufficient; the system must be able to recover dynamically. We need a paradigm where healthy nodes persist while failed ones recover and rejoin.

Introducing PyTorch Monarch

PyTorch Monarch is a distributed programming paradigm allowing a single Python program to orchestrate an entire GPU cluster. It utilizes an actor-based runtime, a process mesh abstraction, and asynchronous execution to unify training, evaluation, and reinforcement learning.

System Architecture

The architecture is decoupled into several layers:

Python API: The high-level interface for distributed GPU execution.
Monarch Runtime: Handles tensor sharding, supervision trees, and actor/mesh management.
Rust Runtime (Tokio): Provides the foundation for memory safety and high performance.
Infrastructure: The bottom layer integrating with SLURM, Kubernetes, SkyPilot, RDMA, and RCCL/NCCL.

Figure 1: PyTorch Monarch architecture decoupling Python API from the Rust runtime and infrastructure.

Fault-Tolerance Model

By separating the parallelism strategy (within replicas) from the fault-tolerance mechanism (across replicas), Monarch ensures:

Isolation: Actors maintain private states; crashes do not propagate.
Hierarchy: Failures are managed at the lowest possible level.
Rapid Recovery: Local restarts take seconds; escalation takes minutes.

Fault Handling Figure 2: Monarch's hierarchical fault-handling model and supervision tree.

Porting Monarch to ROCm: Ecosystem Integration

Bringing Monarch to AMD GPUs required a deep engineering dive into the GPU runtime and communication stacks.

Three Primary Porting Paths

Collective Communications: The team utilized hipify_torch to convert C++ bridge code from CUDA to HIP, linking it against RCCL (which maintains an API similar to NCCL).
GPU Memory Management: The build system was updated to auto-detect the platform, routing CUDA driver API calls to their HIP equivalents.
RDMA Integration: By setting GPU_PLATFORM=rocm, the libibverbs-based RDMA path remains active while swapping CUDA bindings for HIP to enable GPU-direct transfers.

Porting Process Figure 3: Porting Monarch from CUDA to ROCm via hipify_torch and auto-detection.

Engineering Deep-Dives

1. Linking Strategies

A key difference was discovered in how runtimes are shipped:

NVIDIA: Ships libcudart_static.a, allowing for static linking.
AMD: Does not provide a static equivalent for libamdhip64, necessitating dynamic linking.

Despite this, both platforms use dlopen for GPU driver API functions (e.g., hipMemCreate vs cuMemCreate), ensuring the runtime contract remains consistent.

2. The Rust Compatibility Shim

To avoid cluttering the Rust codebase with #ifdef branches after bindgen emitted HIP-specific types (like hipError_t or hipStream_t), the team implemented a rocm_compat module within nccl-sys and rdmaxcel-sys.

This shim re-exports HIP symbols under CUDA names:

// Example of the rocm_compat shim logic
pub type cudaError_t = hipError_t;
pub use hipSetDevice as cudaSetDevice;
pub type cudaDeviceptr_t = hipDeviceptr_t;

These efforts successfully integrated HIP type aliases into the Rust environment, completing the port to ROCm.

Bringing PyTorch Monarch to AMD GPUs