Communication Efficient LLM Pre-training with SparseLoCo

2025-08-21

AI Summary

Communication bottlenecks severely hinder large language model (LLM) pretraining in bandwidth-constrained settings (e.g., cross-data-center training); existing approaches reduce communication frequency but still transmit full-precision gradients and often underperform AdamW-based DDP baselines. Method: We propose SparseLoCo—the first framework enabling efficient LLM pretraining under extreme gradient compression via Top-*k* sparsification and 2-bit quantization. It incorporates error feedback and distributed momentum optimization, theoretically justifying that outer-layer momentum can be well approximated by highly sparse local updates, and introduces a novel sparse gradient aggregation mechanism. Contribution/Results: SparseLoCo significantly reduces communication overhead while improving both convergence speed and downstream task performance over full-precision DiLoCo, establishing a new state-of-the-art for communication-efficient LLM pretraining.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Search and Optimization: Learning to Search

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Communication-efficient distributed training algorithms have received considerable interest recently due to their benefits for training Large Language Models (LLMs) in bandwidth-constrained settings, such as across data centers and over the internet. Despite reducing communication frequency, these methods still typically require communicating a full copy of the model's gradients-resulting in a communication bottleneck even for cross-datacenter links. Furthermore, they can slightly degrade performance compared to a naive AdamW DDP baseline. While quantization and error feedback are often applied to reduce the pseudo-gradient's size, in the context of LLM pre-training, existing approaches have been unable to additionally leverage sparsification and have obtained limited quantization. In this work, we introduce SparseLoCo, a communication-efficient training algorithm for LLMs that effectively leverages Top-k sparsification and quantization to reach extreme compression ratios of up to 1-3% sparsity and 2-bit quantization while outperforming full-precision DiLoCo. Our key observations are that outer momentum can be locally approximated by an error feedback combined with aggressive sparsity and that sparse aggregation can actually improve model performance. We empirically demonstrate in a range of communication-constrained LLM training settings that SparseLoCo provides significant benefits in both performance and communication cost.

Problem

• Reduces communication bottleneck in LLM pre-training
• Leverages sparsification and quantization for compression
• Maintains performance in bandwidth-constrained distributed training

Innovation

• Top-k sparsification for extreme compression
• 2-bit quantization reducing communication size
• Local momentum approximation with error feedback