デカップリングされたディロコ：大規模な分散 AI トレーニングにおける回復力

弊社の新開発した分散アーキテクチャにより、遠隔に位置するデータセンター間において大規模言語モデル（LLM）の学習を、通信帯域幅を削減しつつかつハードウェアへの耐障害性を高めた環境下で行うことが可能になりました。従来の先鋭的な AI モデルの学習には、同型のチップがほぼ完全な同期を保ちながら密結合された大規模システムに依存していましたが、このアプローチは現状の世界最高水準のモデルに対して極めて効果的でした。しかし、未来 generations のスケール拡大において、数千個のチップにわたってこのような同期状態を維持することは、大きな運用上の課題となります。本稿では、この課題に対する新たなアプローチ、「Decoupled DiLoCo（分散型低通信：Distributed Low-Communication）」を発表いたします。同アーキテクチャでは、大規模な学習ジョブを非同期のデータフローによって接続された解像された「compute の島（islands）」に分割し、局所的な障害の影響を他のシステム領域から隔離することで、その他部分は引き続き効率的に学習を継続することができます。その結果、グローバルに分散したデータセンターを跨いで先進的なモデルを学習するための、より耐障害性が高く柔軟な手法が実現されます。とりわけ重要なのは、Decoupled DiLoCo は、以前の分散手法（例：Data-Parallel）において大規模化を不可能にしてきた通信遅延の課題を持たないことです。

先鋭的モデルがスケールと複雑性を増すにつれて、私たちはより多くの計算リソース、場所、および多様なハードウェアを活用したモデル学習のための多種多様なアプローチを探索しています。

図 1: 学習ジョブを個別の「compute の島（learner units）」に分離することで、同等レベルのハードウェア障害が発生しても、ほぼ途切れることなく学習を継続することが可能になります。これは、障害の影響が効果的に隔離されているためです。

スケールにおけるより耐故障性の高い非同期学習の開発

Decoupled DiLoCo は、2 つの prior の進歩に支えられています。一つは、非同期データフローに基づく分散 AI システムを開拓した「Pathways」であり、もう一つは遠隔データセンター間を介する通信に要求される帯域幅を劇的に削減し、遠隔場所間での大規模言語モデル学習を現実的なものにした「DiLoCo」です。Decoupled DiLoCo はこれらのアイデアを組み合わせて、スケールにおけるより柔軟な AI モデル学習を実現します。Pathways を基盤として構築された同システムは、compute の島（learner units として知られる）を超えた非同期学習を可能にし、ある領域におけるチップの故障が他の領域の学習進行に影響を与えることを防ぎます。このインフラストラクチャはさらに自己修復機能も兼ね備えています。試験では、「カオス工学」と呼ばれる手法を用いて学習実行中に人為的なハードウェア障害を導入しました。Decoupled DiLoCo は、learner units の全損失後も学習プロセスを継続し、それらが復旧時にシームレスに再統合されました。

図 2: 左：従来の学習手法と比較して、Decoupled DiLoCo アプローチは桁違いに少ない帯域幅を必要とし、非常に効率的です。中央：ハードウェア障害のレベルが上がるにつれて、Decoupled DiLoCo は高いレベルの「goodput」（実質的な学習成果）を維持し続けます。一方、他のアプローチでは performance が急激に低下します。（最初の 2 つのグラフはシミュレーションされた学習実行に基づいています）。右：現実世界の試験において、Decoupled DiLoCo を用いて学習された Gemma 4 モデルのベンチマーク ML パフォーマンスは、従来の学習手法で達成された性能と同等であることを示しました。

Decoupled DiLoCo を Gemma 4 モデルを用いた検証では、ハードウェア故障が発生した際、システムがより伝統的な学習手法よりも高い学習クラスターの可用性を維持しつつ、最終的には同等のベンチマーク ML パフォーマンスを実現することが確認されました。

Decoupled DiLoCo は障害への耐性が高くなるだけでなく、実運用レベルで完全に分散化された事前トレーニングを実行するためにも実用的です。私たちは、広域ネットワーク帯域幅 2〜5 Gbps（既存のデータセンター間のインターネット接続による比較的実現可能なレベルであり、施設間への新たな専用ネットワークインフラを必要としない）を利用して、4 つの別の米国地域にわたって 120 億パラメータを持つモデルの学習に成功しました。特筆すべきは、このシステムが従来の同期手法よりも 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 様による追加的なサポートと助言にも心より感謝いたします。