解耦式 DiLoCo：韌性分散式 AI 訓練的新前沿

Decoupled DiLoCo：具備韌性且分佈式 AI 訓練的新前沿

Arthur Douillard 與 DiLoCo 團隊

我們全新的分佈式架構有助於跨越遠距離數據中心訓練大型語言模型（LLM）——具備更低的頻寬需求與更強的硬體韌性。

傳統上，訓練前沿 AI 模型依賴於一個大型且緊密耦合的系統，其中相同的晶片必須保持近乎完美的同步。這種方法對於當今最先進的模型非常有效，但隨著我們展望未來更大規模的世代，要在數千個晶片之間維持這種程度的同步，將成為一項重大的物流挑戰。

今天，在一篇新論文中，我們很高興能分享解決此問題的新方法，稱為 Decoupled DiLoCo（解耦式分佈式低通信訓練）。透過將大型訓練任務劃分為多個解耦的運算「孤島」（islands），並讓異步數據在它們之間流動，這種架構隔離了局部故障，使系統的其他部分能夠繼續高效學習。

其結果是提供了一種更具韌性且靈活的方法，能在全球分佈的數據中心訓練先進模型。至關重要的是，Decoupled DiLoCo 不會受到通訊延遲的困擾，而正是這種延遲使得以往的分佈式方法（如數據並行 Data-Parallel）在全球規模下變得不切實際。

隨著前沿模型的規模和複雜性持續增長，我們正在探索多種方法，以便在更多的運算資源、地點和多樣化的硬體上訓練模型。

圖 1：將訓練任務解耦為獨立的運算「孤島」（學習單元），即使在相同程度的硬體故障下，也能讓訓練基本上不受干擾地進行，因為這些故障的影響已被隔離。

開發更具容錯能力的大規模異步訓練

Decoupled DiLoCo 建立在兩項早期進展之上：Pathways（引入了基於異步數據流的分佈式 AI 系統）以及 DiLoCo（大幅降低了分佈式數據中心之間所需的頻寬，使跨遠距離地點訓練大型語言模型變得可行）。

Decoupled DiLoCo 將這些理念結合在一起，以更靈活地進行大規模 AI 模型訓練。它構建於 Pathways 之上，實現了跨獨立運算孤島（稱為學習單元）的異步訓練，因此一個區域的晶片故障不會中斷其他區域的進度。

這種基礎設施還具備自我修復能力。在測試中，我們使用了一種稱為「混沌工程」（chaos engineering）的方法，在訓練過程中人為引入硬體故障。Decoupled DiLoCo 在失去整個學習單元後仍能繼續訓練過程，並在它們重新上線時無縫地重新整合。

使用 Gemma 4 模型測試 Decoupled DiLoCo 的結果顯示，當硬體發生故障時，該系統比傳統訓練方法能維持更高的學習集群可用性，同時最終能達到相同的機器學習（ML）基準性能水平。

圖 2：左圖：Decoupled DiLoCo 方法所需的頻寬比傳統訓練方法低幾個數量級，非常高效。中圖：隨著硬體故障程度增加，Decoupled DiLoCo 仍能提供高水平的「有效吞吐量」（即有用的訓練），而其他方法的表現則直線下降（前兩張圖表基於模擬訓練）。右圖：在實際實驗中，使用 Decoupled DiLoCo 訓練的 Gemma 4 模型的基準 ML 性能與傳統訓練方法達到的性能持平。

Decoupled DiLoCo 不僅對故障更具韌性，而且對於執行生產級別的全分佈式預訓練也非常實用。我們成功地在四個不同的美國區域訓練了一個擁有 120 億參數的模型，使用的廣域網路頻寬僅為 2-5 Gbps（這一水平利用現有的數據中心間互聯網連接即可實現，無需在設施之間建立新的客製化網路基礎設施）。值得注意的是，該系統實現此訓練結果的速度比傳統同步方法快 20 倍以上。這是因為我們的系統將必要的通訊整合到較長的運算週期中，避免了系統某一部分必須等待另一部分的「阻塞」瓶頸。

推動 AI 訓練基礎設施的演進

在 Google，我們對 AI 訓練採取全棧式方法，涵蓋硬體、軟體基礎設施和研究。越來越多的收益來自於重新思考這些層級如何相互配合。

Decoupled 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 提供的額外支持與建議。