通过稀疏化与量化,让大模型训练通信量减少97%以上,性能还更好。
Overcoming the Communication-Performance Tradeoff in LLM Pretraining
- 采用Top-k稀疏化和2比特量化压缩梯度信息。
- 在178M到20亿参数模型上,通信量降低97%-99%,损失更低。
- 适用于多工作节点、长通信间隔等复杂场景,适合分布式大模型训练。
通信高效的分布式训练算法(如DiLoCo)因在带宽受限环境(如跨数据中心、互联网)中训练大语言模型的优势而受到广泛关注。尽管这些本地更新方法通过减少同步实现通信压缩,但仍需传输与模型规模相当的密集伪梯度,导致跨数据中心链路仍存在通信瓶颈。虽然量化常用于减小伪梯度大小,但在大模型预训练中,现有方法难以应用稀疏化而不造成显著性能下降。本文提出SparseLoCo,一种面向大语言模型的通信高效训练算法,可有效结合Top-k稀疏化与2比特量化,使传输的伪梯度达到高达97%-99%的极低稀疏度,同时最终损失低于密集版DiLoCo。在语言模型预训练的实证研究中,SparseLoCo在不同模型规模(178M-2B)、MoE Transformer(645M-A273M)、工作节点数量增加及通信间隔变长等条件下,均保持对DiLoCo的性能与压缩优势。
原文摘要 · Abstract (English)
Communication-efficient distributed training algorithms (e.g., DiLoCo) have received considerable interest due to their benefits for training large language models (LLMs) in bandwidth-constrained settings, such as across datacenters and over the internet. While these local-update methods achieve communication reduction through reduced synchronization, they still require communicating dense model-sized pseudo-gradients, resulting in a communication bottleneck even for cross-datacenter links. While quantization is often applied to reduce the pseudo-gradient's size, in the context of LLM pre-training, existing approaches have not been able to leverage sparsification without incurring significant performance degradation. In this work, we introduce SparseLoCo, a communication-efficient training algorithm for LLMs that can effectively leverage Top-k sparsification and 2-bit quantization to reach extreme sparsity in the communicated pseudo-gradient, as high as 97-99%, while achieving lower final loss than dense DiLoCo. In our empirical study of language model pre-training, we demonstrate that SparseLoCo's improvements over DiLoCo in performance and pseudo-gradient compression are maintained across dense model scales (178M-2B), an MoE transformer (645M-A273M), increasing the number of workers, and increasing communication intervals.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。