M AXIMIZING C OMMUNICATION E
FFICIENCY
FOR L ARGE - SCALE T RAINING VIA 0/1 A DAM
Conglong Li Microsoft Minjia Zhang Microsoft
Christopher De Sa Cornell University Yuxiong He
Microsoft
A BSTRACT 1-bit gradient compression and local steps are two representative techniques that enable drastic communication reduction in distributed SGD.
Their benefits, however, remain an open question on Adam-based large model pre-training (e.g. BERT and GPT).
In this paper, we demonstrate the non-linearity in Adam causes slow convergence even when 1-bit compression or local steps are individually applied.
To alleviate this limitation, we propose 0/1 Adam that linearizes each Adam step via approximating its optimizer states using their stale estimates and linear correlation.
0/1 Adam performs an Adam-like step to preserve the adaptivity, while its linearity allows utilizing 1-bit compression and local steps simultaneously for wall-clock time speed up.
We provide convergence guarantee for 0/1 Adam on smooth non-convex objectives.
On various large-scale benchmarks such as BERT-Base,[-BERT-Large, GPT-
-]
BERT-Large, GPT-2 pre-training and ImageNet, we demonstrate on up to 128 GPUs that 0/1 Adam is able to reduce up to 87% of data volume, 54% of communication rounds, and achieve up to 2 × higher training throughput and end-to-end training time reduction compared to the state-of-the-art baseline 1-bit Adam; while enjoying the same statistical convergence speed and end task model accuracy on GLUE dataset and ImageNet validation set.
I NTRODUCTION
Over the past few years, we have witnessed outstanding performance of foundation models on many applications.
However, these models, including BERT Devlin et al.
and GPT Radford et al. (2019a);
Brown et al.
(2020), usually have hundreds of millions or even billions of parameters and require to be trained on massive GPUs.
For example, the largest dense transformer model, 530B MT-NLG Smith et al.
(2022), was trained over 4000 GPUs in more than a month.
At this scale, the expensive communication overhead across computing processors and servers hinders the scalability (Alistarh et al., 2017).
1-bit gradient compression and local steps are two representative methods to mitigate the communication bottleneck.
1-bit compression drastically reduces the communication volume by quantizing each value in gradients with ultra-low bits (i.e., as low as 1 bit) Seide et al.
; Bernstein et al.
(2018a); and local steps alternatively saves the bandwidth by periodically skipping communication rounds(Stich, 2018).
While these techniques demonstrate tremendous success on distributed SGD, their benefits over large-scale Adam-based model training, such as for BERT and GPT pre-training, remains an open question (Kingma and Ba, 2014; Wang et al., 2019a).
Comparing to SGD where the model parameters are linearly dependent on the gradients, the non-linearity in Adam updates (Kingma and Ba, 2014) limits the direct usage of compression or local steps.
In particular, this non-linearity incurs two challenges: 1) when aggressively compressing the gradient such as with 1-bit quantizer, all the coordinate-wise effect learning rate will become the same value, so that Adam no longer enjoys adaptive and fast convergence; 2) to ensure all parallel workers reach consensus on the optimizer states, which is critical for convergence, the existence of non-linearity incurs the overhead of iteratively synchronizing the states when using local steps.
Tang et al.
(2021) undertook the first investigation of fixing this non-linearity towards compression and proposed 1-bit Adam .
The algorithm follows a two-stage training paradigm: first run Adam with full-precision communication ( full-precision stage 1 ); and then switch to 1 bit when the variance becomes stable ( compression stage ).
While this paradigm avoids compressing non-linear information with a one-time frozen variance, the experimental results from (Tang et al., 2021) indicate the fullprecision stage still incurs non-trivial overhead.
Furthermore, 1-bit Adam is restricted in the scope of gradient compression, and cannot be trivially adapted when other techniques are used, such as local steps.
Besides, the empirical success of (Tang et al., 2021) was not substantiated on GPT-style models, for instance, 175B GPT-3 Brown et al.
(2020), 530B MT-NLG Smith et al.
(2022), etc.
In this paper, we address this gap by proposing 0/1 Adam .
Adam breaks the barrier of non-linearity from two aspects: first it adaptively freezes variance, so that given agreement on a stale variance state, the parallel workers only need to communicate momentum that is linearly dependent on the model update; This technique allows reducing the previous two-stage compression scheme to a unified single stage; 2) it leverages the insight that in adjacent Adam steps, the changes to optimizer states are generally bounded, so that with frozen variance, parallel workers can linearly approximate momentum and parameter updates locally without additional synchronization.
This further pushes the limit of communication reduction towards its extreme, achieving the state-of-the-art speed up for large-scale model training.
To summarize, our contributions are as follows:
• We propose 0/1 Adam , which addresses the limitations of previously proposed 1-bit Adam when applying aggressive 1-bit quantization and local steps (Section 4).
• We provide convergence guarantee of 0/1 Adam on smooth and non-convex objectives (Section 5).
• We conduct experiments on a wide range of large-scale model training tasks, including BERT-Base, BERT-Large, GPT-
pre-training and ImageNet.
We demonstrate on up to 128 GPUs that 0/1 Adam is able to reduce up to 87% of data volume, 54% of communication rounds, and achieve up to 2 × higher throughput and training time reduction compared to the state-of-the-art 1-bit Adam without compromising end-to-end model accuracy (Section 6).
• The 0/1 Adam optimizer and corresponding experimental scripts (e.g. BERT pre-training and GLUE finetuning) have been open sourced in a deep learning optimization library called DeepSpeed 2 .
R ELATED W ORK
Communication-efficient training.
There has been various lines of research focusing on improving communication efficiency in large-scale training, such as using asynchrony (Niu et al., 2011; Lian et al., 2015; Xie et al., 2020), decentralization (Lian et al., 2017; Lu and De Sa, 2021), gradient quantization (Alistarh et al., 2017; Wen et al., 2017), gradient sparsification (Wangni et al., 2017;
Wang et al., 2018a), local steps (Stich, 2018; Lin et al., 2018), etc.
In this paper we study the aggressive 1-bit compression, which was first introduced in (Seide et al., 2014) to speed up speech model training, where an algorithm called 1-bit SGD is proposed.
After that, Wen et al.
(2017) proposes adding 0 as an additional numerical level and Liu et al.
(2018) discusses the use of zero-th order oracle in 1-bit SGD.
Chen et al. (2019a); Balles and Hennig (2018); Xu and Kamilov (2019) study the correlation and combination between 1-bit SGD and other techniques.
Convergence analysis on 1-bit SGD is given in (Bernstein et al., 2018a;
Karimireddy et al., 2019; Safaryan and Richtárik, 2021).
Bernstein et al.
(2018b)
; Sohn et al.
Le Phong and Phuong (2020); Lyu (2021) investigate the robustness of 1-bit SGD.
Among all the variants of 1-bit communication, the design with error feedback mechanism has shown to work best both empirically (Seide et al., 2014) and theoretically (Karimireddy et al., 2019).
Other lines of research applies 1-bit communication to various scenarios such as federated learning (Jin et al., 2020;
Yue et al., 2021), decentralized learning (Lu and De Sa, 2020; Koloskova et al., 2019), meta learning (Fan et al., 2021), etc.
Perhaps the closest works to this paper are (Tang et al., 2021;
Li et al., 2021a), which propose using two-stage training to enable 1-bit Adam and 1-bit Lamb, respectively.
Different from those two work, 0/1 Adam addresses non-linearity challenges in adaptive optimizers by considering both extreme quantization and local steps.
Furthermore, we also study how to apply extreme communication compression on GPT-style models, which to the best our knowledge is still under-explored.
Adaptive learning rate optimizers.
One of the most popular adaptive optimizers is Adam, which was first introduced in (Kingma and Ba, 2014).
It uses both first and second moment information of stochastic gradient to perform optimizer steps and has shown significant benefits on training deep
[Figure]
learning models.
Reddi et al.
(2019) spots the issue of Adam convergence and provides a variant called AMSGrad while Zaheer et al.
(2018) argues the Adam only converges with large batch sizes.
Multiple lines of theoretical study on Adam are given in (Fang and Klabjan, 2019; Alacaoglu et al., 2020; Défossez et al., 2020).
Additionally, Chen et al.
; Zhou et al.
(2018a)
; Lu et al.
Danilova et al.
Zou et al.
(2019) provide more general analysis on Adam-type optimizers.
Subsequently, other variants of Adam are proposed in (Luo et al., 2019;
Chen et al., 2019b; Huang et al., 2018;
Wang et al., 2019b; Zhou et al., 2018b; Zhuang et al., 2021; 2020).
Unlike these methods, which focus on improving the convergence of generic optimizations for DNN models, our work studies how to maximize the communication efficiency of Adam in large-scale distributed training settings.
A C LOSER L OOK AT N ON - LINEARITY IN A DAM
In this section, we provide a more formal description on the problem setting and illustrate the limitations from the original Adam and the state-of-the-art 1-bit Adam (Tang et al., 2021).
Problem Formulation.
In this paper, we consider the following optimization problem:
[Equation]
where x denotes the d -dimensional model.
D denotes the training set and f ( x ; ζ ) is the loss incurred over sample ζ given model parameters x .
The structure of the problem naturally captures many of the model training problems.
The non-linearity in Adam.
At step t ≥ 0 , denote x t and g t as the model parameters and stochastic gradient computed at step t , respectively.
The update formula of SGD and Adam 3 can be summarized as:
[Equation]
where γ is the learning rate, ϵ is a small constant to prevent zero division, β 1 and β 2 are tunable decaying factors.
The linearity in SGD update implies when using compression or local steps, the potential noise from (accumulated) gradients is in the order of O ( γ ) , which approaches zero when learning rate is decaying or set to be small.
By comparison, the two auxiliary optimizer states in Adam, momentum ( m ) and variance ( v ), introduce non-linearity in the model update.
Equation (3) gives the formula of Adam when running it sequentially.
In a distributed setting with n workers, g t in Equation (3) is often computed in parallel on different workers.
Mathematically, if we denote g ( i ) t as the stochastic gradient computed on the i -th worker at step t , then distributed Adam can be written as replacing g t with 1 /n (cid:80) ni
=1 g ( i ) t in Equation (3) as follows:
[Equation]
Issue with non-linearity on 1-bit compression.
The main bottleneck in running distributed Adam is the accumulation of 1 /n (cid:80) ni
=1 g ( i ) t since the gradients are usually high-dimensional.
Based on the profiling results from (Tang et al., 2021; Li et al., 2021a), the communication of gradients could take up to 94% of the total training time on modern clusters.
Gradient compression mitigates this issue by sending and averaging gradients with fewer bits.
However, in Adam this causes the loss on the learning rate adaptivity.
Consider using the aggressive 1-bit compression (Liu et al., 2018), which sends each gradient with only signs and a shared, usually the average over all the coordinates, magnitude.
More specifically, denote C [ · ] as the 1-bit compression, then
[Equation]
It is straightforward to observe that naively applying 1 bit to compress gradients in the original Adam loses coordinate-wise adaptivity since sharing magnitude makes all the coordinates-wise learning rate γ/ √ v t + ϵ the same value.
This makes Adam no difference than momentum SGD.
Issue with non-linearity on local steps.
In SGD (Equation (2)), the model updates are linearly dependent on the gradients and has zero additional states.
It implies with local steps, the parallel workers can entirely reach consensus after a single round of synchronization, even with compression (Basu et al., 2020).
However, in Adam simply synchronizing the model can still leave the momentum and variance out-of-sync.
This makes parallel workers fail to capture the global adaptivity when the system scales up.
To give a more concrete example, we profile a full run of BERT-Large pre-training with original Adam, and summarize different metrics of momentum and variance in Figure 1.
As shown in Figure 1(d) and 1(b), the difference between local and global optimizer states, momentum and variance, remain constants and do not decrease to zero.
1-bit Adam and its limitations.
1-bit Adam (Tang et al., 2021) is a state-of-the-art solution that addresses non-linearity in 1-bit compression.
1-bit Adam adopts a pre-conditioned variance state from running original Adam for T 0 steps first.
The intuition there is that at later stage of training, the variance state becomes stable so that v T 0 can be a good approximation of variance state for the remaining steps.
As paritally illustrated in Section 1, the full-precision stage of 1-bit Adam still presents non-trivial overhead.
For instance: as illustrated in (Tang et al., 2021), when training BERT-Large on 64 GPUs using Ethernet, while the full-precision stage contains 15% of the total steps, it can take more than 50% of the entire training in terms of the wall-clock time 4 .
Additionally, 1-bit Adam is restricted in the scope of compression, how it handles other techniques such as local steps remains open question.
0/1 A DAM
In this section, we give the full description of 0/1 Adam .
To maximize the communication efficiency, ideally we want an algorithm that enables adaptive convergence like Adam, while allowing aggressive compression (e.g. 1 bit) and requires no additional synchronization on the optimizer states when using local steps.
0/1 Adam solves this problem from two aspects.
Adaptive Variance Freezing.
To begin with, 0/1 Adam creates a linear environment that freezes the variance adaptively.
The intuition is leveraged from the observation in Figure 1(a): the change of variance over steps in Adam is generally smooth.
While 1-bit Adam captures a reasonable variance estimate via one-time freezing, it is reasonable to also presume that before its freezing point, the variance within several adjacent steps will stay close due to its smoothness.
This motivates us to extend the one-time freezing policy in 1-bit Adam into an adaptive one, by letting workers agree upon the freezing points from a given step index set T v ⊆{ 0 , ··· ,T − 1 } .
The frozen variance creates multiple intervals over training, during which the workers have agreement on the denominator (Equation (3)) and the only uncertainty is then left in the nominator that is linearly dependent on the model update, just like SGD.
Require: local model on the i -th node x ( i ) 0 , learning rate { γ t } Tt =1 , m
= 0 , v 0 = 0 , auxiliary buffer u 0 = 0 , total number of iterations T , decaying factor β 1 , β 2 from Adam, numerical constant ϵ , variance update step index set T v , synchronization step index set T u , the most recent step with synchronization t ′ =0 .
1: for t =0 , ··· ,T − 1 do 2:
Compute local stochastic gradient g ( i ) t .
Update momentum: m ( i ) t + 12 = β 1 m ( i ) t +(1 − β 1 )
g ( i ) t .
Update model: x ( i ) t + 12 = x ( i ) t − γ t m ( i ) t / √ v t + ϵ .
Update buffer: u ( i ) t + 12 = u
( i ) t
+ γ t m ( i ) t .
6: if t ∈T u then 7:
Perform 1-bit AllReduce: u t + 12 = 1bit-AllReduce (cid:16) u ( i ) t + 12 (cid:17) .
Approximate momentum with compressed buffer: m ( i ) t +1 =
u t + 12 / (cid:80) th
= t ′ γ h .
Update model with compressed buffer: x ( i ) t +1 = x ( i ) t ′ − u t + 12 / √ v t + ϵ .
Reset the auxiliary buffer: u ( i ) t +1 = 0 .
Update the synchronization step: t ′ = t .
12: else 13: x ( i ) t +1 =
x
( i ) t + 12 ; m ( i ) t
+1 = m ( i ) t + 12 ; u ( i ) t
+1 =
u ( i ) t
+ 12 .
14: end if 15: if t ∈T v then 16:
Perform full-precision AllReduce: g t =
AllReduce (cid:16) g ( i ) t (cid:17) .
Update the variance: v t +1 = β 2 v t
+(1 − β 2 )( g t ) 2 . 18: else 19:
Use the stale variance for the next iteration: v t +1 = v t .
20: end if 21: end for 22: return x T .
Including 1-bit Compression and Local Steps.
With frozen variance, we make another observation based on Equation (3) that the model difference on workers will be linearly dependent to the momentum.
So that, the momentum can be approximated locally rather than synchronized additionally based on the communicated model difference, given the premise that the change of momentum is not abrupt within close steps.
Formally, denote x ( i ) t , m ( i ) t , v ( i ) t as the model, momentum, variance on worker i at step t , respectively.
Suppose all the workers are synchronized at step t ′ , then with frozen variance v over all the workers,
[Equation]
Actual sent tensors in the communication.
[Equation]
Sync model parameters with the sent tensors
[Equation]
Approximate momentum via linear estimates via sent tensor.
where we omit the compression part for brevity.
Combined with compression, we provide the full description of 0/1 Adam 5 in Algorithm 1.
Note that here we defer the details of 1-bit compression to Appendix A and treat it as a black-box procedure named 1bit-AllReduce while the original full-precision AllReduce is referred to as AllReduce .
We also remark that although both techniques appear to be natural, to the best of our knowledge, we are the first to apply them to addressing the non-linearity challenges in 1-bit compression and local steps for maximizing the communication efficiency of Adam optimizer. average.
C ONVERGENCE A NALYSIS
In this section, we provide the convergence guarantee for 0/1 Adam (Algorithm 1) under arbitrary freezing policy T v and local steps policy T u .
In the main paper, we provides the convergence rate in the general case.
However, different T v or T u gives us the opportunity to obtain tighter bounds.
We leave these discussion in the appendix.
We start by making the following assumptions.
Assumption 1.
Lipschitzian gradient: f ( · ) is assumed to be with L -Lipschitzian gradients, which means
[Equation]
Assumption 2.
Bounded variance:
The stochastic gradient computed on each worker is unbiased and has bounded variance: E ζ ∼D ∥∇ f ( x ; ζ ) −∇ f ( x ) ∥ 2 ≤ σ 2 , ∀ x .
Assumption 3.
Bounded gradient: The infinity norm of stochastic gradient is bounded by a constant
G ∞ > 0 such that ∥ g t ∥ ∞ ≤ G ∞ , ∀ t .
Assumption 4.
Compression error in Algorithm 1: For arbitrary x ∈ R d , there exists a constant ∆ , such that the output of compressor C [ · ] has the following error bound: E ∥C
[ x ] − x ∥ 2 ≤ ∆ 2 .
Assumption 5.
Given ordered set T u , denote t j as the j -th element in T u , we assume there exists a constant H ≥ 0 , it holds that max 1 ≤ j< |T u | ( t j +1 − t j ) ≤ H .
Remarks on the assumptions.
Assumption 1, 2 and 3 are standard in the domain of non-convex optimization.
Comparing with the 1-bit Adam paper (Tang et al., 2021), we do not explicitly assume the uniform lower bound on the variance coordinate, i.e., e ⊤ j v >v min > 0 , ∀ j for some constant v min .
Instead we assume an infinity-norm bound on the gradient as in Assumption 3 which is more realistic.
Assumption 4 is also assumed in (Tang et al., 2021), in the appendix we discuss the variant of 0/ Adam that converges with weaker condition on the C .
The convergence for Algorithm 1 is then given in the follow theorem.
Theorem 1.
Under Assumption 1 to 5, let m =
|T v | , select β 1 , β 2 ∈ [0 , 1) that fulfills m ≤ log(1 − β 1 ) / log( β 2 ) , if we run Algorithm 1 with a constant learning rate: for all t ≥
[Equation]
then it holds that
[Equation]
where ˜ x t =1 /n (cid:80) ni
=1 x ( i ) t
and we omit f ( 0 ) − inf x ∈ R d f ( x ) , G ∞ , d , ϵ , β 1 , β 2 and L as constants.
Theorem 1 shows that 0/1 Adam Algorithm 1 essentially admits the same convergence rate as distributed SGD in the sense that it achieves linear speed up, at rate 1 /O ( √ nT ) .
The effect of compression ( ∆ ) and local steps ( H ) only appears on a non-dominating term.
E XPERIMENTS
In this section we evaluate the performance of 0/1 Adam over several large-scale model training tasks comparing with baselines (1-bit Adam (Tang et al., 2021) and original Adam (Kingma and Ba, 2014)).
Since Tang et al.
(2021) already demonstrated that 1-bit Adam has similar statistical results to Adam, we omit the comparison of end-to-end model accuracy to Adam for brevity.
Throughout the experiments, we enable FP16 training for all the tasks following (Tang et al., 2021).
That makes the full-precision communication (including Adam, full-precision stage in 1-bit Adam and full-precision AllReduce in 0/1 Adam ) use 16-bit per number.
We use the 1-bit compressor (Equation (4)) in 0/1 Adam .
Experimental details.
We adopt the following tasks for the evaluation: BERT-Base ( L =12 , H =768 ,
A = 12 , 110 M params) and BERT-Large ( L = 24 , H = 1024 , A = 16 , 340 M params) pre-training, training Resnet18 ( 12 M params) on ImageNet (He et al., 2016) and GPT-
pre-training.
For BERT model, we use the same dataset as (Devlin et al., 2018), which is a concatenation of Wikipedia and
BooksCorpus with 2.5B and 800M words respectively.
We use the GLUE fine-tuning benchmark (Wang et al., 2018b) to evaluate the convergence of the BERT models trained by different algorithms. For
ImageNet, we adopt ImageNet-1k dataset, which contains 1.28M images for training and 50K images [Figure] [Figure] And so we test it for 4 to 32 GPUs in Figure (d).
for validation (Deng et al., 2009).
For GPT-2 we adopt the model from its original paper (Radford et al., 2019b) which contains 117M parameters (48 layers, 1600 hidden size, 25 attention heads).
For training data, we adopt the same dataset blend as in (Shoeybi et al., 2019):
Wikipedia (Devlin et al., 2018), CC-Stories (Trinh and Le, 2018), RealNews (Zellers et al., 2019) and OpenWebtext (Radford et al., 2019b).
Other details including learning rate schedules, hyperparameters can be found in Appendix C.
Hardware.
We evaluate two clusters: one with 4 NVIDIA V100 GPUs per node and 40 Gigabit
Ethernet inter-node network (2.7 Gbps effective bandwidth); the other one with 8 V100 GPUs per node and 100 Gigabit InfiniBand EDR inter-node network (close to theoretical peak effective bandwidth).
We use 4 to 128 GPUs for BERT and ImageNet pretraining tasks to measure 0/1 Adam ’s performance gain.
We use 64 GPUs for GPT-2 pre-training.
Additionally, for ImageNet training we apply the accelerated data loading technique from lmdb 6 .
Policy for T v and T u in 0/1 Adam.
We first illustrate our policy on T v .
Observing from our motivation study (Figure 1) that the variance difference in adjacent steps decreases roughly exponentially.
{+Denote k j as the step where j -th variance update takes place, we select T v such that, k j +1 − k j =
 ⌊ j/κ ⌋ , ∀ κ> 0 . We adopt κ =16 for all the three tasks.+} [Table]
