SEPARATE: A Simple Low-rank Projection for Gradient Compression in Modern Large-scale Model Training Process
Authors: Hanzhen Zhao
Xingyu Xie
Cong Fang
Zhouchen Lin
Venue: ICLR 2025 Poster (2025)
Keywords: efficient training
gradient compression
Primary area: optimization

Abstract: Training Large Language Models (LLMs) presents a significant communication bottleneck, predominantly due to the growing scale of the gradient to communicate across multi-device clusters. However, how to mitigate communication overhead in practice remains a formidable challenge due to the weakness of the methodology of the existing compression methods, especially the neglect of the characteristics of the gradient. In this paper, we consider and demonstrate the low-rank properties of gradient and Hessian observed in LLMs training dynamic, and take advantage of such natural properties to design SEPARATE, a simple low-rank projection for gradient compression in modern large-scale model training processes. SEPARATE realizes dimensional reduction by common random Gaussian variables and an improved moving average error-feedback technique. We theoretically demonstrate that SEPARATE-based optimizers maintain the original convergence rate for SGD and Adam-Type optimizers for general non-convex objectives. Experimental results show that SEPARATE accelerates training speed by up to 2× for GPT-2-Medium pre-training, and improves performance on various benchmarks for LLAMA2-7B fine-tuning.

