PyTorch 基金会博客最新发布：@AMD 和 @Meta 贡献者分享了如何将 PyTorch Monarch 移植到搭载 ROCm 的 AMD Instinct GPU 上，以支持大规模容错分布式训练。

本文详细介绍了 Monarch GPU 运行时和分布式通信栈的 ROCm 移植过程，并展示了 Monarch、TorchFT 和 TorchTitan 如何实现健康副本在故障节点恢复并重新加入时继续训练，无需完全从检查点重启。

验证工作包括在 128 块 AMD Instinct MI300 GPU 的 SLURM 集群和 256 块 AMD Instinct MI355 GPU 的 Kubernetes 集群上训练 Llama 3 8B 模型。

阅读完整技术深入解析：https://pytorch.org/blog/bringing-pytorch-monarch-to-amd-gpus-single-controller-distributed-training-on-rocm/

将 PyTorch Monarch 引入 AMD GPU：基于 ROCm 的单控制器分布式训练 – PyTorch

来源：https://pytorch.org/blog/bringing-pytorch-monarch-to-amd-gpus-single-controller-distributed-training-on-rocm/

特色项目

• PyTorch 徽标 (https://pytorch.org/projects/pytorch/)

训练具有数十亿参数的最先进大语言模型（LLM）需要在成百上千块 GPU 上进行分布式训练。在这种规模下，硬件故障不再是例外事件——而是常态。单次 GPU 内存错误、网络分区或节点崩溃都可能导致已运行数天或数周的整个训练任务中断。虽然我们之前的工作（https://pytorch.org/blog/efficient-moe-pre-training-at-scale-with-torchtitan/）展示了 FP8 训练在规模上的近线性扩展（在包含 1024 块 MI325 GPU 的集群上使用 DeepSeekV3-671B 实现了 96.16% 的扩展效率），但关键挑战依然存在：大规模下的可靠性。

为应对这些挑战，我们将 PyTorch Monarch 引入到搭载 ROCm 的 AMD Instinct GPU，将单控制器模型扩展至 CUDA 环境之外，并将这一新兴运行时带入更广泛的硬件生态系统。

在本篇博客中，我们将探讨 PyTorch Monarch 的架构，介绍将 Monarch 的 GPU 运行时和分布式通信栈移植到 ROCm 所需的工程工作，并演示系统如何在节点故障时动态恢复，而无需停止整个训练任务。通过本文，你将了解 Monarch 如何在 AMD GPU 上实现弹性、容错的分布式训练，以及为什么这代表着向稳定、大规模 AI 基础设施迈出的重要一步。

挑战：大规模下的可靠性

传统的容错策略严重依赖周期性检查点：每隔一段时间将完整的模型状态保存到持久化存储中。当发生故障时，整个任务从最后一个检查点重新开始。虽然概念简单，但这种方法存在显著缺点。

对于真正的大规模训练，仅能扩展是不够的——训练还必须能够从故障中恢复。我们需要一种更动态的方法，允许健康节点在故障节点恢复并重新加入时继续训练，从而最小化计算浪费并最大化 GPU 利用率。这正是 PyTorch Monarch 所解决的问题。

什么是 PyTorch Monarch？

PyTorch Monarch 引入了一种新的分布式编程范式，使开发者能够通过单个 Python 程序编排整个 GPU 集群。凭借其基于 actor 的运行时、进程网格抽象和异步执行模型，Monarch 简化了大规模分布式训练，并支持在统一脚本中组合训练、评估和强化学习等复杂工作流。

该架构在多个不同层次上运行：

1. Python API：单一编程接口，开发者编写简单的 Python 代码即可获得分布式 GPU 执行。
2. Monarch 运行时：管理 actor 和网格、监督树和张量分片。
3. Rust 运行时（Tokio）：确保高性能和内存安全。
4. 基础设施：集成 RDMA、RCCL/NCCL、SLURM、Kubernetes 和 SkyPilot。

图 1：PyTorch Monarch 架构将 Python API 与 Rust 运行时和基础设施解耦。

通过将每个训练副本内部使用的并行策略与跨副本的容错机制解耦，Monarch 提供了更清晰的容错模型。故障被隔离（actor 拥有私有状态，崩溃不会传播）、分层处理（在尽可能低的层级处理），并且恢复速度快（本地重启仅需数秒，仅当升级时才需数分钟）。

图 2：Monarch 的分层故障处理模型和监督树。

将 Monarch 移植到 ROCm：生态系统集成

将 Monarch 引入 AMD GPU 需要大量的工程工作来移植 GPU 运行时和分布式通信栈到 ROCm。

我们成功实现了三条主要的移植路径：

1. 集合通信：使用 hipify_torch 将 C++ 桥接代码从 CUDA 转换为 HIP，并链接到与 NCCL API 镜像的 RCCL。
2. GPU 内存管理：扩展构建系统以自动检测平台，并将 CUDA 驱动 API 调用路由到其 HIP 等价调用。
3. RDMA 集成：配置 GPU_PLATFORM=rocm 保持基于 libibverbs 的 RDMA 路径不变，同时将 GPU 端绑定从 CUDA 替换为 HIP 以支持 GPU 直传。

图 3：通过 hipify_torch 和自动检测将 Monarch 从 CUDA 移植到 ROCm。

此外，有两个横切问题影响了移植过程，值得深入探讨：

1. HIP 运行时无静态链接： NVIDIA 提供了 libcudart_static.a，因此 CUDA 路径直接静态链接 cudart_static。ROCm 没有提供 libamdhip64 的静态版本，因此 ROCm 构建动态链接 amdhip64。两个平台还通过 dlopen 加载 GPU 驱动 API 函数，包括 hipMemCreate、cuMemCreate 及相关调用，保持运行时的约定在两侧一致。
2. Rust 兼容性封装而非分叉绑定： 一旦 hipify_torch 重写 C/C++ 头文件，bindgen 会生成 HIP 命名的类型，如 hipError_t、hipDeviceptr_t 和 hipStream_t。我们不在每个 Rust 调用点添加 #ifdef 分支，而是在 nccl-sys 和 rdmaxcel-sys 中添加了 rocm_compat 模块，该模块将 HIP 符号重新导出为 CUDA 名称，例如 pub type cudaError_t = hipError_t 和 pub use hipSetDevice as cudaSetDevice。其余 Rust 代码保持平台无关。

这些努力最终在 Rust 中引入了 HIP 类型别名，所有 1,171 个测试通过，确保完全支持 ROCm 7.0+。我们已将这些贡献上游到开源社区（参见 PR#2393 (https://github.com/meta-pytorch/monarch/pull/2393) 和 PR#2891 (https://github.com/meta-pytorch/monarch/pull/2891)）。

如今，Monarch on ROCm 提供了完整的生态系统支持，包括 Actor 运行时、RDMA、监督和张量分片。它可以在 SLURM（HPC）、Kubernetes（云原生）和 SkyPilot（多云）上无缝运行，为 TorchTitan（训练引擎）和 TorchFT（容错）等下游引擎提供生产负载支持。

案例研究：大规模容错训练

为了展示 Monarch 在 AMD GPU 上的能力，我们将其与 TorchTitan 和 TorchFT 集成，构建了一个弹性、无检查点的分布式训练架构。

架构概览

该架构包含三个层次：

• Monarch：充当编排器，管理进程和集群编排。它生成 ReplicaActor 和 Lighthouse 服务，将 GPU 组织成进程网格。
• TorchFT：在步骤级别处理容错。它联系 Lighthouse 进行法定数量协调，执行 Quorum AllReduce，并跳过故障节点。
• TorchTitan：作为训练引擎，执行前向（FSDP）、反向和优化器步骤，同时管理检查点和指标。

图 4：AMD GPU 上集成了 Monarch、TorchFT 和 TorchTitan 的弹性训练栈。

在此设置中，Monarch 提供了监督树用于细粒度的故障检测和隔离。当训练 actor 中注入故障时，Lighthouse 检测到故障并由 TorchFT 处理。健康副本在同伴故障时继续独立训练，无需全局中断。

动态故障恢复工作流

让我们通过一个包含四个副本组的具体场景来了解恢复工作流。

1. 正常训练：OrchestrationManager 生成 4 个 ReplicaActor（Monarch 监督器）和一个 Lighthouse。每个 ReplicaActor 生成一个包含 8 个 GPU 进程的 Replica，运行 TorchTitan 训练器。所有 4 个副本都就绪（quorum_id=1），DiLoCo 梯度同步每 20 步进行一次。
2. 故障检测：副本 0 中的 GPU 进程崩溃。Monarch 监督器在进程死亡前捕获 report_training_error（包含完整回溯）。副本 1、2 和 3 被标记为未受影响并继续训练。
3. 本地重启：ReplicaActor 0 启动原地重启（_stop_and_restart()），停止旧的进程网格并生成新的进程网格。同时，其他 3 个副本继续同步（quorum_id=2）。
4. 同伴检查点传输：Lighthouse 选择副本 1 作为捐赠者。发起从副本 1 到恢复中的副本 0 的同伴检查点传输（模型、优化器、调度器和训练器状态）。所有副本在法定数量边界短暂暂停，等待新法定数量形成。
5. 恢复训练：副本 0 同步完成后，新的法定数量（quorum_id=3）与所有 4 个副本建立，DiLoCo 同步恢复。

图 5：展示同伴检查点传输的动态故障恢复工作流，无需全局检查点重载。

整个恢复过程无需任何手动干预，无需完全检查点重启，对整体训练吞吐量的影响最小。

性能特征

我们在 SLURM 和 Kubernetes 环境中使用 AMD Instinct MI300 类集群验证了该方法。

SLURM 16 节点 MI300 集群（128 块 GPU）

我们在一个 16 节点 SLURM 集群（总共 128 块 MI300 GPU）上训练了 Llama 3 8B 模型，每 180 秒注入 RCCL 故障，法定数量同步每 20 步进行一次。结果非常显著：

• 活跃工作节点数动态波动（介于 8 到 16 之间），这是由于注入的故障。
• 训练在频繁故障下持续进行——没有完全重启。
• 损失曲线显示稳定收敛，与无故障注入的基线运行非常接近。

图 6：在 16 节点 SLURM MI300 集群上，训练在频繁故障下持续进行。注意没有一个副本宕机超过 30 分钟；恢复速度很快，不同副本被动态杀死并恢复。

Kubernetes 32 节点 MI355 集群（256 块 GPU）

我们将实验扩展到 32 节点 Kubernetes 集群（总共 256 块 MI355 GPU）。参与者数量保持高度稳定（恢复事件期间在 30 到 32 之间轻微波动），全局平均损失从 12 平滑下降到约 4。这表明 Monarch 容错模型在 SLURM 和 Kubernetes 上都能在大规模下可靠工作。

图 7：在 32 节点 Kubernetes MI355 集群上稳定恢复和平滑损失收敛。

总结与未来方向

大规模训练 AI 模型不仅需要原始计算能力——还需要能够优雅处理不可避免硬件故障的弹性基础设施。通过将 PyTorch Monarch 引入搭载 ROCm 的 AMD Instinct GPU，我们展示了一种切实可行的容错分布式训练方法，可最小化计算浪费并最大化 GPU 利用率。

此次集成为 AMD GPU 上的大规模训练带来了多项重要成就：

• 首次在 AMD 硬件上进行大规模验证：我们成功在 AMD GPU 上部署了 Moniterch 与 TorchTitan 和 TorchFT，证明 ROCm 软件栈完全支持先进的容错机制。
• 更清晰的容错模型：Monarch 提供了健壮的监督树和进程网格抽象，能够隔离故障并实现快速本地恢复。
• 生态系统就绪：该方法可在 SLURM 和 Kubernetes 上无缝运行，已为生产负载做好准备。

关键架构洞察在于使用 Monarch 基于 actor 的运行时和监督树来隔离故障，并结合 TorchFT 基于法定数量的同步机制，允许健康节点继续训练。对于在 AMD GPU 上运行大规模训练负载的团队，该集成为更稳定、高效和经济的模型开发提供了路径。

展望未来，我们的下一步包括：

• 扩展 NIC 支持并提高运行时性能。
• 扩展 Monarch 以支持更多基于 ROCm 的预训练和强化学习（RL）框架。
• 进一步优化容错性能，特别是减少重新加入重载延迟，并将恢复与计算重叠。
• 继续与 PyTorch 社区进行开源协作。

附加资源

• PyTorch Monarch GitHub 仓库 (https://github.com/meta-pytorch/monarch)
• TorchTitan 文档 (https://github.com/pytorch/torchtitan)
• TorchFT GitHub 仓库 (https://github.com/pytorch/torchft)
• 弹性大规模训练：将 TorchFT 与 TorchTitan 集成在 AMD GPU 上 (https://rocm.blogs.amd.com/artificial-intelligence/primus-torchft/README.html)

免责声明

第三方内容由第三方所有者直接许可给您，而非由 AMD 许可给您。所有链接的第三方内容均按“原样”提供，不提供任何形式的保证。使用此类第三方内容完全由您自行决定，AMD 在任何情况下均不对您因任何第三方内容承担任何责任。您承担所有风险，并对因使用第三方内容而产生的任何损害负全部责任。

© 2026 Advanced Micro Devices, Inc.