将 PyTorch Monarch 移植到 AMD GPU：在 ROCm 上进行单控制器分布式训练（13 分钟阅读）

# 将 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/) 训练拥有数十亿参数的最先进大型语言模型 (LLMs) 需要在数百或数千个 GPU 上进行分布式训练。在这样的规模下，硬件故障不再是例外事件——它们是可以预期的。单个 GPU 内存错误、网络分区或节点崩溃可能导致已经持续数天或数周的整个训练运行中断。虽然我们之前的工作 (https://pytorch.org/blog/efficient-moe-pre-training-at-scale-with-torchtitan/) 展示了大规模 FP8 训练的近乎线性扩展（在使用 DeepSeekV3-671B 的 1024 GPU MI325 集群上实现了 96.16% 的扩展效率），但关键挑战依然存在：大规模下的可靠性。 为了应对这些挑战，我们将 PyTorch Monarch 带到了搭载 ROCm 的 AMD Instinct GPU 上，将单控制器模型扩展到 CUDA 环境之外，并将这一新兴运行时引入更广泛的硬件生态系统。 在本博客中，我们将探讨 PyTorch Monarch 的架构，详细介绍将 Monarch 的 GPU 运行时和分布式通信栈移植到 ROCm 所需的工程工作，并演示系统如何在不停止整个训练任务的情况下动态地从节点故障中恢复。最后，您将了解 Monarch 如何在 AMD GPU 上实现弹性、容错的分布式训练，以及为什么这代表着向稳定的大规模 AI 基础设施迈出的重要一步。 ## 挑战：大规模下的可靠性 传统的容错策略严重依赖周期性检查点：定期将完整的模型状态保存到持久化存储中。当发生故障时，整个任务从最后一个检查点重新启动。虽然概念简单，但这种方法有明显的缺点。 **挑战** | **影响** ---|--- **检查点开销** | 将数百 GB 的模型状态写入存储会消耗时间和 I/O 带宽。 **计算浪费** | 故障发生时，自上一个检查点以来的所有进度都会丢失。 **集群空闲时间** | 在更换故障节点并重启任务期间，整个集群处于空闲状态。 **可扩展性限制** | 随着集群规模增长，任何检查点间隔内发生故障的概率都会增加。 对于真正的大规模训练，仅有扩展是不够的——训练还必须能够从故障中恢复。我们需要一种更动态的方法，让健康的节点能够继续训练，同时故障节点恢复并重新加入，从而最大限度地减少计算浪费和最大化 GPU 利用率。这就是 PyTorch Monarch 发挥作用的地方。 ## 什么是 PyTorch Monarch？ PyTorch Monarch 引入了一种新的分布式编程范式，使开发者能够通过单个 Python 程序编排整个 GPU 集群。借助其基于角色的运行时、进程网格抽象和异步执行模型，Monarch 简化了大规模分布式训练，并支持在统一脚本内组合训练、评估和强化学习的复杂工作流。 该架构在多个不同级别上运行： 1. **Python API**：单一程序接口，开发者编写简单的 Python 代码即可获得分布式 GPU 执行。 2. **Monarch 运行时**：管理角色和网格、监督树和张量分片。 3. **Rust 运行时 (Tokio)**：确保高性能和内存安全。 4. **基础设施**：与 RDMA、RCCL/NCCL、SLURM、Kubernetes 和 SkyPilot 集成。 *图 1：PyTorch Monarch 架构将 Python API 与 Rust 运行时和基础设施解耦。* 通过将每个训练副本内部使用的并行策略与跨副本使用的容错机制解耦，Monarch 提供了更清晰的容错模型。故障被隔离（角色拥有私有状态，崩溃不会传播）、分层处理（在尽可能低的级别处理），并且恢复速度快（本地重启只需数秒，升级处理也仅需数分钟）。 *图 2：Monarch 的分层故障处理模型和监督树。* ## 将 Monarch 移植到 ROCm：生态系统集成 将 Monarch 带到 AMD GPU 需要大量的工程工作，以将 GPU 运行时和分布式通信栈移植到 ROCm。 我们成功实现了三个主要的移植路径： 1. **集合通信**：我们使用 `hipify_torch` 将 C++ 桥接代码从 CUDA 转换为 HIP，并链接到 RCCL，RCCL 的 API 与 NCCL 类似。 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` 模块，该模块在 CUDA 名称下重新导出 HIP 符号，例如 `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)）。 如今，ROCm 上的 Monarch 提供了完整的生态系统支持，包括 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 提供了一个监督树，用于细粒度的故障检测和隔离。当训练角色中注入故障时，Lighthouse 会检测到该故障，并由 TorchFT 处理。健康的副本在同行失败的情况下继续独立训练，无需全局中断。 ### 动态故障恢复工作流 让我们通过一个包含四个副本组的具体场景来理解恢复工作流。 1. **正常训练**：OrchestrationManager 生成 4 个 ReplicaActor (Monarch 监督者) 和一个 Lighthouse。每个 ReplicaActor 生成一个包含 8 个 GPU 进程的副本，这些进程运行 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：动态故障恢复工作流，展示了无需全局检查点重新加载的对等检查点传输。* 整个恢复过程无需任何手动干预，无需完整的检查点重启，并且对整体训练吞吐量的影响最小。 ## 性能表现 我们使用搭载 AMD Instinct MI300 级别的集群，在 SLURM 和 Kubernetes 环境中验证了这种方法。 ### 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 上部署了集成 TorchTitan 和 TorchFT 的 Monarch，证明了 ROCm 软件栈完全支持高级容错机制。 - **更清晰的容错模型**：Monarch 提供了强大的监督树和进程网格抽象，隔离了故障并支持快速的本地恢复。 - **生态系统就绪**：该方法在 SLURM 和 Kubernetes 上都能无缝工作，可用于生产工作负载。 关键的架构洞察在于使用 Monarch 基于角色的运行时和监督树来隔离故障，并结合 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) - 弹性大规模训练：在 AMD GPU 上集成 TorchFT 与 TorchTitan (https://rocm.blogs.amd.com/artificial-intelligence/primus-torchft/README.html) ## 免责声明 第三方内容由拥有该内容的第三方直接向您授权，而非 AMD 向您授权。所有链接的第三方内容均按“原样”提供，不作任何形式的保证。使用此类第三方内容的风险完全由您自行承担，AMD 在任何情况下均不对任何第三方内容向您承担任何责任。您承担所有风险，并全权负责因使用第三方内容可能导致的任何损害。 *© 2026 超威半导体公司*