DiLoCo: Distributed Low-Communication Training of Language Models | Summary

This article explains the key points of DiLoCo: Distributed Low-Communication Training of Language Models.

• 2023-11-14 (arXiv)
• Douillard, Arthur, Feng, Qixuan, Rusu, Andrei A., Chhaparia, Rachita, Donchev, Yani, Kuncoro, Adhiguna, Ranzato, Marc’Aurelio, Szlam, Arthur, Shen, Jiajun.
• Google DeepMind
• Paper

Read this article in Korean

Summary

• Training language models across poorly connected clusters is difficult because conventional distributed training exchanges information at every optimization step. On C4, a web-text dataset, DiLoCo with eight workers reached validation perplexity, a measure of next-token prediction error where lower is better, of 15.02 versus 15.30 for a fully synchronous eight-times-larger-batch baseline at the same modeled time and compute, while communicating 500 times less.
• Distributed Low-Communication training (DiLoCo) lets each worker train a model replica locally with AdamW, then combines parameter changes using an outer Nesterov momentum optimizer. The main experiment synchronizes every 500 local steps and starts from a model pretrained for 24,000 steps.
• The experiments find similar final performance across different worker data distributions and several changing compute schedules, but test models with only 60 million to 400 million parameters. A baseline using eight times as many updates achieves lower perplexity at eight times the modeled training time, exposing a compute and data efficiency trade-off.

1. Introduction

Synchronous training requires tightly connected accelerators and can be disrupted by a device failure. DiLoCo divides devices into workers that train locally and communicate periodically, reducing the need for high-bandwidth links between clusters while retaining the same model size and inference speed.

2. DiLoCo

Each of k workers starts an outer round from shared parameters, trains independently on its own data shard for H inner steps, and returns the difference between its starting and ending parameters. The outer optimizer applies their averaged difference to the shared model, which is distributed for the next round; each worker performs N = T × H inner steps over T rounds.

• The implementation uses AdamW locally and Nesterov momentum for the outer update. The supplementary materials specify a uniform average for randomly partitioned data and shard-size weighting for non-i.i.d. data, where worker data distributions differ.

Four replicas train independently before their parameter changes are combined and the updated shared model starts another round.

3. Experiments

The evaluation uses decoder-only Chinchilla-style transformers with 60M, 150M, and 400M parameters on C4 (Raffel et al., 2020; Hoffmann et al., 2022). Runs use a sequence length of 1,024 tokens, batch size 512, and 88,000 total steps; by default, DiLoCo starts after 24,000 pretraining steps and uses non-i.i.d. shards.

• For the main 150M experiment, eight workers each perform 128 rounds of 500 local steps. Training steps serve as a proxy for wall-clock time because cross-worker communication is infrequent.
• At equal modeled time and compute, DiLoCo reaches perplexity 15.02 versus 15.30 for the synchronous 8× batch baseline. The 8× updates baseline reaches 14.72 using the same modeled compute and data as DiLoCo but 8× the time.

The three transformer configurations vary in depth and hidden dimension while sharing a 32,000-token vocabulary.

The eight-worker DiLoCo curve ends at perplexity 15.02, below the 15.30 curve for the eight-times-larger batch baseline after the same pretraining and additional training steps.

The comparison separates modeled time from compute, data, and communication: the 8× updates baseline obtains 14.72 perplexity but takes 8× the time assigned to DiLoCo.

3.1. Ablations

Starting DiLoCo from scratch or after 12,000 or 24,000 pretraining steps gives similar final perplexity, although earlier starts use more total compute at a fixed number of steps. For the 150M model, communicating every 1,000 steps raises perplexity by 2.9% relative to every 50 steps while communicating 20× less; the authors select H = 500 for the remaining experiments.

• Randomly partitioned i.i.d. data, where worker shards have the same distribution, converges faster initially than clustered non-i.i.d. data, but their final perplexities are 15.08 and 15.02 with eight workers.
• Increasing replicas improves results with diminishing returns beyond eight, while also increasing total compute and data. Outer Nesterov momentum performs best among the tested outer optimizers, with learning rate 0.7 and momentum 0.9.
• In simulated communication failures, a 50% probability of dropping each worker’s communication increases non-i.i.d. perplexity by 2.1% relative to perfect communication. Changing the replica count over time produces similar results for the tested schedules with similar total compute.

