MuLoCo: Muon is a practical inner optimizer for DiLoCo

要約

Dilocoは、データセンターの設定で並列処理と加速器の利用を増やすための利点を備えたネットワーク制約の下で、大規模な言語モデル（LLM）をトレーニングするための強力なフレームワークです。
ただし、通信頻度が大幅に減少しているにもかかわらず、Dilocoの通信ステップには、モデルのパラメーターの完全なコピーをすべて削減することが依然として含まれています。
既存の作業では、ディロコでの通信を減らす方法を探求しましたが、エラーフィードバック蓄積者の役割と圧縮性に対する内部オプティマイザーの効果は未調査のままです。
この作業では、2つのローカルオプティマイザー（AdamwとMuon）と組み合わせた場合、ディロコの通信オーバーヘッドを減らすためのトップKスパース化と量子化などの標準的な圧縮方法の有効性を調査します。
私たちの実験では、デコーダーのみのデコーダーのみのトランス言語モデル（LMS）が、ミューオンをディロコの内部オプティマイザーとともにエラーフィードバックアキュムレータとともに活用することで、パフォーマンスの低下で通信デルタを2ビットに積極的に圧縮できることが明らかになりました。
重要なことに、Muloco（Muon Inner Optimizer Diloco）は、8倍少ない通信と同一のメモリの複雑さを持っている間、Dilocoを大幅に上回ります。

要約(オリジナル)

DiLoCo is a powerful framework for training large language models (LLMs) under networking constraints with advantages for increasing parallelism and accelerator utilization in data center settings. Despite significantly reducing communication frequency, however, DiLoCo’s communication steps still involve all-reducing a complete copy of the model’s parameters. While existing works have explored ways to reduce communication in DiLoCo, the role of error feedback accumulators and the effect of the inner-optimizer on compressibility remain under-explored. In this work, we investigate the effectiveness of standard compression methods including Top-k sparsification and quantization for reducing the communication overhead of DiLoCo when paired with two local optimizers (AdamW and Muon). Our experiments pre-training decoder-only transformer language models (LMs) reveal that leveraging Muon as the inner optimizer for DiLoCo along with an error-feedback accumulator allows to aggressively compress the communicated delta to 2-bits with next to no performance degradation. Crucially, MuLoCo (Muon inner optimizer DiLoCo) significantly outperforms DiLoCo while communicating 8X less and having identical memory complexity.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google