打通 PyTorch Monarch 与 ROCm：单 Controller 架构的异构算力实战

基础算力前言：

千亿级参数大模型的训练，本质上是几百上千张 GPU 的分布式协同。在这个量级下，我们必须接受一个现实：硬件故障是常态。一张卡的内存错误、一次网络分区或者一个节点的宕机，都足以让持续数天的训练任务直接中断。

把线性扩展效率刷到了 96.16%，算力确实拉满了。但跑得快是一回事，跑得稳是另一回事。现在真正的硬骨头在于：千卡规模下怎么保证系统"不翻车"（可靠性）。

最近一直在折腾 PyTorch Monarch 的跨平台适配，总算在 ROCm 上把它调通了。这相当于把 Monarch 的单控制器模型从 CUDA 舒适区拽了出来，强行安在了 AMD Instinct 上。

过程中踩了不少坑，主要集中在运行时适配和 RCCL 通信层。文章后半部分我会重点讲讲我们是怎么搞定"弹性容错"的------也就是某个节点挂了，整个训练作业还能接着跑。这活儿干完，感觉离稳定的大规模训练基础设施又近了一步。

1：分布式训练的"底线"：规模化容错攻坚

传统的容错套路基本都押宝在周期性 Checkpoint 上：定期把全量模型状态落盘。遇到故障时，整个作业直接回滚到最后一个 Checkpoint 重跑。原理大家都懂，确实简单，但在大模型训练这种场景下，这种方案的代价和瓶颈太明显了。

Checkpoint 开销黑洞：模型状态的落盘（数百 GB 级别）是 I/O 密集型操作，极易造成存储带宽拥堵。
计算资源浪费：故障导致"断崖式"回退，上次 Checkpoint 后的所有迭代成果付诸东流。
集群利用率坍缩：单点故障引发系统性停摆，故障节点替换与作业重启期间，整体算力处于闲置状态。
可扩展性陷阱：随着集群规模的线性扩张，任意 Checkpoint 窗口期内的故障概率呈指数级上升。

在万卡集群里，单纯的横向扩展只是及格线，真正的考验在于容错韧性 。理想的运行时应当支持故障隔离与自愈：健康节点无视异常持续迭代，故障节点修复后动态重连。

这种机制的核心目的是压缩非计算开销，把 GPU 利用率钉在高位。这也是我们引入 PyTorch Monarch 的关键原因------它的架构设计天生就是为了应对这种混乱的分布式环境。

2：浅析 PyTorch Monarch

Monarch 其实是在分布式编程上换了一套思路：它允许你只用一个 Python 进程就把整个 GPU 集群捏在手里。

这背后的核心是一套基于 Actor 的运行时。它通过 Process Mesh（进程网格）把物理资源抽象成一个逻辑整体，配合异步执行引擎，把原本散落在各处的训练、评估和 RL 逻辑，全部揉进了一个统一的脚本里。换句话说，你写的不再是单机脚本，而是集群级的编排指令

整个架构从上到下分成了四层，各司其职：

Python API 层：主打极简封装。开发者只需写几行常规的 Python 代码，就能直接拉起分布式 GPU 任务，完全屏蔽了底层的复杂性。
Monarch Runtime（核心运行时）：调度的"大脑"。负责管理 Actor（执行体）、Mesh 拓扑、监督树（Supervision Tree）以及张量分片（Tensor Sharding）。
Rust 运行时（基于 Tokio）：底层性能与安全的兜底。用 Rust 和 Tokio 异步框架构建，既榨干了并发性能，又保证了内存安全，彻底告别 C++ 时代的内存泄漏噩梦。
基础设施层（Infra）：底层直接打通 RDMA，无缝对接 RCCL/NCCL，并且原生适配 SLURM、K8s 和 SkyPilot 等主流集群调度平台。

2.1剥离控制平面：PyTorch Monarch 的 Python-Rust 异构运行时设计

Monarch 的容错架构主打一个"解耦"：将副本内部的并行策略与副本间的容错机制彻底拆分，让整体模型更加清晰。其核心容错机制体现在三个维度：

故障严格隔离：每个 Actor 维护私有状态，单点崩溃被死死限制在局部，绝不会引发雪崩效应（故障不扩散）。
分级容错处理：故障在尽可能低的层级被拦截和消化，不随意向上层抛异常。
极速恢复机制：本地 Actor 重启只需秒级响应；只有当本地无法解决时，才会触发分钟级的故障上报（Escalation）流程。

2.2 容错基石：Monarch 的监督树与分层熔断机制

