​

​

DiLoCo: Distributed Low-Communication Training of Language Models

Adel Nabli - Johan Obando

Introduction

Related Work

Methodology

Setup

Ablations

Experiments

Limitations and Conclusions

Questions

Agenda

Introduction

Related Work

Methodology

Experiments

Conclusion

Limitations

Questions

Agenda

Introduction

Large language models (LLM) have become a critical component in many applications of machine learning.

​

• … this require a large number 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.

​

Distributed Low-Communication (DiLoCo), that enables training of language models on islands of devices that are poorly connected.

Introduction

Because performance has scaled with model size, practitioners train increasingly larger models on increasingly large data.

• training approach remains standard mini-batch back-propagation of the error.

​

​

At modern scale, training via standard backpropagation poses unprecedented engineering and infrastructure challenges.

• several thousands of devices need to be powered and be placed at the same physical location interconnected with high-bandwidth cables to minimize latency.
• Careful software engineering is required to orchestrate the passage of gradients, parameters and intermediate states between these devices
• The more devices that are used for each synchronous training step, the more chances there are that one of them fails, risking halting training, or introducing subtle numerical issues.

​

Introduction

This work take inspiration from literature on Federated Learning

• In Federated Learning, there are 𝑘 workers, each operating on their own island of devices
• Such workers perform some amount of work locally, and then exchange gradients every 𝐻 steps to get their model replica back in sync.

The current paradigm poorly leverages heterogeneous devices, that might have different speed and topology.

• It is difficult to colocate and tightly synchronize a large number of accelerators.

Introduction

The authos propose a variant of the popular Federated Averaging (FedAvg) algorithm (McMahan et al., 2017)

• the inner optimizer is replaced with AdamW
• the outer optimizer with Nesterov Momentum

​

Address the shortcomings mentioned before

• while each worker requires co-located devices their number is roughly 𝑘 times smaller than the total
• workers need not to communicate at each and every single step but only every 𝐻 steps which can be in the order of hundreds or even thousands
• while devices within an island need to be homogeneous, different islands can operate with different device types.

Introduction

Common Crawl is a publicly-available web archive that provides “web extracted text” by removing markup and other non-text content from the scraped HTML files. “Colossal Clean Crawled Corpus” larger and comprises reasonably clean and natural english text

Introduction

In a large-batch training setting with overtraining, our empirical validation on the C4 dataset (Raffel et al., 2020)

​

• …demonstrates that DiLoCo can achieve better performance
• DiLoCo is capable of effectively utilizing several islands of devices at training time, despite a low bandwidth connectivity among these islands.
• Experiments show that DiLoCo is robust against different data distributions used by local workers and frequency of global parameter updates.
• DiLoCo exhibits robustness to island failure, and nicely leverage islands whenever these become available.

Introduction

Related Work

Methodology

Experiments

Conclusion

Limitations

Questions

Agenda

Related Work

Local SGD and Federated Learning Several communities have proposed and studied local SGD.

• First instantation was in McMahan et al. (2017) who introduced the concept of federated learning and local SGD as a way to enable learning on a network of mobile devices which retain private access to their own data.
• Lin et al. (2020) considered local SGD as a way to improve generalization when learning with large batch sizes.

​

• FedMom (Huo et al., 2020) considers Nesterov as the outer optimizer as the authors did. They also tackle a language modeling task, the setting is much smaller (1-layer LSTM), with only 2 replicas, and rather frequent communication (every 20 inner steps).
• a larger setting with up to a 400M transformer language model, across up to 64 replicas, and up to 100× less communication.
• AdamW as inner optimizer while they used SGD.

​

​

Related Work

• Ortiz et al. (2021) validated on a large-scale setting. They consider ImageNet (Deng et al., 2009) with Resnet50 and Resnet101 (He et al., 2015), and found that local sgd struggles at scale.

​

• The majority of works on linear connectivity considers only averaging once all replicas have been fully finetuned, while we exploit the linear mode connectivity during training.

Related Work

• Linear Mode Connectivity

​

… how to linearly interpolate between several models in parameters space (Frankle et al., 2020; Wortsman et al., 2021).

​

A surprising result from this field is the relative easiness to find a linear interpolation between several models where all intermediary points have a low loss, avoiding any loss barrier.

• Wortsman et al. (2022c) started from a pretrained model, finetuned different replicas on various tasks or choice of hyperparameters (Wortsman et al., 2022b), and then averaged the resulting parameters.

​

Originally proposed in the vision domain, this method has then been used also in NLP (Li et al., 2022), RLHF (Ramé et al., 2023a), noisy data (Rebuffi et al., 2022), and OOD (Ramé et al., 2023b).

​

​

Introduction

Related Work

Methodology

Experiments

Conclusion

