PaperlayerView exact revision on arXiv

Source-linked AI summary

DiLoCo: Distributed Low-Communication Training of Language Models

Arthur Douillard, Qixuan Feng, Andrei A. Rusu, Rachita Chhaparia, Yani Donchev, Adhiguna Kuncoro, Marc'Aurelio Ranzato, Arthur Szlam, Jiajun Shen

arXiv:2311.08105v3cs.LGcs.CL

TL;DR

Training large language models with standard methods requires tightly interconnected accelerators, which are difficult to colocate at modern scale. DiLoCo combines local optimization with infrequent communication across workers. With 500 steps between communications, the paper reports faster wall-clock convergence when communication takes a lot of time, alongside robustness across several training conditions.

  • Problem

    Standard training at modern scale requires thousands of colocated devices and high-bandwidth connections, creating substantial engineering and infrastructure challenges.

    Source: §1 · p. 1

  • Method

    DiLoCo uses AdamW for many local inner steps and Nesterov momentum for outer updates that average worker changes.

    Source: §2 · p. 2, §6 · p. 11

  • Results

    With 500 steps between communications, DiLoCo converges faster in wall-clock time when each communication step takes a lot of time, and experiments report robustness across data distributions, inner-step counts, and changing worker numbers.

    Source: §6 · p. 11

  • Takeaways & Limitations

    DiLoCo is reported as a robust and effective way to distribute transformer language-model training across several poorly connected machines.

    Source: §6 · p. 11

    The presented DiLoCo version assumes homogeneous workers, so waiting for workers with different speeds to complete the same number of steps may be inefficient.

    Source: §5 · p. 10

Abstract

from arXiv · show

Large language models (LLM) have become a critical component in many applications of machine learning. However, standard approaches to training LLM require a large number of tightly interconnected accelerators, with devices exchanging gradients and other intermediate states at each optimization step. While it is difficult to build and maintain a single computing cluster hosting many accelerators, it might be easier to find several computing clusters each hosting a smaller number of devices. In this work, we propose a distributed optimization algorithm, Distributed Low-Communication (DiLoCo), that enables training of language models on islands of devices that are poorly connected. The approach is a variant of federated averaging, where the number of inner steps is large, the inner optimizer is AdamW, and the outer optimizer is Nesterov momentum. On the widely used C4 dataset, we show that DiLoCo on 8 workers performs as well as fully synchronous optimization while communicating 500 times less. DiLoCo exhibits great robustness to the data distribution of each worker. It is also robust to resources becoming unavailable over time, and vice versa, it can seamlessly leverage resources that become available during training.

1. Introduction

DiLoCo addresses the difficulty of colocating and synchronizing many accelerators by training model replicas locally and communicating only after many inner steps. On C4, it achieved better perplexity than fully synchronous optimization while communicating 500 times less, and experiments report robustness to varied worker data and resource availability.

1. Introduction

  • Standard training at modern scale requires thousands of colocated, tightly interconnected devices and careful coordination of gradients, parameters, and intermediate states.Using more devices per synchronous step also increases the chances of failures that risk halting training or introducing numerical issues.Source: §1 · p. 1
  • DiLoCo adapts Federated Averaging with many inner steps, AdamW within workers, and Nesterov momentum as the outer optimizer.Workers exchange information every H steps, which can be hundreds or thousands, and different islands can use different device types.Source: §1 · pp. 1–2
  • 500 times less communication accompanied even better perplexity than a fully synchronous model on C4.The paper reports this comparison in a large-batch training setting with over-training.Source: §1 · p. 2
  • Experiments report robustness to different worker data distributions and global-update frequencies, as well as island failures and newly available resources.The approach can leverage islands as they become available.Source: §1 · p. 2

2. DiLoCo

DiLoCo alternates parallel local optimization on worker-specific data shards with an outer update that averages workers’ parameter changes and redistributes the shared model. Its large inner-step count limits communication, while its optimizer choices connect the method to Federated Averaging and data-parallel training.

2. DiLoCo

  • Workers start from a shared initial model, which may be random or pretrained, and each worker trains a replica using its own data shard.The process repeats for T outer steps across k workers and k shards.Source: §2 · p. 2
  • Because H is much greater than 1, workers can perform hundreds of local steps between communications.This makes communication across workers minimal.Source: §2 · p. 2
  • After H independent inner steps on each worker’s data shard, workers average their parameter deltas, update the shared parameters, and receive the updated copy.This is the only point when workers communicate; a worker performs N = T × H inner steps in total.Source: §2 · p. 2
  • DiLoCo uses AdamW as InnerOpt and Nesterov momentum as OuterOpt; with SGD as OuterOpt it is equivalent to classical Federated Averaging.With H set to 1 and InnerOpt set to SGD, it is equivalent to large-batch data-parallel training.Source: §2 · pp. 2–3
  • DiLoCo is data parallelism with little communication, suitable for poorly connected or geographically distant workers.Workers may use standard data and model parallelism during their inner optimization.Source: §2 · p. 3

