Local AI is not just about inference

Many enterprises already operate compute across multiple offices and data centers worldwide. Taken together, that distributed compute can support AI workloads that wouldn’t fit or scale within a single site. At the largest scale, Google, Microsoft, and OpenAI have reached the limits of single data centers and are moving toward multi-data center training.

Inside a data centre, GPUs are connected over a network bandwidth of hundreds of gigabytes per second, which can be more than 10,000x faster than the Internet. Is it possible to train or fine-tune a high-performant large language model over such a slow connection without taking an unfeasible amount of time?

An outline of the blog

We're going to discuss how to answer this question, starting by introducing the current most popular approaches. We focus on data-parallel techniques, since they are the most applicable for these cases. One of these methods is DiLoCo, which has become state-of-the-art for LLMs and can perform as well as classical data-parallel training (referred to as DDP) inside a data-centre. However, it remains a mystery why it performs so well.

After introducing and explaining how DiLoCo works, we deep dive into DiLoCo to understand where its true value lies. We argue that DiLoCo should not be viewed as a new distributed technique compared to classical federated averaging (FedAvg). Rather, DiLoCo is an amazing method which enhances single-node training and as a consequence, it also improves distributed training. Although this observation can initially be made without empirical evidence, throughout the blog we showcase several empirical implications of it.

We go on to analyze DiLoCo as a single-node technique, to understand why it performs better than regular training. First, we hypothesise that DiLoCo implicitly makes a model take larger steps during training. We find evidence for this by showing that other variants of DiLoCo, which had been thought not to work well, can perform as well as DiLoCo by also ensuring that they take large steps. We also analyze some recent theoretical results to support this. In a second approach, we argue that DiLoCo falls in a general category of techniques exploiting the broader training trajectory of the model. We consider such a method, called Checkpoint averaging, and find that it can outperform DiLoCo - further, we can combine both to achieve even better performance.

Then, we return to distributed training to introduce a new and simple distributed method, OverlapSPARTA, that can achieve zero communication overhead in low-bandwidth environments. This is a purely distributed technique that can achieve zero communication overhead. When combined with the single-node techniques we discuss throughout the blog (like single-node DiLoCo and CheckpointAvg), it outperforms both DiLoCo and DDP per training step — and since there is no communication overhead, training is as fast as it would be inside a data centre. In the conclusion, we discuss further and related work and pose some interesting questions.

An intro to distributed training

We will start by briefly introducing how model training works and go over the most common approaches to distributed training. When training or fine-tuning a model, we perform several consecutive training steps. At each training step, the model is updated using a gradient and an optimizer.

The gradient is computed based on a small batch of training data and the current model. This gradient computation represents the majority of the work done during the training step. The optimizer tells us how to use the gradient to update the model. It is usually a sophisticated version of stochastic gradient descent (like Adam or Muon).

In distributed training, we want to leverage several compute nodes to parallelize the work done during each training step. With optimal parallelization, we can reduce the time a training step takes by a factor equal to the number of nodes.

The slow network problem

To fully utilize the total compute power of the nodes, a large amount of communication between them is required. If the network connection is slow, this becomes the bottleneck, making the parallelization useless: if we spend all our time communicating, we might even take more time than if we had used only one node.

Having a slow network connection is not an unrealistic scenario. For example, a company might have GPU clusters (like an EXO cluster!) in different offices around the world that their employees use. If they want to aggregate all that compute to, for example, fine-tune a large model for a specific domain/style using lots of new documents from all its offices, they will probably have to do this over a really slow network, like the Internet.

Fortunately, a lot of research has been done to address this problem and we will introduce some of the most popular approaches in this section. In order to put these methods in context and so that we can track our progress throughout the blog, it will be good to have a realistic example use-case in mind. Here are estimate numbers of a fine-tuning run on a large amount of data which we will use as our example set-up.

For cases like this, where each node is already a cluster with a large amount of total memory, the popular method is data-parallel training. This means that parallelization will occur along the batch dimension, i.e., the batch of data used at each training step will be split across the nodes.

The common approaches to data-parallel training

The standard data-parallel training strategy, referred to as DDP (for distributed data-parallel) is to have at each training step, each node computing, in parallel, the gradient of an individual batch. In our case, where we have 8 nodes of the same type, each individual batch would have an 1/8-th of the size of the global batch. This means an individual batch size with 12.5k tokens.

After computing their individual gradients, the nodes communicate via an all-reduce, which means that they aggregate these 8 individual gradients and compute their average. Using the resulting average gradient, which is exactly the gradient of the global batch of 100k tokens, the model can then be updated. In this way, we parallelize the global batch across the 8 nodes in order to perfectly mimic the training step we would perform if we had only 1 node.

Inside a data centre, this all-reduce communication can be done lightning quick (e.g., with NVLink) and there is essentially no communication overhead. This implies an 8x training time speed-up compared to using only one node!

The best and worst possible training times

Imagine if we used DDP in our set-up but we had infinite network bandwidth, which means no communication overhead. How much time would training take?

Each one of our nodes is composed of 32 M3 Ultra Mac Studios, connected with very fast networking since they are in the same location. This means that each node can achieve essentially zero communication overhead, that is, essentially full compute utilization, when processing its own individual batch.

32 M3 Ultra Mac Studios will have in total, almost 1 PFLOPS of compute power. As a very rough estimate, each node can process the individual batch of 12.5k tokens in around 6 seconds. With no communication overhead, the whole training run would then take only 6 x (num training steps) = 18000 seconds, which is only about 5 hours.

This is the best possible training time! It's essentially how much time training would take if we were inside a data centre with the compute we have. In our set-up however, we are actually faced with a network bandwidth of 50 Mbps, which is 2000x slower than Thunderbolt 5, the connection type between the M3 Ultra Mac Studios at each location. So, how much time would DDP training take?

In each all-reduce, we communicate each node's individual gradient. This is a vector of length 1 billion (with one element for each one of the 1B trainable parameters). When sent over a 50Mbps network in standard precision (16-bits per element), it will take at least 5 minutes to be received. In practice, an all-reduce between 8 nodes requires more than just one send - at best, it can take 10 minutes to finish.

