技術要約：Factored Gossip DiLoCo

1. 問題提起

大規模言語モデル（LLM）の学習を、分散型かつ低帯域幅の環境（例：消費者向けネットワーク）へとスケールアップさせる際、通信のボトルネックが大きな課題となります。DiLoCo（Douillard et al., 2023）は、グローバルな同期の間に多くのローカル最適化ステップを実行することで通信頻度を減少させますが、その外部同期プロセス自体が依然として重大な制限となっています。具体的には、DiLoCoは正確かつ高ボリュームでブロッキングを伴うグローバル同期（例：リング・オールリデュース）に依存しています。このアプローチは脆弱であり、最も遅いリンクを待つ間、計算を停止させてしまいます。また、単一の一時的なリンク障害によって同期ラウンド全体が中断され、リトライが必要になることで、さらなる効率低下を招きます。

本論文は、数十億パラメータ規模のモデルの学習に必要な安定性を損なうことなく、このブロッキングを伴う通信を削減する必要性に取り組んでいます。

2. 手法：Factored Gossip DiLoCo

著者らは、正確なグローバル同期を、ミキシング／ゴシップアルゴリズムによる近似的な同期へと緩和するフレームワークであるFactored Gossip DiLoCoを提案しています。核心となる革新は、外部同期ステップを以下の2つの異なる操作に分解（ファクタライズ）することです。

1. Mix1（非ブロッキング・パラメータ・ミキシング）:

   • 機能: 前回の外部ステップのパラメータ（$z_{m,r}$）を混合し、次ラウンドの新しい開始点（$x_{m,r}$）を生成します。
   • 特性: このステップは非ブロッキングであり、後続の内部ステップの計算とオーバーラップさせることが可能です。
   • メカニズム: 著者らは、Mix1においてオーバーラップされたグローバル平均化を利用しています。次ラウンドの勾配計算は、$x_{m,r}$（グローバルに混合されたもの）ではなく、$z_{m,r}$（ローカルな更新）から開始されるため、この通信は時間的な鮮度不足（temporal staleness）（古い勾配の使用）を引き起こしませんが、「開始点のノイズ」項を導入します。

2. Mix2（ブロッキング・勾配・ミキシング）:

   • 機能: ローカルステップ完了後に計算された外部勾配（パラメータ差分 $\Delta_{m,r}$）を混合します。
   • 特性: このステップはブロッキングであり、ワーカー間の合意（コンセンサス）を強化して最適化の安定性を確保する役割を果たします。
   • メカニズム: Mix2は、グローバル平均、またはスパースなペアワイズ・ゴシップとして構成できます。これは、Mix1の近似的な性質によって導入された不一致を抑制するための補正ステップとして機能します。

理論的分析

論文では、期待されるコンセンサス誤差（ワーカー間のパラメータの二乗距離）に関する理論的境界を提供しています。$\rho_1$ と $\rho_2$ をそれぞれ Mix1 と Mix2 の収縮係数とすると、ワーカー間の平均距離の境界は以下の通りです：
$$\mathbb{E}[\text{Distance}] \leq 2\eta^2 \sigma^2 H \left(1 + \frac{\rho_2}{1-\rho_1}\right)$$

• 重要な洞察: もし $\rho_1 \to 1$（Mix1での効果的な混合が行われない状態）であり、かつグローバルな Mix2（$\rho_2=0$）がない場合、コンセンサス誤差は発散する可能性があります。
• トレードオフ: Mix1は発散を防ぐために不可ло欠であり、一方で Mix2は境界をタイトにします。このフレームワークにより、計算利用率（ブロッキングの最小化）と安定性（コンセンサス誤差の最小化）の間のトレードオフを調整することが可能になります。

コンセンサス指標