3. Experiments

The C4 experiments compare DiLoCo with synchronous baselines and test its performance across different initializations, communication schedules, data regimes, and resource conditions. The main comparison reports lower perplexity with substantially less communication.

3. Experiments

  • With 8 workers, DiLoCo yields lower perplexity than the 8×-batch baseline, runs 8× faster, and communicates 500 times less.The comparison is on C4 after pretraining a 150M baseline for 24,000 steps.Source: Figure 2 · p. 4, Table 2 · p. 4, §3 · p. 5

3.1. Ablations

  • Starting before 24,000 pretraining steps or from random initialization preserves final PPL; with one replica, DiLoCo also improves convergence and final generalization at no communication cost.Starting DiLoCo before 24,000 steps achieves similar final PPL, and random initialization does not degrade performance.Source: §3.1 · pp. 5, 8, 9
  • At H=1000, perplexity increases only 2.9% relative to H=50 despite 20× less communication; more frequent communication than H=500 has diminishing returns.The experiments use H=500 as a trade-off between generalization and communication cost.Source: §3.1 · p. 6
  • The i.i.d. setting converges faster early, but i.i.d. and non-i.i.d. settings attain comparable final generalization.The comparison uses 8 workers; the authors characterize DiLoCo as robust to data-distribution differences.Source: §3.1 · p. 6, Figure 5 · p. 6
  • More replicas improve generalization in both data regimes with diminishing returns beyond eight; performance also improves monotonically as model size increases.Model-size hyperparameters were tuned on the 150M model and may be suboptimal for other sizes.Source: §3.1 · pp. 6–7
  • Nesterov performs best among the tested outer optimizers; SGD and Adam perform poorly, and outer learning-rate decay yields similar performance.The adopted Nesterov setting uses outer learning rate 0.7 and momentum 0.9; Adam was particularly unstable with a high second-moment norm.Source: §3.1 · p. 7
  • Ultimate generalization depends on total compute, while models with equal compute achieve similar performance across different allocation schedules.Ramping up and ramping down compute yield similar performance, but worse generalization than schedules using more total compute.Source: §3.1 · p. 8, Figure 7 · p. 8
  • At a 50% communication-drop probability in the extreme non-i.i.d. setting, perplexity degrades only 2.1% relative to perfect communication, despite transient spikes at higher drop rates.Dropped workers continue training from their own parameters rather than the shared parameters.Source: §3.1 · p. 8

4. Related Work

Related work places DiLoCo in the local-SGD and federated-learning literature and connects it to linear mode connectivity. It contrasts prior methods by scale, optimizer, communication, and whether parameter averaging occurs during training.

4.1. Local SGD and Federated Learning

  • DiLoCo builds on local SGD and federated learning, using Nesterov outer optimization in a larger language-model setting than FedMom.The comparison covers models up to 400M parameters and 64 replicas, with up to 100× less communication; DiLoCo uses AdamW internally, whereas FedMom used SGD.Source: §4.1 · p. 9
  • Compared with Ortiz et al. (2021), DiLoCo operates with H = 1000, without pretraining, and up to 64 replicas in both i.i.d. and non-i.i.d. settings.Ortiz et al. reported degraded generalization with fewer inner steps, no pretraining, and at least 16 replicas; DiLoCo communicates 125× less in the comparison.Source: §4.1 · p. 10

4.2. Linear Mode Connectivity

  • Linear mode connectivity studies parameter interpolation between models to combine capabilities while keeping intermediate points at low loss.The cited literature extends parameter averaging across vision, NLP, RLHF, noisy-data, and out-of-distribution settings.Source: §4.2 · p. 10
  • Other approaches to alleviate loss barriers were not applied to DiLoCo, though the authors describe them as complementary future work.The passage names methods by Jin et al. (2023), Jordan et al. (2023), and Stoica et al. (2023).Source: §4.2 · p. 10
  • Unlike most linear-connectivity work that averages fully fine-tuned replicas, DiLoCo uses linear mode connectivity during training.BTM and PAPA are described as roughly equivalent frameworks with SGD outer optimization; DiLoCo also merges models trained on different worker data during training.Source: §4.2 · p. 10

