返回

文章详情

将 PyTorch Monarch 引入 AMD GPU

Hacker News2026年7月25日 15:55

特色项目:训练拥有数十亿参数的最先进大型语言模型(LLMs)需要在数百或数千个 GPU 上进行分布式训练。在这个规模下,硬件故障并不是异常事件——它们是可以预期的。单个 GPU 内存错误、网络分区或节点崩溃都可能导致整个训练过程停下来,而这个过程可能已经进行了几天或几周。虽然我们之前的工作展示了 FP8 训练在大规模下的近线性扩展(在一个 1024-GPU MI325 集群上取得 96.16% 的扩展效率,使用 DeepSeekV3-671B),但关键挑战仍然是:大规模下的可靠性。为了应对这些挑战,我们将 PyTorch Monarch 引入基于 ROCm 的 AMD Instinct GPU,扩展了单控制器模型超越 CUDA 环境,并将这一新兴运行时带入更广泛的硬件生态系统。在这篇博客中,我们将探索 PyTorch Monarch 的架构,全面介绍将 Monarch 的 GPU 运行时和分布式通信栈移植到 ROCm 所需的工程工作,并展示系统如何在不停止整个训练工作的情况下动态恢复节点故障。到最后,您将理解 Monarch 如何使在 AMD GPU 上实现弹性、容错的分布式训练成为可能,以及为什么这代表了向稳定的大规模人工智能基础设施迈出的重要一步。 挑战:大规模下的可靠性 传统的容错策略在很大程度上依赖于周期性检查点:在规定间隔将整个模型状态保存到持久存储中。当发生故障时,整个作业从最后的检查点重新启动。虽然在概念上很简单,但这种方法有显著缺点。 挑战影响 检查点开销:写入数百千兆的模型状态到存储中消耗时间和 I/O 带宽。 浪费的计算:自上次检查点以来的所有进展在故障时丢失。 集群空闲时间:在替换故障节点并重新启动作业时,整个集群处于空闲状态。 扩展限制:随着集群规模的增长,在任何检查点间隔内发生故障的概率增加。对于真正的大规模训练,单纯的扩展是不够的——训练还必须能够从故障中恢复。我们需要一种更动态的方法,允许健康节点在失败节点恢复和重新加入时继续训练,从而最小化计算浪费并最大化 GPU 利用率。这正是 PyTorch Monarch 的用武之地。 什么是 PyTorch Monarch? PyTorch Monarch 引入了一种新的分布式编程范式,使开发者能够通过单个 Python 程序协调整个 GPU 集群。借助基于演员的运行时、进程网格抽象和异步执行模型,Monarch 简化了大规模分布式训练,并支持将训练、评估和强化学习组合在一个统一的脚本中进行的复杂工作流。该架构在多个不同的层次上运行: Python API:开发者编写简单的 Python 代码以实现分布式 GPU 执行的单程序接口。 Monarch 运行时:管理演员和网格、监督树和张量分片。 Rust 运行时(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 将来自 CUDA 的 C++ 桥接代码转换为 HIP,并链接 RCCL,其 API 与 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。两个平台还额外使用 dlopen 加载 GPU 驱动 API 函数,包括 hipMemCreate、cuMemCreate 和重新启动。

赞助内容

NordVPN Next-gen Antivirus

本站免费、广告极少。如果觉得有帮助,可以请我们喝杯咖啡 —— 任何金额都对持续运营有实际帮助。

请我喝杯咖啡