of Resnet
at the end of epoch 90 from different algorithms.
The original accuracy is provided by Pytorch pretrained model library (Pytorch, 2014).
For the other two algorithms, the accuracy is the highest score over 3 runs.
The other two columns shows zero-shot evaluation of the trained GPT-2 on WikiText-103 and LAMBADA datasets, the evaluation methodology follows (Shoeybi et al., 2019).
The number for Adam is obtained from (Li et al., 2021b).
[Table]
[-Denote k j as the step where j -th variance update takes place, we select T v such that, k j +1 − k j = 
-]
Then we move on to discuss the policy for T u .
Based on the derivation in Section 4, the approximation noise from local step is proportional to the learning rate.
And so if we denote t j as the step where j -th synchronization takes place, then our intuition is to increase t j +1 − t j roughly inversely proportional to thelearningrateat t j soastomaketheapproximationnoisebounded.
ForBERT-Base/Largepretraining, as illustrated before, the learning rate exponentially decreases by 0.99 every 520 steps after 12.5K linear increase warmup steps.
So that we set t j +1 − t j =1 for the first 12.5K steps and after that let it multiply by 2 every 32678 steps based on the calculation that the learning rate will decrease by half.
Similarly, for ImageNet we set t j +1 − t j =1 for the first 50050 steps (10 epochs) and after that let it multiply byevery 50050 steps (10 epochs).
We clip the interval at 16 in all the tasks.
This corresponds to H =16 in Assumption 5.Finally, since our theory in Section 5 indicates that approximation will be more accurate when the variance is frozen.
So that we additionally stop updating variance
when
t j +1 − t j > 1 .
Remarks on the selected policy.
As described, the policies for T v and T u generally follow the learning rate schedule adopted.
This is favored in practice for three reasons: (1) It does not require too much hyperparameter tuning.
Consider adopting a constant or decaying policy based on some utility function, capturing the dynamics of the training to reach a useful interval would require tedious searching on the hyperparameters and retraining; (2) Learning rate policy is well-motivated.
As naturally, communication is a way of eliminating the difference among workers.
If workers adopt a larger learning rate, more frequent communication will be needed since the workers are taking large steps in the weight space
and so we require higher-frequent communication for them to reach consensus; (3) Learning rate schedule is universally used in all the ML/DL applications, and thus the method can be easily adapted to other applications.
In fact, if we consider the scope of large model training, a typical learning rate schedule is a linear warm-up followed by a decaying phase, which is very similar to our test cases here.
6. C ONVERGENCE S PEED AND Q UALITY A NALYSIS
Figure 2 presents the sample-wise and time-wise convergence results for different algorithms with 128 GPUs on the Ethernet cluster.
We find that 0/1 Adam provides the same sample-wise convergence speed compared to the baseline, with up to 2 × time-wise speed up.
Table 1 summarizes the GLUE results using the checkpoints from our BERT pretraining experiments.
0/1 Adam achieves similar end task accuracy compared to the numbers reported in previous work.
0/1 Adam achieves faster training time than prior work because it reduces the communication overhead in distributed training by using both 1-bit quantizer to compress the communication volume (up to 32 × reduction) and 1-bit AllReduce to reduce the expensive synchronization overhead for local steps in both warmup and non-warmup phases.
Table 2 provides the ImageNet validation accuracy of trained models from different algorithms, and we find the final accuracy can achieve the reported accuracy from Pytorch library (Pytorch, 2014).
For brevity, convergence comparison on GPT-2 is given in the Appendix C.
6. T RAINING T HROUGHPUT A NALYSIS
Figure 3 summarizes the throughput results on different tasks and different clusters.
We observe that 0/1 Adam can consistently outperform baselines in all settings.
It is also worth mentioning that 0/
Adam on Ethernet (2.7 Gbps effective bandwidth, 4 GPUs per node) is able to achieve comparable throughput as 1-bit Adam on InfiniBand (near 100 Gbps effective bandwidth, 8 GPUs per node), as shown in the red line in Figure 3(b) and the blue line in Figure 3(c), which demonstrates 0/1 Adam further removes the redundancy in communication effectively that exceeds the hardware barrier.
Communication reduction and the role of local steps.
To better understand the importance and effect of local steps, we additionally run a special case of 0/1 Adam where we keep the same policy of T v but use T u = { 0 , ··· ,T − 1 } .
This special version of 0/1 Adam does not skip rounds but use the same variance freezing policy.
We plot the data volume usage and throughput results in Figure 4 and 5, respectively.
We see that although no local steps suffice to reduce the data volume overhead from 1-bit Adam towards 1-bit-per-parameter in general, the throughput improvement is limited compared to Figure 2.
C ONCLUSION
In this paper, we study the challenges of using 1-bit communication on Adam, and limitations of the state-of-the-art 1-bit Adam algorithm.
We propose an algorithm named 0/1 Adam that adopts two novel design: adaptive variance state freezing and 1-bit sync.
We provide convergence proof for 0/1 Adam and measure its effectiveness over baseline Adam and 1-bit Adam on various benchmarks, including BERT-Base/Large, GPT-2 pretraining and ImageNet. A CKNOWLEDGMENTS
Yucheng Lu is supported by Meta PhD Fellowship.
The authors would like to thank anonymous reviewers from ICLR2023 for providing valuable feedback.
A F ULL D ESCRIPTION TO A LL R EDUCE
As introduced in Section 2, the error feedback based 1bit-AllReduce works best both in theory and in practice.
In fact, the original 1-bit Adam also adopts the error-feedback design (Tang et al., 2021).
We give the full description of this 1bit-AllReduce in Algorithm 2, to replace the 1bit-AllReduce in Algorithm 1.
In the theoretical analysis, our proofs will also rely on this algorithm.
Note that this algorithm does not require any additional assumptions for our theory to hold, since this fits the black-box procedure in Algorithm 4 and Algorithm 1.
Require: communication buffer z ( i ) t , worker error δ ( i ) t , server error δ t , 1-bit compressor C [ · ] .
Both worker and server errors will be initialized at 0 at t =0 .
1: (On i -th node) 2:
Compress z ( i ) t into ˆ z ( i ) t = C {+[ z ( i ) t + δ ( i ) t ] , and update the compression error by δ ( i ) t +
 =+}
