ACCO: Accumulate while you Communicate, Hiding Communications in Distributed LLM Training

We thank reviewer by2S for their thorough analysis of our paper, for acknowledging the strengths of our method, and for the suggested literature.

W1:

”PipeSGD is not cited or discussed.”

We thank the reviewer for this relevant citation that we missed, and will add it to our related work section in the next iteration of our paper. We will also make clear in the different parts of our paper that “DPU” is not the only name of this particular method, but that it also appeared in other papers under other names.

W2:

”PipeSGD + delay compensation has been proposed in SAPipe”

We thank again reviewer by2S for this highly relevant citation! We will make sure to discuss it in the next version of our paper. SAPipe introduces 3 strategies (see section 3.2 of SAPipe):

1. WP with local gradient in the current step,
2. WP with the latest synchronized gradient,
3. WP with all the above combined and the delay compensation technique from [Zheng et al., 2017].

However, if we partition the optimizer’s states, i.e. split the All-Reduce operation in two parts All-Reduce = Reduce-Scatter + All-Gather (see NCCL doc for visuals on collective communication primitives), methods (1) and (3) are not implementable, which is a major drawback to train LLMs due to memory constraints (see line 60).

Indeed, see in Fig.1 that the “Opt step” is sandwiched between the two communication parts Reduce-Scatter and All-Gather in the communication stream, which is run in parallel to the computation stream. Thus, when the Opt. step is done, only half of the communication load is finished, meaning that the computation is not finished either (the computation continues until the communication finishes) and the “local gradient in the current step” is thus not available because only half of it has been computed so far. This means that methods (1) and (3) cannot be implemented with sharded optimizers as they require access to the current gradient in the optimizer update. Moreover, sharding the optimizer leads to an update that can only be made in in conjunction with a communication: line 11 in Algorithm 3 of SAPipe cannot be done locally but only while communicating.

However, method (2) re-uses the last averaged gradient to make a prediction, which is exactly the method we experimented with under the name “ACCO-WP” (and is inspired by [Chen et al., 2019] ) - however, we show in our ablations line 391 that ACCO completely outperforms ACCO-WP thanks to the novel interleaved update method we introduce. We will add the reference to SAPipe when talking about ACCO-WP in the final paper, and how ACCO allows us to overcome its limitations.

Questions:

1. The main reason for this interleaved update method using half-accumulated gradients is to make the method compatible with sharded optimizers, as the algorithms using locally accumulated gradients are not compatible with this (see the remark above on SAPipe). Thus the idea of “using the locally accumulated gradient without averaging across workers” is impossible in practice with sharded optimizers.
2. We agree that our theoretical analysis is not the most precise as it is, and the ones in SAPipe could help us in this regard. However, as the theoretical conclusions for standard 1-step delayed SGD for smooth functions do not seem to apply in practice for using AdamW to train Transformers (see our experiments with DPU), we preferred to rather focus on practical and empirical results for our paper.

We thank reviewer by2S for their suggestion.

We agree that it is possible to slice a local gradient vector and use it in the opt step. But we would also remark that our goal is to run computations in parallel to communication and use a sharded optimizer. Whatever strategy we may want to use, this means that the opt. step is done in the communication stream after the reduce-scatter and before the parameter all-gather while some gradient is computed in the computation stream. Thus, during the opt. step at time $t$, we have access to the following information:

• The gradient computed locally at the last time step $t-1$ that we can slice.
• The averaged gradient (after reduce-scatter) computed at the last time step $t-1$.

But we do not have access to the gradient that we are still computing locally at time $t$ because this one will finish to be computed after the all-gather operation that is done after the opt. step. Thus, strategies (1) and (3) of SAPipe are not implementable.

If the optimizer states were not sharded, then the opt. step could be done in the computation stream right after a gradient computation and the SAPipe strategies could be implemented, but this is not the case here.

We thank reviewer by2S for precising their concerns, and hope to clear our explanations in the following.