This means that we would spend almost all our time communicating, more than 99% of the time, and training would take more than 100 times longer than it would inside a data centre - going from 5 hours to 20 days!

Low bandwidth methods - FedAvg and DiLoCo

How can we avoid having to train for 20 days? How do we avoid communicating so much? The state-of-the-art method for distributed training over low-bandwidth connections is DiLoCo (for Distributed Low Communication), introduced by Douillard et. al at Google DeepMind^[1]. DiLoCo is an extension of a simple idea:

What if instead of communicating at every training step like in DDP, each node performs several training steps before an all-reduce between the nodes happens?

This is commonly referred to as federated averaging.

FedAvg

Here, instead of having a synchronization at every training step, it only happens every, say, $H$ training steps. During these $H$ steps, the nodes train independently using their own individual batches without ever communicating.

At every $H$-th step, the nodes sychronize by performing an all-reduce of their models. This means that they aggregate their models and compute their average. All nodes are then updated to this average model.

Federated averaging is great since it reduces the communication overhead, as we are only communicating every $H$ steps. However, the drawback is that the lower synchronization frequency will imply worse performance than DDP.

To illustrate this, below are the loss curves of training a 50M parameter nanoGPT-style model with 4 nodes using the two techniques. FedAvg is performing 100 training steps between synchronizations and the individual training steps are using the Adam optimizer.

DiLoCo

DiLoCo resolves this degradation of training performance by substituting the all-reduce with a special type of synchronization, a global outer optimization step. This general idea had been previously introduced, and the authors of DiLoCo discovered that a specific optimizer, stochastic gradient descent (SGD) with Nesterov momentum, is the best choice for this global outer optimizer. They observed that it allows for the same, or even better, performance as DDP, while having, e.g., 100 times less communication overhead!

DiLoCo (H = 100) decreases the communication overhead by 100x compared to DDP. Since in our set-up we saw that DDP would give 1% compute utilization, DiLoCo (H = 100) increases this to 50%, with seemingly no loss in performance. This reduces the training time from 20 days to 10 hours.

How DiLoCo works

DiLoCo is an amazing solution to the slow network problem, allowing us to train/fine-tune models as well and only twice as slow as we would inside a data centre with DDP. But what precisely does it do? For those less interested in the details, feel free to skip to the next section.

As we mentioned in the previous section, DiLoCo is an extension of FedAvg. In both, synchronization between nodes only occurs every, say, $H$ steps. During these $H$ steps, each node trains independently.

It is at the synchronization that the two methods differ. DiLoCo treats each sequence of $H$ steps as one big outer training step, using it to perform an outer update of a global model. At every $H$ steps, the global model is updated and the synchronization between nodes happens by setting all nodes to this updated version of the global model.

To perform this outer update of the global model, there is an outer optimizer and an outer gradient. The outer optimizer is, like in usual training, a sophisticated version of stochastic gradient descent, namely, SGD with Nesterov momentum. As an outer gradient, DiLoCo uses the average change in the node's models over the last $H$ steps. More precisely, each node computes its individual pseudo-gradient — which represents how much that node's model has changed during the sequence of $H$ steps. These pseudo-gradients are then aggregated across all nodes, and the outer gradient is the average of the individual pseudo-gradients.

DiLoCo with $N = 1$ node is different from standard training

If we take a look at the FedAvg algorithm we described before, we can see that if we have only one compute node, it is the same algorithm as standard training. DiLoCo is not.

With only one compute node in the DiLoCo algorithm we described above, there is only one pseudo-gradient, so the all-reduce step does not happen. But, there is still a global model which is updated every $H$ steps and an outer optimizer step which uses the single-node's pseudo-gradient. This can be a different algorithm from standard training on one node, since additionally to the normal training steps, there is an outer update happening at every $H$ steps.

In fact, in the original DiLoCo paper^[1], it was shown that even single-node DiLoCo can perform better than standard single-node training.

Should DiLoCo be viewed as a new distributed algorithm?

DiLoCo is highly regarded as the state-of-the-art distributed technique for low-bandwidth training. Despite this, it is poorly understood and it is sort of a mystery why it works. In the following sections, we will try to unpack it, beginning with understanding where the true value of DiLoCo lies.

A thought experiment

To start, let's consider a thought experiment where we have two different distributed training algorithms. The first is DiLoCo, which we broadly reiterate.

1st algorithm: DiLoCo

• $N$ compute nodes which start with the same model.
• Every $H$ steps, the nodes sychronize with a global outer optimization step - this uses SGD with Nesterov momentum as an outer optimizer.
• Outside of these training steps, the nodes train independently on their own individual batches of data.
• Each synchronization step works by first aggregating the pseudo-gradients of the nodes with an all-reduce, and then performing an outer update of the global model using the average pseudo-gradient. All nodes are then set to the updated global model.

The second algorithm uses DiLoCo in the non-distributed setting along with FedAvg. In the previous section, we noted that DiLoCo with only one node is a different method than standard single-node training, and so this is not just classical FedAvg.

2nd algorithm: FedAvg with single-node DiLoCo

• $N$ compute nodes which start with the same model.
• Every $H$ training steps, the nodes sychronize exactly like in FedAvg, by performing an all-reduce of their models.
• Outside of these training steps, the nodes train independently on their own individual batches of data.
• (single-node DiLoCo outer updates) Before each synchronization, each node performs its own independent DiLoCo outer update using only its own pseudo-gradient.

This second algorithm only uses DiLoCo locally at each node. It is running FedAvg, with synchronization every $H$ steps, on $N$ nodes which are each individually and independently running the single-node DiLoCo training strategy (with outer updates also every $H$ steps, using independent and individual pseudo-gradients).

Which algorithm is better?

Since we showed that DiLoCo performs much better than FedAvg, you might at first bet that the first algorithm does better than the second, with the motivating argument that DiLoCo's global outer optimizer is a superior synchronization to the simple all-reduce in FedAvg. However, these two algorithms are actually mathematically equivalent!

