将 PyTorch Monarch 引入 AMD GPU

特色项目 训练具有数十亿参数的最先进的大型语言模型 (LLM) 需要在数百或数千个 GPU 上进行分布式训练。在这种规模上，硬件故障并不是异常事件——而是预料之中的。单个 GPU 内存错误、网络分区或节点崩溃可能会导致已经进行数天或数周的整个训练运行瘫痪。虽然我们之前的工作展示了大规模 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 简化了大规模分布式训练，并实现了将训练、评估和强化学习结合在一个统一脚本中的复杂工作流程。该架构在多个不同的级别上运行： Python API：一种单程序接口，开发人员可以在其中编写简单的 Python 代码来获得分布式 GPU 执行。 Monarch Runtime：管理参与者和网格、监督树和张量分片。 Rust Runtime (Tokio)：确保高性能和内存安全。基础设施：与 RDMA、RCCL/NCCL、SLURM、Kubernetes 和 SkyPilot 集成。图 1：PyTorch Monarch 架构将 Python API 与 Rust 运行时和基础设施解耦。通过将每个训练副本中使用的并行策略与跨副本使用的容错机制解耦，Monarch 提供了更清晰的容错模型。故障是隔离的（参与者具有私有状态，崩溃不会传播）、分层（在尽可能最低的级别处理），并且恢复速度很快（本地重新启动只需几秒，升级时只需几分钟）。图 2：Monarch 的分层故障处理模型和监督树。将 Monarch 移植到 ROCm：生态系统集成 将 Monarch 移植到 AMD GPU 需要大量的工程工作才能将 GPU 运行时和分布式通信堆栈移植到 ROCm。我们成功实现了三个主要移植路径： 集体通信：我们使用 hipify_torch 将 C++ 桥接代码从 CUDA 转换为 HIP，并链接到 RCCL，后者反映了 NCCL 的 API。 GPU 内存管理：我们扩展了构建系统，以自动检测平台并通过其 HIP 等效项路由 CUDA 驱动程序 API 调用。 RDMA 集成：配置 GPU_PLATFORM=rocm 可以保持基于 libibverbs 的 RDMA 路径完整，同时将 GPU 端绑定从 CUDA 交换到 HIP 以进行 GPU 直接传输。图 3：通过 hipify_torch 和自动检测将 Monarch 从 CUDA 移植到 ROCm。此外，两个横切问题塑造了端口，值得仔细研究： HIP 运行时没有静态链接：NVIDIA 附带 libcudart_static.a ，因此 CUDA 路径直接链接 cudart_static 。 ROCm 没有提供 libamdhip64 的静态等效项，因此 ROCm 构建动态链接 amdhip64。两个平台还额外开放了 GPU 驱动程序 API 函数，包括 hipMemCreate 、 cuMemCreate 和 re