Decoupled DiLoCo：レジリエントな分散AIトレーニングの新たなフロンティア
Arthur Douillard および DiLoCo チーム
新しい分散アーキテクチャにより、離れたデータセンター間でLLM（大規模言語モデル）を訓練できます。帯域幅を抑えつつ、ハードウェアのレジリエンシー（耐障害性）も高められます。
最先端のAIモデルの訓練は、伝統的には、同一のチップをほぼ完全に同期させ続ける必要がある、大規模で強く結合されたシステムに依存してきました。このアプローチは、今日の最先端モデルには非常に効果的です。しかし、将来の世代に向けてスケールをさらに大きくしていくにつれ、何千ものチップにわたって同じレベルの同期を維持することは、重大なロジスティクス上の課題になります。
そこで今日、新しい論文にて、この問題に対する新しいアプローチ、Decoupled DiLoCo（Distributed Low-Communication）を紹介できることを嬉しく思います。大規模な訓練を、計算の「島（islands）」へと分割し、それらの間で非同期のデータが流れるようにすることで、このアーキテクチャはローカルでの中断や支障を切り分け、システムの他の部分が効率的に学習を継続できるようにします。
その結果、世界中に分散したデータセンターにまたがって先進的なモデルを訓練するための、よりレジリエントで柔軟な方法が得られます。さらに重要な点として、Decoupled DiLoCo は、Data-Parallel のような従来の分散手法がグローバル規模では非現実的にならせた通信遅延の問題を抱えていません。
最前線のモデルが、規模と複雑さの両面で引き続き拡大していく中で、より多くの計算資源、より多様なロケーション、そして多様なハードウェアにまたがってモデルを訓練するための、さまざまなアプローチを検討しています。
図1：学習実行を計算の「島」（学習ユニット）ごとに分離することで、同じレベルのハードウェア障害が起きていても、影響が隔離されるため、学習は概ね中断されることなく継続できる。
大規模でのより耐障害性の高い非同期学習の開発
分離型のDiLoCoは、2つの先行する進歩に基づいています。Pathwaysは、非同期データフローに基づく分散AIシステムを導入し、そしてDiLoCoは、分散データセンター間で必要とされる帯域幅を大幅に削減し、離れた場所にまたがって大規模言語モデルを学習することを現実的にしました。
分離型DiLoCoは、これらの考え方をまとめて、大規模においてより柔軟にAIモデルを学習するためのものです。Pathwaysの上に構築されており、計算の別々の島（学習ユニットとして知られる）にまたがる非同期学習を可能にします。これにより、ある領域でチップ障害が起きても、他の領域の進捗を中断しません。
このインフラストラクチャは自己修復型でもあります。テストでは、「カオスエンジニアリング」と呼ばれる手法を用いて、学習実行中に人工的なハードウェア障害を導入しました。分離型DiLoCoは、学習ユニット全体の損失が発生しても学習プロセスを継続し、その後それらがオンラインに戻ったときにシームレスに再統合しました。
Gemma 4モデルで分離型DiLoCoを検証したところ、ハードウェアが故障しても、より従来型の学習手法に比べて学習クラスタの可用性が高く維持されることが示されました。その一方で、最終的には機械学習（ML）のベンチマーク上の同等の性能レベルも達成しています。
図2：左：デカップルドDiLoCoのアプローチでは、従来の学習手法に比べて必要な帯域幅が桁違いに少なくて済むため、非常に効率的です。中央：ハードウェア障害のレベルが高まっても、デカップルドDiLoCoは高いレベルの「グッドプット」（有用な学習）を継続して提供しますが、他のアプローチは急激に性能が落ち込みます。（最初の2つのグラフは、シミュレーションによる学習実行に基づいています。）右：実環境での実験では、デカップルドDiLoCoで学習したGemma 4モデルのベンチマークされたML性能は、従来の学習アプローチで到達した性能と同等でした。
デカップルドDiLoCoは、障害への耐性が高いだけでなく、本番レベルの、完全に分散された事前学習を実行するためにも現実的です。私たちは、広域ネットワーキングの2〜5 Gbps（施設間で新しいカスタムネットワーク基盤を要求するのではなく、データセンター施設間の既存のインターネット接続で比較的実現可能なレベル）を用いて、4つの米国の別々のリージョンにまたがる形で、120億パラメータのモデルを正常に学習しました。特筆すべき点として、本システムは、従来の同期方式よりも20倍以上速く、この学習結果を達成しました。これは、システムの一部が別の部分を待つことで発生する「ブロッキング」のボトルネックを回避し、必要な通信をより長い計算期間に組み込むよう設計しているためです。
AI学習インフラの進化を牽引する
Googleでは、AI学習をハードウェア、ソフトウェア基盤、研究にまたがるフルスタックのアプローチで捉えています。ますます、こうしたレイヤー同士がどのように組み合わさるかを見直すことで得られる改善が大きくなっています。
デカップルドDiLoCoはその一例です。インターネット規模の帯域幅で学習ジョブを可能にすることで、システムのどこにあっても未使用の計算資源を活用でき、使われずに滞留していたリソースを有用な処理能力へと変えられます。
効率性や耐障害性にとどまらず、この学習パラダイムは、TPU v6eやTPU v5pのように異なる世代のハードウェアを、単一の学習実行の中で混在させることも可能にします。この方式は、既存のハードウェアの有効寿命を延ばすだけでなく、モデル学習に利用できる総計算量も増やします。私たちの実験では、異なる世代のチップが異なる速度で動作していても、単一のチップ種別のみで行う学習実行と同等のML性能が得られました。これにより、より古いハードウェアでもAI学習を実質的に加速できることが保証されます。
さらに、新しい世代のハードウェアはどこにでも一斉に導入されるわけではないため、世代をまたいだ学習が可能であれば、繰り返し発生する運用上の調整や能力（キャパシティ）のボトルネックを緩和できます。
私たちは今日、AIインフラの最前線を押し進める中で、次世代のAIを解き放つために必要となる、レジリエントなシステムのアプローチを引き続き探求しています。
謝辞
この取り組みは、Google DeepMindおよびGoogle Researchのメンバーからなるチームによって行われました。
デカップルド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による追加サポートおよび助言にも、心より感謝します。
