解耦DiLoCo:弹性分布式AI训练的新前沿
谷歌DeepMind提出解耦DiLoCo架构,通过异步计算岛实现低带宽、高容错的分布式大模型训练,显著提升硬件故障恢复能力和跨数据中心训练效率。
解耦DiLoCo:弹性、分布式AI训练的新前沿我们的新分布式架构有助于跨远程数据中心训练LLM,具有更低的带宽和更高的硬件弹性。训练前沿AI模型传统上依赖于一个大型、紧密耦合的系统,其中相同的芯片必须保持近乎完美的同步。
这种方法对于当今最先进的模型非常有效,但随着我们展望未来几代的规模,在数千个芯片之间保持这种水平的同步将成为一个重大的后勤挑战。今天,在一篇新论文中,我们很高兴分享一种解决这个问题的新方法,称为Decoupled DiLoCo(分布式低通信)。
通过将大型训练运行划分为解耦的计算“岛”,并在它们之间异步流动数据,该架构隔离了局部中断,以便系统的其他部分可以高效地继续学习。其结果是,一种更有弹性、更灵活的方式来训练全球分布式数据中心的先进模型。
至关重要的是,解耦DiLoCo不会遭受通信延迟,而通信延迟使以前的数据并行等分布式方法在全球范围内不切实际。随着前沿模型的规模和复杂性不断增长,我们正在探索多种方法来跨更多计算、位置和各种硬件训练模型。
大规模开发更容错的异步训练解耦DiLoCo建立在两项早期进步的基础上:Pathways引入了基于异步数据流的分布式AI系统,而DiLoCo则显著减少了分布式数据中心之间所需的带宽,使跨远程位置训练大型语言模型变得可行。
解耦DiLoCo将这些想法结合在一起,以更灵活地大规模训练AI模型。它建立在Pathways之上,可以跨独立的计算岛(称为学习器单元)进行异步训练,以便一个区域的芯片故障不会中断其他区域的进程。
这个基础设施也是自愈的。在测试中,我们使用了一种名为“混沌工程”的方法,在训练运行期间引入人为硬件故障。解耦DiLoCo在丢失整个学习器单元后继续训练过程,然后在它们重新上线时无缝地重新整合它们。
使用Gemma 4模型测试解耦DiLoCo表明,当硬件出现故障时,系统比更传统的训练方法保持更高的学习集群可用性,同时最终提供相同基准水平的机器学习(ML)性能。解耦DiLoCo不仅对故障更具弹性,而且对于执行生产级、完全分布式预训练也很实用。
我们成功地在四个不同的美国地区训练了一个120亿参数模型。
使用2-5 Gbps的广域网络(利用数据中心设施之间的现有互联网连接相对可实现的水平,而不需要设施之间的新定制网络基础设施)。值得注意的是,该系统实现这一训练结果的速度比传统同步方法快20倍以上。
这是因为我们的系统将所需的通信融入到更长时间的计算中,避免了系统的一个部分必须等待另一个部分的“阻塞”瓶颈。推动AI训练基础设施的发展在谷歌,我们采取全栈方法进行AI训练,涵盖硬件、软件基础设施和研究。
越来越多的收益来自于重新思考这些层如何组合在一起。解耦DiLoCo就是一个例子。通过支持互联网规模带宽的训练作业,它可以利用任何未使用的计算资源,无论其位于何处,将闲置资源转化为有用容量。
除了效率和弹性之外,这种训练范式还解锁了在一次训练中混合不同硬件代的能力,例如TPU v6e和TPU v5p。这种方法不仅延长了现有硬件的使用寿命,还增加了模型训练的总计算量。
在我们的实验中,以不同速度运行的不同代芯片仍然与单芯片类型训练运行的ML性能相匹配,确保即使是较旧的硬件也可以有意义地加速AI训练。
更重要的是,由于新一代硬件不会同时到达所有地方,因此能够跨代训练可以缓解反复出现的物流和容量瓶颈。今天,当我们推动AI基础设施的前沿时,我们正在继续探索解锁下一代AI所需的弹性系统方法。
致谢这项工作由Google DeepMind和Google Research的一个成员团队完成。
Decoupled DiLoCo背后的领导者和核心贡献者是Arthur Douillard、Keith Rush、Yani Donchev、Zachary Charles、Ayush Dubey、Blake Woodworth、Ionel Gog、Josef Dean、
Nova Fallen、Zachary Garrett。运营支持由Nate Keating和Jenny Bishop提供。
我们还感谢Jeff Dean、Marc'Aurelio Ranzato、Raia Hadsell、Arthur Szlam、Edouard Yvinec、Henry Prior、Paul Barham、Michael Isard、Daniel Ramage、
Brendan McMahan、Chase Hensel和Zoltan Egyed的额外支持和建议。
本文为机器翻译辅以 AI 润色,仅供参考。原始事实以原文为准。