If we run both algorithms with the same starting model as well as the same choice of individual batches at all training steps and nodes, they will produce the exact same models at all training steps.

What does this imply?

Before explaining why the algorithms are equivalent, note that this observation has a strong implication for DiLoCo. It tells us that DiLoCo's value might not be as a novel distributed training technique compared to FedAvg, but rather as a technique which improves single-node training, that is, a training optimizer improvement.

In the second algorithm, the single-node DiLoCo updates before synchronization use pseudo-gradients which represent changes from the independent and individual training runs of each node. If doing these updates improves individual training performance at each node, then the performance of the global averaged model should also improve. In fact, we already noted that single-node DiLoCo can outperform regular single-node training.

Why: the details...

To better understand what is going on and explain why the two algorithms above are indeed equivalent, we will now have to go into the details of the DiLoCo algorithm. If you are not too interested in the details, feel free to skip down to where we start discussing some empirical implications of this observation.

Let's rewrite the DiLoCo algorithm in a much more explicit form. What does it actually mean in DiLoCo outer update to perform the outer step $\theta_{j+1} := \text{OuterOpt}(\theta_j, \Delta_j)$ with a pseudo-gradient $\Delta_j$? If the $\text{OuterOpt}$ is plain stochastic gradient descent (SGD) with learning rate $\lambda$ and no momentum, it is clear that $\theta_{j+1} := \theta_j - \lambda \Delta_j$. But what if it is SGD with Nesterov momentum, as in DiLoCo?

Classical Nesterov momentum is a well-known convex optimization technique where at each training step, the gradient is not calculated at the current model, but at a prediction of where it would be based on the previous trajectory. So, how should it be applied here, where we are instead given a pseudo-gradient $\Delta_j$ to use?

Since the exact implementation is not detailed in the initial DiLoCo paper^[1] we can look at open-source implementations of DiLoCo and its variants - for example, AsyncDiLoCo by Google DeepMind and OpenDiLoCo by Prime Intellect. They use PyTorch's SGD implementation, which has a boolean Nesterov variable and allows for using a given gradient.

Following this, the update $\theta_{j+1} := \text{OuterOpt}(\theta_j, \Delta_j)$ when the OuterOpt is SGD with Nesterov momentum is defined as follows:

• (Velocity term) $v_{j+1} = \eta v_j + \Delta_j$
• (Update) $\theta_{j+1} = \theta_j - \lambda (\eta v_{j+1} + \Delta_j)$

for some hyperparameters $\lambda$ (outer learning rate) and $\eta$ (outer momentum). Notice that even with a gradient (and not a pseudo-gradient), PyTorch's Nesterov implementation does not calculate a gradient at an accelerated step - this seems to be intended as an approximation of true Nesterov.†

We can now expand the velocity term over several steps as $$v_{j+1} = \Delta_j + \eta \Delta_{j-1} + \eta^2 \Delta_{j-2} + \ldots + \eta^{j} \Delta_0,$$which lets us rewrite the DiLoCo outer update as

In particular, the outer update is linear in the pseudo-gradients. This means that in DiLoCo, we can commute this outer update with the all-reduce of the pseudo-gradients $\Delta_j = \frac{1}{N} \sum_i \Delta^{(i)}_j$, without changing the result of the algorithm. Instead of first aggregating the pseudo-gradients and then performing the outer update, as is written in DiLoCo, we can first perform the outer update on each node using its individual pseudo-gradient and then aggregate and do an all-reduce of the nodes' models, like in FedAvg.

Doing so creates the second algorithm we described above! Thus, the two algorithms are equivalent. DiLoCo is an instance of FedAvg running with a specific local training strategy at each node (single-node DiLoCo). For an even more detailed explanation, we have written a small proof of the equivalence of the two algorithms in the appendix.

The empirical implications

To illustrate some first implications of the observation we made, we compare DiLoCo (with H = 100) with an extended version of the second algorithm from before: we still do an all-reduce of models every 100 steps, but instead of doing single-node DiLoCo updates only every 100 steps, we do them every 50.

This results in the same communication overhead as DiLoCo (H = 100), but we should now expect better performance, since it benefits from more frequent (single-node) DiLoCo updates. Indeed, it has been shown^[1] that decreasing H in DiLoCo until a certain point can improve performance.

We see that the performance of having more frequent single-node DiLoCo updates is better, despite having no additional communication time. All we did was change the single-node updates and we did not touch the distributed part of the algorithm, which is still FedAvg.

A rigorous comparison

Given what we've learned, we can now also argue that comparing DiLoCo to DDP, as we did in Figure 2, is not a rigorous comparison of distributed training methods. Since we realized that DiLoCo can be rewritten as FedAvg with individual nodes running single-node DiLoCo, what we should do is compare DiLoCo to a version of DDP where we are also running single-node DiLoCo, which adds no communication overhead. Unsurprisingly, we see a loss gap similar to the loss gap between DDP and FedAvg in Figure 1.

Maybe DiLoCo should not be viewed as a distributed technique?

In simple terms, DiLoCo's value comes primarily as a single-node technique, since, under the hood, DiLoCo is the same as FedAvg with independent single-node DiLoCo updates. We will give more empirical implications of this throughout the blog. In the next section, we will look into why DiLoCo improves performance, even if primarily as a single-node training technique or in other words, as a new training optimizer.

DiLoCo as a single-node technique

In the beginning of the last section, we discussed how DiLoCo can be rewritten as FedAvg + single-node DiLoCo. This indicated that DiLoCo should not be interpreted as a new distributed training technique compared to FedAvg, but rather as a technique which improves single-node training. However, it is still unclear why it can improve training performance - the aim of this section is to understand this. Since we will only really be discussing single-node DiLoCo throughout, let's reiterate it so that it is clear.

Reiterating the DiLoCo update in the single-node setting

In the single-node setting, DiLoCo differs from standard training in the addition of an outer update which updates the model every $H$ steps. This uses an outer optimizer and pseudo-gradients, which are computed as the change in the model over each sequence of uniterrupted $H$ training steps.

