我们的新分布式架构帮助在远程数据中心训练大型语言模型(LLMs),并且降低带宽需求,提高硬件弹性。传统的前沿AI模型训练依赖于大型紧耦合系统,其中相同芯片必须保持近乎完美的同步。这种方法对当前的最先进模型非常有效,但随着未来规模的增长,保持数千个芯片的同步成为一项重大后勤挑战。
在新论文中,我们兴奋地介绍了一种名为解耦式DiLoCo(Distributed Low-Communication)的方法。通过将大型训练任务分配到解耦的“计算岛”上,并在其间异步传输数据,这种架构能够隔离局部干扰,使系统的其他部分能够高效学习。结果是一个更具弹性和灵活性的方式来训练跨全球分布的数据中心的先进模型。至关重要的是,解耦式DiLoCo不受以往分布式方法(如数据并行)在全球规模下造成的通信延迟影响。
图1:将训练运行解耦为单独的“计算岛”(学习单元)允许在发生硬件故障时,保持大部分训练不受影响,因为这些故障的影响是隔离的。
解耦式DiLoCo建立在两个早期进展之上:Pathways和DiLoCo。Pathways引入了基于异步数据流的分布式AI系统,而DiLoCo显著减少了分布式数据中心之间所需的带宽,使得在远程位置训练大型语言模型成为可能。解耦式DiLoCo将这些思想结合在一起,使得在规模上更灵活地训练AI模型。
它使得异步训练在不同的计算岛(学习单元)之间进行,从而确保一个区域的芯片故障不会中断其他区域的进展。该基础设施也是自愈的。在测试中,我们使用了一种称为“混沌工程”的方法在训练过程中引入人工硬件故障。即使在整个学习单元丢失的情况下,解耦式DiLoCo仍能继续训练过程,并在它们重新上线时无缝重新集成。
对Gemma 4模型的测试表明,当硬件故障发生时,该系统保持更高的学习集群可用性,相较于传统训练方法,最终提供了相同水平的机器学习(ML)性能。
图2:左:解耦式DiLoCo方法所需的带宽比传统训练方法少几个数量级,效率极高。中:随着硬件故障级别的增加,解耦式DiLoCo继续交付高水平的“有效输出”,而其他方法则急剧下降。右:在实际实验中,使用解耦式DiLoCo训练的Gemma 4模型的基准ML性能与传统训练方法相当。
解耦式DiLoCo不仅对故障更具弹性,而且适用于执行生产级的完全分布式预训练。我们成功地在美国四个地区使用2-5 Gbps的广域网连接训练了一个拥有120亿参数的模型,这一水平相对容易通过现有的数据中心设施之间的互联网连接实现,而无需在设施之间建立新的定制网络基础设施。值得注意的是,该系统以比传统同步方法快20倍以上的速度实现了训练结果。这是因为我们的系统将所需的通信整合到更长的计算周期中,避免了“阻塞”瓶颈,即系统的一个部分必须等待另一个部分。
在Google,我们采取全面的AI训练方法,涵盖硬件、软件基础设施和研究。越来越多的收益来自重新思考这些层次如何结合在一起。解耦式DiLoCo就是一个例子。通过在互联网规模带宽下启用训练作业,它可以利用任何未使用的计算资源,将闲置资源转化为有用的能力。除了效率和弹性外,这种训练范式还解锁了在单次训练运行中混合不同硬件代的能力,例如TPU v6e和TPU v5p。这种方法不仅延长了现有硬件的有效使用寿命,还增加了可用于模型训练的总计算能力。在我们的实验中,不同代的芯片以不同速度运行,仍能匹配单一芯片类型训练运行的ML性能,确保即使是老旧硬件也能有效加速AI训练。此外,由于新一代硬件并不会同时到达所有地方,能够跨代训练可以缓解反复出现的后勤和容量瓶颈。
随着我们推动AI基础设施的前沿,我们将继续探索解锁下一代AI所需的弹性系统的方法。
博主点评: 解耦式DiLoCo方法无疑是应对未来AI训练挑战的创新解决方案,通过降低带宽需求和提高故障容忍度,开辟了分布式训练的新天地。这种灵活的架构不仅能有效利用现有资源,还能在硬件更新换代时延长旧设备的使用寿命,真正做到资源的优化配置与高效利用。