5. Limitations

The reported evidence is limited to language modeling with transformers and models of 60–400 million parameters, with larger-scale performance still unvalidated. The authors also identify homogeneous-worker assumptions, diminishing returns beyond eight workers, and a wall-clock versus compute/data-efficiency trade-off.

5. Limitations

  • The evaluation covers only language modeling with transformers; other domains and architectures, including CNNs, remain untested.The authors note that CNNs are known to be more sensitive to linear mode connectivity.Source: §5 · p. 10
  • Results cover 60–400 million parameters, while any benefit at larger scales is an unvalidated hypothesis.The authors state that contemporary state-of-the-art models use three orders of magnitude more parameters.Source: §5 · p. 10
  • This version assumes homogeneous workers; waiting for every worker to complete the same number of steps can be inefficient when worker speeds differ.An asynchronous extension is proposed as future work.Source: §5 · p. 10
  • DiLoCo exhibits diminishing returns beyond eight workers, limiting the benefit of adding further compute under the reported setup.The authors identify improving use of additional compute as a future research direction.Source: §5 · p. 11
  • DiLoCo converges quickly in wall-clock time, but its distributed computation reduces FLOP and data efficiency.The authors attribute this partly to an effectively large outer-update batch; naively reducing it risks destabilizing workers because their batches are too small.Source: §5 · p. 11

6. Conclusion

DiLoCo distributes transformer language-model training across poorly connected machines using long local runs and infrequent communication. The authors report robustness across data distributions, inner-step counts, and changing worker numbers, while noting that broader scale, domain, and architecture generalization remains untested.

6. Conclusion

  • DiLoCo uses AdamW inner optimization, Nesterov momentum outside, and many local steps to reduce communication between poorly connected machines.The default number of inner steps is 500.Source: §6 · p. 11
  • Sending data only once every 500 steps can make DiLoCo converge faster in wall-clock time when communication steps take a long time.The paper contrasts this cadence with standard mini-batch methods that communicate every few hundred milliseconds.Source: §6 · p. 11
  • The authors report robustness to worker data distributions, inner optimization steps, and worker counts that change over time.The conclusion describes the approach as robust and effective for transformer training on poorly connected machines.Source: §6 · p. 11
  • Whether the findings extend to larger-scale models, other domains, or other architecture types remains to be seen.Source: §6 · p. 11

Supplementary Materials

The supplementary material reports implementation choices and ablations, including gradient averaging and pruning, then examines how outer-gradient similarity varies with data regime, inner steps, and worker count.

6.1. Implementation Details

  • Main-experiment hyperparameter choices are identified in Table 5, and the implementation uses AdamW inner optimization with separate optimizer states for each replica.Synchronizing optimizer states gave no significant improvement and increased communication by 3×.Source: §6.1 · p. 14, Table 5 · p. 14
  • For i.i.d. shards, replicas’ outer gradients are uniformly averaged; for non-i.i.d. shards, gradients are scaled by shard example count.The authors report that disjoint merge performed slightly worse, while greedy soup was too time-costly in this setting.Source: §6.1 · p. 14
  • The experiments used machines hosting 16 A100 GPUs, not necessarily in the same geographic region, with outer optimization performed on a connected CPU server.Source: §6.1 · p. 15

6.2. Experiments & Ablations

  • Pruning up to 50% of individual outer-gradient values caused negligible performance loss (+0.39% perplexity), offering another way to reduce communication demands.The study used sign-based pruning; the table caption identifies Table 6 as per-neuron sign pruning.Source: §6.2 · p. 15, Table 6 · p. 15
  • Outer gradients have near-zero similarity in the i.i.d. case; across both data regimes, similarity is inversely related to communication frequency.The authors surmise that more inner steps can move replicas toward a similar general direction while averaging stochastic-gradient noise.Source: §6.2 · p. 15, Figure 10 · p. 16
  • As learning rates anneal, i.i.d. gradient similarity increases, whereas non-i.i.d. variance increases and averaging the more orthogonal gradients is associated with better late-training generalization.Source: §6.2 · p. 15
  • In the non-i.i.d. setting, increasing replicas from 4 to 8 produces less-correlated outer gradients, and the averaged gradient norm varies inversely with the square root of replica count.Figure 11 depicts the similarity trend as replica count increases.Source: §6.2 · p. 15, Figure 11 · p. 16

Related papers

Loading 2311.08105v3…