Limitations

Questions

Agenda

Methodology

• model architecture (e.g., a transformer) with parameters 𝜃.
• training dataset as D = {(x, y), ...} with x and y being respectively the input data and target.

In language modeling (Vaswani et al., 2017), the input is a sequence of tokens and the target is the input sequence shifted by one.

​

1. First, we start from an initial model with parameters 𝜃(0)
2. We also have 𝑘 workers each capable of training a model replica and 𝑘 shards of data, one for each worker.
3. There is an outer optimization (line 1, 12, and 14 in Algorithm 1), which consists of 𝑇 outer steps.

*At each outer step 𝑡, gradients from each worker are gathered, averaged and sent to an outer optimizer (OuterOpt) to update the shared copy of the parameters. Afterwards, this shared copy of the parameters is re-dispatched to each local worker (line 3).

Methodology

4. Within each phase, each worker (line 3) performs independently and in parallel its own inner optimization (lines 4 to 9) for 𝐻 steps using an inner optimizer, denoted by InnerOpt.
5. Each worker samples data from its own shard (line 5), and updates its own local copy of the parameters (line 8).
6. Once all workers have completed their inner optimization step, the cumulative gradients of each worker are averaged (line 12), and the resulting outer gradient is used to update the shared copy of the parameters (line 14), which is then used as starting point for the next round of inner optimizations.

​

Methodology

They use as inner optimizer (InnerOpt) AdamW (Kingma and Ba, 2014; Loshchilov and Hutter, 2019), which is the most widely used optimizer for transformer language models.

​

As for the outer optimizer (OuterOpt) we use Nesterov momentum (Sutskever et al., 2013) because it gave the best convergence em- pirically.

Methodology

Methodology

Introduction

Related Work

Methodology

Experiments

Conclusion

Limitations

Questions

Agenda

Experiments

Setup:

chinchilla-style (Hoffmann et al., 2022)

Experiments

Setup:

• We consider a language modeling task on the C4 dataset, a dataset derived from Common Crawl (Raffel et al., 2020).

​

• We report perplexity on the validation set against number of steps used at training time.

​

• The total number of steps is set to 88,000. We consider three model sizes, all decoder-only transformers adapted from the Chinchilla architecture (Hoffmann et al., 2022).
• ​
• We perform experiments both in the i.i.d. and non-i.i.d. settings.

Experiments

Setup:

Experiments

Experiments

Ablations:

​

• Number of pretraining steps.
• Communication frequency.
• i.i.d vs non-i.i.d data.
• Number of replicas.
• Model size.
• Outer optimizer.
• Adaptive compute pool.
• Dropping communications.
• With a single worker.

Experiments

Experiments

Experiments

Experiments

Experiments

Experiments

Experiments

Limitations

In the paper:

• Small model size (60M -> 400M).
• Diminishing returns with .
• No experiments with heterogeneous devices.

​

Additional limitations:

• Name of the algorithm block.
• No study of the impact of “overtraining regime”.
• Unfair comparison with post-local SGD (different settings).
• No comparison with established “distributed adaptive opt.”

Limitations

Copy-pasted from [Reddi et al., 2020].

Limitations

“ DiLoCo presented here assumes that all workers are homogeneous” p.10

DiLoCo is synchronous:

-> if heterogeneous device used, wait for slowest at each outer step.

Limitations

“in an overtrained setting with large amount of steps” p.7

​

-> for 60M model, k=8:

• “Chinchilla optimal” [Hoffmann et al., 2022]: 256 steps/worker (1.08B tokens).
• DiLoCo: 88 000 steps/worker (369.1B tokens).

​

​

-> Theoretical result from [Reddi et al., 2020]:

​

“the number of local iterations can be large when T is large”

​

​

What is the impact on DiLoCo? Is overtraining necessary to effectively use large H?

Limitations

• From FL: local SGD with adaptive opt. & non-i.i.d data is hard:�-> naive algo can diverge [Chen et al., 2021, Wu et al., 2023].
• Specific optimizer have been crafted for this problem:�Mime Adam, Local AMSGrad, FedAMS, FedAdam, FAFED
• DiLoCo “heterogeneous federated learning [...] we use this setting by default”:
• no comparison with existing methods.
• baseline used (adaptive opt + SGD) may diverge.

Conclusion

• Distributed training of LLMs with poorly connected machines seems possible :
• Use AdamW inner + SGD w/ Nesterov outer!�
• Important questions still unanswered:
• Possible/efficient to use heterogeneous devices?
• Previous adaptive opt. from FL do not work for LLMs?
• Large H only possible with “overtraining”?
• Conclusions hold for larger models/number of replicas?�
• Still preprint: correct attribution of Algo 1 + fair comparison with previous work + missing literature.

Questions