DiPaCo: 分散型パス構成

DiPaCo: Distributed Path Composition

March 15, 2024

著者: Arthur Douillard, Qixuan Feng, Andrei A. Rusu, Adhiguna Kuncoro, Yani Donchev, Rachita Chhaparia, Ionel Gog, Marc'Aurelio Ranzato, Jiajun Shen, Arthur Szlam

cs.AI

要旨

機械学習（ML）の進歩は、ニューラルネットワークモデルのスケーリングによって推進されてきました。このスケーリングは、並列動作するデバイス間の高帯域幅通信を必要とするMLアプローチに対応するための、ますます高度なエンジニアリングの成果によって可能となりました。本研究では、MLモデルのための共同設計されたモジュール式アーキテクチャとトレーニングアプローチを提案します。これを「DIstributed PAth COmposition（DiPaCo）」と呼びます。トレーニング中、DiPaCoは計算を共有モジュールのセットを通るパスに分散します。Local-SGDにインスパイアされた最適化手法（DiLoCo）と組み合わせることで、通信を大幅に削減しながらモジュールを同期させます。このアプローチは、接続性が低く異種混在のワーカー間でのトレーニングを容易にし、ワーカーの故障やプリエンプションに対する堅牢性を確保する設計となっています。推論時には、各入力に対して単一のパスを実行するだけでよく、モデルの圧縮は必要ありません。このアプローチは、より非同期でモジュール化された大規模学習の新しいパラダイムに向けた最初のプロトタイプと位置付けています。広く使用されているC4ベンチマークでの実験では、同じトレーニングステップ数で、より短い実時間において、DiPaCoは1億5000万パラメータの256の可能なパスのうち1つを選択することで、10億パラメータの密なトランスフォーマー言語モデルの性能を上回りました。

English

Progress in machine learning (ML) has been fueled by scaling neural network models. This scaling has been enabled by ever more heroic feats of engineering, necessary for accommodating ML approaches that require high bandwidth communication between devices working in parallel. In this work, we propose a co-designed modular architecture and training approach for ML models, dubbed DIstributed PAth COmposition (DiPaCo). During training, DiPaCo distributes computation by paths through a set of shared modules. Together with a Local-SGD inspired optimization (DiLoCo) that keeps modules in sync with drastically reduced communication, Our approach facilitates training across poorly connected and heterogeneous workers, with a design that ensures robustness to worker failures and preemptions. At inference time, only a single path needs to be executed for each input, without the need for any model compression. We consider this approach as a first prototype towards a new paradigm of large-scale learning, one that is less synchronous and more modular. Our experiments on the widely used C4 benchmark show that, for the same amount of training steps but less wall-clock time, DiPaCo exceeds the performance of a 1 billion-parameter dense transformer language model by choosing one of 256 possible paths, each with a size of 150 million parameters.