Making the outer update explicit

In the last section, we saw that in DiLoCo, with the outer optimizer being SGD with Nesterov momentum, the outer update $\theta_{j+1} = \text{OuterOpt}(\theta_j,\Delta)$ can be rewritten as

where the $\Delta_j, \Delta_{j-1}, \ldots$ are the pseudo-gradients from the different runs of $H$ uniterrupted training steps, with $\Delta_j$ being the most recent one. Moreover, using the standard DiLoCo hyperparameters◊ of $\lambda = 0.7$ and $\eta = 0.9$ (as proposed in the initial DiLoCo paper), the outer update is

Given this rewriting of the DiLoCo update, there are two natural questions you probably have:

(1) Larger step and other outer optimizers

We start by discussing the first question. As an example to give intuition for why the numbers attached to the pseudo-gradients must be important, let's look at what would happen if the DiLoCo outer update was instead simply $\theta_{j+1} = \theta_j - \Delta_j$. As explained above, the pseudo-gradient $\Delta_j$ is the difference between the current global model $\theta_j$ and the model we have just before the outer update happens, call it $M$. Therefore, the update $\theta_{j+1} = \theta_j - \Delta_i$ is actually $\theta_{j+1} = \theta_j - (\theta_j - M) = M.$

This means the outer update does nothing and using it is the same as running standard single-node training. So, there must be something important about the weights DiLoCo uses, since it performs significantly better than standard single-node training. To understand what, let's first compare DiLoCo to using different outer optimizers instead of SGD with Nesterov momentum - thus having a different outer update. With other outer optimizers, the outer updates become:

In the original DiLoCo paper^[1], the authors perform a comparison between DiLoCo and the result of having these two SGD strategies as outer optimizers (as well as Adam). They show that DiLoCo with SGD + Nesterov momentum as an outer optimizer performs significantly better - below is a plot from their paper showcasing this while training a Chinchilla-style language model (from an already given starting point).

Do they really perform worse than DiLoCo?

In their comparison above, they evaluated the other outer optimizers with values of $\lambda < 1$ and $\eta = 0.9$. However, if we look at the explicit definition of these outer updates, doing this implies that, compared to DiLoCo (with outer SGD + Nesterov), the other outer updates will get a lower contribution from the pseudo-gradients - not only lower from the most recent pseudo-gradient $\Delta_j$, but likely the others as well.

Could this be bad? Well, the converse of each pseudo-gradient, $- \Delta_i$, tells you the change in the model over a sequence $H$ training steps. So, very informally, if this change is a good training direction (towards models of low loss), an outer update with a higher contribution from $-\Delta_i$ will make the model take a larger step in that direction.

To do this, we should use $\lambda > 1$ in the other outer optimizers! We test this and plot their performance, as well as of standard training and DiLoCo, during the single-node training of a 50M nanoGPT-style model.

We now get essentially the same performance as DiLoCo, specially with SGD with momentum. Further, to show that this also translates to the distributed setting, we also test the performance of SGD with momentum as an outer optimizer (with outer learning rate $\lambda > 1$) when we have 4 nodes, showing that it can do as well as DiLoCo, if we choose the correct hyperparameters.§

So, is the key larger steps?

We might actually be able to use other outer optimizers to achieve the same performance as DiLoCo. But we can think more generally - all the previous outer updates we considered are of the following form:

where the $w_i$'s are weights that come from a choice of hyperparameters (in the previous cases, a learning rate $\lambda$ and possibly a momentum $\eta$).

What really matters in the choice of these weights for good performance?

• Is it the size of $w_0$?

  Since we observed that plain SGD as an outer optimizer with $\lambda = 1.5$ performs better than standard training, a large value of $w_0$, that is, increasing the step size in the direction of the most recent pseudo-gradient, seems to have a speed-up effect on training.

• Is it the values of the first few weights, say, $w_0, \ldots, w_4$?

  DiLoCo and SGD with momentum (as an OuterOpt) seem to perform even better in Figure 7, and these have a lower value of $w_0$ than SGD ($\lambda = 1.5$) but non-zero values for the other weights $w_i$ with $i > 0$. So, the contribution from pseudo-gradients other than the most recent one, might be additionally important.

• Is it their sum $\sum_i w_i$?

  Mimicking the first two weights in DiLoCo with SGD with standard momentum as an OuterOpt showed very similar performance (sidenote §). However, they have a very different sum of the weights $\sum_i w_i$, so this might be less important than the value of the initial weights.

• Is an exponential-like weighting the best?

  Although both DiLoCo and SGD with momentum (as an outer optimizer) make the weights $w_0,w_1,w_2, \ldots$ behave exponential-like, it is not clear if this is the best type of weighting. To assess this, a more general sweep through different weighting schemes is necessary - the best choice of weights may end up being very dependent on factors like the model and other training hyperparameters.

Some theoretical insights - what happens in convex problems?

While writing this blog, the work of Khaled et. al ^[10] came out, where the aim is to understand the role of outer optimizers in data-parallel optimization algorithms for convex problems. Unlike with LLMs, with convex problems we are theoretically guaranteed to achieve the best model.◊ ◊ Note: This is the model that fits best with the training data — the global minimum of the training loss function.

Because of this, a lot of the research focuses on proving that certain optimization algorithms converge to this best model faster than others. Further, because of the theoretical guarantees, researchers often use convex problems to sketch out, or confirm, their hypotheses for more complicated problems, such as LLMs.

The authors of ^[10] performed an analysis of distributed training algorithms which use outer optimizers, like DiLoCo, providing reasonable convergence speeds for these when used in convex problems. They considered three options of outer optimizers: SGD, SGD with momentum and SGD with accelerated momentum.

In the previous section, we discussed at length how these outer optimizers have value as single-node methods and their inherint purely distributed method is FedAvg, for which convergence rates already exist. However, their convergence rate results offer incredibly valuable insights, even if we think of these techniques primarily as single-node methods.