z ( i ) t
+ δ ( i ) t − ˆ z
( i ) t .
3: Send ˆ
z
( i ) t to the server.
4: (On server) 5: Take the average over all the ˆ z ( i ) t and compress it into z t =
C [ 1 n (cid:80) n i =
1 ˆ z
( i ) t +1 + δ t ] , and update the compression error by δ t +1 = 1 n (cid:80) ni
{+=
+}
ˆ z ( i ) t
+ δ t − z t . 6: Send z t to all the workers.
7: (On i -th node) 8: return z t , δ ( i ) t +1 , δ t +1 .
{+Algorithm 
+}
The full description of AllReduce
Require communication z i t .
B P ROFILING R ESULTS FOR F IXED C OST OF C OMMUNICATION
We profile the time taken in computation and others (including initialization of a communication round and compression) during one 1-bit AllReduce round at different scales on Ethernet cluster in Table [Table]
[Table]
C A DDITIONAL E XPERIMENTAL D ETAILS
Training Parameters.
For BERT pretraining, we follow the settings from (Devlin et al., 2018) and let the learning rate linearly increases to 4 × 10 − 4 as a warmup in the first 12.5K steps, then decays into 0.99 of the original after every 520 steps.
We set β 1 =0 .
9 and β 2 =0 .
999 for all the algorithms.
We adopt the batch size of 4096.
For 1-bit Adam, we follow the hyperparameters given in (Tang et al., 2021) and set the full-precision stage for 1-bit Adam as 16K and 23K on BERT-Base and BERT-Large, respectively.
All the hyperparameters used here (e.g. learning rate) strictly follow (Tang et al., 2021) for fair comparison.
For ImageNet, we follow the example script from Pytorch 7 and use batch size of 256 and a milestone decay learning rate schedule: starting at 1e-4 and decay by a factor of 10 at epoch 30 and 60, with 90 epochs in total.
We set 10 epochs (50050 steps) as the full-precision stage for 1-bit Adam.
For GPT-2 we set batch size to be 512, and use 300K training steps (158B tokens).
The learning rate schedule follows a linear warmup of 3K steps and a single cycle consine decay over the remaining 297K steps ( 1 × 10 − 5 min ).
For 1-bit Adam, we set its full-precision stage length to be 80K steps, and for the 0/1 Adam , we follow the same learning rate based policy from BERT on T v and T u .
For GLUE benchmarks we use Adam optimizer and perform single-task training on the dev set.
Following the setup in the BERT paper (Devlin et al., 2018) and 1-bit Adam paper (Tang et al., 2021), we search over the hyperparameter space with batch sizes ∈ { 8 , 16 , 32 } and learning rates { 1 × 10 − 5 , 3 × 10 − 5 , 5 × 10 − 5 , 8 × 10 − 5 } .
The convergence plots for GPT-2 pre-training are given in Figure 6.
D T HEORETICAL A NALYSIS D. A NALYSIS TO AN INTERMEDIATE VERSION OF 0/1 A DAM
Algorithm 4 Generic framework of applying 1-bit communication to Adam with frozen variance state.
1-bit Adam can be viewed as a special case of setting T v = { 0 , ··· ,T 0 − 1 } where T 0 denotes its total number of steps in the full-precision stage. [Table] [Table] [Table]
[Table] [Table] [Table] [Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
v
9: end Update model 12: end for 13: return We start from a special case of 0/1 Adam that compresses gradients without local steps.
This is given in Algorithm 4.
Note that the following proof will use Algorithm 2 to replace Compressed-AllReduce in Algorithm 4, as introduced in Section A.
Algorithm 4 allows us to work with weaker assumption as given in the following
Assumption 6.
Compression error in Algorithm 4: For arbitrary x ∈ R d , there exists a constant 0 ≤ ω< 1 , such that the output of compressor C [ · ] has the following error bound:
[Equation]
Theorem 2.
Under Assumption 1, 2, 3, and 6, let m =
|T v | , and select β 1 ,β 2 ∈ [0 , 1) such that m ≤ log(1 − β 1 ) / log( β 2 ) .
If we run Algorithm 4 with a constant learning rate: for all t ≥γ
[Equation]
min
[Equation]
then it holds that
[Equation]
where we omit f ( 0 ) − inf x ∈ R d f ( x ) , G ∞ , d , ϵ , β 1 , β 2 and L as constants.
Proof.
The main update of Algorithm 4 (with constant learning rate) can be summarized as: for every t =0 , ··· ,T − 1 ,
[Equation]
where the g t is the output of the 1-bit AllReduce algorithm 8 .
Note that based on Algorithm 2, the gradient approximation term follows:
[Equation]
[Equation] (cid:33)
[Equation]
i [Equation]
( i
[Equation] (cid:33) =
g [Equation] where we denote
[Equation]
To prove the convergence, we now define the following auxiliary sequence: for any t ≥ 0 ,
[Equation]
The rest of the proof is to use this auxiliary sequence to bound two types of steps separately.
We call a step t as reuse step if t ̸∈T v and update step otherwise.
We see for all the update steps, v t ̸ = v t +1 while for all the reuse steps v t =
v t +1 .
The bounds on two different types of steps are provides by Lemmaand Lemma 6.
Specifically, denoting V 1 = (cid:13)(cid:13)(cid:13) 1 √ v 1 + ϵ (cid:13)(cid:13)(cid:13) 1 , from Lemma 5 we obtain for all the reuse steps, (cid:88) t ̸∈T
[Equation]
≤ (cid:88) t
[Equation]
v . This, however, does not affect our analysis, since our proof holds for a noisier case. The original Algorithm 4 is mainly for practical concern while from Lemma 6 we obtain for all the update steps, (cid:88)
[Equation]
Note that the two inequalities above hold when the learning rate fulfills
[Equation]
Combine them together,
[Equation]
Dropping the constants
[Equation]
we finally obtain
[Equation]
where in the last step we use the condition in the theorem that β m 2 ≥ 1 − β 1 .
To meet the requirement of learning rate we set
[Equation]
then it holds that
[Equation]
That completes the proof.
D. P ROOF TO T HEOREM
Note that the following proof will use Algorithm 2 to replace 1bit-AllReduce in Algorithm 1, as introduced in Section A.
Theorem 1.
Under Assumption 1 to 5, let m =
|T v | , select β 1 , β 2 ∈ [0 , 1) that fulfills m ≤ log(1 − β 1 ) / log( β 2 ) , if we run Algorithm 1 with a constant learning rate: for all t ≥
[Equation]
then it holds that
[Equation]
where ˜ x t =1 /n (cid:80) ni
=1 x ( i ) t
and we omit f ( 0 ) − inf x ∈ R d f ( x ) , G ∞ , d , ϵ , β 1 , β 2 and L as constants.
Proof.
We now prove Theorem 1.
Similar to the proof to Theorem 2, in this proof we discuss the case of t ∈T v and t ̸∈T v separately.
Following the proof of Theorem 2, we define the following auxiliary sequence
[Equation]
where
[Equation]
[Equation]
And we additionally define that
[Equation]
Note that the definition of ˜ g t is different from the g t in Theorem 2 since the former is computed on local models which potentially can be different before the sync step.
To expect a compression error bound to scale in the order of O ( γ 2 ) , we slightly modify the update of line 5, 8, 9 of Algorithm 1 into
[Equation]
Note that since Theorem 1 states the convergence results for constant learning rate, such modification does not change the semantics of the original Algorithm 1.
Based on Algorithm 2, we know that
[Equation]
Based on Lemma 10, we know that for all the t ∈T v , we have the following bound,
[Equation]
On the other hand
[Equation]
for all the t ̸∈T v
[Equation]
we have the following bound,
[Equation]
Note that they hold if learning rate is set to be
[Equation]
Combine them together
[Equation]
we obtain
[Equation]
[Equation]
Omitting constants:
[Equation]
where in the last step we use the condition in the theorem that β m 2 ≥ 1 − β 1 .
To meet the requirement of learning rate we set
[Equation]
then it holds that
[Equation]
And that completes the proof.
D. T ECHNICAL L EMMA
Lemma 1.
Consider running Algorithm 2 over a communication buffer
z (same notation in Algorithm 2) under Assumption [Equation] let δ t denote:
[Equation]
then based on Assumption 6 and 3, it holds that t ≥ 0 , if E ∥ z ( i ) t ∥ 2 ≤ C for some constant C > 0 ,
[Equation]
Proof.
Note that the error is initialized by 0 , so that when t = 0 the bound trivially holds.
We next prove the case for t ≥ 1 .
For any i ∈{ 1 , ··· ,n } and t ≥ 1 , by the definition of the sequence δ ( i ) t ,
[Equation]
Selecting η = 1 − ω 2 ω , we obtain
[Equation]
Similarly, we can show that for any t ≥ 1 ,
[Equation]
[Equation]
where in the last step we apply the Jensen Inequality.
Since we do not assume a bound on the [Table] ˆ z [Table]
i [Table] [Table] we need to bound it in terms of
[Equation]
where we apply the results from the bound on E (cid:13)(cid:13)(cid:13)
δ
( i ) t (cid:13)(cid:13)(cid:13) 2 .
Given this bound, and following the analysis for E (cid:13)(cid:13)(cid:13) δ ( i ) t (cid:13)(cid:13)(cid:13) 2 , we can now bound the E (cid:13)(cid:13) δ t (cid:13)(cid:13) 2 as follows
[Equation]
Finally, we obtain t ≥ 1 ,
[Equation]
n
[Equation]
That completes the proof.
Lemma 2.
For the variance term, we have the following upper and lower bound: for any t ≥ 1 ,
[Equation]
where the inequality holds element-wise.
Proof.
On one hand, for any t j ≤ t<t j +1 , where t j denotes an update step, we obtain element-wise:
[Equation]
so that
[Equation]
On the other hand, for any t ≥ 1 and j ∈{ 1 , ··· ,d } ,
[Equation]
so that
[Equation]
That completes the proof.
Lemma 3.
In Algorithm 4, for any t ≥ 0 ,
[Equation]
Proof.
For any t ≥ 0 ,
[Equation]
(cid:88)
[Equation]
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
[Equation]
≤
[Equation]
[Equation]
[Equation]
[Equation]
k
[Equation]
[Equation]
≤ (1 −
[Equation]
[Equation]
(cid:88) k
[Equation]
β
[Equation]
k
[Equation]
E (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) 1 n n (cid:88) i = g
( i ) k
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) +3 E
[Table]
δ k ∥
[Table]
3 E
[Table]
  ≤ (1 − β 1 ) t (cid:88) k
= β t
[Table]
k(cid:32)n n (cid:88) i =
E (cid:13)(cid:13)(cid:13) g
[Table]
i
[Table]
+
[Table]
∥
[Table]
k
[Table]
+3 E
[Table]
+
[Table]
∥(cid:33) ( i ) ≤ (1 − β 1 )
[Table]
k = β
[Table]
(cid:18) 3 G
[Table]
d + 192 ω (1+ ω )
[Table]
G
[Table]
d (
[Table]
[Table]
(cid:19) ≤ (cid:18) 3(1+ ω )
[Table]
G
[Table]
∞ d (1 − ω ) + 192(1+ ω ) [Table] G
[Table]
∞ d (1 − ω ) [Table] (cid:19) [Table]
[Table] [Table] β
[Table]
≤ 195(1+ ω ) 3 G 2 ∞ d (1 − ω ), where in the step (i) we use Lemma 1. That completes the proof.
Lemma 4.
For any a , b ∈ R d , the following bound holds:
[Equation]
Proof.
Denote the subscript j as the index of the coordinate.
[Equation]

   [Equation]
(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13) .
Note that the second step holds not because Cauchy-Schwarz Inequality but due to the fact that a 2 j , b j > 0 (since √ b would implicitly assume so).
Lemma 5.
In Algorithm 4, for all the t ≥ 1 that fulfills v t = v t +1 , i.e., ∀ t such that t ̸∈T v , if we let
[Equation]
the following bound holds,
[Equation]
Proof.
Recall the auxiliary sequence
[Equation]
For all the steps t ≥ 0 that fulfills v t +1 = v t
[Equation]
we obtain
[Equation]
From Assumption
[Equation]
we have
[Equation]
ϵ (cid:13)(cid:13)(cid:13)(cid:13) = − γ E (cid:28) ∇ f ( y t ) , ∇ f ( x t ) √ v t + ϵ (cid:29) + Lγ2 E (cid:13)(cid:13)(cid:13)(cid:13) g t √ v t + ϵ cid:13)(cid:13)(cid:13)(cid:13)
= − γ E (cid:28) ∇ f ( x t ) , ∇ f ( x t ) √ v t + ϵ (cid:29) + γ E (cid:28) ∇ f ( x t ) −∇ f ( y t ) , ∇ f ( x t ) √ v t + ϵ (cid:29) + Lγ2 E (cid:13)(cid:13)(cid:13)(cid:13) g t √ v t + ϵ
cid:13)(cid:13)(cid:13)(cid:13) = − γ E (cid:28)
∇ f ( x t ) , ∇ f ( x t ) √ v t + ϵ (cid:29) + γ E (cid:28) ∇ f ( x t ) −∇ f ( y t ) √ v t + ϵ , ∇ f ( x t ) (cid:29) + Lγ2 E (cid:13)(cid:13)(cid:13)(cid:13)
g t √ v t + ϵ (cid:13)(cid:13)(cid:13)(cid:13) ≤− γ E ∥∇ f ( x t ) ∥(cid:112)
G 2 ∞ + ϵ + γ 2 η E (cid:13)(cid:13)(cid:13)(cid:13) ∇ f ( x t ) −∇ f ( y t ) √ v t + ϵ (cid:13)(cid:13)(cid:13)(cid:13) 2 + γη 2 E ∥∇ f
( x t ) ∥
+ Lγ 2 2 E (cid:13)(cid:13)(cid:13)(cid:13)
g t √ v t + ϵ
cid:13)(cid:13)(cid:13)(cid:13) , where in the last step we use Lemma 2 and the fact that for any a , b and constant η> 0 , ⟨ a , b ⟩≤ η 2 ∥ a ∥ 2 + 1 2 η ∥ b ∥ 2 .
Set η =( (cid:112) G 2 ∞ + ϵ ) − 1 , with Assumption 1 and Lemma 4,
[Equation]
f
[Equation] ≤− [Equation]
(cid:13)(cid:13)(cid:13)(cid:13) g [Equation] v t
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)
= −
[Equation]
E
[Equation]
√ [Equation] E (cid:13)(cid:13)(cid:13)(cid:13) g
[Equation]
v
[Equation]
ϵ (cid:13)(cid:13)(cid:13)(cid:13) ≤−
[Equation]
E
[Equation]
(cid:13) m
[Equation]
√ v
[Equation]
+ ϵ (cid:13)(cid:13)(cid:13) (cid:13) + γ 3 L 2 V 1 (cid:112) G 2 ∞ + ϵ β m
E (cid:13)(cid:13)(cid:13) (cid:13) δ t √ v t + ϵ (cid:13)(cid:13)(cid:13) (cid:13) + Lγ2 E (cid:13)(cid:13)(cid:13)(cid:13)
g t √ v t + ϵ (cid:13)(cid:13)(cid:13)(cid:13) ≤− γ E ∥∇ f ( x t ) ∥2 (cid:112)
[Equation]
21 (cid:112)
[Equation]
ϵ
[Equation]
[Equation]
[Equation]
ϵ β
[Equation]
E ∥ δ t
[-cid:13)(cid:13)(cid:13)(cid:13) + Lγ
 E (cid:13)(cid:13)(cid:13)(cid:13)-][Equation]