Starting from scratch and starting after 24,000 pretraining steps both end at 15.02 perplexity, while the 12,000-step start ends at 14.91.

Less frequent synchronization progressively worsens final perplexity, with a larger deterioration at 2,000 steps than at 1,000 steps.

The random-shard run improves faster early in training, but its final perplexity is close to the clustered-shard run.

At fixed local steps, non-i.i.d. perplexity changes from 15.02 at eight replicas to 14.96 at 64, while additional replicas consume more compute and data.

Against single-worker baselines, eight-worker DiLoCo improves perplexity by 1.01, 1.21, and 1.01 for the 60M, 150M, and 400M models respectively.

Nesterov ends at 15.02 perplexity, compared with 15.45 for Adam, 15.5 for SGDM, and 16.81 for SGD in the plotted comparison.

Doubling and halving the replica count yield similar final perplexities when their total compute is matched; ramping up and down also produce similar outcomes.

Increasing the probability of dropped outer-gradient communication makes learning curves noisier in both data regimes, while training continues.

With one replica, periodic outer optimization reaches 16.03 perplexity versus 16.23 for the baseline, without cross-worker communication.

4. Related Work

The related-work discussion connects DiLoCo to local stochastic gradient descent (SGD), which makes multiple local updates between exchanges, federated learning, and linear mode connectivity. These lines of work motivate local updates, outer optimization, and combining model parameters during training.

4.1. Local SGD and Federated Learning

Federated averaging combines locally trained models (McMahan et al., 2017), and later work introduces stronger outer optimizers (Reddi et al., 2021; Wang et al., 2020). The authors compare DiLoCo with FedMom (Huo et al., 2020), which also uses outer Nesterov momentum but studies a one-layer long short-term memory network with two replicas and communication every 20 steps; DiLoCo tests transformers up to 400M parameters and 64 replicas.

4.2. Linear Mode Connectivity

Linear mode connectivity concerns interpolation between model parameters along paths with low loss. DiLoCo uses repeated parameter updates during training, whereas much of the cited model-merging work averages replicas after fine-tuning; the paper identifies branch-train-merge (BTM) and population parameter averaging (PAPA) as related exceptions (Li et al., 2022; Jolicoeur-Martineau et al., 2023).

5. Limitations

The study covers one task, language modeling, and one architecture family, transformers, at 60M–400M parameters. Its standard procedure waits for workers to complete the same number of steps, so different worker speeds may reduce efficiency; simulated dropped communication does not establish a fully asynchronous implementation.

• Returns diminish beyond eight workers, and the 8× updates baseline shows that DiLoCo sacrifices some compute and data efficiency for modeled time efficiency. The paper’s suggested improvements at larger model scales remain a hypothesis.

6. Conclusion

The authors conclude that infrequent synchronization with local AdamW and outer Nesterov momentum can train transformer language models across poorly connected workers. Their evidence covers variations in communication interval, data distribution, and worker availability, while leaving larger scales and other domains untested.

Appendix

• Supplementary Materials, 6.1. Implementation Details: The inner learning rate is 4e−4, warmup lasts 1,000 steps, and weight decay is 0.1. Each replica keeps its own AdamW optimizer state: synchronizing those states gave no significant improvement while tripling transmitted data. Experiments used machines with 16 A100 GPUs and a CPU server for outer optimization.
• 6.2. Experiments & Ablations: Pruning 50% of individual outer-gradient values increases perplexity from 15.02 to 15.08, or 0.39%; pruning 75% raises it to 15.27. Cosine similarity, a measure of directional agreement, shows greater variability among outer gradients for non-i.i.d. shards and lower similarity when the number of non-i.i.d. replicas increases from four to eight.

Brief Thoughts

The equal-compute comparison makes the communication reduction concrete, and the additional 8× updates baseline shows the cost of faster modeled convergence. An open deployment question is how the synchronization scheme behaves with substantially larger models and workers of unequal speed.