Indeed, let's take a look at the convergence speeds they obtained for when different outer optimizers are used on top of standard training. (To keep things simple, we show their results assuming zero noise, $\sigma^2 = 0$.) To illustrate these results, we will compare how close to the best model each algorithm gets at a given step $t$. As a baseline, we use the convergence speed that has been proved for standard training - we use $d$ to denote the distance to the best model with standard training at step $t$.

What do these results tell us?

These results, particularly the second and third row, are a good fit with what we saw empirically in Figure 7 and Figure 8. Having SGD as an outer optimizer with $\lambda > 1$ performed better than standard training, and even better did SGD with momentum and $\lambda > 1, \eta < 1$.

They also prove a convergence speed for when the outer optimizer is SGD with an accelerated momentum. They show an improvement which grows with the number of training steps, which seems like a really good sign for DiLoCo. However, their definition of the outer update for this case is actually very different from DiLoCo.◆

The sum of the weights is crucial for convex problems

When considering the three outer optimizers they use in the general form of equation (1), the improvement factors in the convergence speeds they obtained, are of the same order as the inverse of the sum of the weights, $1/\sum_{i\leq R} w_i$, in all three cases.◇ ◇ Note: This is not difficult to derive in the first two cases (SGD and SGD with momentum), but requires several algebraic manipulations in their third case (SGD with accelerated momentum) by looking at their exact outer update in their Theorem 3.6 in ^[10]. This indicates that for convex problems, a larger value of $\sum_i w_i$ is crucial for quick convergence!

In practice, for LLMs, the sum of the weights might not be the most important parameter though, as evidenced by Figure 7 and Figure 8 - there, DiLoCo performed very similarly to using SGD + momentum as an OuterOpt (lr = 1.33, mom = 0.43) but it has a much larger sum of weights $w_i$.★ ★ Note: This sum is roughly 7 for DiLoCo (with the standard lr = 0.7, mom = 0.9) and 2.3 for SGD + momentum as an OuterOpt.

Nonetheless, these theoretical results confirm that adjusting the hyperparameters of other outer optimizers, specially using outer learning rates larger than 1 can give better performance. By having equation (1) in mind, this reiterates that larger steps in the direction of the pseudo-gradients could be the value of (single-node) DiLoCo.

(2) But why use pseudo-gradients at all?

There seems to be a clear value in appropriately choosing the weights attached to the pseudo-gradients in the outer update. But why use pseudo-gradients and a linear combination of them at all?

The pseudo-gradient tell us how much the model has changed over a sequence of H consecutive training steps, so they represent general trends in the training trajectory of the model. The concept of exploiting the overall trajectory of the model during training is not novel. One motivation for it is the observation that in several interesting learning problems, for example in LLM training, there is a stage of training after which the model remains in what looks like a convex-like basin - see, for example, Izmailov et al.^[9].

Models end-up at a convex-like basin

For an illustration of this, we perform a standard training run on a 50M nanoGPT-style model and consider three checkpoints (steps 22,000, 23,000 and 24,000) during the training run. This gives us three models and we can consider the 2-dimensional plane they define containing all the models which are a linear combination of them. We can also then calculate the losses at several models in that plane, to plot a visualization of the loss landscape of the plane.

Looking at the landscape, the loss function clearly seems convex in this plane. The reason for this is not 100% clear, specially since this is only a 2-dimensional plane in the space of models. A natural hypothesis is that maybe the loss function (in the space of all possible models) is actually composed of several isolated convex-like basins, where "good" models live - once the model finds one of these basins it remains there, since it is convex-like.

Even if we do not fully understand why this happens at the moment, we can use of the empirical observation. In particular, notice in the landscape that the average of these checkpoints should have significantly better validation loss than the current model itself! Indeed, this is marked in the plot, where the average of the checkpoints has significantly lower loss (roughly $3.55$) than any of the checkpoints themselves (at best, $3.62$).

Checkpoint averaging can perform better than DiLoCo

Using the average of the checkpoints is reminiscent of a classical technique in convex optimisation, called Polyak averaging, where a speed-up in convergence is achieved by evaluating the loss at an average of the models from all the optimization steps. A similar approach has been considered for language models, in the work of Kaddour^[7] and Sanyal et. al^[8].

The idea is that while a model is being trained, we are also evaluating an average of checkpoints, which, with the evidence from the loss landscape shown before, should have lower loss than the model itself.

Note that this evaluation of the checkpoint average model introduces essentially no compute overhead compared to standard training, since it is evaluated only every thousand steps, and if necessary, this evaluation can be performed on CPUs while the GPUs are working.

How does CheckpointAvg compare to DiLoCo? We plot below a comparison where the companion model in CheckpointAvg is evaluated every 1000 steps and 4 checkpoints are averaged (T = 1000, k = 4).

We can achieve better performance with this technique than we get with DiLoCo!♦ We can also try to join both techniques to see if we achieve even better performance. Below we plot the loss curve of using DiLoCo together with CheckpointAvg (that is, every 1000 steps, we are evaluating the average of four checkpoints of the DiLoCo run). That performs even better.

DiLoCo is also an average of checkpoints

Both DiLoCo and CheckpointAvg exploit, through pseudo-gradients and checkpoint averages, the overall trajectory of the model to perform an update or an evaluation. The main difference is that DiLoCo uses pseudo-gradients and updates the model, whereas CheckpointAvg uses model checkpoints and only evaluates at the average of the checkpoints. However, it is actually not too difficult to observe that the DiLoCo update can be written without pseudo-gradients and also as a weighted average of checkpoints.

Roughly, this is the case since the pseudo-gradients themselves are the difference between two model checkpoints.◆ Knowing this, we can ask a similar question to the one we asked before concerning what are the best weights to put in equation (1): what is the best weighted average of checkpoints to use for an update or an evaluation?

What is coming next?

