将 PyTorch Monarch 引入 AMD GPU实现弹性容错分布式训练迈向稳定大规模 AI 基建2026 年 7 月 6 日消息AMD 的侯朝军、Liz Li、Zachary Streeter 等以及 Meta 的 Matthias Reso、Hamid Shojanazeri、Monarch 团队将 PyTorch Monarch 引入了搭载 ROCm 的 AMD Instinct GPU。训练具有数十亿参数的最先进大语言模型LLM需要在数百或数千个 GPU 上进行分布式训练。在这种规模下硬件故障是预料之中的事单个 GPU 内存错误等都可能导致整个训练运行中断。此前的工作虽展示了 FP8 训练在大规模场景下的近线性扩展能力但大规模训练的可靠性仍是关键挑战。为应对这些挑战研究团队将 PyTorch Monarch 引入搭载 ROCm 的 AMD Instinct GPU把单控制器模型扩展到了 CUDA 环境之外让这个新兴运行时适配更广泛的硬件生态系统。下面我们来扒一扒这项研究都透露了哪些重点——挑战大规模训练的可靠性传统容错策略依赖定期检查点机制即按固定间隔将完整的模型状态保存到持久存储中故障时整个作业从上一个检查点重新启动。但这种方法有显著缺点挑战影响检查点开销将数百 GB 的模型状态写入存储设备会消耗大量时间和 I/O 带宽。计算资源浪费故障发生时自上一个检查点以来的所有进度都会丢失。集群空闲时间在替换故障节点并重新启动作业时整个集群会处于空闲状态。可扩展性限制随着集群规模的增大在任何检查点间隔内发生故障的概率也会增加。真正的大规模训练不仅要实现扩展训练过程还得能从故障中恢复。这就需要一种更动态的方法允许健康节点在故障节点恢复并重新加入时继续训练以减少计算资源浪费提高 GPU 利用率而 PyTorch Monarch 就能发挥这样的作用。什么是 PyTorch MonarchPyTorch Monarch 引入了一种新的分布式编程范式让开发者通过单个 Python 程序协调整个 GPU 集群。借助基于 actor 的运行时、进程网格抽象和异步执行模型Monarch 简化了大规模分布式训练还支持在一个统一脚本中组合训练、评估和强化学习等复杂工作流程。该架构在多个层面运行Python API一个单程序接口开发者编写简单 Python 代码就能实现分布式 GPU 执行。Monarch 运行时管理 actor、网格、监督树和张量分片。Rust 运行时Tokio确保高性能和内存安全。基础设施与 RDMA、RCCL/NCCL、SLURM、Kubernetes 和 SkyPilot 集成。Monarch 通过将每个训练副本内使用的并行策略与副本间使用的容错机制解耦提供了更清晰的容错模型。故障被隔离、分层处理恢复速度快。将 Monarch 移植到 ROCm生态系统集成将 Monarch 引入 AMD GPU需要大量工程工作把 GPU 运行时和分布式通信栈移植到 ROCm。研究团队成功实现了三条主要移植路径集体通信使用 hipify_torch 将 C 桥接代码从 CUDA 转换为 HIP并链接到 RCCL其 API 与 NCCL 类似。GPU 内存管理扩展构建系统以自动检测平台并通过其 HIP 等效项路由 CUDA 驱动 API 调用。RDMA 集成配置 GPU_PLATFORMrocm 可保持基于 libibverbs 的 RDMA 路径不变同时将 GPU 端绑定从 CUDA 替换为 HIP 以实现 GPU 直接传输。此外有两个跨领域问题影响了移植工作HIP 运行时无静态链接NVIDIA 提供 libcudart_static.aCUDA 路径可直接链接 cudart_static而 ROCm 没为 libamdhip64 提供静态等效项所以 ROCm 构建动态链接 amdhip64。两个平台都会额外使用 dlopen 加载 GPU 驱动 API 函数确保两侧运行时契约一致。使用 Rust 兼容性垫片而非分叉绑定hipify_torch 重写 C/C 头文件后bindgen 会生成 HIP 命名的类型。为避免在每个 Rust 调用点添加 #ifdef 分支研究团队在 nccl - sys 和 rdmaxcel - sys 中添加了 rocm_compat 模块将 HIP 符号重新导出为 CUDA 名称其余 Rust 代码保持平台无关性。这些努力最终在 Rust 中引入了 HIP 类型别名1171 个测试全部通过确保了对 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 提供了完整的生态系统支持可在 SLURM、Kubernetes 和 SkyPilot 上无缝运行为 TorchTitan 和 TorchFT 等下游引擎支持生产工作负载。案例研究大规模容错训练为展示 Monarch 在 AMD GPU 上的强大功能研究团队将其与 TorchTitan 和 TorchFT 集成构建了一个无需检查点的弹性分布式训练架构。架构概述该架构由三层组成Monarch作为协调器管理进程和集群编排。它生成 ReplicaActors 和 Lighthouse 服务将 GPU 组织成进程网格。TorchFT在步骤级别处理容错。它与 Lighthouse 联系以进行仲裁协调执行仲裁 AllReduce 操作并跳过故障节点。TorchTitan作为训练引擎执行前向传播、反向传播和优化器步骤同时管理检查点和指标。在这种设置下Monarch 提供了一个监督树用于细粒度的故障检测和隔离。当训练 actor 中注入故障时Lighthouse 会检测到故障并由 TorchFT 处理。即使出现对等节点故障健康的副本也能继续独立训练无需全局中断。动态故障恢复工作流程下面通过一个包含四个副本组的具体场景了解恢复工作流程正常训练OrchestrationManager 生成 4 个 ReplicaActorsMonarch 监督器和一个 Lighthouse。每个 ReplicaActor 生成一个包含 8 个 GPU 进程的副本运行 TorchTitan 训练器。所有 4 个副本就绪quorum_id 1每 20 步进行一次 DiLoCo 梯度同步。故障检测副本 0 中的一个 GPU 进程崩溃。Monarch 监督器在进程死亡前捕获 report_training_error包含完整的回溯信息。副本 1、2 和 3 被标记为未受影响并继续训练。本地重启ReplicaActor 0 发起原地重启_stop_and_restart()停止旧的进程网格并生成一个新的。同时其他 3 个副本继续同步quorum_id 2。对等检查点传输Lighthouse 选择副本 1 作为捐赠者。发起从副本 1 到恢复中的副本 0 的对等检查点传输模型、优化器、调度器和训练器状态。在新仲裁组形成时所有副本在仲裁边界处短暂暂停。恢复训练副本 0 同步完成后新的仲裁组quorum_id 3建立所有 4 个副本恢复 DiLoCo 同步。整个恢复过程无需人工干预无需完整的检查点重启对整体训练吞吐量影响最小。性能特征研究团队在 SLURM 和 Kubernetes 环境中使用 AMD Instinct MI300 系列集群验证了这一方法。SLURM 16 节点 MI300 集群128 个 GPU在 16 节点的 SLURM 集群共 128 个 MI300 GPU上训练 Llama 3 8B 模型每 180 秒注入一次 RCCL 故障每 20 步进行一次仲裁同步。结果出色由于注入的故障活跃工作节点数量在 8 到 16 之间动态波动。尽管频繁出现故障训练仍能无缝继续没有进行完整的重启。损失曲线显示出稳定的收敛与未注入故障的基线运行结果非常接近。Kubernetes 32 节点 MI355 集群256 个 GPU实验扩展到 32 节点的 Kubernetes 集群共 256 个 MI355 GPU。参与节点数量在恢复事件期间在 30 到 32 之间略有波动保持高度稳定全局平均损失从 12 平稳下降到约 4。这表明 Monarch 容错模型在 SLURM 和 Kubernetes 大规模集群上都能可靠工作。总结与未来方向大规模训练大型人工智能模型需要强大计算能力也需要能应对硬件故障的弹性基础设施。通过将 PyTorch Monarch 引入搭载 ROCm 的 AMD Instinct GPU研究团队展示了一种实用的容错分布式训练方法减少了计算资源浪费提高了 GPU 利用率。这种集成在 AMD GPU 大规模训练方面取得多项重要成果首次在 AMD 硬件上进行大规模验证成功在 AMD GPU 上部署 Monarch 与 TorchTitan 和 TorchFT证明 ROCm 软件栈支持先进容错机制。更清晰的容错模型Monarch 提供强大监督树和进程网格抽象隔离故障并实现快速本地恢复。生态系统就绪该方法可在 SLURM 和 Kubernetes 上无缝运行适用于生产工作负载。关键架构见解是使用 Monarch 基于 actor 的运行时和监督树隔离故障结合 TorchFT 基于仲裁的同步机制让健康节点继续训练。对于在 AMD GPU 上运行大规模训练工作负载的团队来说这种集成提供了更稳定、高效和经济的模型开发之路。展望未来研究团队的下一步计划包括扩展 NIC 支持并提高运行时性能。扩展 Monarch 以支持 ROCm 上更多的预训练和强化学习RL框架。进一步优化容错性能特别是减少重新加入的重新加载延迟并使恢复过程与计算过程重叠。继续与 PyTorch 社区进行开源合作。更多资源PyTorch Monarch GitHub 仓库TorchTitan 文档TorchFT GitHub 仓库弹性大规模训练在 AMD GPU 上集成 TorchFT 与 TorchTitan文档访问 PyTorch 全面的开发者文档查看文档 ›教程获取面向初学者和高级开发者的深入教程查看教程 ›资源查找开发资源并获取问题解答查看资源 ›保持联系获取更新、活动信息和最新消息提交表单即同意接收 Linux 基金会LF及其项目关于其活动、培训、研究、发展和相关公告的营销电子邮件可随时使用收到的电子邮件页脚中的链接取消订阅。隐私政策。x - 推特脸书领英YouTubeGitHubSlackDiscord