著者らは、ワーカー間のトークンごとの出力分布間のイェンセン・シャノン（JS）距離を、機能的な不一致の指標として導入しています。L2パラメータ距離が理論と相関している一方で、JS距離は学習の不安定性や損失のスパイクをより良く追跡できることを示しており、いつブロッキング通信（Mix2）が必要になるかを知るための実用的なシグナルとして機能します。

3. 主な貢献

• 同期の分解（Factorization of Synchronization）: DiLoCoの同期を、非ブロッキングなパラメータ混合（Mix1）とブロッキングな勾配混合（Mix2）に分解することで、時間的な鮮度不足を導入することなく、通信と計算のオーバーラップを可能にしました。
• 近似的な同期: 脆い正確なグローバルバリアの代わりに、遅延や一時的な障害に対して緩やかに劣化するゴシップベースのミキシングを採用しました。
• JS距離による最適化: 機能的なJS距離を用いて、敏感なパラメータブロック（例：トークン埋め込み、LM-Head、初期層）を特定し、これらのサブセットに対してのみブロッキングMix2を適用することで、安定性を維持しながら通信量を大幅に削減しました。
• 実証的検証: 本手法が、低帯域幅の設定において、DiLoCoと同等の学習進捗を維持しつつ、計算利用率を大幅に向上させることを実証しました。

4. 実験結果

実験は、1.5BパラメータのLlama 3モデルを用い、FineWebデータセット上で、模擬的な低帯域幅接続（100Mbps–200Mbps）を備えた8つのワーカー（A100 GPU）を使用して行われました。

• 計算利用率:
  • DiLoCo: 100Mbpsで36%、200Mbpsで53%の利用率を達成（$H=100$ の場合）。
  • GlobalM1 (Mix1のみ、Mix2なし): 100Mbpsで56%、200Mbpsで**100%**の利用率を達成。
  • GlobalM1LocalM2 (Mix1はグローバル + Mix2はローカル): 100Mbpsで36%、200Mbpsで**66%**の利用率を達成し、強力なバランスを実現。
• 学習パフォーマンス（パープレキシティ）:
  • GlobalM1LocalM2 は、10Bトークンにおいて 18.80 の検証パープレキシティを達成しました。これは DiLoCo の 18.65 に近く、純粋なローカルミキシング（LocalM1M2: 19.37）を大きく上回っています。
  • GlobalM1LocalM2-subset（パラメータの43%にのみMix2を適用）は、19.15 のパープレキシティを達成しつつ、さらに高い利用率（100Mbpsで45%）を実現しました。
• 堅牢性: 5%の通信失敗率をシミュレートした場合、DiLoCoの利用率は再起動により大幅に低下しましたが、Factored Gossip DiLoCoの手法は、より弱いコンセンサスでも混合が進むため、高い利用率を維持しました。
• スケーリング: これらの利点は、異なるワーカー数（$M=8, 16$）および内部ステップ数（$H=100, 500, 1000$）においても観察されました。

5. 意義と主張

本論文は、Factored Gossip DiLoCoが、大規模な分散学習を高帯域幅のデータセンターの外側でも実用的なものにすると主張しています。ブロッキング通信を削減することで、本フレームワークは以下のことを実現します：

1. 標準的なDiLoCoと比較して、低帯域幅かつ分散型の設定において、計算利用率を大幅に向上させます。
2. DiLoCoと同等の学習進捗を維持し、実時間あたりのパープレキシティを低減します。
3. 一時的なネットワーク障害に対して緩やかな劣化を提供し、グローバルバリアに伴う壊滅的な中断を回避します。
4. 計算利用率と最適化の安定性の間の調整可能なトレードオフを提供し、実務者が自身の帯域幅制約に基づいて最適な構成（例：サブセットベースのMix2）を選択できるようにします。

著者らは、彼らのアプローチが、遅延勾配に依存する非同期手法とは異なり、時間的な鮮度不足を導入することなく、DiLoCoベースの設定において通信と計算をオーバーラップさせる初めての手法であることを強調しています。