So, having an update or an evaluation which is a linear combination of pseudo-gradients or of model checkpoints can have a big impact on training performance. The real gain and the reason why DiLoCo was performing better than using other outer optimizers in Figure 6 from ^[1] was due to the weights chosen in this linear combination - indeed, we were able to exploit this to achieve similar performance using SGD with standard momentum as an outer optimizer by appropriately adjusting the hyperparameters.

As we discussed before, the question of what the best weights are for these linear combinations would be great to understand, although its answer might have high variance, as it might be dependent on several training components (model, data, inner optimizer hyperparameters,...). The correct answer might even change based on what stage of training we are in. We discuss several other things that are left to be understood further about DiLoCo in the conclusion.

With the improved understanding we have of DiLoCo and adjacent techniques, like CheckpointAvg, we will now go back to distributed training. We'll introduce a new purely distributed technique to achieve full compute utilization in low bandwidth environments. In our set-up it achieves the best possible training time (half of DiLoCo) and when paired alongside the single-node methods discussed here (which add no communication), can significantly outperform DiLoCo in both training performance as well as time.

A new distributed training technique for no communication overhead

Having uncovered where the value of DiLoCo really lies, we can now return to our distributed training set-up. As we saw in the introduction, performing DDP, a classical training technique, would yield less than 1% compute utilization and require at least 20 days of training. In contrast, DiLoCo can achieve the same or better performance while communicating 100 times less.

We also noted that DiLoCo should be interpreted as an improvement to the training optimizer and not a new distributed method. We motivated this by showing that it is equivalent to the product of independent single-node DiLoCo runs and FedAvg. Under this perspective, we argued that the right comparison for DiLoCo is not against DDP alone, but against DDP also equipped with single-node DiLoCo updates. This was done in Figure 5, where we observed a performance gap similar to that between DDP and FedAvg.

At the start we said that DiLoCo ideally allows you to achieve the same performance as we would have inside a data centre but in low-bandwidth environments. However, taking into account that DDP + single-node DiLoCo can be performed inside a data-centre and performs significantly better than DiLoCo (see Figure 5), this is not really true.

How can we achieve full compute utilization and maintain good performance?

Currently in our setting, with DiLoCo or FedAvg (H = 100) we are achieving 50% compute utilization - increasing it to 100% will halve the training time to the best possible 5 hours! A first naive approach to do this is to further reduce the frequency of the DiLoCo updates (i.e., increase H). However, as demonstrated in the original DiLoCo paper^[1], this comes at the cost of performance.

We need a different idea. In fact, for zero communication overhead, we need to be hiding communication with computation. That is, the communication should be overlapped and happen while the devices are performing computations◆. How should we do that?

As we outlined at the start, in our set-up, a training step (i.e., processing the individual batches in parallel) takes 6 seconds and our network bandwidth of 50Mbps implies that an all-reduce of the 1B trainable parameters takes roughly 10 minutes. So, an all-reduce of 1% of the trainable parameters should take only 6 seconds (1% of 10 minutes), the same time as a training step.

So, why not simply do an all-reduce of 1% of the trainable parameters while the training step is being performed? This would go as follows.

The algorithm - OverlapSPARTA

The algorithm is very simple. Each node is initialized with the same model. At each training step, two processes are done in parallel:

• 1) Individual training steps: Each node performs their own training step on their own individual batch using their current model.
• 2) All-reduce of a small portion of the trainable parameters: For a set value of p (in our case p is 1% or slightly less), a p-fraction of the trainable parameters is chosen and the nodes aggregate their current values for these parameters and compute the average value for each.

The choice of p is precisely so that the all reduce, in which communication occurs, can be done while the individual training steps are happening so that it is hidden and incurs zero additional time.§ Once these are finished, we can set each node's model for the next step as follows: each node updates its model by keeping its locally trained values in 1) for most parameters, but replacing the chosen p-fraction with the averaged values from the all-reduce in 2).

To summarize, instead of sharing parameters infrequently, as in FedAvg, we share at every step, but only a small (random) portion of the parameters. This allows the communication to be performed while the training steps are being performed, resulting in no communication overhead and training will now take the best possible 5 hours!

What is the performance like?

The obvious question is whether this achieves good performance. At each step we share a delayed version of 1% of the parameters and this delay could hurt how well the model trains: these parameters are shared in the all-reduce in 2) but do not benefit from the training step in 1), as their locally trained values get overwritten by the averaged values from the all-reduce.

Fortunately, we find in practice that this does not seem to hurt performance! To showcase this, we start by comparing the performance of FedAvg with our new method, as well as a version of it where communication is not overlapped with computation. The latter had already been introduced by us as SPARTA - there, the partial all-reduce in step 2) happens only after the training step 1) has finished, thus not requiring overlapped communication; this implies also that the shared parameters no longer suffer from a training delay.

Method	Training Time	Theoretical Wall-Clock	Compute Utilization
DDP	20 days	~50×	<1%
FedAvg with $H{=}100$	10 hrs	1×	50%
SPARTA with $p{=}0.01$	~10 hrs	~1×	~50%
OverlapSPARTA with $p{=}0.01$	~5 hrs	~0.5×	100%

Looking at the loss curves, the methods seem to essentially perform the same per training step. However, the key insight is that OverlapSPARTA achieves 100% compute utilization, and thus, in the setting we are considering, reduces wall-clock time by half.

To do more in comparing FedAvg vs OverlapSPARTA per wall-clock time, we plot also below the loss curves of FedAvg with other values of H. In particular, FedAvg with H = 1000 would have at least 90% compute utilization in our example setting, but clearly underperforms OverlapSPARTA.

What about DiLoCo?

Because of our previous discussion that DiLoCo should not be interpreted as a new distributed algorithm, it was more rigorous to first compare our new method, a new and purely distributed technique, directly to FedAvg. We did that above, so we can now consider DiLoCo.

Just like how we argued that DiLoCo is essentially FedAvg with individual nodes performing single-node DiLoCo training, we should do the same to OverlapSPARTA. That is, we should try having the individual nodes in OverlapSPARTA performing single-node DiLoCo updates, which incurs no additional communication overhead.