Monarch 的 ROCm 适配与底层改造

将 Monarch 引入 AMD GPU 平台是一项硬核的底层工程。核心难点在于对 GPU 运行时和分布式通信栈进行深度改造，使其能够无缝跑在 ROCm 之上。

在底层适配层面，我们重点完成了以下三条核心链路的移植：

集合通信适配 ：利用 hipify_torch 将 C++ 桥接代码由 CUDA 转换为 HIP，并链接至 RCCL。由于 RCCL 与 NCCL 的 API 高度一致，这保证了集合通信的平滑过渡。
显存管理改造：扩展了原有的构建系统，增加平台自动检测能力。在运行时，将 CUDA Driver API 的调用请求，统一映射并替换为等效的 HIP 接口。
RDMA 深度集成 ：通过配置 GPU_PLATFORM=rocm 编译选项，维持了基于 libibverbs 的 RDMA 通信路径。同时，将 GPU 端的底层绑定（Bindings）从 CUDA 切换至 HIP，以支持 GPU Direct 显存直传。

2.3异构算力互通：Monarch 的 HIP 化改造与平台自适应加载

在底层移植时，有两个跨领域的关键问题直接影响了架构设计：

HIP 运行时的静态链接问题 ：NVIDIA 提供 libcudart_static.a，允许 CUDA 路径直接静态链接。而 ROCm 缺乏 libamdhip64 的静态等效库，因此 ROCm 构建只能采用动态链接。为了保证一致性，两个平台均通过 dlopen 动态加载 GPU 驱动 API（如 hipMemCreate / cuMemCreate），确保了底层运行时逻辑的统一。
Rust 兼容层设计（替代条件编译） ：hipify_torch 重写头文件后，bindgen 会生成 hipError_t 等 HIP 专属类型。为了避免在 Rust 调用点堆砌 #ifdef 分支，我们在 nccl-sys 和 rdmaxcel-sys 中引入了 rocm_compat 模块。该模块作为兼容层（Shim），将 HIP 符号以 CUDA 的名称重新导出（例如 pub type cudaError_t = hipError_t），从而让核心 Rust 代码保持平台无关。

通过引入 Rust 层面的 HIP 类型别名，我们成功跑通了全部 1171 个测试用例，实现了对 ROCm 7.0+ 的完整支持。相关代码已合入上游开源仓库（PR #2393 和 PR #2891）。

目前，Monarch on ROCm 已具备完整的生态支撑能力，涵盖 Actor 运行时、RDMA、监督树及张量分片。它能够无缝运行在 SLURM（HPC）、Kubernetes（云原生）和 SkyPilot（多云）环境中，为 TorchTitan（训练引擎）和 TorchFT（容错）等下游引擎在生产环境跑真实负载提供底层支撑。

容错训练：从理论到千卡落地

在 ROCm 平台上搞了一套组合拳：用 Monarch 做编排，TorchFT 做容错，TorchTitan 做训练引擎。目的很明确------打造一个不怕节点挂掉的弹性训练环境。

拆开来看这三层：

Monarch 层 ：扮演"包工头"角色。它拉起 ReplicaActors 进程，启动 Lighthouse 作为仲裁者，然后把物理 GPU 抽象成好管理的 Process Mesh。
TorchFT 层：这是容错的核心。它在每个训练 Step 介入，通过 Lighthouse 协调集群状态。一旦 AllReduce 过程中有节点掉链子，它立马执行剔除操作，保证计算继续。
TorchTitan 层：纯粹的执行单元。负责 FSDP 切分、梯度回传和参数更新，同时兼顾 Checkpoint 落盘和 Metrics 采集。

2.4 节点宕机不停训：拆解 Monarch 的"秒级感知-毫秒级跳过"机制

为了更直观地理解 Monarch 的容错恢复机制，我们以一个 4 副本组 的具体场景为例进行拆解：

1. 正常运行阶段 OrchestrationManager 负责编排，生成 4 个 ReplicaActor（Monarch Supervisor）及一个 Lighthouse 服务。每个 ReplicaActor 进一步 spawn 出一个包含 8 个 GPU 进程（运行 TorchTitan 训练器）的副本。当 4 个副本全部就绪后，法定人数组 quorum_id=1 建立，DiLoCo 梯度同步按每 20 step 的周期执行。

2. 故障检测与隔离 假设 Replica 0 中的某个 GPU 进程发生 Crash。Monarch Supervisor 在进程终止前成功捕获 report_training_error（包含完整堆栈信息）。此时，Replica 1、2、3 被标记为"未受影响"，训练任务继续执行。

