Heterogeneous Low-Bandwidth Pre-Training of LLMs

TL;DR

This work addresses the challenge of scaling LLM pretraining under limited bandwidth by introducing Heterogeneous SparseLoCo, a framework that blends sparse, infrequent synchronization with low-bandwidth activation- and gradient-compression. It enables a mix of full replicas on high-bandwidth hardware and groups of resource-limited participants to form a replica via pipeline parallelism with subspace communication, maintained by outer SparseLoCo synchronization. The study shows that activation compression can be combined with SparseLoCo at modest cost and that selectively applying compression to bandwidth-constrained replicas improves the loss-communication tradeoff, with gains growing at higher compression. These results demonstrate a practical path for low-bandwidth model parallelism and heterogeneous participation in LLM pretraining, potentially enabling broader access to large-scale language modeling. The approach is validated on decoder-only LLaMA-2 models (178M and 512M), with observations that heterogeneity helps under SparseLoCo but not with standard AdamW, underscoring the interplay between optimization dynamics and communication strategies.

Abstract

Pre-training large language models (LLMs) increasingly requires distributed compute, yet bandwidth constraints make it difficult to scale beyond well-provisioned datacenters-especially when model parallelism forces frequent, large inter-device communications. We study whether SparseLoCo, a low-communication data parallel method based on infrequent synchronization and sparse pseudo-gradient exchange, can be combined with low-bandwidth pipeline model parallelism via activation and activation-gradient compression. We introduce a heterogeneous distributed training framework where some participants host full replicas on high-bandwidth interconnects, while resource-limited participants are grouped to jointly instantiate a replica using pipeline parallelism with subspace-projected inter-stage communication. To make the recently introduced subspace pipeline compression compatible with SparseLoCo, we study a number of adaptations. Across large-scale language modeling experiments (178M-1B parameters) on standard pretraining corpora, we find that activation compression composes with SparseLoCo at modest cost, while selective (heterogeneous) compression consistently improves the loss-communication tradeoff relative to compressing all replicas-especially at aggressive compression ratios. These results suggest a practical path to incorporating low-bandwidth model parallelism and heterogeneous participants into LLM pre-training.

Figures (3)

• Figure 1: Illustration of heterogeneous training proposed in this work. Across the data parallel axis, we train with $M$ SparseLoCo replicas that perform $H$ local steps before synchronizing using compressed pseudo-gradient communication. A replica can be hosted by a single high-bandwidth cluster (M0 and M49), or by clustering a group of resource-limited participants that collectively form a replica via pipeline parallelism with compressed inter-stage communication (M1). Within each participant, any standard parallelism may be used. This framework allows effective training of large language models with low-bandwidth, heterogeneous resources.
• Figure 2: Compute utilization versus inter-stage bandwidth for a 70B-parameter model partitioned into 4 pipeline stages. We compare SparseLoCo with $H{=}50$ inner steps under different PP-compression ratios. X-axis denotes the bandwidth between adjacent stages. We observe that PP-compression significantly improves compute utilization in limited-bandwidth regimes.
• Figure 3: Extended token budget results (left) and wall-clock training time under bandwidth constraints (right) for a 512M model.