AIが爆発的に発展している現代において、大規模言語モデル(LLM)は機械学習アプリケーションを推進する強力なエンジンとなっています。しかし、これらの巨大なモデルを訓練するには膨大な計算資源が必要です。世界中に分散したデバイス上でこれらのモデルを効率的に訓練できたらどうなるでしょうか?OpenDiLoCoはその驚きの答えを提供します。
従来の分散型訓練方法では、頻繁な通信と大量の帯域幅が必要であり、訓練の規模と効率が制限されていました。一方、DiLoCo(分散型低通信)訓練方法は、通信の必要性を減らすことで、LLMの世界規模での訓練を可能にします。
OpenDiLoCoは、DiLoCo訓練方法を実装し、Hivemindライブラリを使用して拡張可能で分散型の訓練フレームワークを提供するオープンソースフレームワークです。このフレームワークの優れた点は、世界規模で、2大陸3カ国にまたがってモデルの訓練を行いながら、90~95%の計算利用率を維持していることです。
主な特徴:
リソースの動的増減:訓練中に計算リソースを動的に調整でき、新しいデバイスやクラスタは訓練途中で参加したり退出したりできます。
フォールトトレランス:分散型訓練では、一部のデバイスが信頼できない可能性があります。Hivemindのフォールトトレラントな訓練により、デバイスが使用できなくなっても訓練プロセスは停止しません。
ピアツーピア通信:マスターノードはなく、すべての通信はピアツーピアで行われます。
研究者たちは、DiLoCoの実験を再現しただけでなく、それを10億パラメータのモデル規模にまで拡張しました。彼らはアブレーションスタディを通じて、DiLoCoアルゴリズムの計算効率と拡張性における優位性を示しました。さらに印象的なのは、DiLoCoの勾配を、性能を低下させることなくFP16を使用して全還元できることを証明したことでしょう。
主な貢献:
再現と拡張:DiLoCoの元の研究結果を再現し、それを10億パラメータのモデル規模に拡張しました。
オープンソース実装:Hivemindライブラリに基づいた拡張可能な実装により、分散型訓練を幅広い開発者や研究者に提供します。
世界規模の分散型訓練:2大陸3カ国間のモデル訓練を通じて、OpenDiLoCoの実用的な可能性を示し、同時に90~95%の計算利用率を維持しました。
効率に関する洞察:アブレーションスタディを通じて、アルゴリズムの拡張性と計算効率に関する貴重な知見を提供します。
実験結果:
Prime Intellect社は、DiLoCoの主要な実験結果を再現することで、その方法の有効性を証明しました。C4データセットを使用して1億5000万パラメータのモデルで言語モデリングタスクの訓練を行うと、DiLoCoは通信要求を500分の1に削減しながら、ベースライン性能と同等の性能を示しました。
DeepMindの元のDiLoCo論文では、最大4億パラメータのモデルに対してのみ実験が行われました。この研究では、Prime Intellect社は、この方法を1億1000万パラメータのモデルに拡張し、TinyLlamaと同じハイパーパラメータを使用しました。
OpenDiLoCoの異なる大陸における分散型訓練機能を示すために、カナダ、フィンランド、アメリカの2つの州にある4つのDiLoCoワーカーノードを使用しました。各ノードには8つのH100 GPUが搭載されています。
Prime Intellect社は、DiLoCoの主要な実験結果を再現し、この方法を元の研究のパラメータサイズの3倍に拡張し、現実世界の分散型訓練環境におけるその応用を示しました。
今後、同社はDiLoCoをより多くの分散型ワーカーノード上のより大きなモデルに拡張し、安定性と収束速度を向上させる可能性のあるモデル集約技術、および非同期重み平均通信方法の実装による計算アイドル時間の削減を検討する予定です。
論文アドレス:https://arxiv.org/pdf/2407.07852