"how could the computation of the local gradient be finished after the parameter all-gather?".

When training Transformers with a sharded optimizer, we have two versions of the parameter vectors $x$ and $x_{\text{slice } i}$ on each worker, one solely to compute gradients in half precision, the other to update the model parameters in full precision:

• $x$ is the full vector of model parameters $\in\mathbb R^d$, stored in half precision (brain float 16) used to compute gradients.
• $x_{\text{slice } i}$ is a full precision version (float32), but part of the optimizer's state, i.e a $1/N$ slice $\in \mathbb R^{d/N}$ (we assume that $d$ is divisible by $N$ for simplicity) used to update the parameters.

If we have a communication/computation overlap and use a sharded optimizer, then we have 2 parallel processes at step $t$ on worker $i$:

• $\nabla$ computation: while the communication is not finished, use $x_t$ to compute gradients $g_i^{(t)}$.
• communication: Reduce-Scatter the gradients computed at last step $( \frac 1 N \sum_i g_i^{(t-1)} )\_{\text{slice } i}$ $\rightarrow$ use the sharded optimizer to update the $1/N$ slice $x_{\text{slice } i}^{(t)}$ with an Opt. step and thus create $x_{\text{slice } i}^{(t+1)}$ $\rightarrow$ All-Gather a half-precision version of parameters $x_{\text{slice } i}^{(t+1)}$ to create the full parameters (in bf16) $x_{t+1}$ for the next step.

We hope that our explanation makes clear how the computation of local gradients is finished after the parameter all-gather.

"you suddenly lose the access to the local gradient?"

We must emphasize that the optimizer being sharded (local worker $i$ has only access to a $1/N$ slice of the optimizer states), an optimizer step can only be made in conjunction with a communication. Without a communication, given a local gradient $g_i \in \mathbb R^d$, each worker $i$ could only update a slice of size $\mathbb R^{d/N}$.

What we are saying is that for its weight prediction methods (1) and (3), SAPipe requires an update spanning over the whole local model $\in \mathbb R^d$ using local gradients $g_i \in \mathbb R^d$ but without communications (line 11 of Algo.3 of the SAPipe paper). This is not possible with sharded optimizers, as the local optimizer shard can only update a slice in $\mathbb R^{d/N}$.

Note that ACCO only uses gradients averaged over the whole $N$ workers in its parameter updates and weight predictions (equations (1)-(2)), meaning that the update and weight predictions are done with a communication, allowing updates spanning the whole model $\in \mathbb R^d$.

We thank reviewer j9YN for their kind comments, and are happy that they find our method very novel and recognize its strengths.

”Could you provide model performance v.s. number of workers like Figure 3?”