Below is the comparison of doing this with DiLoCo itself. Again, we see the same performance per training step, but OverlapSPARTA + single-node DiLoCo updates reduces wall-clock time by half and has 100% compute utilization. We also add DDP to the plot.

Furthermore, we can do like we did in Figure 4 with FedAvg and again exploit the single-node nature of DiLoCo by increasing the frequency of single-node DiLoCo updates - as expected, this also gives a slight improvement in performance, without adding communication overhead.

What about CheckpointAvg?

Recall that previously, we introduced the CheckpointAvg method and showed how it can improve single-node training performance, possibly even better than DiLoCo. Interestingly, putting the two together gave even better performance (see Figure 12). Adding both to OverlapSPARTA as single-node methods (and thus no further communication) shows a similar improvement in the performance.

This means that OverlapSPARTA + single-node DiLoCo updates + CheckpointAvg achieves significantly better performance than DDP (and DiLoCo) but with 100% compute utilization, even at a bandwidth as low as 50Mbps!

But, just as we argued in the case of DiLoCo, for a truly rigorous comparison, we should compare this improvement not against DDP alone, but also equipped with these single-node techniques. Indeed, doing that recovers a clear performance gap, much like the one we saw in Figure 1 and Figure 5.

Conclusion

Distributed training in low-bandwidth environments is still an open problem. Although OverlapSPARTA combined with CheckpointAvg and single-node DiLoCo significantly outperforms DDP alone, Figure 19 shows that DDP (which can be done inside a data-centre) equipped with the same single-node techniques still performs better. We are not yet able to match data-centre performance.

We're not alone in tackling this — some amazing teams at Apple, Prime Intellect, Google DeepMind, Pluralis, etc. are actively researching solutions. We hope this blog contributes a useful perspective and sparks further interest in this area! For those interested in more details, we discuss future and related work and some interesting questions below.

Future extensions and related work

There are many questions about OverlapSPARTA which we chose not to dive deep into in this blog. For instance, is the method tolerant to node failures? How does it scale for a larger number of nodes, or for larger models? How should it be adjusted for heterogeneous devices? We are also currently choosing the set of parameters to be shared randomly with each parameter chosen independently with probability p, but it is not clear if this is the best approach — a better informed choice, maybe dependent on the inner optimizer, is worth exploring. These are questions we're actively thinking about — we'd love to hear your thoughts or ideas on any of them!

Several recent works have also been aimed at increasing the compute utilization of data-parallel training in low-bandwidth settings, without hurting or even improving performance. We discuss some of these below.

DiLoCo with overlapped communication

Two examples are works from teams at Google Deepmind on extensions of DiLoCo (Eager Updates For Overlapped Communication and Computation in DiLoCo^[14] and Streaming DiLoCo with overlapping communication: Towards a Distributed Free Lunch^[13]). The base idea is the same as in OverlapSPARTA - hide communication with computation. However, the approach is quite different.

Both of these works adapt DiLoCo to allow for overlapped communication. This means that the all-reduce of pseudo-gradients in DiLoCo happens while individual training steps are being performed, and so, when they are applied they suffer from a delay. To mitigate this delay they do an outer update for only a portion of the parameters at a time, so that the all-reduce takes less time.

A key difference between this and OverlapSPARTA is that in this blog, we have argued that DiLoCo should be viewed as a single-node technique. This made it so that when thinking of how to overlap communication, we did not have to think of DiLoCo at all, making the algorithm much simpler (no need to think of pseudo-gradients and outer updates and how to take care of these while overlapping). We just after added on top single-node DiLoCo at each node, which recovered the same performance as DiLoCo per training step at 100% compute utilization. Adding CheckpointAvg gave even better performance.

Secondly, the way the authors of Streaming DiLoCo choose the small portion of parameters to update is by taking a group of layers. Almost all models have less than 100 layers, which implies always sharing more than 1% of the parameters pseudo-gradients. At significantly low bandwidths, this implies that the all-reduce is finished only after many steps. It is interesting to understand if this large delay in parameter training can hurt performance, compared to using OverlapSPARTA which avoids this by choosing the p-portion randomly and the parameter p so that we always have at most a 1-step delay.

An interesting idea in these works which makes sense to try with OverlapSPARTA is the idea of not completely updating the model when it receives a delayed portion of the parameters/pseudo-gradients. Instead, as is done in Streaming DiLoCo, one might be better off updating the model to a weighted average of the current model and the one we would update to - maybe this helps mitigate a small performance degradation coming from using 1-step delayed parameters.

SparseLoCo

Another example of a very recent work increasing compute utilization is SparseLoCo by Covenant AI. Among other changes, this method adds to DiLoCo the idea of sparsification of the shared pseudo-gradients. This means that at each round only a "Top-k" portion of the parameters share their pseudo-gradients. The authors show that it can perform better than DiLoCo in certain sparsification scenarios and has lower communication overhead.

As we did with DiLoCo in this blog, SparseLoCo can also be rewritten as FedAvg + single-node SparseLoCo updates, revealing that it might also be more appropriately viewed as a single-node technique. The ultimate test for this is if OverlapSPARTA + single-node SparseLoCo achieves similar performance per training step as SparseLoCo itself, which would be interesting to know. Furthermore, SparseLoCo's outer update can also be rewritten as a variant of the general outer update equation (taking into account the sparsification). This could be interesting to look into to understand if SparseLoCo outperforms DiLoCo because of its weights in this equation or because of the sparsification itself (which leads to the non-"Top-k" parameters not being updated, unlike in DiLoCo).

Finally, it would be interesting to incorporate the idea of sparsification in OverlapSPARTA to further improve its performance - does this give insight into how we should choose the p-fraction of parameters shared at every step?

There is more to understand with DiLoCo

Throughout this blog, we've gained a deeper understanding of DiLoCo. We argued that it should be interpreted as an improvement to the training optimizer, and so not as a new distributed technique. We then explored why it is an improvement over standard training. In this exploration, one of the key insights was to re-write the outer update as a linear combination of the pseudo-gradients. We noticed that correctly choosing the weights (in general, making them larger) can ensure better training performance - this allowed us to achieve similar performance to DiLoCo with other outer optimizers, like SGD with standard momentum, and was further supported by the theoretical results of Khaled et. al^[10].

