Communication-Efficient Language Model Training Scales Reliably and Robustly: Scaling Laws for DiLoCo

TL;DR

This work develops and validates scaling laws for DiLoCo, a distributed low-communication training method, and contrasts it with data-parallel training under a fixed token budget. Through extensive experiments from 35M to 2.4B parameters and extrapolations to 4B and 10B, DiLoCo demonstrates predictable scaling, robustness to synchronization cadence, and often superior evaluation loss, particularly when using multiple replicas. The analysis reveals that DiLoCo enables larger optimal batch sizes, significantly reduced bandwidth, and faster wall-clock times, while maintaining or improving generalization as models scale. The findings offer a principled framework to predict hyperparameters and performance for larger-scale deployments and highlight promising directions for extending scaling laws and deployment systems for communication-efficient training.

Abstract

As we scale to more massive machine learning models, the frequent synchronization demands inherent in data-parallel approaches create significant slowdowns, posing a critical challenge to further scaling. Recent work develops an approach (DiLoCo) that relaxes synchronization demands without compromising model quality. However, these works do not carefully analyze how DiLoCo's behavior changes with model size. In this work, we study the scaling law behavior of DiLoCo when training LLMs under a fixed compute budget. We focus on how algorithmic factors, including number of model replicas, hyperparameters, and token budget affect training in ways that can be accurately predicted via scaling laws. We find that DiLoCo scales both predictably and robustly with model size. When well-tuned, DiLoCo scales better than data-parallel training with model size, and can outperform data-parallel training even at small model sizes. Our results showcase a more general set of benefits of DiLoCo than previously documented, including increased optimal batch sizes, improved downstream generalization with scale, and improved evaluation loss for a fixed token budget.

Figures (19)

• Figure 1: DiLoCo. Each DiLoCo model replica trains independently for $H$ inner optimization steps. These models are synchronized via an outer optimization step, usually involving momentum across outer optimization steps. In this figure, there are $M=4$ replicas.
• Figure 2: DiLoCo does better with scale. We compare Data-Parallel to DiLoCo for varying model sizes $N$. For all $M$, DiLoCo improves monotonically wrt Data-Parallel as $N$ increases.
• Figure 3: DiLoCo with $M = 1$ generalizes better than Data-Parallel. We present the evaluation loss and downstream accuracy of Data-Parallel and DiLoCo with $M = 1$ for varying model and global batch sizes (measured in tokens). In all settings, DiLoCo with $M = 1$ does better than Data-Parallel, and the gap between them increases with batch size. We see similar results for other model sizes, but omit for the sake of brevity.
• Figure 4: DiLoCo increases optimal batch size, part 1. Evaluation loss of Data-Parallel and DiLoCo as a function of global batch size (in tokens). For all $M$, DiLoCo exhibits larger optimal batch size than Data-Parallel. Moreover, the optimal batch size increases as a function of $M$. We see similar results for other model sizes, but omit for conciseness.
• Figure 5: DiLoCo increases optimal batch size, part 2. Zero-shot accuracy on HellaSwag of Data-Parallel and DiLoCo as a function of global batch size (in tokens). Even at smaller model sizes, DiLoCo with $M = 2$ attains higher accuracy for larger global batch sizes. We see similar results for other model sizes, but omit for conciseness.