V
[Equation]
m E ∥ g t ∥ 2 , where in the last step we apply Lemma 2 and 4.
Using the bound on the error from Lemma 1, Lemmaand the assumption on the stochastic gradient, we obtain (cid:32)
[Equation]
Based on the learning rate bound
[Equation]
and summing over all the reuse steps
[Equation]
we obtain
[Equation]
That completes the proof.
Lemma 6.
In Algorithm 4, for all the t ≥ 0 that fulfills v t ̸ = v t +1 , i.e. t ∈T v , if the learning rate fulfills γ < 1, the following bound holds
[Equation]
v
[Equation]
Proof.
For all the steps t that fulfills v t ̸ = v t +1 , [Equation] [Equation] −
[Equation]
ϵ (cid:19)
[Equation]
γ (cid:18) δ t √
[Equation]
t
[Equation]
ϵ −
[Equation]
{+t +1 √ v t +
+}
+ ϵ (cid:19)
[Equation]
− γ
[Equation]
γ
[Equation]
−
[Equation]
(cid:18)
[Equation]
ϵ
[Equation]
m
[Equation]
v
[Equation]
ϵ (cid:19)
[Equation]
(cid:18)
δ
[Equation]
v
[Equation]
v
[Equation]
ϵ (cid:19)
[Equation]
m
[Equation]
v
[Equation] Based on the smoothness assumption [Equation] for constant η> 0 that will be assigned later,
[Equation]
[Equation]
Now we can bound the last two terms as follows, note that
[Equation]
where in the last step we apply Lemma 1.
On the other hand,
[Equation]
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
where we again apply Lemma 1 and Lemma 3.
Put everything together,
[Equation]
Set η =34 , and considering γ < 1125 , we get
[Equation]
Summing over all the update steps, we obtain 0 ≤ (cid:88)
[Equation]
Adding
[Equation]
x t ) ∥ 2 on both sides, and note that
[Equation]
[Equation]
we finally obtain (cid:88)
[Equation]
That completes the proof.
Lemma 7.
Under Assumption 4, for any t ≥ [Equation] it holds that
[Equation]
Proof.
Based on the definition of the compression error
[Equation]
we obtain
[Equation]
4∆ 2 .
That completes the proof.
Lemma 8.
In Algorithm 1, for any t ≥ 0 , the momentum term is uniformly bounded by the following:
[Equation]
Proof.
We prove this lemma via induction.
Note that when t = [Equation] the inequality trivially holds due to initialization at 0 and Jensen Inequality.
Now suppose the inequality holds up to step t ≥
[Equation]
then for t +1 , if t ∈T u
[Equation]
then
[Equation]
[Table]
[Equation]
On the other hand
[Equation]
if t ̸∈T u
[Equation]
then
[Equation]
For all the t + 12 case, the inequality holds trivially due to Jensen Inequality.
Finally, all the ˜ · bound can also be obtained via Jensen Inequality.
And that completes the proof.
Lemma 9.
In Algorithm 1, for all the t such that t ̸∈T v , it holds that if we set learning rate
[Equation]
then, (cid:88)
[Equation]
t
[Equation]
Proof.
Since when t ̸∈T v [Equation] it can either belongs to T u or not
[Equation]
We first prove the case for t ∈T u .
From the definition of the auxiliary sequence
[Equation]
we obtain,
[Equation]
[Equation]
From Assumption
[Equation]
we have
[Equation]
We now bound A 1 to A 4 separately.
Note that from Lemma 8, the momentum term can be uniformly bounded by a constant.
For brevity of the derivation, we use M to denote such constant bound [Equation] and fit in its value at the end of the proof [Equation]
For A 1 ,
[Equation]
[Table]
[Equation]
where in the last step we use Assumption 1, Lemma 2 and Lemma 4.
For the second term, denote the last sync step before t is k , then we have:
[Equation]
where the first step holds because Lemma 4, Lemma 2, and the fact that at the sync step k , ˜ x k = x ( i ) k .
For the third term [Equation] we have
[Equation]
where we again apply the Lemma 2 and Lemma 4.
Then we can get
[Equation]
where in the second step we reuse Equation (5).
Next we can bound A 2 as follows
[Equation]
where in the sixth step we reuse Equation (5) [Equation]
For A 3 ,
[Equation] [Equation]
where in the last step we reuse Equation (6).
Combine the bound of A 1 to A
[Equation]
we obtain
[Equation]
We set the two constants η 1 ,η 2 as
[Equation]
then we have,
[Equation]
Finally
[Equation]
we need to bound the norm of q t .
If we denote the last sync step was k steps before t , then,
[Equation]
based on which we obtain
[Equation]
[Equation]
Put everything together
[Equation]
and let γ fulfills
[Equation]
we finally obtain
[Equation]
To this end
[Equation]
we have provided bound to all the sync steps t with ( t ̸∈T v and t ∈T u ).
For all the t with ( t ̸∈ T v and t ̸∈ T u ) [Equation] they can be seen as a special case of q t = 0 .
Since A 3 + A 4 > 0 , this bound will continue to hold for them, so that to sum over all the t with t ̸∈T v , we obtain (cid:88)
[Equation]
where we replace M with Lemma
[Equation]
That completes the proof
[Equation]
Lemma 10.
In Algorithm
[Equation]
For all the t ≥ 0 that fulfills v t ̸ = v t +1 , i.e. t ∈T v , if the learning rate fulfills [Equation] , the following bound holds
[Equation]
Proof.
From the definition of the auxiliary sequence
[Equation]
we obtain,
[-+
 + ϵ − ˜ m t √ v t + ϵ (cid:19) − (cid:18) γ δ t +1 √ v t +
 + ϵ − γ δ t √ v t + ϵ (cid:19) .-][Equation]
Based on Assumption 1,
[Equation]
[Equation]
We now bound the three norm terms separately.
From Equation (5), we obtain for the first term,
[Equation]
where we again use M to denote the constant bound from Lemma 8 for brevity.
On the other hand, based on a similar derivation to Equation (6) [Equation] we obtain [Equation]
Finally
[Equation]
for the last norm [Equation] it’s possible that the update towards t +1 step contains synchronization on the buffer.
So that we need to discuss the two cases separately.
First, for all the t ∈T u , denote the last sync step before t is k , then we have
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
[Table]
[Equation]
where in the last step we use Lemma 7, 8 and 4.
It is straightforward to verify that this bound also holds for t ̸∈T u (since there will be no noise from the sync step).
Combine the three norm term bounds
[Equation]
we obtain
[Equation]
where in the last step we set η =1 and use the requirement that γ < 1 / 6 .
Summing over all the t ∈T v , we get 0 ≤ (cid:88)
[Equation]
Adding
[Equation]
on both sides, and note that
[Equation]
We finally obtain (cid:88)
[Equation]
That completes the proof.