DiLoCo: Distributed Low-Communication Training of Language Models

3 months ago by cm0002 to c/Aii

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.

Training LLMs requires thousands of accelerators that need to be wired together with ultra fast cables because they constantly share intermediate states and gradients during every single step of the optimization process. Building and maintaining a single compute cluster of that size is something that only large corps can afford to do.

Instead, we could piece together spare computing resources that happen to be scattered across the world. A new distributed optimization algorithm called DiLoCo, which stands for Distributed Low Communication, allows training language model across multiple smaller islands of devices that are poorly connected.

Rather than forcing the chips to sync up after every single calculation they let each island train independently for hundreds of steps before updating the global parameters. This approach takes heavy inspiration from a concept called Federated Learning where devices do a bunch of work locally before syncing their model replicas.

The technical breakthrough here is in the way they structure the math behind the inner and outer learning steps using an optimizer called AdamW for the local training steps on each island and then apply Nesterov momentum when gathering and blending all those outer gradients together.

Their approach was tested on the popular C4 dataset and proved that an 8 worker setup performs just as well as a fully synchronized model while communicating 500 times less. It is also remarkably resilient since the model still converges even if the data distribution is wildly different on every worker or if compute islands randomly drop offline during training.

This is a massive step towards being able to train models completely outside corporate control, paving the way towards genuine community driven projects developing this technology in open source fashion. It would be possible to build something akin to SETI@Home where people contribute some of their idle computing power towards training that's done using a distributed network of regular computers across the world.

ExperimentalGuy 2 points 3 months ago

It's crazy how its robust to boxes being taken offline during training. That's honestly magic.

save

path: 0 24882424, hotness: undefined, score: 2, children: 0