DeepMind 突破:异步训练带来更稳健性
谷歌 DeepMind 推出解耦式 DiLoCo:提升异步训练架构,容忍硬件故障
谷歌 DeepMind 近期发布了一项名为“解耦式 DiLoCo”的创新分布式训练架构,旨在显著提升大规模人工智能模型训练的效率,并增强其面对硬件故障时的鲁棒性。
传统的模型训练方法通常要求所有计算单元在执行梯度更新时进行严格同步。这种同步机制使得整个训练过程极易受到单个硬件故障的影响而中断。为了克服这一挑战,解耦式 DiLoCo 采取了将训练过程分散到多个异步、且能进行故障隔离的“计算孤岛”中的策略。这种设计允许每个计算单元独立进行训练,而无需等待其他单元的完成,从而提高了训练的灵活性和效率。

该架构的核心理念是将训练任务分配到多个被称为“学习单元”的集群中。每个学习单元可以在本地独立进行多次梯度计算,并将压缩后的梯度信息传递给外部优化器进行汇总。由于这一过程是异步进行的,即使某个学习单元发生故障,其余单元仍可继续执行训练任务,有效避免了传统方法中因单点故障导致整体训练停滞的问题。
实验结果显示,在经历高硬件故障率的环境下,解耦式 DiLoCo 架构仍能保持高达 88% 的计算资源利用率,而采用标准数据并行训练方法的对照组仅为 27%。此外,新架构大幅降低了跨数据中心的带宽需求,从原先的 198 Gbps 降低至仅 0.84 Gbps,这使得在现有的商业互联网基础设施上实现全球分布式训练成为可能。
解耦式 DiLoCo 还具备一项重要的“自愈”能力。在进行了混沌工程测试时,该系统展现出在整个学习单元失效后仍能持续训练的能力,并在故障单元恢复后能够无缝地重新整合到训练流程中。这种灵活性同样体现在对多种硬件平台的支持上,它能够让不同代数的 TPU 芯片在同一次训练中协同工作,这不仅有助于延长旧有硬件设备的使用寿命,还能有效缓解因硬件更新换代过程中可能出现的算力瓶颈问题。
核心亮点:
- 解耦式 DiLoCo 通过将训练任务分配给多个异步的“学习单元”,显著提升了大规模模型训练的鲁棒性。
- 该架构将跨数据中心的带宽需求降低至 0.84 Gbps,为全球分布式训练的实现铺平了道路。
- 具备自愈能力的解耦式 DiLoCo 即使在面临硬件故障时,也能维持高效的训练状态,并支持不同硬件平台的混合使用。
谷歌 DeepMind 发布的解耦式 DiLoCo 架构无疑为解决当前大规模 AI 模型训练面临的效率和稳定性挑战提供了一个创新的思路,其异步设计和故障容忍能力预示着更可靠、更经济的分布式训练新时代。
谷歌 DeepMind 的解耦式 DiLoCo 架构在 AI 模型训练领域展现了巨大的潜力,其异步设计和强大的容错能力不仅提高了训练效率,也为未来构建更稳定、更具成本效益的分布式训练系统奠定了基础。