We hinted that maybe, making the weights in this linear combination larger (specially for the most recent pseudo-gradients) is beneficial because the pseudo-gradients reveal correct training directions and we want to push the model into those directions. But is this really what is happening?

To understand this better, there is a simple experiment we can perform. While running the DiLoCo algorithm (single-node) we can plot the loss landscape that shows what occurs during an outer update. Since the loss landscape will probably be convex-like, we might expect that the DiLoCo outer update will move the model closer to the centre (the lower loss region), moving it in the "correct" direction.

The loss landscape we need to plot is of the two-dimensional plane containing the models:

• The model that resulted from the previous outer update, that is, the first "Global model" in the single-node DiLoCo diagram.
• The model we have just before the outer update, that is, the "Model" in the diagram.
• The model after the outer update, that is, the final "Global model" in the diagram.

Something unexpected happens - even though DiLoCo's validation loss curve clearly shows faster convergence than standard training, this DiLoCo outer update seems to actually pull the model to a point of worse loss! We do indeed see a larger step taking place, but it seems to be either too large or in the wrong direction. However, it still performs better in the long run.

Understanding this would be very interesting - clearly, the general assumption that DiLoCo performs better and so we should see the model at a place of lower loss after the outer update is performed is wrong. Why?

References

1. Douillard, Arthur, et al. "Diloco: Distributed low-communication training of language models." arXiv preprint arXiv:2311.08105, 2023.
2. McMahan, Brendan, et al. "Communication-efficient learning of deep networks from decentralized data." Artificial intelligence and statistics, PMLR, 2017.
3. Reddi, Sashank J., et al. "Adaptive Federated Optimization." International Conference on Learning Representations, 2021.
4. PyTorch Documentation: SGD.
5. EXO Labs Blog.
6. Charles, Zachary, et al. "Communication-Efficient Language Model Training Scales Reliably and Robustly: Scaling Laws for DiLoCo." arXiv preprint arXiv:2503.09799, 2025.
7. Kaddour, Jean. "Stop wasting my time! saving days of imagenet and bert training with latest weight averaging." arXiv preprint arXiv:2209.14981, 2022.
8. Sanyal, Sunny, et al. "Early weight averaging meets high learning rates for llm pre-training." arXiv preprint arXiv:2306.03241, 2023.
9. Izmailov, P., et al. "Averaging weights leads to wider optima and better generalization." 34th Conference on Uncertainty in Artificial Intelligence, UAI 2018.
10. Khaled, Ahmed, et al. "Understanding Outer Optimizers in Local SGD: Learning Rates, Momentum, and Acceleration." arXiv preprint arXiv:2509.10439, 2025.
11. Apple. Partial Parameter Updates for Efficient Distributed Training. arXiv preprint arXiv:2509.22418, 2025.
12. Prime Intellect. OpenDiLoCo: An Open-Source Framework for Globally Distributed Low-Communication Training. arXiv preprint arXiv:2407.07852, 2024.
13. Douillard, Arthur, et al. "Streaming DiLoCo with overlapping communication: Towards a Distributed Free Lunch." arXiv preprint arXiv:2501.18512, 2025.
14. Kale, Satyen, et al. "Eager Updates For Overlapped Communication and Computation in DiLoCo." arXiv preprint arXiv:2502.12996, 2025.
15. Ramasinghe, Sameera, et al. "Protocol Models: Scaling Decentralized Training with Communication-Efficient Model Parallelism." arXiv preprint arXiv:2506.01260, 2025.

Appendix

A. Equivalence of DiLoCo and FedAvg with single-node DiLoCo updates

Why both experiments are the same. Let's reiterate the second algorithm. During a sequence of consecutive $H$ steps, each node has the same starting model $\theta_j$ and then does $H$ independent training steps without communicating. After that, each node performs their own single-node DiLoCo outer update. Then they do an all-reduce of their models to create the model $\theta_{j+1}$ that each node will start the next sequence of $H$ steps with.

More precisely, let $\theta^{(i)}_j$ denote the model node $i$ has after the $H$ independent training steps, and let $\Delta^{(i)}_j := \theta_j - \theta^{(i)}_j$ be its own individual pseudo-gradient. Let $w^{i}$ denote the model that node $i$ is updated to after its own single-node DiLoCo outer update, i.e., $w^{i} = \text{OuterOpt}(\theta_j, \Delta^{(i)}_j)$. Then, as we saw in the previous section, we have:

• (Individual update) $w^{i} = \theta_j - \lambda(1+\eta)\Delta^{(i)}_j - \lambda \eta^2 \Delta^{(i)}_{j-1} - \lambda \eta^3 \Delta^{(i)}_{j-2} - \ldots$.

• (All-reduce) $\theta_{j+1} = \frac{1}{N}\sum_i w^{i}$.

B. DiLoCo update as a weighted average of checkpoints

Rewriting the DiLoCo update in terms of checkpoints. The DiLoCo outer update can be expressed as a weighted average of model checkpoints rather than using pseudo-gradients.

For example, if the OuterOpt in DiLoCo were SGD with learning rate $\lambda = 1.5$, the update $\theta_{j+1} = \theta_j - 1.5 \Delta_j$ can be rewritten as an exponential average: $$\theta_{j+1} = 1.5 \theta'_{j+1} - 0.75 \theta'_j + 0.375 \theta'_{j-1} - \ldots$$ where each $\theta'_i$ denotes the model checkpoint just before the $i$-th outer update.

For DiLoCo with SGD with Nesterov momentum (learning rate $\lambda = 0.7$ and momentum $\eta = 0.9$), the rewriting is: $$\theta_{j+1} = 1.33 \theta'_{j+1} + 0.13 \theta'_j - 0.29 \theta'_{j-1} - 0.2\theta'_{j-2} - 0.03\theta'_{j-3} + \ldots$$ where the remaining terms have negligible weight.