Published as a conference paper at ICLR 2025
SEPARATE: A SIMPLE LOW-RANK PROJECTION FOR GRADIENT COMPRESSION IN MODERN LARGE-SCALE
MODEL TRAINING PROCESS
Hanzhen Zhao', Xingyu Xie², Cong Fang Zhouchen Lin
(hzzhao, xyxie, fangcong, zl mnnnn @pkup.u.eee edu..
1 State Key Lab of General AI, School of Intelligence Science and Technology, Peking University
2 Department of Mathematics, National University of Singapore
Institute for Artificial Intelligence, Peking University
4 Pazhou Laboratory (Huangpu), Guangzhou, Guangdong, China
ABSTRACT
Training Large Language Models (LLMs) presents a significant communication bottleneck, predominantly due the growing scale of the gradient communicate across multi-device clusters. However, how to mitigate communication overhead practice remains formidable challenge due to the weakness of the methodology of the existing compression methods, especially the neglect of the characteristics of the gradient. this paper, we consider and demonstrate the low-rank properties of gradient and Hessian observed in LLMs training dynamic, and take advantage such natural properties to design SEPARATE, a simple low-rank projection for gradient compression in modern large-scale model training processes. SEPARATE realizes dimensional reduction by common random Gaussian variables and an improved moving average frodbbce technique. We theoretically demonstrate that SEPARATE basedd optimizers maintain the original convergence rate for SGD and Adam-Tyyn optimizers for general non-convex objectives. Experimental results show that SEPARATE accelerates training speed by up to 2x for GPT-2-Medium pre-training, and improves performance various benchmarks for LLAMA2-7B fine-tuning.
INTRODUCTION
With massive amounts of data, billions of parameters, and multi-device clusters, the remarkable strides of Large Language Models (LLMs) across multiple disciplines (9 such scalable intrinsic characteristics and training paradigms. However, corresponding challenges erupt, as multiple rounds of communication massive model parameters across multi-device clusters have created significant communication bottleneck in training process. For example, pre-training a BERT-Large model (340M parameters) with a single batch size on 16 GPUs shows that 92% of the training time is spent on gradient all-reduce in backward propagation (49). Therefore, a critical problem urgently to be addressed is how to relieve the communication burden while ensuring the quality of model training.
48
5 8 are attributed to
In response to the aforementioned problem, a variety of gradient compression strategies have been proposed reduce the communication overhead of gradients between devices. Such strategies can be broadly categorized into two compression techniques, including low-precision compressors (e.g..
SignSGD [6]. 1-bit Adam [49] and 1-bit LAMB (29]) and low-rank compressors (e.g., Atomo (57) and PowerSGD (55) ). Moreover, freob-feedback technique [46] 49 29] focuses compensating for
errors accumulated during compression by changing the objects to compress from the gradient to
the gradient and historical error summation 49 29 Error-feedback based compression algorithms
have been theoretically proven to have lower communication complexity when applied to SGD-Type
optimizers
[53 30
"Corresponding author

---

Published as a conference paper at ICLR 2025
However, from the system perspective, how to mitigate communication overhead in practice remains a core problem and a formidable challenge. For example, low-bit gradients (c.g., less than 8-bit) not supported by typical hardware. Thus, quantizing gradients into less than 8 bits can lead to significant precision loss and numerical instability. For biased low-precision compressors with error- feedback techniques, which do not support communication primitive all-reduce due to their inherent structures, they must use all-gather for aggregation. It limits them only suitable for master-server communication patterns, incompatible with advanced ring-based and tree-based patterns in current large-scale model training (38) and extremely slows down the communication speed (T.. Therefore, there remains a significant gap between the performance of these methods in training practice and their theoretical results. Considering the system-level speed, recent studies [T] 62 indicate that when training representative L.MMs with off-the-shelf lastr (DDP), most of gradient compressors show longer wall-clock training time than vanilla Adam. Furthermore, integrating these gradient compressors into commonly used optimizers often requires significant modifications, necessitating additional effort to ensure compatibility and maintain effective.
These system-level limitations drive us to consider improvements at the methodological level. The ineffectiveness of these methods stems from inherent flaws in their methodology. Specifically, these algorithms seldom take the characteristics the gradient of LLMs into account, leading to additional computation communication rounds to compensate for the errors introduced by compression. For low-precision compressors, the compression ratio is upper bounded by floating-poin digit number 32. Worse, vector-wise quantization is independent of the properties of the gradient and is commonly computationally heavy. Several low-rank compressors (55) consider the low-rank approximation of gradient, but the compression and decompression are so complex that the time cost in extra computation is close to or even larger than the saved communication time cost. Therefore, the design concept of modern "workable" gradient compression algorithms is that
We should design compression algorithms by leveraging the properties of the model, ensuring the suitability for modern LLMs efficient training ecology, and lightweighting compression and decompression computation to save time.
Contribution. In this work, we focus on how to design a universal gradient compression technique in line with the above concept. We propose SEPARATE, SimplE low-rank Projection for grAdient compRession in modern IArge-scale model Training procEss, which compresses the gradient to arbitrary low-dimension one before communication and then reconstructs after, no matter what optimizers and training frameworks are used. Instead of low-rank approximation estimate of weight matrix for parameter-efficien fine-tuning like LoRA (22). stand on the low-rank structure of the gradient itself.
The motivation of SEPARATE originally comes from the observation of geometric properties of the Hessian spectrum during training dynamic. Especially, many studies reveal that the
of Hessian is often "top-heavy" [55 43 44 61], indicating that several top eigenvalues are dominant the trace of Hessian. This fact intuitively suggests the theoretical feasibility of low-rank gradient compression with restricted variance through careful designed compression strategy. We show the existence of such low-rank properties in Section 3 and the theoretical validity of SEPARATE in
Section 5
Let us give a general introduction to SEPARATE to illustrate its effectiveness. Considering gradient g E Id to communicate, we generate a common Gaussian random matrix M ddxm on each node and do p g x M to compress g. After communication, we use the same random matrix for decompression g = p x M Considering that EI pMM = p MM = p. we show the compressor is unbiased, and the compression ratio is arbitrary. Moreover, the variance analysis seems to be more important. We show theoretically that the variance is bounded in Section 4 and the
gradient complexity for SGD and Adam- Type optimizers maintains the same order as vanilla SGD and Adam in Section 5
We demonstrate that SEPARATE works well in both LLMs pre-training and fine-tuning tasks. To reduce compression eтror in training process. we design a novel crro--feddbcck mechanism for SEPARATE. We estimate more stable compression errors using moving average of historical errors and incorporate these eтrors back into the gradient before compression. We also reset the error periodically eliminate bad historical information and select update directions preciously. For
2

---

Published as a conference paper at ICLR 2025
2 125M
GPT2 345M
Hessian Spectrum
Hessian Spectrum
Layer
Layer
Density (Log Scoae
0.23
10 35
Spectrum of Gradient
Spectnum of Gradiem
Egenvvaue
Eigenvalae
Figure I: Low-rank gradient and Hessian observations in practice. The two pictures on the left show the spectrum of gradient matrices of chosen layers in GPT-2-125M and GPT-2-345M respectively. The X-axis represents the order of eigenvalues from largest to smallest, and the Y-axis represents the magnitude of eigenvalues. The two pictures on the right show the spectrum of Hessian in ResNet-20 and ResNet-1 10 trained on CIFAR-10 respectively. The X-axis represents the eigenvalue, and the Y-axis represents the corresponding density.
GPT-2 40 pre-training. SEPARATE can improve the training speed of Adam by 2 times. For LLAMA2 51] fine-tuning, SEPARATE can obtain better performance on downstream tasks.
As gradient projection method, SEPARATE is independent of the choice of optimizers and can be regarded as plug-in, which can be simply implemented with few lines of code in off-the shelf frameworks shown in Algorithm 2 In addition, the method is robust to some extent, and its performance insensitive to the selection of hyperparameters. As long as the selection of compression ratio meets the random direction sampling quantity required in theoretical analysis, the algorithm can show effective performance.
2 RELATED WORK
Communication Efficient Training. As the scale of models explosively grows in recent years [48
36 25 52], the redundant communication of gradient has become the main bottleneck of training.
Gradient compression 22
56 34) is promising solution for communication-effiienn training among GPUs clusters. In general, the compression techniques can be divided into low- precision methods [49 29
55
49
46
46
34) and low-rank methods 32 2 Besides, several theoretical and
empirical analyses of more general unbiased compressor C 53
21 are also proposed to show the 56 has been proposed to apply
effect of gradient compression theoretically. Moreover, ZeRO
quantization techniques to sharding strategy (41).
Error-feedback Technique. Gradient compression techniques bring information loss during the training process and impair the training accuracy and performance of the model. The accumulation of error may cause algorithm divergence [46]. To solve such problem, the error-feedback technique was first proposed by [46). which adds the compression error to the gradient before compression in this iteration. Afterward, various works study the effect of the error feedback technique from theory to practical methods. Recent works [18 533 study the theoretical convergence of unbiased compressor C with error-feedbace technique on convex and non-convex settings. Moreover, many induce the feebback technique to adaptive gradient algorithms (27 15) which is commonly used in training of large models, such as I-bit Adam (49) and 0/1 Adam (34). LoCo [60] first presents the prror-feedbakk-eee based efficient training method for the sharding strategy.
Low-rank Approximation. Many studies focus on the low-rank approximation of weight or gradient matrix in neural network training dynamic and propose the error bound (55] [2]. However, such a low-rank approximation of weight (c.g.. LoRA (22) may not reach a comparable performance as full-rank pre-training (31) and fine-tuning [58]. Some studies show that the Hessian in deep
33 44], which means the eigenvalue drops fast and the Hessian maintains low-rank, and some recent studies theoretically illustrate that the gradient is naturally low-rank 13 63). In fact, some classical methods have considered low-rank estimation of trace, like the Hutchinson estimator (23). Such related works provide the potential to design a simple gradient compression by low-rank projection with restricted variance. Some recent research about training or gradient accumulation consider the low-rank property of gradient or Hessian, and use random projection for memory reduction 67 19), but lack deep analysis on bounded variance.
neural networks is "top-heavy" 43
3

---

Published as a conference paper at ICLR 2025
Algorithm 1 SEPARATE: A Simple Low-rank Projection for Gradient Compression
Require: Initialization model parameters, nodes, one-round communication budget m, error reset frequency T' 3 E [o, 1I, a common Gaussian random number generator, initialize e" E ((0,cc) c1)
while K do
STEP 1. In each node n compute stochastic gradient g" and h = g.. + ch;
STEP Generate fresh i.i.d. common random Gaussian vectors Eι, , Sm (0, Ia) and compute [pl,n. -mm,. with pi,n (h E.) as the low-dimension projection of h:: STEP 3. Do all-reduce and obtain global projected gradient [pI,... pmll
STEP 4. Compute h 1 E P: E; and use h for model weight update node n:;
STEP 5. Update error e/
I (1 B)e + B(hf - h if KfT * 0, or e ++1 0f k%T 0
(error reset); end while
LOW-RANK PROPERTY OF GRADIENT AND HESSIAN IN TRAINING
Through some phenomena observed in practice and simple theoretical analysis of Transformer [54] we further illustrate the widespread existence of the low-rank property of gradient. Moreover, we also demonstrate that the eigen- gggen-spectrum of Hessian is "top-heavy", consistent with some previous work about Hessian spectrum of neural networks [43 44). We study these observations practically in this work. As shown in Figure T. the landscape
gradient in representative I.L.Ss training possesses low-rank property. Specificaly, We observe the spectrum of GPT-2-Small (125M, with 12 Transformer-bass observe the gradient spectrums of attention and MLP in middle layer and last layer, respectively, The eigenvalues of grad drop fast, especially in MLP layers. This phenomenon is more pronounced in larger models. Based on recent study on Transformer's dynamic (50), theoretically prove that the gradient of MLP layer in Transformer shows low-rank property over time in Appendix [
layers) and GPT-2-Medium (345M, with 24 Transformer-basee layers). We
Moreover, study the spectrum of Hessian of trained deep neural networks. In fact, early studies have discovered that the eigenvalues of simple networks (e.g., 3-layer networks) drop fast. As shown in Figure [ study the spectrums of Hessian of different scales of ResNet (20) trained on CIFAR-10 (3. Except several extremely large eigenvalues, most of eigenvalues of the Hessian are close to zero. is more obvious in larger network structures. This means that lot of the directions along the loss are almost flat, and the trace of Hessian is bounded much less than (d). which ensures our theoretically analysis in Section 5
4 SEPARATE: A SIMPLE LOW-RANK PROJECTION
Considering low-rank property of gradient and "top-heavy" observation of Hessian, natural idea of gradient compression is randomly project the gradient to dindemmnnional subspace. As the sample is up to a certain level (usually much less than dimension d), the dominant information of gradient can be exactly recovered. This leads our SEPARATE method. The core idea of SEPARATE is composed with two parts. First, we design Gaussian random projection compressor for dimensionality reduction of gradient. Hessiann-rrace- tressiantracee-boudee bounded proposes the theoretical guarantee for random projection. Second, to solve the critical challenge that the compression error in gradient compression would accumulate during the training process, we propose a novel moving average error feedback technique to ensure practical utility. We introduce SEPARATE method in Algorithm I
4.1 COMMON RANDOM PROJECTION COMPRESSOR
For classic multi-node parallel training, each node computes its local stochastic gradient g with their mini-bated data. Then through communication like all-reduce each node obtains the global gradient for model update.
As mentioned early, we take the low-rank property of gradient and "top-heavy" property of Hessian into account to design bouunddd low-rank projection for gradient compression, applicable to most of practical scenarios in large-scale model training. Specifically, for the gradient vector g, each node generates m common random Gaussian variables e.,'.. Sm fm N(O, L2). Then each node
4

---

Published as a conference paper at ICLR 2025
computes
P: = Ei) Vi E [m].
(1)
After compression, each node constructs [pr.... Pmn| as the compressed gradient for communication like all-reduce. Then they obtain the global [p Pm land reconstruct the gradient as
E
Then we theoretically illustrate the unbiasedness and variance- variance-bounded property of our common random projection in Lemma 4.1 and Lemma 4.2 as below. The proof details are shown in Appendix
B
Lemma 4.1. g is unbiased estimator of g.'g
E g.
Lemma 4.2. The variance of under norm matrix, can be bounded by Atr(A)
|| where a given positive semi-definite symmetric 1 Ilg.l..
3tr(A)
EE
Bll
gll
lg Igll-
m
m
Remark 4.3. Lemma 4.1 directly indicates the unbiasedness of common random projection, which standard for low-rank compression techniques (55) 67 197. What we really highlight is Lemma which illustrate that the variance of our compression strategy under the Mahalanobis norm can be bounded by the trace of given matrix. Considering the standard convergence analysis of SGD or Adam-Type optimizers, the second order factor of Taylor expansion of the objective function can be written as Mahalanobis norm under Hessian matrix, and the "top-heavy" property ensures the trace of Hessian is much smaller than dL, where d is the dimension and L is Lipschitz constant defined in Assumption 5.2 This provides the potential for SEPARATE-baa based algorithms to surpass other compressors in the order of in convergence speed. We show the theoretical results in Section
4.2
4.2 MOVING AVERAGE ERROR FEEDBACK
Error-feedback technique widely used in gradient compression methods 46
mitigate the information loss during compression. It is worth considering how error-feedback works for our common random projection compressor, even though Lemma 4.1 and Lemma 4.2 show the unbiasedness and trace-bounded variance, because the accumulated error in continuous updates may cause wide deviation. One straight solution is to use the gradient compression error from the previous
42 60
= h - gh where g": is the original gradient on iteration as e = arg min.erd e - h -
the n-th node and h is the reconstruction of compressed gradient. However, we empirically discover the instability of this estimate, because for single-step iteration, random vectors for projection may have large deviation from the direction of true gradient. This means potential abrupt fluctuations of c,, especially for several continuous iterations with series of random vectors with very different directions. This may lead to the whole training process converging to another suboptimal region. This phenomenon particularly prominent when training from scratch, changing the shape of loss curve and the convergence result of loss (see Section 6.3). To solve this problem, inspired by the idea of momentum technique in the analyses of accelerated gradient descent (37), take the moving average of the historical compression error and the current one to maintain the continuity to some degree:
* argmin e
B 2
h
gm
1
1-
3 |le
1122 (1
e + (h 8t) (5)
where 8 E 10 represents the trade-off between two factors. We demonstrate that the moving average error feedback can reduce the total accumulated error when applying to SGD Adam-Type optimizers in Lemma C.2 and Lemma
C.3 In brief, we simply formulate the iteration of SGD and n g where 0 means the model's weight and means
Adam-Type optimizers 0 0º
5

---

Published as a conference paper at ICLR 2025
the Hadamard product of two vectors. Our analysis shows that moving average error-feedback effectively reduces the accumulated gap of weights updated by g and by g as
Σ
n n'o''
Σ n' og'
S o (nle
(6)
which means as the training goes on, the accumulated compression error of the top k rounds only of eo-order, instead of the accumulation from e e Compared with vanilla error feedback technique
[6]. taking moving average of the historical error maintains stability to some extent and reduces the variance of accumulated error.
Moreover, we consider the compression error in the early training process cannot guide the later iteration, so we reset the compression error after given number of iterations like 128. Combining the techniques above we propose SEPARATE method in Algorithm Our SEPARATH can be applied to all the gradient-basee optimizers with few lines of codes, as shown in Algorithm 2
CONVERGENCE GUARANTEE
In this section, we propose the convergence guarantee of our SEPARATE method. We consider the following non-convex optimization problem
min f(0) BS-D [F(0,C)),
where F is the non-convex optimization objective and 0 is the model weight to update. Data is sampled from a given distribution D. We focus on the convergence rate of SEPARATE-bast SGD and Adam-Type Optimizer to find an e-approximate first-order stationary point of the objective (()). For such SEPARATE-bbsed optimizers, they update the model weight as below:
m (1 m 3,h
SGD 0 0 - -
Adam-Type
n7* 1 h h
(8)
@k+1 0 n m
where $1 E (0, 1), and v(.) computes series of pre-conditioners based on different AdamType optimizers. For example, v(-) computes the inverse of the second order moment of Adam's gradient
like v (ho..... & = 1/vv* + where k (1
x2) + (2 Adam-Typp Type optimizers We analyze them uniformly.
differ from various definitions of v(-) 27
15
47
Like the common analysis of the convergence rate of general non-convex objective define the e-approximate first-order stationary point as below.
24
59
we
Definition 5.1 (e-stationary point). The model weight 0 is an e-appriximate first-order stationary point of J [xx]] <.
Definition 5.1 means for SGD and Adam-Type optimizers, we need to find the iteration round K such that [ E f < E, or more tightly
K-1
1
V || S2
K
Moreover, formally present some assumptions to constrain the objective function and the op- timization problem. Assumption 5.2 and [33 are commonly used to describe the properties of the objective in stochastic setting [24 16). Assumption 5.4 ensures the trace of Hessian is bounded
globally, which is also common in practical applications 43
Assumption 5.2 (L-smoothness). The function f is L-smooth if satisfies
[[(()) f(02)) < L| 02l1. V01. 02 Rdd.
(10)
Assumption 5.3 (Stochastic Gradient Boundness). The stochastic gradient g* on each device is unbiased and its infinite norm and variance are bounded as:
E 00 Coc.
E. (((*) 111 - <0²
4k: [K].
(11)

---

Published as a conference paper at ICLR 2025
GPT-2 Single Step Cost
GFT-2-345M
Adamil 26-80
SOPAMITEICR-G
130
1000
Trahking Loss
Traiking LoGa
60-
32-
Steps
Timmist
Madel Sizw
(a) Loss against Steps
(b) Loss against Time
(c) Single Step Time
Figure 2: Loss curve and single step time cost of vanilla AdamW and SEPARATE with different compression ratios (CR) GPT-2-345M trained with 10B tokens Open OpenWettext dataset. (a) shows the loss against the iteration steps, and (b) shows the loss against the wall-clock time. (c) shows the single-step time cost of vanilla A and SEPARATE with different compression ratios.
Assumption 5.4 .omiition)ominniion).... The function f is A-Hessian dominated if there exists positive semi-definite symmetric matrix A such that V f(0) 3 A, E R
Now we present the main theorem of the convergence rate of SEPARATE-basee SGD and Adam-Tyy Type optimizers in general non-convex setting. The proof of two Theorems can be seen in Appendix
Theorem 5.5 (SEPARATE- based SGD convergence). Suppose Assumption 5.2 5.3 and 5.4 hold.
Let e (800g)) e)), namely Vx c))l|| S e1- Let m ut(A) and Atr(A)
With
L
7) O (d-1/22) in SEPARATE-basee SGD, after K Ω (2 iterations in Algorithm 1 holds
K
1
Σ K=O
E V
(()))) so(2).
(12)
K
This improved convergence rate under standard non-convex stochastic setting is derived from our analysis of trace-bounded variance in Lemma In the common analysis of compression techniques, compression reduces the amount of information through single step communication by Old) times, or with ((d) O(d)-variance [2]. Moreover, considering error feedback technique brings bias, the analysis is commonly element-wise (2 46). Thus the convergence rate is often dimensional dependent O (de-)., but they assume is defauli and usually ignore this item [46]. Considering modern I.LMs training, the dimension of parameters is often extremely large and cannot be ignored. Fortunately, our analysis provides an improved rate on the order of d from (d)) to O (d)/2) under natural defined Assumption from "top-heavy" Hessian observation. It is the first gradient compression method to achieve such rate to the best of our knowledge.
Remark 5.6.
Theorem 5.7 (SEPARATE-basee Adam- Type convergence). Suppose Assumption 5.2 5.3 hold.
Assume < |m(())lm & Cu and ||n - I 1131 (1 B1) K-) &u.
Let e
B00,00 c1),
< min
{
a Bic 2224 (1-)) With o (d- and B1 = 0 (d-12) in SEPARATE- based
Adam-Type optimizers, after K = S (de 11) iterations in Algorithm 7 it holds
K
+ 1m so(2).
(13)
K=O
Remark 5.8. Because different coordinates of the gradient are updated with different step sizes in Adam-Tya optimizers, the convergence rate is inevitably dependent dimension d, and combined with compression and error feedback, Adam-Typp optimizers are often O(d) at the order of Theorem 5.7 shows that integrating SEPARATE with various Adam-Type optimizers does affect their convergence speed. Thus SEPARATE based Adam Type optimizers share the same stochastic gradient complexity as O (2- when ignoring the dimensional factor d.
7

---

Published as a conference paper at ICLR 2025
Table 1: Performance comparison with representative communicatio-effiiien methods. We fine- tune LLAMA2-7B on the alpaca-gpt4 dataset and evaluate the performance on downstream tasks, including commonsense reasoning, world knowledge, math, and code.
Method
TriQA
GSM8K
MBPP
NQ
WinoG
Arc-e
Arc-c
PIQA
HellaS
Avg.S
Avg.R
Adam
49.39
15.69
19.40
3.07
46.09
72.31
53.90
57.07
26.45
38.15
3.22
PowerSGD
50.78
18.65
22.40
2.52
40.57
62.43
42.71
52.88
30.12
35.90
3.00
1-hit Adam
62.08
16.53
16.80
5.15
49.57
49.21
37.97
52.88
22.04
34.69
3.22
ZeRO++
$7.49
18.42
22.20
1.72
49.88
43.03
30.85
47.61
33,97
33.91
3.22
SEPARATE
57.18
20.17
21.40
4.18
49.25
72,66
49.83
52,99
29.73
39.71
2.22
6
EXPERIMENTS
6.1 SOTA COMPARISON WITH COMMUNICATION EFFICIENT METHODS
We evaluate SEPARATE on both pre-training and fine-tuning tasks LLMs. The experiment setting details shown in Appendix
Pre-Training GPT-2 from Scratch. We train GPT-2-345M with vanilla AdamW optimizer as a baseline and compare the convergence speed of our SEPARATE-base AdamW optimizer with baseline and PowerSGI [55]. We set different compression ratios of SEPARATE and compare the convergence speed of iteration steps and wall-clock time, respectively in Figure 2 (a) and (b). Figure (a) demonstrates that though high compression ratio results in more steps to reach the same loss, SEPARATE-based optimizer with the proper compression ratio even has a faster convergence speed of iteration steps. Figure 2 (b) shows that SEPARATE can accelerate the training speed up to 2x for GPT-2-345M, pre-tratee tre-training by trade-off the compression ratio and one-step cost, compared with baseline and PowerSGD. Thus we think effective designing algorithms based on low-rank gradient and Hessian of model ensures accurate training with lower cost. Moreover, we test the single-step time cost of vanilla AdamW and SEPARATE. with different compression ratios GPT-2-125M and GPT-2-345M. We find that when model scale increases, the effect SEPARATE is more remarkable, So SEPARATE is significant in training large-scale models.
Fine-tuning LLAMA2 and Evaluating on Downstream Tasks. We follow Llama-Accessory [66] and fine-tune LLAMA2-78B 7 (51) model on pilacca-gp44 (39) dataset for three epochs, and evaluate the ability of fine-tuned model on several downstream tasks including commonsense reasoning, world knowledge, math and code. We report the scores obtained on these evaluation benchmarks in Table We compute the average score (Avg.S, f) and average rank (Avg. R, 4, of each method. The results indicate that SEPARATE outperforms other low-rank or low-bit optimizers with error feedback and ZeRO++, even beyond the performance of vanilla Adam optimizer without compression on the comprehensive performance of all the downstream tasks. Moreover, other communication- efficient optimizers have much worse performance than SEPARATE and vanilla Adam, especially on commonsense reasoning tasks. The possible reason for this result that their iast-iteration-bate error feedback cannot deal with the accumulated error during the whole process, while SEPARATI is applied with the moving average the historical error to stabilize the error and error reset mechanism to remove useless information.
6.2 ABLATION EXPERIMENTS
We conduct ablation experiments to explore the effect of each component in SEPARATE, including 1 Derror-feeddad technique (Frr.Fed)), 2)moving average update of error (Err.Avg.) and 3)error reset mechanism (Err.Re.). For training from scratch, without frror-feedback and moving average techniques, the training process becomes unstable and cannot converge. We show the results in Section 6.3 For fine-tuning tasks, we fine-tune the LLAMA2-7B model alpaca-gpt4 dataset, applying SEPARATE on Adam optimizer with b6loat precision. The results are shown in Table
Error-feedback Technique. We first study the effect of directly applying the vanilla aroo-feedback technique on the common random projection compressor SEPARATE. Comparing SEPARATEJ with SEPARATE2,
observe that directly applying ennor-eeedbac technique slightly
8

---

Published as a conference paper at ICLR 2025
Table 2: Effects of components in SEPARATE to the performance of I.LAMA27B fine-tuned alpaca-gpt4 dataset, including error-feedback technique (Err.Fed.), moving average on error (Err.Avg.) and error reset mechanism (Err.Re.).
Method
Err.Fed.
Err.Avg.
Er..Re.
GSM8K
MBPP
NQ
Arc-e
Arc-c
PIQA
Avg.
SEPARATEI
False
N/A
N/A
21.53
22.00
4.57
71.25
49.15
52.61
36,86
SEPARATE2
True
False
N/A
20.17
21.80
4.35
69.14
51.19
52.29
36.49
SEPARATE3
True
True
N/A
20.92
22.00
4.49
69.49
51.53
52.23
36.77
SEPARATE4
True
True
20.39
21.40
4.46
70.02
49.49
52.45
36.37
SEPARATES
True
True
128
20.17
21.40
4.18
72.66
49.83
52.99
36.87
Table 3: Effect of different compression ratios of SEPARATE the performance of LLAMA2-7B fine-tuned on alpaca-gpt4 dataset, where SEPARATE-X means we set the compression ratio at X.
Method
TriQA
GSM8K
MBPP
NQ
WinoG
Arc-e
Arc-c
PIQA
HellaS
Avg.
SEPARATE-8
55.43
18.88
21.00
4.13
45.78
74.07
50.85
52.83
32.43
39.49
SEPARATE-16
57.18
20.17
21,40
4.18
49.25
72.66
49.83
52.99
29.73
39.71
SEPARATE-64
59.00
19.03
20.40
4.54
49.8
68.43
48.47
52.88
28.59
39.02
SEPARATE-128
58.93
18.65
22.40
4.79
50.28
65.78
47.46
52.67
27.79
38.75
impairs the whole fine-tuning performance, especially on commonsense reasoning tasks such as WinoG, Are and PIQA. This may be due to the randomness of projection directions. When the random projection directions are far from the dominant directions of Hessian in several continuous iterations, the variance of error will become extremely large and misguide the next iteration. Thus we need another technique smooth errors and drop the misleading information arising from randomness.
Moving Average Update of Error. We use the moving average of the historical error and the current to replace the direct from the last iteration, reduce the instability of vanilla error-feedbad technique. Comparing SEPARATE2 and SEPARATE3 we can find that the moving average update of error improves the whole fine-tuning performance, especially Arc-c and MBPP tasks. However, because of the randomness of the compressor, the accumulated misleading information of inexact directions still impairs the performance compared with SEPARATEI.
Error Reset Mechanism. We notice that random projection may generate bad directions far away from the dominant ones of Hessian even with unbiased estimate. Especially compared with quantization methods some bad projections may cause extremely large error. To solve the accumulated bad random information arising from the random projection in the compression error, we design error reset mechanism to reset the accumulated error frequently. Comparing SEPARATE3, SEPARATE4 and SEPARATES, we find that the appropriate set of error reset frequency T can improve the performance, and too large set of eтror reset frequency still causes the misleading information accumulation and worse accuracy.
6.3
CHOICE OF SEPARATE
Random Seed Choice of Projection. In the SEPARATE method, introduce extra randomness due to the random projection. Thus in practice we need to set an extra random seed for commonly generating the same random Gaussian vectors on different devices. In order clarify that our method works not depending the choice of seed, we select different common random seeds for projection to repeat training GPT-2-350M with compression ratio 16. The results are shown in Figure (a), where we randomly set seed as 3407, 4396 and 37, and the results make diffeerence.
in Moving Average Error. SEPARATE method introduces an extra paramm-areetter for moving average update of error. We tested the effect of selection on the convergence speed. We test different 3 in the first 10000 steps of GPT-2-345M training. The results in Figure 3 (b) show that SEPARATE is sensitive the choice of B. If we set 3 1, which means do not moving average update of error, the convergence speed shows down remarkably in the first few steps, and in the subsequent iterations it seriously deviates from the optimal dynamic. Setting 0.95 seems

---

Published as a conference paper at ICLR 2025
Peak Memory
GPT:2-345M
GPT-2-345M
BASELINE
Intenn
SEPARATE
02ED-57
hete-33) 30
Paak McmoryyGa
Tranigg 2.
Tranigg
1900e
20038
10000
2000
2003
4000
3006e
Model
(a) Choice of Seed
(b) Choice of Error 8
(c) Peak Memory
Figure 3: Hyper-parameter choice of SEPARATE and memory cost report. (a) shows the effect
of different common random seeds to SEPARATE training with compression ratio 16. (b) shows the effect of different | choices of moving average update of error to SEPARATE training with compression ratio 16. (c) shows the peak memory Adam and SEPARATE training on different models with compression ratio 16.
be acceptable, and when decreasing 3, the convergence speed slows down again. When 8 < 0.80, the training dynamic becomes unstable with loss value increasing to NaN, which we do not present.
Compression Ratio. We conduct experiments performance of models. One advantage of SEPARATE that SEPARATE can set the compression ratio by users themselves instead of or (at most 32) like quantization methods. We the compression ratio as 8,16,64, and 128, respectively, and tested the performance of the fine-tuned model on downstream tasks. The results are shown in Table 3 The results show that increasing the compression ratio to certain extent enhances the performance of the model, due to the certain noise improving the generalization of models. Besides, an overlarge compression ratio reduces the performance of models, but the accuracy loss is within the acceptable range because if we compress the gradient 128 times, the communication cost is quarter of 1-bit quantization methods and the performance is better than it. We also find that with different compression ratios, SEPARATE focuses on different types of tasks, so users can set the ratio adaptively.
the effect of the scale of compression ratio the
6.4 MEMORY COST OF SEPARATE
We report the peak memory in Figure (c). Results show the peak memory when training or fine- tuning corresponding models, which demonstrates that the extra memory of SEPARATE is negligible in comparison with the communication cost it saves. Moreover, by expanding the random vector buffer, we can save time in generating random Gaussian vectors, which means we can use some memory cost to exchange computational time in compression and decompression to accelerate the training process.
7 CONCLUSION
In this paper, we propose a simple low-rank projection for gradient compression in modern large- scale model training process named SEPARATE. We have carried out theoretical analysis and lot of experimental verification to illustrate SEPARATE provides an easy way to apply gradient compression to all types of gradient- based optimizers across various training frameworks, while maintaining high-quality training performance and achieving significant compression. SEPARATE holds great potential for future research, particularly in the reduction of the memory cost in the training process. Intuitively, by updating the optimizer state in the low-dimension subspace, it has the potential to save optimizer state memory, and this potential similar to its ability to save communication overhead. In addition, SEPARATE with adaptive selection of the compression ratio also worth studying, providing more general framework. In summary, we think SEPARATE is helpful for the further study and popularization of large-scale models.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work is supported by National Key R&D Program of China (2022DD0160300) and the NSF
China (No.s 62376008 and 62276004).
REFERENCES
[1] Saurabh Agarwal et al. "On the utility of gradient compression in distributed training systems".
In: Proceedings of Machine Learning and Systems (2022). pp. 652-672.
[2] Alham Fikri Aji and Kenneth Heafield. "Sparse communication for distributed gradient de- scent". In: arXiv preprint arXiv:1704 05021 (2017).
[3] Krizhevsky Alex. "Learning multiple layers of features from tiny images". In: https://www. cs.
toronto. Shau/riic/lariing-ffatrre features-2009- TR. pdf (2009).
[4] Jacob Austin et al. "Program synthesis with large language models". In: arXiv preprint arXiv:2 07732 (2021).
[5] Yutong Bai al. "Sequential modeling enables scalable learning for large vision models". In: arXiv preprint 1rXiv:23/2.00722 (2023).
[6] Jeremy Bernstein et al. "signSGD: Compressed optimisation for non-convex problems". In: International Conference on Machine Learning- PMLR. 2018, pp- 560-569.
[7] Yonatan Bisk et al. "Piqa: Reasoning about physical commonsense in natural language". In: Proceedings of the AAAI conference on artificial intelligence. Vol. 34. 05. 2020, pp. 7432- 7439.
[8] Andreas Blattmann et al. "Align your latents: High-resolution video synthesis with latent diffusion models", In: Proceedings of the IEEE/CVF Conference on Computer Vision and Patlern Recognition. 2023, pp. 22563-22575.
[9] Tom Brown et al. "Language models few-shot learners". In: Advances in neural information processing systems 33 (2020), pp. 1877-1901.
[10] Peter Clark et al. "Think you have solved question answering? try arc, the ai2 reasoning challenge". In: arXiv preprint arXiv:1803.05457 (2018).
[11] Karl Cobbe et al. "Training verifiers solve math word problems". In: arXiv preprint arXiv:21 14188 (2021).
[12] OpenCompass Contributors. "Opencompass: A universal evaluation platform for foundation models". In: GirHub repository (2023).
[13] Romain Cosson et al. "Low-Rank Gradient Descent". In: IEEE Open Journal of Control
Systems (2023).
[14] Ning Ding et al. "Delta tuning: comprehensive study of parameter efficient methods for pre-trained language models". In: arXiv preprint arXiv:2203 06904 (2022).
[15] John Duchi, Elad Hazan, and Yoram Singer. "Adaptive subgradient methods for online learning and stochastic optimization." " In: Journal of machine learning research 12.7 (2011).
[16] Cong Fang, Zhouchen Lin, and Tong Zhang. "Sharp analysis for nonconvex sgd escaping from saddle points". : In: Conference Learning T'heory. PMLR. 2019, pp. 1192-1234.
[17] Aaron Gokaslan and Vanya Cohen. Open WexTett Corpus. http
Skylion007
github.
2019.
[18] Eduard Gorbunov et al. "MARINA: Faster non-convex distributed learning with compression".
In: International Conference on Machine Learning. PMLR. 2021, pp. 3788-3798. Yongchang Hao, Yanshuai Cao, and Lili Mou. Flora: Low-Rank Adapters Are Secretly Gradient
[19]
Compressors. 2024. arXiv: 2402. 032.9 [cS... LG] URI.: https::/aaxii..ttt
[20] Kaiming He et al. "Deep residual learning for image recognition". In: Proceedings of the IEEE conference on computer vision and pattern recognition. 2016, pp. 770-778.
[21] Yutong He, Xinmeng Huang, and Kun Yuan. "Unbiased Compression Saves Communication in Distributed Optimization: When and How Much?" In: Advances in Neural Information Processing Systems (2024).
[22] Edward Hu et al. "Lora: Low-rank adaptation of large language models". In: arXiv preprint a 106.000.0.60 096.9965 (2021).
11

---

Published as a conference paper at ICLR 2025
[23] M.F. Hutchinson. "A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines". In: Communication in Statistics- Simulation and Computation 18 (Jan. 1ww.110000101100011111
1989), pp. 1059-1076.
[24] Prateek Jain, Purushottam Kar, et al. "Non-convex optimization for machine learning". In:
Foundations and Trends@ in Machine Learning 10.3-4 (2017), pp. 142-363.
[25] Albert Q Jiang et al. "Mixtral of experts". In: arXiv preprint arXiv:2401.0444 0444088.888888 (2024).
[26] Mandar Joshi al. "Triviaqa: A large scale distantly supervised challenge dataset for reading
comprehension". In: arXiv preprint /rXiv:/70s 770.03355 (2017).
[27] Diederik P Kingma and Jimmy Ba. "Adam: A method for stochastic optimization". In: arXiv
preprint arXiv: 1ee.i441266888 (2014).
[28]
Tom Kwiatkowski et al. "Natural questions: a benchmark for question answering research". In: Transactions of the Association for Computational Linguistics 7 (2019), pp. 453-466.
[29] Conglong Li et al. "1-bit LAMB: communication efficient large-scale large-batch training with LAMB's convergence speed". In: 2022 IEEE 29th International Conference on High Performance Computing, Data, and Analytics (HiPC). IEEE, 2022, pp. 272-281.
[30] Zhize Li et al. "Acceleration for compressed gradient descent in distributed and federated optimization". In: arXiv preprint arXiv:2002. 1334 (2020).
[31] Vladislav Lialin al. "Relora: High-rank training through low-rank updates". In: The Twelfth International Conference on Learning Representations: 2023.
[32] Yujun Lin al. "Deep gradient compression: Reducing the communication bandwidth for distributed training". In: arXiv preprint 171..0:712.018880 (2017).
[33] Yuanshi Liu et al. Accelerated Gradient Algorithms with Adaptive Subspace Search for Sittance-Faste. Optimization, 2023. Xiv: 2312. 03218 [es. LG]
[34] Yucheng Lu et al. "Maximizing communication efficiency for large-scale training via 0/1 adam". In: arXiv preprint arXi::2202.06000 (2022).
[35] David Luebke. "CUDA: Sealable parallel programming for perforeeformance scientific com- puting". In: 2008 5th IEEE international symposium on biomedical imaging: from nano to macro. IEEE. 2008, pp. 836-838.
[36] Deepak Narayanan et al. "Efficient large-scale language model training on gpu clusters us- ing megatron-Im".. In: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. 2021, pp. 1-15.
[37] Yurii Nesterov. Introductory lectures on convex optimization: A basic course. Vol. 87. Springer Science & Business Media, 2013.
[38] Pitch Patarasuk and Xin Yuan. "Bandwidth optimal all-reduce algorithms for clusters of workstations". In: Journal of Parallel and Distributed Computing 69.2 (2009), pp. 117-124. Baolin Peng al. "Instruction tuning with gpt-4". In: arXiv preprint arXiv:2304.03227 (2023). Alec Radford et al. "Language models are unsupervised multitask learners". In: OpenAl blog 1.8 (2019), p. 9.
[39]
[40]
[41] Samyam Rajbhandari et al. "Zero: Memory optimizations toward training trillion parameter models". In: SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE. 2020, pp. 1-16.
[42] Peter Richtárik, Igor Sokolov, and Ilyas Fatkhullin. "EF21: A new, simpler, theoretically better, and practically faster error feedback". In: Advances in Neural Information Processing Systems 34 (2021), pp. 4384-4396.
[43] Levent Sagun, Leon Bottou, and Yann LeCun. "Eigenvalues of the hessian in deep learning: Singularity and beyond". In: arXiv preprint 16Xiv:1611.0744 (2016).
[44] Levent Sagun al. Empirical Analysis of the Hessian of Over-Parametrized Neural Networks.
2018. arXiv: 1706. 04454 [cs. LG]
URL: https arxiv org abs / 1706
04454
[45]
Keisuke Sakaguchi et al. "Winogrande: An adversarial winograd schema challenge at scale". In: Communications of the ACM 64.9 (2021), pp. 99-106.
[46] Frank Seide et al. "1-bit stochastic gradient descent and its application to data-pan aata-parallel dis- tributed training of speech dnns". In: Fifteenth annual conference of the international speech communication association. 2014.
12

---

Published as a conference paper at ICLR 2025
[47] Noam Shazeer and Mitchell Stemn. "Adafactor: Adaptive learning rates with sublinear memory
cost". In: International Conference on Machine Learning. PMLR. 2018, pp. 4596-4660
[48]
Mohammad Shoeybi et al. "Megatron-Im: Training multi-billion parameter language models using model parallelism". In: arXiv preprint arXiv:1909.0000 000900803 (2019).
[49] Hanlin Tang et al. "1-bit adam: Communication efficient large-scale training with adam's con- vergence speed". In: International Conference Machine Learning. PMLR. 2021, pp. 10118- 10129.
[50] Yuandong Tian et al. "Joma: Demystifying multilayer transformers via joint dynamics of mlp and attention". In: arXiv preprint arXiv:2310.00535 (2023).
[51] Hugo Touvron et al. "Llama 2: Open foundation and fine-tuned chat models". In: arXiv preprint
arXiv:2307.099 0008809888 (2023).
[52] Hugo Touvron et al. "Lllama Open and efficient foundation language models". In: arXiv preprint arXiv:2302. 13971 (2023).
[53] Alexander Tyurin and Peter Richtárik. "DASHA: Distributed nonconvex optimization with
communication compression, optimal oracle complexity, and no client synchronization". In: arXiv preprint arXiv:2202.01268 (2022).
[54] Ashish Vaswani al. "Attention is all you need". In: Advances in neural information process- ing systems 30 (2017).
[55] Thijs Vogels, Sai Praneeth Karimireddy, and Martin Jaggi. "PowerSGD: Practical low-rank gradient compression for distributed optimization". In: Advances in Neural Information Pro- cessing Systems 32 (2019).
[56] Guanhua Wang et al. "Zero++: Extremely efficient collective communication for giant model
training". In: arXiv preprini arXiv:2306. 10209 (2023).
[57] Hongyi Wang et al. "Atomo: Communication-effiiee learning via atomic sparsification". In: Advances in neural information processing systems 31 (2018).
[58] Wenhan Xia, Chengwei Qin, and Elad Hazan. "Chain of lora: Efficient fine-tuning of language models via residual learning". In: arXiv preprint arXiv:2401.044
(2024).
[59] Xingyu Xie et al. "Adan: Adaptive Nesterov momentum algorithm for faster optimizing deep
models". In: IEEE Transactions on Pattern Analysis and Machine Intelligence (2024).
[60] Xingyu Xie et al. "LoCo: Low-Bit Communication Adaptor for Large-scale Model Training". In: IEEE Transactions on Pattern Analysis and Machine Intelligence (2025).
[61] Xingyu Xie et al. "Optimization prper-parameter laws for large language models". In: arXiv preprint WxXiv:240904477 (2024).
[62] Hang Xu et al. "Compressed communication for distributed deep learning: Survey and quanti- tative evaluation". In: (2020).
[63] Greg Yang. James B Simon, and Jeremy Bernstein. "A spectral condition for feature learning". In: arXiv preprint arXiv:2310.1783 (2023).
[64] Pengyun Yue et al. CORE: Common Random Reconstruction for Distributed Optimization
with Provable Low Communicution Complexity. 2023. arXiv: 2309. 13307 [cs. LG]
[65] Rowan Zellers et al. "Hellaswag: Can a machine really finish your sentence?" In: arXiv preprint
ar05v: 07030 (2019).
[66] Renrui Zhang et al. "Llama-adapter: Efficient fine-tuning of language models with zero-init attention".. In: arXiv preprint arXiv:2303. 16199 (2023).
[67] Jiawei Zhao et al. "Galore: Memory-effiiien Ilm training by gradient low-rank projection". In: arXiv preprint 2xXiv:2403.0350) (2024).
[68] Yanli Zhao et al. "Pytorch fsdp: experiences on scaling fully sharded data paralle!"' In: arXiv preprint arXiv:2304 /277 (2023).
13

---

Published as a conference paper at ICLR 2025
Algorithm 2 SEPARATE: PyTorch-like
random_ random_variable bufferr Inttaaiizaii) $ initialize the variable buffer
def Communication hook (model. (model.weight): i definition hook function for overlap computation projected _aradgradda = Project (grad, random_ vaaioobeearaaeeeeeaaaaaaa borbble_bffffrr random project the grad grad)
Communication(proeeete
grad (eeprojectprrojeteee grad, random_varian cariaale_bbufees
fandom_variall__ bufferr U Update( update variable buffer with new random variables
for model. weight in moddl..wegght. model. weight.register (Communiciaiiinn hook)
f register hook for each layer's weight of the model
A DISCUSSION
In this section, we make a discussion about some related methods for memory- training, the details of our method application, and the robustness. Such a discussion provides us with a clearer understanding of the domain, practicability, and extensibility of our method.
Compared with Memory-Efficien Training. Memory--fficient training is another area of efficient training, focusing on how to reduce the memory cost during training process. We notice that random projection is also used for memory reduction, such as Flora for gradient accumulation [19] and GaLore for memory-efficien update in compact space [67] These approaches seem to show similarity and potential in combination with our method, but we consider different sides of efficient training. It is usually a trade-off because there are key bottlenecks about the extra computational cost for memory reduction. Communication in modern LLMs training framework (e.g.. Mcgatron LM (48) overlaps between local computation and global all-reduce. This allows us to tolerate only minimal extra computational preprocessing overhead, lest long synchronization times cause the wall-clock training time to rise. One matrix projection is almost at the limit of tolerance, and SVD is almost intolerable. To the best of our knowledge, only our "simple but efficient operation" can manage reduce the training wall-clock time.
Details of SEPARATE Application. For the common random projection compressor shown in Algorithm D we introduce common random vector generators on different nodes of the distributed training clusters. This setup leads our analysis, but it does not incur any additional overhead in actual application, because we can easily set dedicated common random seed to generate the same random variables on each node during initialization. Moreover, we could easily apply SEPARATE all kinds of gradient-based optimizers and can be regarded plug-in, which can be simply implemented with a few lines of code off-the-shelf frameworks as shown in Algorithm 2 This means that our method is also the first gradient compression technique that can be scamlessly integrated with FSDP rather than DDP to the best of our knowledge.
Robustness with Few Parameter Tuning. We expect to propose simple and effective gradient compression method that can be efficiently adapted to existing LLMs training frameworks. Therefore, expect the method have certain degree of robustness, that is, does not depend on the resetting of hyperparameters. In our experiments, all hyperparameters are derived from the default settings of the corresponding model trained under the corresponding framework. In other words, do not adjust the hyperparameters individually for SEPARATE. We show the hyperparameters setting in Appendix E In addition, the method introduces the compression ratio. Our variance analysis and the main theorem demonstrate that, provided the selection of the compression ratio aligns with the conditions in our theoretical analysis (Theorem 5.5 and 5.7]. our method can exhibit effective performance. This theoretical underpinning ensures that our method is robust within the specified compression ratio. Second, in practical applications, especially for models with a substantial number of parameters, such those in the millions or billions, we can defaultly set the compression ratio as 16, 32 or 64. As illustrated Table 3, the performance differences between compression ratios of 16. 32, 64 times for gradient information ratio is more influenced by the user's s device constraints rather than the intrinsic characteristics the model.
minimal. Consequently. the choice of compression
14

---

Published as a conference paper at ICLR 2025
B PROPERTIES OF SEPARATE COMPRESSOR
In this section, we propose several properties of the expectation and variance of the common random projection compressor in SEPARATE. The first analysis of such properties is proposed in the research of distributed optimization theory (64). Let represent the vector for communication and à represent the estimate one generated by the compressor. In Lemma 4.1 and Lemma 4.2 we show that is an unbiased estimator, and the variance of à can be tr( ( b-o-bbdnded under arbitrary matrix A-norms, respectively.
Lemma B.1. â is an unbiased estimator of a,
Eg.... fm a.
(14)
Proof of Lemma 4.1
E ā Egi..... -m
a, Ex) Ei
m 1=1
(15)
= -Eg E15 a =
Ia
a
The variance of à under norm symmetric matrix, can be bounded by 3tr(A)
Lemma B.2.
|| ll. where A is a given positive semi-definite
m |a| m ||a|| A
3tr(A)
E Em
la
a 1 |al
(16)
Proof of Lemma 4.2 For the simplicity of notation, we use Ee to denote Eg,... S..
Ee ll
all2
EE
(a, E) ----) a)
m
m
(17)
Ee
a E AC: a - a Aa)
m
Eg, al E AE 1
|all
Let A UTUU be the eigenvalue decomposition of A where D diag (bi. ba) is a diagonal matrix, and Ç - Ući be a linear transformation of the random variable & We have
Eg, [ AE
_ E, [ SC
DCC
d
= =U
Σ
b. SC
U
(18)
b TT
be I+ 2D
- ir( ir(A) I+ 2A
- 3tr(A) I.
In =, we use Ç (0, La) based on the rotational invariance of the standard Gaussian distribution.
In b
we use the second and forth moment of standard Gaussian variables: EC? = and ECC = In
we use tr(U DU) tr(UT T UD) = tr(D). The last inequality of 18] is due to tr(A) I A.
Combining (18) and (17). we have
31r(A)
E S... la - all2 <
m |||
|||
(19)
15

---

Published as a conference paper at ICLR 2025
C DEFERRED PROOF IN SECTION 5
C.1 USEFUL LEMMAS
In this Section, we propose several useful lemmas for the proof of our main theorem in Section
Lemma C.1. For the gradient h; E Pi & where i; N En=1 Pi,u is the global average
of as in each machine after communication in Algorithm 4
we have
h = + B (et -eht)).
(20)
where h²
En=
g
N E gh and e
N En= e
Proof of Lemma
C.7
For the estimate h defined in Lemma C.I we have
N
N
h E1 E
N
m
hh {.) E:
(21)
N
==1
1
m
h
h Ei)
E: h
m
8
R=1
where 6t
= h3
1
Em=
(h <i)
Taking average of ot. we have
m
N
1
1 Σ (h (. Es
8²
h
N
n=1
h
ht
(22)
(e - -11 B)et),
= 3
where the last equality uses e ++1 (1 B) + +13hh* htt which is the average of the iteration of k+1 in Algorithm Combining the two equalities above, we have
h = + 8h = g" + B (e" -.
(23)
We complete the proof of Lemma
Lemma C.2. Suppose the setting in Theorem 5.5 hold, and we use h in Algorithm 1 to update the model weight 0 in SGD
0 g++ 8¹ nnh
(24)
Then we have
< 1 (1 3)h+1
E
Σ
h
(25)
Consider the expectation of ||* about the random Gaussian variables E ... Em and the stochastic mini-batch as below
Proof of Lemma C.2
E El
k+1
||
E
[
Be Em. ||e
(26)
16

---

Published as a conference paper at ICLR 2025
First we have
Es,... Em em.le ||e+t|1 BEL.... "...... Em ||[1 -3e² + 3(h -hf)ll S (1 E &m ||o*|| + Bg.... &. 1..11.11. 11 h0*
(27)
(1 3)1
where the first inequality uses h is the unbiased estimate of h Thus have
E E
|| - (1 e) - (1 3y* || S (1 B) 3)k+la c1.
(28)
Considering the accumulated bias of h and h based on Lemma C.I we have
1 (1 3)++1
h g¹)
=
Ee' 1=0
i+1
S [I
B
c1
(29)
1=0
Lemma C.3. Suppose the setting in Theorem model weight 0 in Adam-Type optimizers {m²) ) Jk=o and {mh K
5.7 hold, and we use h in Algorithm / to update Then consider the following two sequences
Jk=n as:
m * (1 (1- 31) k + Big" k-1 3
(30)
m k (1 -1)mm
Assume that for the sequence
[ Ik= each element of TH satisfies
na v S 1/us
Inf n 1 m1 (1 31) KK
cns vi E id], k E [K].
(31)
Then we have
K -m') S
2ccennVa
Σ n (m
(32)
Proof of Lemma
C.3
First we extend the iterations of [m* J Jk=o and {m*) k
Jk=o
k
m (1 - 31)m k-1 + +Big" = =(1 -11-8,) k-1 m + B1 Σ (1 31) g''
*g';
(33)
1=1
and
m (1- )m 3,h (1 (1-3,)** 1-1 mm + 31 Σ
31 h
(34)
Based on Lemma C.I we have
Σ 1=0
m
3
Σ
k
n
1u-
i
(1) -
(35)
3 ==
Then element-wisely analyse the upper bound of the accumulated error below. We the non-blackbody letters to represent the element vectors for convenience with some symbolic abuse.
8
(1-81)² '(" el42)
3
(36)
k-
Σ
(n' - 1/¹ (1 - 3 -/ 11)
8
17

---

Published as a conference paper at ICLR 2025
Then using Lemma
C.2
take the expectation of the accumulated error as below:
k
E
Σ
n
o
(m²
m²)
1=
I Vd. E
n
(m'
m
i=0
I Vd. sup ne(d]
B1
&
k
mn
Th (1- '1 - Blyb-s cn
(37)
3
3Iaicuvva
k-1
E
31
t(1 - 8(1-81)+
+ t=1
Σ (1- 31)²
B
20caaaV B
For the first inequality we use the property of vector norm that -| S Ilso. For the second inequality we use the definition of vector's infinite norm. For the third inequality we use the assumption that " n (1 ( f Vi [K] to obtain the upper bound of each element, and use Lemma C.2 to estimate the upper bound of Ele;, that Ele",1 I E|ch I For the forth
inequality the summation of series to obtain that B1 EE= 1(1 B1)² It=1 (1 B1) B, Thus have
2acannVa
Σn n (m' m')
S
(38)
Lemma C.4. Consider the moving average iteration like AdamType as m = (1 33m k-1 + +gg
(39)
where g f((()) + e is the stochastic gradient with Eg and E
E S 0²2. Then we have
E @ - 8* 11 2 + 82.. )( < (1 B) 1 (1 3)
E m
B
(40)
Proof of Lemma C.4 Based on the iteration of m we have
m ( = (1-3) (m k-1
-V(00))))(((0))+++ -1000)))---00010000"--00=01)1111 (1)(0)))))(
vF(("))
(41)
Then taking expectation of both sides we have
E I
\1( (1 -3 I| (S(0 +(1 "E (0 v)(00~1)11
=3 2(1 8 1 -
V5 (((0))1)))
(1 (1 + aa)
Im k- [/( 2
+(1 B)3 1 (1+ E. [ <S(0k=1)1 8202
(1 3)E I k-1 vf(0t=1)|1 1|
(1 3)2
| - -1(00-1111 + 3202
(1- L
(1 8EE
E 0h 0-112 + 8202
3
(42)
18

---

Published as a conference paper at ICLR 2025
C.2 PROOF OF THEOREM 5.5
Proof. Consider the two SGD-based iteration {0*}
kk=1 as:
@a+1 0² "/8"
0
(43)
and { sk=1 as:
gh+1 0² 1/hh
0
g
+
g
h²)
(44)
i=0
1=0
For & E [K], due to the L-smooth assumption of the objective f and Lemma C.1 we have
E V v
LIE 0z
0²|
LLE
Σ g' nii)
A--1
[1+(1 8) ML
(45)
Next, we write the second-order Taylor expansion of f(0k+1) 1) at 0h as below:
s(Ok+1) s 5(0*) (VS(("), 0 0², + (A 0++1 0²).
(46)
Taking expectation of (46) with condition on the iteration before h and random Gaussian variables to both sides of (46). using Lemma 4.1 Lemma 4.2 and Assumption 5.4 we have
3tr 32 (2)((k)))
V V )| + (
Ef +1
< E f((h) + ((th)) @k+1 0h) +1²
2m
$ <f((0* - TE ( V)(0*)) +nB (V V -V)(0)))((0))))
+1² (3tr(A)
2m +L) E | |
< Ef((t)) - 11 /( 1|² + +E (TS(O*). V5 Vs(ek))
+ ant (1)14)) ( V)(@E)) )) +1² 5tr(A) E |
2m
< <Ef(0*) // / 1(0 00
(1 - (1 - B) b vdcce
+ 122 (lr(A)
(|
o²
+
2m
N
(47)
where in 2 we use m S tr(A) L
Then, using 7) S Atr
we have
K Σ
2 W
8 -(()))
+0 (
BygLer vdcco ) +0
V
(48)
K k=
3mK
B
N
By letting 17 0 1/2 k= (2 we have
Σ V (+t) || 50(2 (((0)) (0((1100)) .(0-) )++++ (2-
coo N
= ( (49)
3
Thus finish the proof of Theorem 5.5
19

---

Published as a conference paper at ICLR 2025
C.3 PROOF OF THEOREM
Proof. We also consider two AdamType-based iteration (0* [0= as:
g=+l = 0z -1 mm = 00 Σ
n' o m'.
(50)
1=0
where m = =(1- (1-31)mm -1)mm² --1 + 31g", and {0 k==1 as:
gh+1 g - -nn m 0 Σn n' m 00
Σ
n'
m +
Σ>n
n'
a (m' - -m'), (51)
==0
1=0
1=0
where r (1 B1 ) k- B1 Based on Lemma C.3
have
2e
E Vf(0h)
LE
@h
8
LE Σ n
(m²
m²)
1
(52)
and
4cf
E 0h
E
Σ n
(m'
m
(53)
82
Based on the L-smooth assumption of the objective S. we have
1f( S = @++1 0*)
L
12|000
1
1
2
L
-E((0h) E (V m
E
g+++1
044
2
4
-ES(0*) E V I(0)), 1/ m²
E
0++ 041
2
+R +R(V/(0') (/(0"), " mm
L
g 2e
< Ef(ee) E V f((t),
17 m
+
2
B
Ef(e*)
E
VnF
(V
-m
||
v
2
2
(54)
E V1t² m 2
+
E ||17**
x|
2L
Ef(e0) + Gill E
V((0t)
m k - - E
| ||?
2
Le cm
2cc1cucoo11
E. 2||m* R|| +
S E.((0)
+ Gill E
VS(@") m²|
-
E
|
2
cill E
m m1
CL
For the first inequality we use (52) For the second inequality we use the upper and lower bound of r7.
For the third inequality we set 17 ELe?
Then based Lemma C.4 we have
E - - r < 3. -
(1
3,
E
Ø6-
@
0
3f02
31
N
(55)
For the second factor we have
E @
< E
0 -
gh
+++ F @+ gz| 00(
+ R
g -h 03-11(3)
(56)
20

---

Published as a conference paper at ICLR 2025
Thus, have
B,)
m h - v/(0t)|
- (1 (1-3,))B m k- |
* --1
m
81
0300 + 24L22 24L22g76,7774
B
(57)
By adding equation
54 at 0++1 and ax equation 57 at @h++, we have
Es((
aE
m ++1
(
2
S <E + Cull + + - Bi) E
V
m
m
-
cm
3ad 11 - --,)²
E |m
2
4
B1
apf0²
2L
24aT? 1e
N
B
B²
(58)
V s|(
Letting
nc B1
and G(Ot) Ef(0*)
neu
E
m
we have
31
G( G((
S
QM
V 11²
3
E
m
2
39
2
24L2
+
N
B
3133
(59)
S G(
|5
||2
2
8
24L
N
8,132
3,
where we set 17 S
Then we sum the inequality above from k = 0 K - and obtain
K
|/
||m
2(G - -G(OK)
2cuB102
4lc
E
+
<
+
K
nar
.N
k=0
(60)
48T. 29
0,3/32
Considering that
G(0") - GG(k))(k) < R0 - Ef((K) + 7JC
m
31
(61)
$ Ef Ef(0)) + 1
N
and letting (0" 10 = A, we have
|"
+ *
2A
20.
20023002
ALcrcucoond
K E
+
qarK K
GK3NN
aN
(62)
01,1,
Finally setting K 9 (d)--)). 1 O (d (d) and 31 0 (d- we have
K K-1
T || + m | s0(2).
(63)
K
We end the proof of Theorem 5.7
21

---

1 Published as a conference paper at ICLR 2025
D LOW-RANK PROPERTY OF GRADIENT IN TRAINING DYNAMIC
In this section, we give an simple theoretical analysis to demonstrate that the gradient of project-up layers MLP in Transformer-based models becomes low-rank with training time. Our analysis is based on recent study on the framework of Transformer dynamic, JoMA [50]. and fix an error in related study [67].
We follow the definitions and assumptions of Transformer in JoMA 50 and compute the dynamic of project-up layer in Transformer MLP as below.
Lemma D.1. Suppose the embedding matrix U E &RRxMM is fixed, where M represents the vocabulary size. The active function is linear and the back propagation gradient is stationary [50]. The weight
matrix is W E gmxn Let A -AA,.... An). where A Eg[g,x] RM 9j is the back propagated gradient of hidden node j in MLP, Eg represents the conditional expectation given the query q, and x represents the token distribution the previous layer, for example, the activation output of the previous layer. Following the same definitions and assumptions in JoMA [50], the project-up up matrix V =UW RRMxx satisfis
Vo V
V exp
A
(64)
2
where A is the normalization of softmax.
Proof. Let V [Vi, vn. and W [WI,'.. wyl. Following the Theorem 2 in study of JoMA
framework 50). for each hidden node j. vj = U w satisfies
A o exp (v]
(65)
where v? vj o V, is the element-wise product of vector. Then notice that Lemma D.1 is the matrix form of equation (65).
With Lemma D.I we analyse the training dynamic of V. Row-wisely write V as V [u]
M and A as A- .... T Mr Based Lemma D.I for E [M], we have
1 ú₁ A exp (u//2) o A diag (exp (uz/2)) 8;
(66)
Based on the assumption that the back propagation gradient is stationary and x is the output of previous layer, which means it given. Thus the A, is constant, and the direction u, is controlled by diagonal matrix. Let ú, = D2({)ō,, where D,(1) is diagonal matrix with initialization (0 = O. Then we have
D3 ))
1
11(1)= = diag exp
(67)
A
2
Let a² = maxxe[1] 0% and b3 = min 83.-
Then we have
))
DJ (0
diag exp
= DD,(t) = diag
exp
(68)
Element-wisely analyze the diagnoal matrix (t diag (d,1 ... (t).))). For i [n]. we
have
(1)3
df,(()ag
exp
< dj. S A
exp
(69)
A
2
2
For dynamic like è = Ce 2a
we have
/23C
xxl)= = erf
(70)
3
22

---

Published as a conference paper at ICLR 2025
Table 4: Comparison of SEPARATE with previous common gradient compression algorithms in standard application metrics, including gradient complexity real communication time, memory cost, ring-based communication support, and sharding support. Here, V represents the scale of model parameters, .N represents the number of nodes in distributed cluster, B is the communication bandwidth (bytes/s), r the dimension of low-rank subspace of gradient PowerSGD, and m the random Gaussian samples in SEPARATE in Algorithm [ Here we only consider the order of in gradient complexity, ignoring others especially the dimension parameters d. We discuss the order We consider the mixed -exeedprecision setting for memory overhead.
of d in detail in Section 5
Gradient
Ring-bake ding-based
Gradient Sharding
Methods
Communication Cost
Memory Cost
Complexity
Comm.
SGD
o (c +)
4V(N (BN)
2V 6M/N
Adam
27
O (c +)
4 1)/(BN)
2 + 14 /N
v
1-bit Adam [49]
O (E +)
0.625W.N/B
18V
X
x
1-bit LAMB
29
0 (E-a)
A,625VN/B
22
X
PowerSGD 55)
x
Tr (BN(
14t + 2r
v
SEPARATE-SGD
O (c-
4mm( - L)/(BN)
4 6M/N
v
SEPARATE-Adan
O (c-
Am(N - 1)/(BN)
AV + 14V/N
where erf(z) is Gaussian error function. Then we have
V2
ert (V (V2%4)
12 A
V2
-2eer-
1 /20 (V22))
< dz:(t) S
(71)
by
aj
A
Since the bound is independent of i, for each j we have
V2 eri bj
2 b (v π A
(V D50
v2 erl
2 aj (v
I.
(72)
π A
It demonstrates that for each line of V, the direction of u, is near to the direction of 8, with error dependent to aj and by. Next we show that the change of direction from one line to another is very different resulting in low-rank property of V. Let
/2 ((,aa) = e
2 a (V2
222)) π A
(73)
Then we have h(t,b;)I = (t) 3 a(t,a))L Considering lim, -/a//// h(t,a)
+ toc, let
j = arg max, by be the row with the largest entry of by. Then if by- 2 aj for all j + j", when t t- /.TA/22br we have (t) = hit, by- I + too but (t( = h(t, az I is still finite because 1' /TA//22 1"
Such analysis demonstrates that the magnitude of each row of V changes drastically. The dominant direction of V will become extremely large when others are still in small range over time. V will become near rank-1 I over time and V has lower rank because D have D on exponents. Such analysis is consistent with what we observed in Figure I in Section B
E EXPEIMENT DETAILS
E. MODEL CONFIGURATIONS
We pre-train the GPT-2-345M [40] model OpenWebtext Webtext dataset [7 of 10B tokens from scratch. We show the results in Figure 2 baselines to fine-tune the I.LAMA2-BB model and evaluate on downstream tasks. The baselines
Second, we compare the SEPARATE with several representative
23

---

Published as a conference paper at ICLR 2025
Table 5: Compare SEPARATE with LoRA. Suppose that the weight matrix W E TRxx", the low-rank reparameterized size of LoRA is r, and the compression ratio of SEPARATE is k.
SEPARATE
LoRA
Weight
Grad Comm.
ma/k
Pre-train
X
Fine-tune
Table 6: Performance comparison with LoRA. We fine-tune LLAMA2-7B on the alpaca-gpt4 dataset and evaluate the performance downstream tasks, including commonsense reasoning, world knowledge, math, and code.
Method
TriQA
GSM8K
MBPP
NQ
WinoG
Arc-e
Arc-c
PIQA
Avg.
Adam
49.39
15.69
19.40
3.07
46.09
72.31
53.90
57.07
39.62
Adam-LoRA
62.08
16.53
16.80
5.15
49.57
49.21
37.97
52.88
36.27
SEPARATE
57.18
20.17
21.40
4.18
49.25
72.66
49.83
52.99
40.96
include typical low-rank optimizers with error feedhack technique like PowerSGI) [55 representative low-bit optimizers like 1-bit Adam [49], and recent quantization efficient training strategy ZeRO++ [56]. For fairness, we use SEPARATE-basee Adam for comparison. In this section, we propose the experiment details including model and hyper-parameters settings. Without special emphasis all experiments are under bfloat16 6 precision.
GPT-2. We conducted the training of GPT-2-345M model [40] from scratch on the Open Webtext [77] dataset on 8x NVIDIA 3090 24G GPUs cluster. We ran total of 50000 iterations and processed a total of 10B tokens. The experiment was configured with PyTorch DDP training framework without model and pipeline parallelism to simplify the experimental setting and fit existing devices. Our global batch size was 8x 512. We the learning rate at 6.0e-4 with the cosine decay down the minimum 6.0e-5 after 2000 iterations of warm-up. We also used the gradient accumulation and set the gradient accumulation step at We used the global gradient norm clipping of and set the AdamW with 31 = 0.9. 32 0.95 and weight decay 0.1.
We conducted the fine-tuning of LLAMA2-7B 51 model on alpaca-gpt4 dataset 39 within the PyTorch FSDP [68] framework. We followed protocols from LLAMA-Accessory 66 and fine-tuned LLAMAA-7H for three epochs on 8x NVIDIA A6000 48G GPUs cluster. To fit the devices set the data parallelism at 4 and model parallelism 2 without pipeline parallelism.
LLAMA2-7B.
We set the global batch size at 32, the learning rate of 2e-5, gradient clipping of 2 and gradient accumulation step 8.
E.2 BENCHMARKS
Following the previous work [66], we conducted evaluation on chosen benchmarks on the popular OpenCompass [12] platform. We chose the benchmarks including commonsense reasoning, world knowledge, math and code as follows:
Commonsense Reasoning: Hellaswag
665
Winogrande
45
ARC-Easy
10
ARC-
Challenge 110 and PIQA (7).
World Knowledge: NaturalQuestions
28 and TriviaQA 26
Math: GSM8K [
Code: MBPP
E.3 APPLICATION TO TRAINING PROCESS AS PLUG-IN
SEPARATE can be easily applied to real large-scale model training tasks. SEPARATE only needs small amount of extra computation costs, which can be formulated tensor multiplication form that
24

---

Published as a conference paper at ICLR 2025
Algorithm 3 G-SEPARATE: General Version of Simple Low-rank Projection
Require: Initialization model parameters with L layers, N nodes, layer-wise communication ratio {mi layer-wise {B P = 0.95, AA E [L], error reset frequency T., adaptive update
frequency Ta, a common Gaussian random number generator, initialize cº c))
while h S K do
STEP In each node n compute stochastic gradient g*. and h - Bn. + " STEP 2. Generate fresh i.id. common random Gaussian vectors E1, Smp ~ ( La)
and compute [p1, [pi,n,t""" "Ppt:"" in.) with Pi,n, (b,, 1. Et) as the low-dimension
projection of b STEP 3. Do all-reduce and obtain global projected gradient [i...... ... Pmnt .ll
1 STEP 4. Compute h m Em Pa & and use h for model weight update
in node n: STEP 5. Update error
k+1 = (1 - )e + 8h (H.) if 1%T. + O.
n.
n.l
k+1 0 if k.TT
(error reset);
STEP Layer-wisely update compression ratio and 3
k+1 int
1 1 1
(1
(h (h", ..s.......
m
(b) (.0...55. Ba.)
0 max
min Bf
0.99
0.90
if KTTT 0,
m m
S gh+1 +1_ 33,, K%TT #0;
end while
well supported by modern computational device architectures like CUDA [35]. By using hook mechanism in PyTorch, we can compress the latter layer's gradient and do communication while the optimizer computes the former layer's gradient in backpropagation process. Thus we properly arrange the overlapped computation and communication order, the time cost of compression can be hidden in the backpropagation computation, and SEPARATE can realize faster training time. Moreover, SEPARATE can be applied to ring-based or tree-based communication. Different from quantization methods, SEPARATE is not troubled by numerical anomalies like overflow or underflow, so it well supports reduce-scatter operation. SEPARATE also supports sharding strategy which is commonly used for training large language models [41]. Even though the gradient separated into different nodes, the linearity of random projection ensures the correctness of SEPARATE'S results. We propose an example that how to use SEPARATE as plug-in by hook mechanism Algorithm 2
Besides, SEPARATE brings extra memory costs for training because of the store of error and random variables. Setting up random variable buffer, we can reduce the memory costs by reusing this part of memory store the stream of random variables. In summary, from the results in Table 4 it is evident that SEPARATE demonstrates superior properties compared to other methods.
E.4 ADDITIONAL EXPERIMENTS
Parameter-efficient fine-tuning (PEFT) techniques accelerate the application of pre-trained language models different downstream tasks without the need to fine-tune all of the model's parameters [14). Among them, the popular Low-Rank Adaptation (LoRA (22)) reparameterizes the weight matrix to low-rank approximations and fine-tuning the reparameterized ones. We compare SEPARATE with popular LoRA on scopes of application and performance on downstream tasks in Table and Table 6 The results demonstrate that SEPARATE has better versatility and effect than LoRA.
E.5 GENERAL VERSION OF SIMPLE LOW-RANK PROJECTION
25

---

Published as a conference paper at ICLR 2025
In this section, we study how to adaptively select the compression ratio and 3 for moving average error feedback. We expect to dynamically adjust these hyperparameters by combining the character- istics of the parameters at each layer of the model and the training dynamic. Our research focuses two aspects. First, the adaptive strategy needs to remain simple and efficient to ensure the over- all training wall-clock time is reduced. Thus, still do not consider the strategies that require pe- riodic heavy computations, such periodic SVD. Second, considering the robustness of our method shown Appendix[A] we think the dynamic adjust- Figure 4: Loss curve of vanilla Adam, SEP- SEP-ment ment of hyperparameters should not be too sharp ARATE with compression ratio 8, and G-to avoid instability in training. Under this consid- SEPARATE GPT-2-345M trained with eration, we get the general version of SEPARATE 10B tokens OpenWebtext dataset.
GPT-2-345M
10.0
Adamw(16 bit)
SEPARATE
Training Loss
G-SEPARATE
Steps
in Algorithm 3
We layer-wisely compute the cosine similarity of the gradient estimate and the gradient itself fre- quently to evaluate the estimate accuracy, and update the compression ratio and / in moving average error feedback dynamically. As shown in Algorithm STEP 6, if the estimate is accurate, we try to use a more aggressive compression ratio and a larger 3 in the nest stage, and vice versa. All steps in Algorithm are layer-wise, and we use the subscript to represent the operation on the I-th layer. The initialization my can be arbitrary such as 16, 32 or 64. The adaptive mechanism can adjust to the proper region of the corresponding layer through several dynamic update processes.
We pre-train GPT-2-345M on 10B tokens Open Webtext dataset from scratch to verify the effectiveness of G-SEPARATE. We follow the same hyperparameter setting of our pre-training experiment in Appendix E.I For G-SEPARATE, we set the adaptive update frequency T. = 2000 to ensure the stability of training. The results shown in Figure 4 indicate that the adaptive extension more slightly fits the baseline, but shares the similar performance of the original algorithm in total.
26


Reviews:
[{"review_id": "UlodIrYnQ8", "reviewer": "Reviewer_Kod5", "summary": "This work proposes an efficient way to communicate gradients between nodes for training large models. The method works by communicating random projections of the gradients, where the projections have zero mean and identity covariance. The projections can be regenerated from the random seed, so a shared random number generator between nodes is assumed. Thus, since the projections $P$ have identity covariance, an unbiased estimate of the original gradient $g$ can be found by noting $\\mathbb{E}[gPP^\\top]=g$. The authors further improve this method by incorporating an error-feedback mechanism that works like momentum. They derive an Adam variant of this method and show that it is effective in practice. They also provide convergence analysis of this method on non-convex objectives.", "questions": "None.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Using the random seed for reconstructing the random projections is a smart trick for efficient communication and saving space.\n- This method has the benefit of being an unbiased estimator of the true gradient, which is not the case with low-bit and some low-rank methods. Even though the variance might be bad (i.e., worse than the distance between the true gradient and its truncated SVD), error feedback (with restarts) seems to mitigate this issue during training, allowing SEPARATE to be stable and perform better than low-rank projection methods, such as PowerSGD. This implies that old gradient directions can still be relevant even after many iterations.\n- The authors additionally show an Adam variant of SEPARATE and derive the convergence rate of both the normal variant and the Adam variant. \n- The low-rankedness in LLMs have been (empirically) demonstrated before, but the authors here show a theoretical analysis of this phenomenon in the appendix.", "weaknesses": "- Sharing random generator between nodes might introduce a layer of complexity that could potentially make debugging difficult when things go wrong. For example, this might prevent a real-life implementation to use non-deterministic approaches for accelerating training. Elaborating on this part would be great. For example, the authors can emphasize that this is not a problem and discuss how to ensure a proper regeneration of the random variable across nodes.\n- This trick for getting an unbiased estimator from a sketch matrix M is actually not new as far as I'm concerned. The authors did not claim its novelty, but I just wanted to mention this since they did not cite relevant works to this trick. For example, the Hutchinson estimator does a very similar thing to calculate the trace (or diagonal) of the Hessian. Citing this estimator can help future readers to look for relevant tricks in the literature.\n- Table 1: the authors should explain what \"performance\" mean here, at least in the main text. I also do not think the use of average performance across datasets is sound. Perhaps the average rank or quantile per dataset is a better aggregated metric.\n- Regarding regularization by projection, an elaboration of the claim in line 407 would be great. If the authors believe that some sort of regularization is happening, then it would be great to provide some experiments or analysis to corroborate this hypothesis, which would be interesting to see. Otherwise, this might feel like a handwavy explanation of the results.\n- The discussion regarding the hessian spectrum is already known and have been demonstrated in quite a few works, e.g. [1], which the authors duly noted, but I'm mentioning to emphasize that this is a motivation rather than a novel discovery. This also applies to studies regarding low-rank properties in LLMs. Thus, the originality of the contribution comes from applying random low-rank projections with an error-feedback mechanism along with the analysis in the appendix. Perhaps stating this clearly in the contributions section can help understand this work better.\n- Some explanations are provided without support, such as line 458: “When the random projection directions are far from the dominant directions of Hessian in several continuous iterations, the variance of error will become extremely large and misguide the next iteration.” It is easy to see how this holds intuitively, but it might not necessarily be the case that error feedback mitigates this phenomenon, or that this phenomenon occurs in the first place in practice. For example, SEPARATE1 in Table 2 seems to be working fine.\n- I don't see how the Ablation study in Table 2 is conclusive. It would be better to provide the average + std for a few seeds, say 5. That would make it clearer whether the difference is significant or not.\n\n[1] Empirical Analysis of the Hessian of Over-Parametrized Neural Networks. Sagun et al. ICLR 2018", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "e6vboQkmDA", "reviewer": "Reviewer_Notb", "summary": "This paper introduces SEPARATE, a gradient compression technique designed to address communication bottlenecks in large-scale model training. SEPARATE leverages low-rank properties observed in the gradients and Hessians of large language models (LLMs), compressing gradients using random Gaussian projections. The method also includes an error-feedback mechanism to counteract compression-induced errors by averaging historical errors over time, which aims to stabilize training dynamics. Experimental results on models such as GPT-2 and LLAMA2 demonstrate SEPARATE’s potential for up to 2x speedup compared to baselines while maintaining similar accuracy across several downstream tasks.", "questions": "Can SEPARATE generalize to architectures without low-rank gradient properties? \nCould a more adaptive error-feedback mechanism improve stability? \nCould SEPARATE benefit from dynamically adjustable compression ratios?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "SEPARATE introduces an innovative error-feedback mechanism that effectively addresses the inherent variability associated with random projections by smoothing out compression errors, leading to more stable gradient updates and robust training dynamics. To support its effectiveness, the authors provide a thorough theoretical analysis demonstrating that SEPARATE maintains the convergence rates for non-convex objectives across both SGD and Adam-type optimizers, ensuring that the compression process does not compromise the underlying optimization goals. Additionally, SEPARATE is designed as a flexible, plug-in module that operates independently of specific optimizers or model architectures, making it a versatile tool that can be seamlessly integrated into existing training frameworks. This adaptability simplifies its deployment in diverse training setups, enhancing its practical utility for large-scale distributed model training environments.", "weaknesses": "The random Gaussian projection approach employed by SEPARATE, though theoretically sound, introduces variance that is only partially mitigated by its error-feedback mechanism. This variance arises from SEPARATE’s use of fixed Gaussian random matrices in each round, which can sometimes yield suboptimal projections that distort the gradients, affecting convergence stability. The variance analysis indicates that stable convergence relies heavily on precise tuning of the error reset frequency and compression ratio, which compromises SEPARATE’s robustness, particularly in scenarios with limited tuning flexibility due to computational constraints. Moreover, while SEPARATE offers flexibility with varying compression ratios, it lacks a mechanism for dynamically adjusting the compression ratio based on factors like gradient sparsity or model variance. This limitation restricts SEPARATE’s adaptability to evolving model sizes or training dynamics, which may lead to suboptimal performance in cases where manual tuning is impractical. Additionally, the experimental results shown in Table 1 do not provide strong evidence of competitive performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Gri0Ii12jF", "reviewer": "Reviewer_acpN", "summary": "The paper proposes a gradient compression technique called SEPARATE that aims to reduce communication overhead  across multi-device clusters in LLM training. SEPARATE leverages the natural low-rank properties of gradient and Hessian matrices by projecting gradients onto a low-dimensional subspace using common Gaussian random matrices. The accumulated compression errors are then handled by an error-feedback mechanism. This paper shows a 2x speedup in training time for tasks like GPT-2 pre-training and improved performance in fine-tuning LLMs like LLAMA2-7B.", "questions": "Refer to weaknesses section", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper tackles the a very relevant issue of communication overhead in large-scale model training. The writing is clear and well presented. SEPARATE’s design as a simple, plug-and-play gradient compression method makes it highly practical. The paper also presents theoretical proof showing that the convergence rates of SGD and Adam are maintained while using SEPARATE and also shows relevant experiments.", "weaknesses": "The method GaLore (Zhao et.al) has proven that gradients are low rank during training. This work can be cited in this context.\n\nMinor:\nThe experiments compare the training time of SEPARATE with other efficient communication techniques. This is a relevant indicator of communication overhead, but training time also includes computation time which could different for other methods. Therefore reporting bandwidth usage or total data transfer volume in addition to time could help demonstrate  reduction in data exchanged.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ejQzTDOxMk", "reviewer": "Reviewer_nVYQ", "summary": "The paper proposes a random projection compression scheme for speeding up gradients' all-reduce operations in distributed training setting (e.g. DDP) via random Gaussian projection. By relying on the unbiased down project and up project and that the variance is bounded, convergence of stochastic optimizer like SGD and Adam are still guaranteed, albeit the bounds suffer a bit from the increased noise from the compression scheme. The authors provide experiments showing that the compression schemes speed up training compared to baselines while does not suffer too much performance degredation.", "questions": "• I'm sure I understand the authors' connection between the gradient being low rank and the eigenvalues of the Hessian. For gradient compression, low rank could imply that there is some low-error decomposition but I'm not sure what the top-heavy eigenvalues have to do with gradient compression? \n\n• In defining the error feedback, what is the domain for the argmin of e? Isn't that just the orthogonal component of the gradient relative to the random projection subspace? Also, how do you justify the claim “taking moving average of the historical error maintains the stability and variance of accumulated error?” I am not familiar with \n\n• The authors should consider other sketching scheme like fast-JL and importance sampling scheme that are much faster than dense Gaussian projection. \n\n• Theory requires fresh i.i.d. Gaussian matrix to be generated across all nodes identically every round to ensure the compression and decompression is unbiased. What is the additional cost to communicate this?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "• Experimental speedup seems useful.", "weaknesses": "• Contributions: The compression scheme is not too new and have been recently explored in a variety of papers like GaLore and Flora (also uses dense Gaussian projection). While the focus of this paper is on communication, I do not see any new idea since the main ideas in GaLore and Flora are also to compress the gradients via some form of projection and then perform optimization in the compressed space before projecting it back up for the update. Extending the ideas from GaLore and Flora to the distributed setting seems straightforward to me. Furthermore, the convergence guarantees are just simple adaptation of well-known proofs i.e. convergence for unbiased and light-tailed noise (bounded variance).\n\n• The moving average error feedback section is not well written. The authors make several claims that are not obvious to me. For example, “[due to the] instability and discontinuity of random projection, the compression error may fluctuate when the random directions are far from the dominant directions of Hessian or the random directions change acutely” is not obvious since the Hessian describes properties of the optimization landscape and the compression scheme impacts (primarily) the direction of the gradient. I think the authors should rewrite the section to compare with the previous error-feedback work and argue better for why their use of EMA and reset are better. \n\n• Experiments and reproducibility concerns: The authors should disclose the tuning effort and hyperparameters for all methods being compared. I couldn't find how the experiments are conducted to determine the learning rate and what parameters are tuned for what amount. Providing code and reproduction steps should be standard. An improvement is not significant if one method is tuned (i.e. given a lot more compute) significantly more than others. Also, experimental improvement for the error feedback seems insignificant and noise prone. From Table 2, I couldn't tell if the error-feedback helps or not.\n\nReferences:\n• Zhao, Jiawei, et al. \"Galore: Memory-efficient llm training by gradient low-rank projection.\" arXiv preprint arXiv:2403.03507 (2024).\n\n• Hao, Yongchang, Yanshuai Cao, and Lili Mou. \"Flora: Low-Rank Adapters Are Secretly Gradient Compressors.\" arXiv preprint arXiv:2402.03293 (2024).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)