3. 本地重启（Local Restart） ReplicaActor 0 触发就地重启机制（_stop_and_restart()），停止旧的进程网格并 spawn 新实例。在此期间，其余 3 个副本维持同步状态，法定人数组演进为 quorum_id=2。

4. 对等检查点传输（Peer Checkpoint Transfer） Lighthouse 服务选举 Replica 1 作为数据源（Donor），发起从 Replica 1 到恢复中 Replica 0 的检查点传输（涵盖模型、优化器、调度器及 Trainer 状态）。所有副本在法定人数边界处短暂暂停，等待新组构建完成。

5. 训练恢复 当 Replica 0 完成状态同步后，包含全部 4 个副本的新法定人数组 quorum_id=3 正式建立，DiLoCo 同步机制恢复运行。

2.5 性能基准测试：ROCm 环境下的吞吐表现与开销剖析

为了验证这套容错机制的可靠性，我们在 AMD Instinct MI300 级集群上，分别基于 SLURM 和 Kubernetes 环境进行了真实的分布式训练验证。

SLURM 16节点 MI300 集群实测（128 GPU） 测试场景设定为在 16 节点（128张 MI300）SLURM 集群上训练 Llama 3 8B 模型。测试期间，我们模拟了恶劣的硬件环境：每 180 秒主动注入一次 RCCL 故障，并维持每 20 步一次的 Quorum 同步。实测表现如下：

活跃节点动态波动：受注入故障影响，集群内活跃的 Worker 节点数在 8 到 16 之间实时变化。
避免全量重启：面对高频故障，训练任务保持了平滑运行，成功避免了"一挂全挂、从头再跑"的灾难。
收敛效果高度一致：Loss 曲线呈现平稳收敛态势，与无故障注入的基线（Baseline）结果高度吻合。

2.6 SLURM MI300集群压力测试------故障频发场景下训练持续运行，副本快速轮换恢复

跨运行时验证：从 SLURM 到 K8s，Monarch 在 256 卡集群下的表现一致性

为了验证架构的通用性，我们在 K8s 上拉起了一个 32 节点（256 MI355）的巨无霸集群。结果显示，即便遭遇节点抖动，参与计算的 Worker 数量依然维持在 30-32 的高水位，波动微乎其微。

更有说服力的是 Loss 曲线：从初始值 12 平稳收敛至 4。这证明了 Monarch 的容错模型具备极强的环境适应性，无论是在传统的 HPC 调度器（SLURM）还是云原生环境（Kubernetes）中，其弹性伸缩与恢复能力都保持了高度一致。

2.7 32节点 K8s MI355 集群压测：恢复机制靠谱，Loss 收敛丝滑

未来的展望：把 Monarch 跑在 ROCm 上的几点思考

搞大模型训练，光堆算力是没用的，硬件迟早会崩，关键是谁能优雅地收拾残局。我们把 PyTorch Monarch 搬到 AMD Instinct + ROCm 这套栈上，算是趟出了一条路子：既能把废掉的算力捡回来，又能把 GPU 的利用率钉死在高位。

这次移植有几个里程碑式的节点：

首次大规模亮机：在 AMD 卡上把 Monarch、TorchTitan 和 TorchFT 串起来了。实测证明，ROCm 这套软件栈完全撑得住这种高级别的容错逻辑，不是玩具。
容错逻辑更清爽：Monarch 那套 Supervision Tree（监督树）和 Process Mesh（进程网格）确实好用。故障被关进了笼子里，重启也是秒级的局部操作，不用惊动全局。
环境通吃：不管你是蹲在机房用 SLURM，还是上云用 Kubernetes，这套逻辑都能无缝切换。生产环境落地，算是有了底气。

核心心法其实就两条：一是靠 Monarch 的 Actor 模型把故障"隔离"；二是靠 TorchFT 的 Quorum（多数派）机制，只要活着的节点够数，训练就不停。这对于那些想在 AMD 平台上省钱又省心的团队来说，应该是个不错的信号

后续演进方向：

网络层：扩展 NIC 兼容性，深度调优运行时性能。
框架层：支持更多预训练/RL 框架，构建开放生态。
容错层：降低 Rejoin 延迟，实现 Recovery-Compute Overlap。

GitHub - meta-pytorch/monarch: PyTorch Single Controller · GitHub

GitHub - meta-pytorch/torchft: Fault tolerance for PyTorch (HSDP, LocalSGD, DiLoCo, Streaming DiLoCo) · GitHub