In Tab.2, the pre-training experiments with 8 and 32 GPUs were done on different hardware (see line 421): the 8 GPUs experiment was performed on a single node of H100, and the 32 GPUs one was on the A100 cluster described line 307. Moreover, due to hardware availability constraints, we stopped the experiments on the 8 GPUs sooner than the ones on 32 GPUs (see the #tokens entry in Tab.3). So we are not able to fairly compare the two to this date. However, given more time (to secure the access to the necessary hardware for the corresponding length of time) we can re-run our experiments on 8 and 16 GPUs on our A100 cluster for the same number of tokens as our 32 GPUs experiment and display the results as in Fig.3. We will provide this experiment in our final version of the paper.

”Lacking in theoretical analysis. It will be better to see how $\tilde \theta$ affects the convergence.”

We agree that our theoretical analysis is not the most precise as it is. However, due to the interleaved update of ACCO, our attempted derivations do not appear to be that straightforward compared to classical “delayed SGD methods” either. Moreover, as the theoretical conclusions for standard 1-step delayed SGD for smooth-convex functions do not seem to apply in practice for using AdamW to train Transformers (see our experiments with DPU), we preferred to rather focus on practical and empirical results for our paper.

”It does not seem appropriate to claim "as a form of plain SGD with no delay"”

We agree with the reviewer on that point and will re-formulate the sentence in the next version of our paper.

”Do you update optimizer states in both Eq(1) and (2), or only Eq(2)?”

This is a great question, thank you! We only update the optimizer state for Eq (2). In our implementation, for Eq (1), we checkpoint the optimizer states right before taking the optimizer step and reverse to them right after, so that they are only updated for Eq (2).

We thank reviewer tJ8M for their questions, and hope to clarify our contributions in the following.

”why the ACCO algorithm is particularly beneficial for training large language models (LLMs)?”

LLM training is particularly challenging because of the number of parameters in the model (typically in the tens of billions or more). This poses two challenges for distributed training with data parallel:

1. the GPU memories are maxed out before fully hosting the models and optimizers states copies (see the memory footprints in our Tab.1),
2. the communication load necessary to average the large gradient vectors over the distributed workers at each step can be so large that it takes more time to perform than actually computing them (see our Fig.3).

Thus, for LLM training, due to the models’ sizes and from a memory footprint viewpoint, it is necessary to use a sharding method (see line 60). But problem (2) is not tackled by partitioning methods such as ZeRO-1, leading to GPUs remaining idle for most of the time due to communication bottlenecks (see our Fig.3). ACCO is the first method allowing to tackle both problems (1) and (2) effectively, leading to significant time speedups (see Tab.3) compared to methods only tackling (1).

” Are there specific cases or experimental results that support the unique advantages of the ACCO algorithm for LLMs?”

Yes, in the two main LLM training tasks (LLMs pretraining and LLM finetuning), we showed in Tab.3 that ACCO allows to train LLM models as good as standard DDP with ZeRO methods, with a similar training dynamic and memory footprint, but significantly faster thanks to the parallel execution of computations and communications.

”does your approach offer direct contributions to memory optimization?”

We do not claim that ACCO introduces a new memory splitting method (it is the same as in ZeRO-1). However, ACCO is the first method that does both memory splitting and parallel execution of communication and computations, which is not trivial (as shown with our ablations and experiments with DPU). As shown in Fig.3 and Tab.3, memory optimization methods such as ZeRO are greatly impacted by communication overheads. On the other hand, previous methods parallelizing computations and communications such as CO2 cannot be seriously considered for LLMs training at scale due to the severe memory overhead they incur (see Tab.1). ACCO efficiently solves the two separate problems with a new approach.

”Memory cost is more than ZeRO-1 actually.”

This is indeed the case, but a natural by-product of the parallel execution of gradient computations and averaging: the two processes editing a “gradient” vector, they must refer to different memory slots. Thus all methods parallelizing computations and communications need an additional communication buffer, and ACCO is not different in that regard. However, the buffer being in half precision (see our Tab.1 and Fig.4 for an illustration), the memory overhead compared to ZeRO-1 is arguably negligible compared to the significant time speedups that it brings.

We thank reviewer a1EP for finding our method highly relevant and our experiments compelling.

Weaknesses:

• Numerical stability issues when slower devices are involved: first, we must emphasize that ACCO does not use "locally computed minibatch gradients" in its optimizer updates but rather averaged gradients over all workers (see Fig.1 and equations (1)-(2)). Thus if one worker is slow and do not have time to accumulate for multiple steps to reach a large batch-size locally, it does not affect the stability of the training procedure as every worker has access to the averaged gradient when updating parameters. However, we agree that the presence of arbitrarily slow workers can negatively impact the training speed for our method. In a typical data center, GPUs can crash for multiple reasons when training LLMs (see the Llama 3 paper). Our method does not handle this extreme case, and akin to what is done today when this happens, the solution is to remove the faulty GPU from the pool as it would only slow down the training.
• Explanation of sharding optimizers: we agree that a more detailed explanation of the sharding optimization would help clarify the paper for readers not familiar with this literature and thank the reviewer for this comment. We will add a paragraph detailing this in the final version of our paper.

Question:

The sharding method we follow is the same as ZeRO-1, but contrary to previous literature, ACCO allows to shard optimizer’s states across workers while running computation and communications in parallel and implementing an effective delay compensation method.

In a typical data parallel method for LLM training, each distributed worker needs to host in its local memory a copy of the model’s parameter in bf16, a gradient buffer in bf16, and for the optimizer: a float32 version of the parameters, and two float32 moment vectors. Taking the notations of our Tab.1, this means that for a model with $\Phi$ parameters, each of the $N$ workers hosts $2\times 2 \Phi$ Bytes for the model and gradient buffer, and $4 \times 3 \Phi$ Bytes for the optimizer states. As shown in our figure 4, here the sharding strategy allows each distributed worker to only host $\frac{4 \times 3} N \Phi$ Bytes for the optimizer states (the bf16 parts are untouched). This is done by splitting the All-Reduce communication primitive into two halves: Reduce-Scatter and All-Gather. In standard DDP, each worker computes its gradients, they are then averaged through all-reduce and each worker has thus access to the full averaged gradient before taking an optimizer step over the full $\Phi$ parameters. Here, each worker stores only $1/N$ of the optimizer states, and can thus only update a slice of the parameters. But as All-Reduce = Reduce-Scatter + All-Gather (see NCCL doc for visuals on collective communication primitives), instead of doing All-Reduce $\rightarrow$ Opt step, one can do Reduce-Scatter $\rightarrow$ Opt step over the $1/N$ slice $\rightarrow$ All-Gather for the exact same result and communication cost, but with a reduced memory cost, which we do (see our Fig.1).

In practice in our code, we just split the optimizer states in $N$ equal slices, each distributed worker only hosting its slice (which requires knowing the world size and the rank of the worker at initialization, which is common for distributed methods).

We thank reviewer aeVa for the suggested reference, we will make sure to add it to our literature review as it is relevant to our work. The “Elastic Scheduling” paper introduces a theoretical framework to analyze the convergence property of a relaxed version of distributed SGD, showing that some degree of computation/communication overlap can theoretically be possible for smooth functions, and experiments with a 20% computation/communication overlap with ResNet28 on CIFAR-100.

However, we respectfully disagree that our method is incremental compared to it : Elastic Scheduling requires 80% of a non-delayed gradient, making the computations and communications sequential for the most part whereas ACCO allows a perfect overlap of gradient computations and communications, which by design is a major difference and a huge advantage in terms of speed. By contrast, the purposes of ACCO are three-fold:

1. Running gradient computations and communications in parallel, hiding 100% of the communication costs (as opposed to some).
2. Allowing sharding the optimizer states.
3. Matching the training dynamic of standard backpropagation when training LLMs with the AdamW optimizer.

We believe that taken together and in consideration that our method is feasible for LLM is a significant step: “Elastic Scheduling” only partially address (1) (but does not achieve full overlap) for SGD with ResNets. However, ACCO is the first method perfectly meeting the three goals together and evaluated on LLMs.

We emphasize that prior work has shown evidence that delayed methods work fine in theory for SGD on smooth functions, and in practice in the case of using SGD for training ResNets (cf e.g. “Elastic Scheduling”). However, we demonstrate empirically that for AdamW on Transformers, this is not the case (cf our experiments with DPU). This calls for mitigating methods. Given the nonlinear update rule of AdamW and the regularity of Transformers’ loss function, our purpose is not to provide a theoretical answer (that would be harder to derive than the analysis for “SGD with smooth functions”), but rather a practical method that is demonstrated empirically. We show in our ablation (line 392) that not all mitigating methods are equally good at compensating delays in this context, but ACCO does it perfectly.

While “Elastic Scheduling” provides a theoretical analysis for SGD with smooth functions, it does not experiment with training Transformers with AdamW, only hides 20% of the communication cost ($\beta=0.8$, use of the $L0$ norm in their experiments page 8), and does not try to implement (2), which is not straightforward nor possible for all methods designed for (1) and (3) (see line 123, or our answer to reviewer by2S). ACCO does all that.

”There is no insight into why the performance gain happens in some cases and not in multiple other cases”

Would it be possible to point us to the precise inconsistencies you see to help us clarify our paper please? We demonstrated in Fig.3 that communication is a bottleneck, our goal is to design a method matching the training loss and memory requirement of standard backpropagation (DDP with a sharded optimizer) while being faster. In terms of loss/sample seen, all our Figures show that DDP and ACCO behave in the same way. The difference in the time gains reported in Tab.3 depend on the size of the model, the number of GPU nodes, and the communication bandwidth between the GPUs: all of these variables make the communication bottleneck more or less severe and are reported in Tab.3.

”Can you please explain why ACCO does not outperform DDP in Figure 6 results?”

Figure 6 shows the evolution of the training loss with respect to the number of training samples for 3 methods: DDP (baseline), ACCO-wp (ablation) and ACCO (our final method). The results show that ACCO-wp is not as good as our final method, ACCO, at countering the negative effect of the delayed gradient. The ACCO curve (blue cross) perfectly overlaps the DDP curve (plain purple), demonstrating that it perfectly compensates the delay. Tab.3 shows that it outperforms DDP in terms of training time.

We thank reviewer aeVa for their answer, and hope to completely address their concerns in the following.

There are other existing works that derive it for non-smooth objectives

We must precise that we are talking about the theoretical study of optimizers with non-linear update rules such as Adam (not SGD) on less regular functions than the classical setting of "convex-smooth" that is far from truthfully depicting the loss landscape of Transformers. This is still an open area of research, see e.g. the conclusion of "Advances in Asynchronous Parallel and Distributed Optimization, Assran et al."

Are you suggesting that experiments with AdamW is a contribution?

Not exactly: we are affirming that algorithms using a delayed update may work easily with SGD on convolutional neural networks in practice but not when using Adam on Transformers. See for example Fig.4 of the SAPipe paper where Pipe-SGD (other name for the "1-step delayed update" or DPU) is shown to work perfectly well for SGD with VGG16 and ResNet50 on ImageNet/CIFAR10, but does not work at all for Adam on Transformers/GPT2. This proves there is a discrepancy between the optimistic conclusions of both the theoretical analysis of "delayed update methods for SGD on smooth functions" and simple experiments with SGD on CNNs, and what is observed in practice for Adam on Transformers.

Our contribution is thus a offering a method that does work for AdamW and Transformers.

Sharding in itself is not a novelty, and computation-communication overlap is not a novelty

We agree on this statement. However, having both at the same time is a novelty. We must insist that while sharding is straightforward to apply for standard backprogation, it is highly non-trivial to manage to design a method that allows to shard the optimizer's parameters in a "delay-compensating" framework (see our answer to reviewer by2S). The main reason for that is because many "weight prediction" methods use "local gradient approximates" in their update rule, which means that each worker must host a full optimizer and not a sharded one.

We summarize in the following our contributions:

Dear reviewer aeVa,

We thank you for your feedback and believe that your primary concern regarding the novelty/importance of our work may stem from some misunderstanding.While we agree that the working implementation of ACCO is a contribution in itself, we must emphasize that ACCO is a novel algorithm with a design justified through a comprehensive ablation study. ACCO is indeed motivated by performance issues (here, to tackle communication and memory bottlenecks when training LLMs) in the same fashion as many recent advances in optimization for deep learning (e.g. local SGD for the Federated Learning use case), but the solution we propose is an algorithmic one and not a mere implementation trick.

Our novel algorithm vastly differs from previous ones in terms of design (see the unique interleaved update in Fig.2), and allows for implementation advantages impossible to obtain from previous methods (100% communication overlap, optimizer sharding).