MQSP: Micro-Query Sequence Parallelism for Linearly Scaling Long Sequence Transformer
Authors: Ying Zhong
Jianjiang Zhu
PengCheng Yang
Xiaoming Zhang
Ke Zhang
Venue: Submitted to ICLR 2023 (2023)
Keywords: Sequence parallelism
Long Sequence Transformer
Distributed training

Abstract: Long sequence modeling of Transformer gains prevalence in fields involving long texts and high-resolution images and videos but suffers from quadratic memory complexity. Existing work investigates low-complexity variants or parallel methods to handle it. The former attempts to approximate full attention and is limited by a single device's capacity. The latter struggles to manage quadratic memory of attention maps, leading to insufficient sequence scalability. In this work, we propose a novel parallel method named $\textbf{M}$icro-$\textbf{Q}$uery $\textbf{S}$equence $\textbf{P}$arallelism. MQSP slices sequences across devices and projects local queries, keys, and values in self-attention. For communication and memory efficiency, MQSP all-gathers the queries while keys and values remain locally to acquire the local attention map, on which a distributed softmax gets conducted to amortize memory by column. Meanwhile, the queries get further partitioned as Micro-Q to divide the computation and recycle the attention map by row, jointly decomposing the quadratic memory to achieve linear scalability. The evaluation result shows that MQSP scales up sequence length linearly and achieves 4.5$\times$ sequence length of ColossalAI's sequence parallelism and 4.3$\times$ of Megatron-LM3, enabling training BERT-large of 78848 sequence length on 32 A100 GPUs. MQSP can reduce up to 78.6$\%$ of memory occupation and achieve up to 3.3$\times$ throughput when training on 17408 sequence length. The convergence quality experiment proves that MQSP provides the means for long sequences with guaranteed convergence, bringing the potential for the Transformer to explore longer sequences.

Under review as a conference paper at ICLR 2023
MQSP: MICRO-QUERY SEQUENCE PARALLELISM FOR LINEARLY SCALING LONG SEQUENCE TRANSFORMER
Anonymous authors Paper under double-blind review
ABSTRACT
Long sequence modeling of Transformer gains prevalence in fields involving long texts and high-resolution images and videos but suffers from quadratic memory complexity. Existing work investigates low-complexity variants or parallel methods to handle it. The former attempts to approximate full attention and is limited by a single device's ' capacity. The latter struggles to manage quadratic memory of attention maps, leading to insufficient sequence scalability. In this work, we propose novel parallel method named Micro- Micro-Query Sequence Parallelism. MQSP slices sequences across devices and projects local queries, keys, and values in self- attention. For communication and memory efficiency, MQSP all-gathers the queries while keys and values remain locally acquire the local attention map, which distributed softmax gets conducted to amortize memory by column. Mcanwhile, the queries get further partitioned as Micro-Q to divide the computation and recycle the attention map by row, jointly decomposing the quadratic memory to achieve linear scalability. The evaluation result shows that MQSP scales up sequence length linearly and achieves 4.5 x sequence length of ColossalAI's sequence parallelism and 4.3x x of Megatron-LM33, enabling training BERT-large of 78848 sequence length on 32 A 100 GPUs. MQSP can reduce up to 78.6% of memory occupation and achieve up to 3.3x throughput when training on 17408 sequence length. The convergence quality experiment proves that MQSP provides the means for long sequences with guaranteed convergence, bringing the potential for the Transformer explore longer sequences.
INTRODUCTION
Transformer Vaswani et al. 2017), an attention-base -baed model initially proposed for natural language processing (NLP), shows its promising potential in computer vision (CV), multi-modality, and more
(Carion et al. 2020 Dosovitskiy et al. 2021 Liu et al. 2022, Amnb et al. 2021 Neimark et al. 2021 Radford et al.
al. 2022). The self-attention associations between arbitrary pairs of tokens enable the Transformer to learn global context-aware sequence representation for many modalities. Furthermore, there is emerging trend toward long-range modeling, which scales up the sequence length of the Transformer. Long-range modeling is essential for the long texts in question answering, document classification, and other NLP tasks, as well the high-resolution pictures and the series of video frames in image modality.
As the sequence gets extended, memory consumption increases rapidly due to the quadratic com- plexity of self-attention, inevitably exceeding the limit of single device (e.g., GPU). This problem obstructs the exploration of the Transformer for modeling longer sequences. Recently researchers focused sparse mechanisms in self-attention (Child et al.) 2019: Beltagy et al., 2020 Zaheer et al. 2020) or low-complexity substitutes [Wang et al. 2020 Xiong et al.] 202 Shorromhho et al. a 202T Qin et al. 2022) approximate full attention, boosting sequence length. However, besides concern about performance influence, the memory upper bound of a single device still limits those methods from scaling up further.
Much existing work investigates parallel methods to distribute the Transformer model across the devices, such as tensor parallelism (Shoeybi et al.) (209)), pipeline parallelism (Harlap al. 2018 Huang et al. 2209:; Narayanan et al.
2021; Yang et al. (2022,, and zero redundancy optimizer (ZcRO) (Rajbhandari et al. 2020 Ren et al.
2021 However, these methods separate the model parameters

---

Under review as a conference paper at ICLR 2023
across different dimensions, not alleviating the enormous intermediate activations introduced by long sequence uuffaattention.Teeefoor.. several more recent methods are devoted to partitioning Trans- former along the sequence dimension. ColossalAIIs sequence parallelism (Li et al. 202I) transfers keys and values in ring-style to compute the partial attention map. Megatron-LMM Korthikanti et al. 2022) modifies the conjugate operators parallel layer-norms and dropouts in sequence dimension. Despite amortizing some intermediate activations, these approaches still consume local memory proportional global sequence length, making it difficult to scale up a longer sequence.
In this paper, we propose MQSP, a novel sequence parallel method, to diminish the quadratic memory overhead and efficiently scale up the long sequence Transformer. MQSP splits the input sequence to devices for parallel computation and projects local queries, keys, and values in self-attention. We design a distributed self-attention to handle the global context-awarenes in self-attention. Specifically, MQSP synchronizes the queries across the devices through an all-gather operation while the keys and values remain local, acquiring local attention maps amortized along the column dimension. Since the rows for softmax are incomplete locally, we conduct distributed softmax with hierarchical reduction across the devices, which introduces negligible communication. More importantly, we divide local queries into m finer- aeee-gg queries, called Micro-Q, and process them step by step to get the corresponding attention outputs. Each micro step's 's memory space for the attention map would be shared along the row dimension, jointly decomposing the quadratic memory for linearly scaling up the sequence length.
The proposed MQSP shows significant sequence scaling ability compared with the existing leading approaches. Our evaluation indicates that for the Transformer BERT-large (Kenton & Toutanova 2019), MQSP could scale up to 78848 sequence length with 32 100 GPUs, 4.5x of ColossalAI's sequence parallelism, and 38912 sequence length with 16 A100 GPUs, 4.3x of Megatron-LM3, In memory usage and throughput comparison, MQSP saves 78.6% memory and achieves 3.3 speedup, demonstrating the superiority of its distributed attention and communication method. Furthermore, the convergence quality experiments on wikitext, SQuAD, QQP and MRPC datasets prove that MQSP maintains the convergence quality for scaling up sequence length, bringing the prospect of exploring longer sequences to Transformers.
2 RELATED WORK
This section briefly introduces the self-attention complexity and the parallel methods for Transformer.
Self-Attention Complexity. In the Transformer proposed by Vaswani al. (2017), self-attention is the vital module for the global dependencies modeling. Omit the batch and multi-head dimensions for brevity. For the input hidden states x e F4xdm, where L is the sequence length and dm the model hidden size, the attention mechanism can be formulated as:
QKT Vare VV
Q, K, V = Lykto (x), C = softmax (SVV = softmax
(1)
Where Lyke is the linear layer projecting the tokens to Q, K E Lxde, and V E 4 in the query, key, and value embedding spaces. The scaled dot product of Q and K produces the attention scores map S E I L which incurs the quadratic complexity O(L²). Subsequently, the softmax operation along the rows converts S to attention probabilities P. which reweights V the context output C.
Many recent approaches introduce sparse mechanisms, such as Sparse Trans. (Child et al., 2019). Longformer (Beltagy et al., 2020), and BigBird (Zaheer et al. 2020), or low-complexity substitutes. such as Linformer (Wang et al. 2020), Nyströmformer (Xiong et al. 2021), Performer (Choromanski al., 2021), and Cosformer (Qin et al.) 2022). These methods algorithmically approximate the full attention with sparsity or low-rank assumption, reducing complexity to O ([Llog.) O(L). However, unsatisfied assumptions could meet performance deficiency in a broad task spectrum, and accommodating the entire Transformer within one device still limits the further expansion of the sequence length. Thus we set our sights on scaling up standard Transformer through more devices.
Parallel Methods for Transformer. Parallelism approaches have been the innovative techniques for training the large Transformer. Pipeline parallelisms (Harlap et al., 2018 Huang et al. 2019 Narayanan al. 2021: Yang et al. (002)) split the model layerwise without handling self-attention
within layers. ZeRO (Rajbhandari et al. 2020) spreads the model's parameters and the optimizer
states and conducts the same computation. Tensor parallelism in Megatron-L.M (Shoeybi al. 2019)
2

---

Under review as a conference paper at ICLR 2023
1//1
L/mmn
L/n
X2
(a) Megatron-LMMe
(b) ColAISP
(c) MQSP
Figure I: The attention maps for (a) Megatron- -M3, (b) ColossAI sequence parallelism, and (c) the proposed MQSP, with memory complexity of O # I,),, O(H 122 and O L.
man
decomposes the continuous linear layers and divides self-attention along multi-hean dimension (Fig. Tay. The methods above deal with model parameters but not the tremendous intermediate activations self-attention that are quadratic to the sequence length, insufficient scale up the sequence.
Therefore, slicing along the sequence dimension comes into mind for scaling long sequence Trans- former. Intuitively, the input sequence x could be split into n chunks, x - {zig #1, .., Jn-1 ..,-.-n
where Iy E IR. xdm, and fed into n devices to compute in parallel. With similar insight, ColossalAT (Bian al. 2021) proposed sequence parallelism (Li et al. 2021), ColAISP for short. Considering the global associations of local query/key/valle embeddings, which are projected in self-attention as as: Qi K2, V; Lyke (z:), ColAISP designs Ring RingQK
RingAV
C = softmax V softmax Qi K K
VJ. VL
(2)
Vd
By ring-style transferring K; and V,, ColAISP circularly computes Q collect the partial
attention scores map S IR (Fig. lb.. It indicates that ColAISP requires quadratic device resources n to scale up L. Additionally, the efficiency of ring communication suffers from the weakest link, e.g., inter-node bandwidth. Megatron-LM3 (Korthikanti al., 2022) modifies its conjugate operators to all- ll-ggabeerreduc--caattee to introduce sequence parallelism in layer-norms and dropouts, yet leaves the self-attention tensor parallel mode, resulting in the same quadratic attention maps.
METHOD
This section introduces the proposed Micro-Query Se- quence Parallelism. We analyze the communication pat- terns in self-attention and propose Micro-Q for reused memory. Then the distributed softmax is described. Fur- thermore, we compare memory usage with other meth- ods and analyze scalability.
Ase to Nerm
Add Nom
Adi & Nomm
MLP
MLP
MLP
Soll-Attemion
Atsention Map
3.1 MICRO-QUER SEQUENCE PARALLELISM
Instead of dividing model parameters along different dimensions previous parallel methods, we focus on the sequence dimension that affects intermediate activa- tions. When the input sequences are partitioned and fed devices in parallel, most modules, such as multi-layer perceptron (MLP), dropout, and layer normalization, are computed independently
softmax
the sequence dimension. However, as illustrated in Fig. 2 self-attention across devices does not natively support sequential parallelism. In MQSP, we dive into the self-atee attention mechanism to attain distributed attention with efficient communication and memory usage.
Figure Illustration of sequence paral- lelism. Besides self-attention with global associations and softmax, other modules natively support sequence parallelism.

---

Under review as a conference paper at ICLR 2023
X1
X m
Q
Micro Q
rank i
other ranks
S
softmax
k=1
Ckzi
X
Ckei
Figure 3: Overview of the MQSP self-attention. f and f are the conjugate all- gellgatherrrduceeseeee communication operations. The dotted box for m times would reuse memory space. The dotted arrow lines represent the aligned row for distributed softmax.
Attention Analysis. Recalling the general form of self-attention in Eq. the self-attention for each query can be analyzed independently. Assuming the k-th query in Q is qu E IR lxd, the corresponding context Cg is calculated as:
T&KT VTL
softmax(S,owe.)) = softmax
V
(3)
Where Srowe.: E RIXL represents the k-th row the attention scores map. ColAISP transfers K, as ring and multiplies with qu to collect the complete Srowy.: which encounters efficiency deficiency bandwidth- hettiihhh-ibbllaneed environment and includes the L factor of memory overhead. contrast, MQSP chooses to share the 9 across the devices, remaining K; and V; locally. In this way, each rank only needs to compute Srowz, RIx its corresponding columns of attention scores:
d-softmax CanttmmaxSSron..... (Sroun, Vi f(d-softmax( ((xx) KT
(4)
Vdl
Where d-softmax means the distributed softmax operation acquire the corresponding columns of attention probabilities, further described in section 3.2 For the communications required here, we define as qi synchronization across devices and as locally reweighted Vú sumumten, e.g., broadcast/reduce or all all- wwagggatreeuceeeeeeeee scatter, as shown in Fig- [4] and f are conjugate, which means the forward pass of one equals the backward pass of the other. Considering the symmetric workload across devices and the efficiency of the duplex collective communication, we adopt all-gather/reduuc- scatter as the conjugate operators to demonstrate our method (Fig.
The complete attention includes the whole sequence 9kE(o.1) in Eq. leading to the equally significant attention scores S. IR 1 4. To this end, propose Micro- Q, the finer-grained query, to reduce memory consumption orthogonally. For number yn of Micro-Q, the local query Q: gets chunked as {4p .... q,"" 1), as depicted in the red box of Fig, where qf IR 4. xdy would be the j-th micro-step's 's input the i-th device. The following f all- aall-gahhers ke to form the shared
Micro-Q.
d
device
device
device
2 Backward: redure.
Forward:
Forward
Bac owendt
broudkan
al g-ggahher
attention d-softmax + reweight
attention d-softmax reweight
Fonward:
Backward: broadcast
Fument
reduce
gat-gr
Ille
Figure 4: Conjugate pair comparisons. Broadcast/reduce: asymmetric simplex one-to-all communica- tion for times. tlggahherrrduue-scateee symmetric duplex collective communication for once.
4

---

Under review as a conference paper at ICLR 2023
and concatenated Q' E R xds Hence the distributed attention in minor range could be conducted as:
Q¹K. V
C d- = d-softmax
(5)
Where the C/ the sub-context matching Qi on the i-th device. The memory complexity of each micro-step's attention scores map Sy would be R x illustrated as in Fig. Ic The memory space could be reused across the micro-steps, implemented through the checkpointing technique (Chen et al. 2016). The existing methods typically employ layerwise checkpointing, not addressing the excess intermediate activations within single layer, while MQSP fractionizes it finer-grained partitions. At last, I conducts reduce-scatter to sum up the sub-contexts O back to their ranks, the reduced micro-context e. Each rank concatenates micro-contexts to produce its local context C.-
memory space for the attention map. Coefficient n along column axis and coefficient along row axis make joint efforts to disintegrate the quadratic memory of the Transformer, enhancing the scalability of sequence length.
Comparison. Compared with R x in ColAISP, MQSP requires only IR d th
In concern of communication, ColAISP transfers queries and keys in rings while MQSP collectively transfers queries and contexts, incurring comparable communication volumes. However, in the heterogeneous network environment, e.g., a cluster with 4 nodes x 8 GPUs with Nvlink, the ring- style ColAISP suffers from the bottleneck of the low inter-node bandwidth. In contrast, the duplex collective MQSP benefits from the hierarchical NCCL (Nvidia). Additionally. if ColAISP adopts the Micro-Q method to reduce memory overhead, the ring communication would be repeated m times, unlike amortized in MQSP. Since the inputs and outputs are uncorrelated among the micro-steps, we could alleviate communication overhead by overlapping with the other step's 's computation, further boosting efficient distributed self-attention (Appx. A.3)
3.2 DISTRIBUTED SOFTMAX
This section introduces the distributed softmax with low cost in MQSP. similar technique could be found in ArcFace (Deng et al.) 2019) recognize large-scale faces. ArcFace sums up the local denominator in the forward pass, while the backward formula is simplified with cross-entropy loss. Here we analyze the general softmax. Defining the whole sequence of scores as [So
SL-1 (a max (ie(o,)))) and the probabilities as [Po. P1. ..... 11, the original form is formulated as:
P: exp(s; 8: - o)
V == Vp. x p1(1 p.) + Σ Vp X (-Pip,
(6)
=== exp(s
In MQSP, the scores distributed across devices, which incurs O 10 T.) communications cost and ((L) memory overhead in the form of Eq. 6 To this end, we convert the arithmetic form and hierarchically reduce the maximum or summation, attain an efficient distributed softmax:
0:
P
Vs, X P: Tsum (A,)
Fauma
Where Fsum means the reduce-sumt operation. 0, = exp(s: .r) and its local summation is θ;
A; = Vp x p, and its local summation is A, Please refer to Appx. A.2 for details of mathematical derivation. It introduces only Onn 1) complexity to allreduce-sum the scalar of local summation, negligible compared with Eq. 6 The pseudo-code of the distributed softmax is given in Algo. !
Algorithm The forward and backward pass of the distributed softmax.
attention scores
p: local cttention prohabilitie P_grad: gradi ent of p
[..., gh/m, Toda me
def [_eoftmax__erwardd13)
both I..., quim, Localsequim] p_grad):
4 max loccl = a.max 1-11 # [..., gDim)
def
5
max_g max_giobal obal all__redueee reduce local, op-Max)
5 P -_ggrad x
6 s_exp (s max_globel (..., Nonel) .exp ()
sum_local (-1um(-1) # I..., gDim)
sum_lacal s_exp_sum(-1 (-1) [...r q0im]
sum_global all_reduce (sum lsca_ cp=suM)
aum_gi obal all_reduce local, op=SUM)
s_grad
sum_glohal [..., Nane]
s_exp sum_glcbal(... None]
return s_grad
return
5

---

Under review as a conference paper at ICLR 2023
Table Memory consumption comparison of model parameters and intermediate activations for the vanilla Transformer and the different parallel methods for the Transformer.
METHOD
MODEL PARAMETERS
INTERMEDIATE ACTIVATIONS
VANILLA TRANSFORMER
(8 + 4
)
10BDL + 2BHL
H
MEGATRON-LM3
(8 + 4 H
D
(2 + n ) + 2BH Li
COLAISP
(8 H ))
11BD +2H L.?
MQSP
(8 + 4 D
10BD + BDL + L
H
3.3 MEMORY USAGE ANALYSIS
This section analyzes the memory usage of the model parameters and intermediate activations in single Transformer layer. We omit the gradients and optimizer states proportional to the model parameters and the statistical buffers or masks in dropout and layer normalization. B, H, I,, and m represent the batch size, multi-head size, sequence length, device number, and Micro-Q number, respectively. To be concise, we assume D = dm Hd, H Hdg, consistent with most implementations.
As for the model parameters, the linear layers in MLP take 8D2, and the Lyke and output layer in self-attention take fT while Megatron-LMMs divides it by n. For the intermediate activation, the MLP takes 5BDL, LL, and the self-attention takes LD + 2BHL². MQSP divides the MLP part by n and introduces the min factor in the attention map and an all- aal-eeahered Micro-Q buffer BD L We apply the same memory m. analysis on Megatron-LMS and ColAiSP (Appx. [A.4), as shown Tbl.
4D2 It is the same for the vanilla transformer and sequence parallel methods,
According to the analysis, tensor parallelism has memory superiority in model parameters, yet the intermediate activations consume most memory the long sequence Transformer. The activations of previous methods include L or L2 factors, which limit their sequence scalability. Contrastively, MQSP adjusts the granularity of Micro-Q m, attaining efficient memory usage. Assuming the remained upper bound memory for intermediate activations as M and adjusting m equivalent to n:
L 10BD + BD +2H SABH-MM- - LS m
110 2H -
M
11D )n Cn (8)
mi
2BH
2H
The maximum L grows proportionally MQSP. It meets > 1 to ensure Micro-Q includes at least one query, which indicates Eq, in the condition of C. L C². Moreover, owing to the query granularity flexibility, we can set larger m to obtain finer-grained Micro-Q to save memory, attaining further sequence length scaling.
with constant C, demonstrating the linear scalability of
EVALUATIONS
This section evaluates the proposed MQSP, verifying its convergence quality and comparing it with other parallel methods in sequence length scalability, memory footprint, and throughput. We further investigate the influence of the Micro-Q setting. Specifically, we implement MQSP with Pytorch-1.9 (Paszke et al. 2211)), referencing Bert (Kenton & Toutanoval 2019) implemented in the HuggingFacc transformers (Wolf al. 2020). The experimental hardware a private cluster, each node containing Intel(R) Xeon(R) Platinum 8369B CPU, 760-GB of RAM, and eight 80-GB A100 GPUs with Nvl ink, resulting in 96 times intra-node bandwidth compared with inter-node (300 GBps v.s. 25 Gbps).
4.1 QUALITY OF CONVERGENCE
MQSP distributes Transformer in sequence dimension to scale up long sequences with the full attention in the vanilla Transformer. To verify the convergence of MOSP, we experiment with our MQSP and distributed data parallel (DDP) on the datasets, including WikiText- 103 (Stephen et al. 2017) SQuAD (Rajpurkar al. 2016), QQP, and MRPC(Wang et al. 2018).

---

Under review as a conference paper at ICLR 2023
TASKS(MENIC))
DDP
MQSP
SQUAD(FI)
86.37
87.17
QQP(FI)
72.10
73.49
MRPC(FL)
86.38
87.99
INkiTTcct
SQUAD
Figure 5: Part of the training loss.
Table 2: The convergence results.
We set the same parallel size as data parallel, e.g., n = in one node, and identical training hyperparameters. The convergence results are shown in Tbl. - and the training loss curves depicted Fig. [ The similar convergence quality demonstrates that MQSP provides a means for long sequences under the guarantee of convergence, which brings exploration space for Transformer toward longer sequences.
4.2 SEQUENCE LENGTH SCALABILITY
This subsection demonstrates the superior sequence length scalability of the proposed MQSP. We compare the maximum sequence lengths achieved by the Transformer using dif- ferent distributed approaches as the number devices increases.
13
MQSP-eq
12
MQSP-mem
11
ColAISP
Megatron-LM
Max Sequence Length(x10n
10
We evaluate the methods by training the BERT- large model (Kenton & Toutanova] 2019). We use the batch size of 16 and the Adam op- timizer following Kenton Toutanoval No- tably, because the checkpointing technique in the Mícro-Q Q mechanism erases layers' stacked self-attention activations, we apply layerwise checkpointing eliminate the accumulated ac- tivations across the Transformer layers for a fair comparison. Furthermore, ular number of devices n, configure MQSP in two ways: a) MQSP-eq, where m = n for linear scalability analyzed in section (..3 and b) MQSP-mem, where m = L/n as the finest- grained Micro-Q investigate the largest scaling potential. MQSP-mem serves as an upper bound reference, and for more about the Micro-Q setting, please refer to section 4.5
each partic-
012
Device Number
Figure 6: Comparison of maximum sequence lengths for training the Transformer as the GPU scale increases.
As shown in Fig- 6. the maximum sequence lengths of these methods measured 2n GPUs, ranging from to 32. The metric data range for LMM3 is limited to 2 GPUs because the tensor parallel size must be able divide multi-head dimension, which is 16 in BERT-large. In comparison, the scalability of the sequence parallel methods is less constrained.
We observe that MQSP-eq acquires 3.2x and 4.3x longer sequence than ColAISP and Megatron- LM3 (38912 12 v.s. 12160/8960) when 2 GPUs, and 4.5x than ColAISP (78848 v.s. 17408) when 2 GPUs. MQSP-eq requires only quarter of GPUs achieve the maximum sequence of ColAISP and Megatron-IMM LM3, denoted by the dotted arrow lines.
As the GPUs increase, Megatron-LM3 and ColAISP climb at a consistently decreasing rate, while MQSP-eq scales up almost linearly with slope of C 2464. As analyzed in section 3.3 MQSP maintains linear scalability under the condition of n < 2464, L < 24642, which is practically a hardly attainable upper bound. It demonstrates that our MQSP resolves the quadratic memory overhead in the long sequence Transformer and achieves superior sequence length scalability. Moreover, single device
MQSP-mem achieves further scalability, about 2x compared with MQSP-eg, (6144 v.s. 2688). It proves that the flexibility of the Micro-Q brings the potential to a more extended sequence, orthogonal to the sequence parallelism.
7

---

Under review as a conference paper at ICLR 2023
PLAUTH
2.4
TO No
1-3
MaP
38
HELP
C4097
259
50.0
1500
27500
320
10000
13,00
JERC J000 22000
1000
2000
20000
43.03
Saaia 1030
Senarre levelt
Mid
HOLE
HOSE Cidde
Calle
24
AM
SACARTA lango
(a) 1n8g
(b) 2:1 6g
(c) 4n32g
Figure Comparison of the memory footprint and the token throughput with sequence length scaling up. w/o 1_ckpt means without layer checkpointing.
4.3 MEMORY FOOTPRINT
In the sequence length experiment, MQSP achieves longer sequences with superior scalability, reflecting the advantage MQSP's low memory footprint. Here, we conduct a more specific memory footprint comparison. Inheriting the model and training settings from the previous section, arrange the experimental environments as 1r.8g, 2n16g, and 4n32g, where XnYg represents Y YGPUs of X nodes in cluster. MQSP below represents MQSP-eq for linear scalability setting. In addition, since the checkpointing in Micro-Q erases most stacked activations, we also evaluate MQSP without layer checkpointing, w/o 1_ckpt for short, to compare the memory and throughput. With the consistent hardware resources, we scales up the sequence length and compare their maximum allocated memory during training, shown in the top row of Fig. 7
Except for Megatron-LMP3 advantage in short sequences with fewer model parameters, MQSP occupies less memory footprint than other methods on most configurations, saving up to 78.6% memory when 17408 on 4n32g, as denoted by the dotted arrow line. Even with stacked activations, MQSP w/o 1.ckpt also requires less memory in long sequences. indicates that MQSP expects less memory to support the training of long sequence Transformer. Furthermore, the memory advantage of MQSP grows with longer sequences, benefiting from its advantageous memory efficiency in self-attention.
4.4 THROUGHPUT
Token Throughput. In addition to the capability for training longer sequence Transformer, training efficiency also deserves attention. We also conduct the token throughput comparison, as shown in the bottom row of Fig. 7 The throughputs generally decline with sequence length scaling up due the quadratic computation complexity self-attention. For the maximum sequence in different environments, MQSP achieves similar token throughput, demonstrating that MQSP could scale up Nx sequence length with Nx devices and N time consumption.
In 1n8g, Megatron-LMM3 has throughput advantages, and MQSP w/o l_ckpt is comparable with ColAISP, while MQSP scales up to 18432 sequence without significant drop throughput. When training longer sequences, which are 8192 for Mcgatron-LM3 and 17408 for ColAISP, their insufficient scalability incurs the inter-node parallel group, marked as and B in Fig. 7 Our MQSP supports the same length within node, achieving 2.1 x and 3.3x throughput per device to Megatron-LMS and ColAISP, respectively. MQSP also has throughput advantages in modti-node environments, benefiting from efficient collective communications, In addition, MQSP w/o ckpt
8

---

Under review as a conference paper at ICLR 2023
gains further throughput with less recomputation while maintaining better sequence scalability than other methods.
Time Ratios. To analyze the throughput advantage of MQSP, measure the time consumption ratios in a Transformer layer for MQSP and ColAISP. Configuring two environments as 4n and 2n1 6g, train these sequence parallel methods on their maximum sequence and profile the execution timeline of the forward passes. Here we adopt no overlapping to directly show the intra-node and inter-node communication costs, Fig. exhibits the ratios of each part.
Loky
Satmaax
S_comm
(a) MQSP-1n4g
(b) ColAISP-In4g
(c) MQSP-2n166
(d) ColAISP-2n1 6g
Figure Pie charts of time consumption with the communication part emphasized. The numeric suffix represents the number of nodes and GPUs. It indicates MQSP's efficient inter-node communication.
The order parts are the projection of Lykes dot-production of and k, softmax or d-softmax, reweighting, d-softmax communications, ring or conjugate communications, and the MLP. demon- strates that d-softmax introduces negligible communication costs. The communication costs occupy acceptably low ratios in intra-node scope for both MQSP and ColAISP, 11% and 12%, respectively. In inter-node environment, the inadequate bandwidth between nodes results in different increased communication ratios, 19% for MQSP but 65% for ColAISP, as a sharper performance drop. It proves that the duplex collective communication adopted by MQSP brings an advantage heterogeneous network environment, compared with ColAISP's ring-styla one restricted by the lowest link.
4.5 MICRO-Q SETTING
Micro-Q serves as the core part of the proposed MQSP, introducing the number m of Micro-Q as new hyperparameter. Thus this subsection explores the effect how we set m. Training the BERT- large 8192 sequence measure the maximum allocated memory and the throughput. As shown in Fig. the memory gain saturates as m reaches a certain level, for Micro- Q's attention computation is no longer where the maximum memory allocation occurs. Moreover, the throughput drops slowly as m increases and then drops off rapidly after the memory gain saturates. The reason could be that excessive segmentation results in inefficient tiny computations and communications. The result indicates that we could make the trade-off between speed and memory, benefiting from the flexibility of the Micro-Q mechanism.
(akensss)
Memory (GB)
12
20
10
2n16g, 6g, vary to
15
3
m Figure 9: Memory (bar) and throughput (line) for different Micro-Q settings.
5 CONCLUSION
This paper presents the Micro-Quer, sequence parallelism, an efficient distributed method for linearly scaling long sequence Transformer. MQSP achieves distributed self-attention through all- latheanh queries, maintaining only partial columns of attention map with low-cost distributed softmax. Furtherly, MQSP introduces the finer-grained query, Micro-Q, to reuse memory among the rows of attention map. jointly decomposing the quadratic memory. MQSP attains 4.5 x sequence length compared with ColAISP and 4.3x with Megatron-LMB. achieving up to 78848 sequence on 32 A100 GPUs. The flexibility of Micro-Q boosts further scalability orthogonally, MQSP saves 78.6% memory and achieves 3.3x speedup in memory and throughput evaluations. With guaranteed convergence, MQSP facilitates scaling longer sequence Transformer.
on single device.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Anurag Arnab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Lučié, and Cordelia Schmid. Vivit: A video vision transformer. In Proceedings of the IEEE/CVE International Conference on Computer Vision, pp. 0836-6846, 2021.
Iz. Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05/500 2020.
Zhengda Bian, Hongxin Liu, Boxiang Wang, Haichen Huang, Yonghin Li, Chuanrui Wang, Fan Cui, and Yang You. Colossal-ai: A unified deep learning system for large-scale parallel training. arXiv preprint arXiv:2110. 2021.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European conference on computer vision, pp. 213-229. Springer, 2020.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint 104.066644.0674, 2016.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 1rXXiv:904..0 1900099994000,,,.. 2019.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. In International Conference on Learning Representations, 2021.
Jiankang Deng, Jia Guo, Niannan Xue, and Stefanos Zafeiriou. Areface: Additive angular margin loss for deep face recogmition. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 4690-4699, 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference Learning Representations, 2021.
Aaron Harlap, Deepak Narayanan, Amar Phanishayee, Vivek Seshadri, Nikhil R. Devanur, Gregory Ganger, and Phillip B. Gibbons. Pipedream: Fast and efficient pipeline parallel dnn training. ArXiv, fbbs/100.003777 2018
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoe V Le, Yonghui Wu, et Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32, 2019.
Jacob Devlin Ming-Was Chang Kenton and Lee Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACI-HLI,, pp. 4171- 4186. 2019.
Vijay Korthikanti, Jared Casper, Sangkug Lym. Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. arXiv preprint arXiv:2205.05) /98, 2022.
Shenggui Li, Fuzhao Xue, Chaitanya Baranwal, Yongbin Li, and Yang You. Sequence parallelism: Long sequence training from system perspective. arXiv e-prints, pp. arXiv-2105, 2021.
Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, al. Swin transformer v2: Scaling up capacity and resolution. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12009-12019, 2022.
Deepak Narayanan, Amar Phanishayee, Kaiyu Shi, Xie Chen, and Matei Zaharia. Memory--fficient hapeline-parally dan training. In International Conference on Machine Learning, pp. 7937-7947. PMLR, 2021.
10

---

Under review as a conference paper at ICLR 2023
Daniel Neimark, Omri Bar, Maya Zohar, and Dotan Asselmann. Video transformer network. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 3163-3172, 2021.
Nvidia. Nvidia collective communication library (necl), 2022. http:://gtthubb.cmm/vvdddd nec1
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Zhen Qin, Weixuan Sun. Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan. Lingpeng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. In International Conference on Learning Representations, 2022.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. In ICMI., 2021.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking. Storage and Analysis, pp. 1-16. IEEE, 2020.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100, 000+ questions for machine comprchension of text. In EMNLP, 2016.
Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He. (ZeRO-Offoad): Democratizing (Billion-Scale) model training. In 2021 USENIX Annual Technical Conference (USENIX ATC 21), pp. 551-564, 2021.
Mohammad Shoeybi. Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catan- zaro. Megatron-Img Training arXiv preprint arXiv: 10X.v0990000005 2019.
parameter language models using model parallelism.
Merity Stephen, Xiong Caiming, Bradbury James, and Richard Socher. Pointer sentinel mixture models. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv: 2018
Peng Wang, An Yang, Rui Men, Junyang Lin, Shuai Bai, Zhikang Li, Jianxin Ma, Chang Zhou, Jingren Zhou, and Hongxia Yang. Unifying architectures, tasks, and modalities through a simple sequence-to-seen tocennceett--eeeeeee so-sequence learning framework. In ICML, 2022.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04700 04768, 2020.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 conference on empirical methods in natural language processing: system demonstrations, pp. 38 45, 2020.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn M. Fung, Yin Li, and Vikas Singh. Nyströmformer: A nyström-base algorithm for approximating self-attention. Proceedings of the AAAI Conference on Artificial Intelligence. AAAI Conference on Artificial 2021.
Intelligence, 35
11

---

Under review as a conference paper at ICLR 2023
PengCheng Yang, Xiaoming Zhang, Wenpeng Zhang, Ming Yang, and Hong Wei. Group Sroup-based interleaved pipeline parallelism for large-scale dnn training. In International Conference on Learning Representations, 2022.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, al. Big bird: Transformers for longer sequences. Advances in Neural Information Processing Systems, 33:17283- 17297, 2020.
A APPENDIX
A.1 PSEUDO-CODE
Communication Operators.
Algorithm 2 The pseudo-code of the conjugate communication operators.
class A.1GaaherMMicre (astograd. Function]
class ReduceScatterd iatt ann
2
def forward g_micre):
def forward(etx, c):
[1l_gatherrgg.mice q
s_micro reduca_scatter(c)
return c_miere
def backwardictxx, g_gred):
def backward (ctx, c_micro_gradi:
6
g_micro_grad = veduce_scatterr/g_graa
e_grad = all__ather (c_microo_grd)
7
zelurn miicro__grad
zelurn e_grad
Micro-Q Implementation.
Algorithm 3 The pseudo-code of the implementation of Micro-Q in self-attention.
class DietribiteeSeeeAAtte deff forward(sel1 x)
(nn. (nn. Xedule)
2.
3
0, K, V T_gkv 1x)
4
s_micro_ist = for e_miaro G.chunk (m)
5
e_mtoro checkpoint seelt.aoftaax_reeeehht e_mioro,
Lmmier__ltt.agpanicc.meeer
8
return sot ic_microl
9 def softmax_reweighe (self, q_micro, X, V):
Q btgiaabher@iiiroooolly......
actn_scores motmul 1_0, K.T()) sqrt (d_k)
attn_probs bostrtootedsartta...oo
ylattr_scores)
matnul (actn__prbrb e_micro #eduesstaatterc.aplly..
return q_mtora
A.2 MATHEMATICAL DERIVATION OF DISTRIBUTED SOFTMAX
Consider arbitrary of the attention map of the scores and probabilities Srow Prowe, PL-1].
RIXL To be specific, they include (Su, $1. $L-1 and [Pu, PL--)- For numerically stable softmax, 8; would minus their maximum value = maxx($;ejo..)) before exponentiating:
exp(s; -o)
Pi
(9)
ビ exp(s; o)
Mathematically, according the derivative Jacobian matrix of softmax: Op:
Op:
p1 - ps).
- + j)
(10)
Os,
For the training error e, the gradients for s; could be backpropagated as: Σ 1--1
Öp,
Ves
Vep;
(11)
V.P pp(1 - p.) + Σ L-1
VeP; (-P.p))
jæi
12

---

Under review as a conference paper at ICLR 2023
Specifically, get each device's s local maximum scores o; max( before commu- nicating to collect Tmaz(a;) where maar means the alfreduce-max operation. Defining the exponent of the normalized scores as 02 exp(s, σ), we sum them up stages:
θ Σ j=Hl
(++1)-1
Oj. P: Tsum
0,
(12)
Where sum means the reduce-sum operation. Similarly, for the backward pass, define X VeP; x pj and its local summation A, changed as:
-(4 Xg. The form of backpropagation could be
Σ 1-1 Vepj x p.P;)
Ves; VePi x
j=0
(13)
Vepi P pi 1=
Vep x Pj
A Pr Tsum (A:)
A.3 OVERLAPPING
As shown in Fig. 10 the computation of softmax on Micro-Q and values reweighting can overlap with other micro steps' communication of all-gather and reduce-scatter, offsetting the time overhead.
All-Gather
Reduce-Scatter
softmax_reweigh
(a) w/o Micro-Q
All Gather 0
Reduce Scatter
All Gather Bether_
Reduce Scatter
softmax reweight 0
softmax reweight 1
(b) Micro- Q w/o overlapping
All Gather
All Gather 1
Reduce Scatter_0
Reduce Scatter_1
softmax reweight 0
softmax reweight 1
(c) Mícro-Q with overlapping
Figure 10: The overlapping of computation and communication.
A.4 DETAILS ON MEMORY FOOTPRINT
B, H, L, n, and m represent the batch size, multi-head size, sequence length, device number, and Micro-Q number, respectively. Assume D = dm Hd, = Hdd.
Vanilla Transformer.
Model parameters:
MLP: The first linear layer d'm X 4dm, and the second 4dm X dm-
Self-attention: Lyku layer dm Hdy, and output layer Hd, x dm Summation: sdZ, 3/H + H dry = (8 + 4 D2
Intermediate activations:
MI.P: Input BLd.r and intermediate BLAdd
13

---

Under review as a conference paper at ICLR 2023
Self-attention: Input BLdm- Lyku produces 3 x BHLd,, attention scores and probabilities 2 BHT?, and reweighted value BHLd,
Summation:
5BLdm + BLdm + BH BHHL(3d, + de) + 2BHL HUBD++BBIH
Megatron-LM3.
Model parameters:
MLP:
The first linear layer dm dmm, and the second 4 x dm-
Self-attention:
Lyku layer 3 dm H dg, and output layer d, dm- Summation: 8d 3Hdmd, Hdddd (8 + 4 1
Intermediate activations:
MLP: Input BLdm, and intermediate BL4 dms.
Self-attention: Input BLdma fyke produces 3 B H Ldk, attention scores and probabilities B L2, and reweighted value B Ldy-
Summation:
BLdm + 4BLdm BLdm + BHL(3d; d,) + 2BHL2
L
= BDL 2BH
n
n2
n
ColAISP.
Model parameters: Same as the Vanilla Transformer.
Intermediate activations:
MLP: Input B dm. and intermediate B 4dm
Self-attention: Input B Bd Laku produces BH _dg, attention scores and probabilities
2 x BH LL L. the ring buffer BH d... and reweighted value BH Ldg.
Summation:
BBLdm
BLd.. + BHL(4ds + dw) + 2BH1²
11 BDD + +2H L
MQSP.
Model parameters: Same as the Vanilla Transformer.
Intermediate activations:
MLP: Input B dd and intermediate B BL 4d
Self-attention: Input B dm Lyke produces 3 x BH 4 d2, attention scores and probabilities
BH 12
the all-gather Micro-Q buffer BH L dg, and reweighted value BH Ld..
mun
Summation:
5dLdm
L BLdm + BHL(3dd + d.). BHLds 2BHL 10BD BBD mn
+2BH L
m
n
m
mn
14


Reviews:
[{"review_id": "lBoyLW_3dTd", "reviewer": "Reviewer_h8ct", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This work studies long sequence problem from the system perspective. The self-attention part of a Transformer requires quadratic memory cost, limiting the applicability of full-attention for thousands of tokens. There are major two lines of directions for tacking with the high complexity due to long sequence: 1) full-attention approximation such as sparse attention, and 2) model parallelism approach. This work focuses on the 2nd line. The existing works such as Megatron and ColAISP reduce memory complexity by n times, where n is the tensor parallel size. This paper proposed to further reduce the memory cost by splitting the attention computation into m micro-steps, and in each micro-step, a micro-context output is computed. In this way, the memory buffer can be reused and the memory cost of attention map is reduced by nm times. The experimental results show that the proposed method, dubbed MQSP, supports training a BERT-large model with 78848 sequence length on 32 A100 GPUs.", "summary_of_the_review": "The paper studies an important problem since long sequence has a wide variety of applications such as question-answering and summarization. The empirical results are also promising and interesting. However, I am concerned with the novelty. The paper is lack of discussion and empirical comparison with FlashAttention. FlashAttention adopts very similar technique and fuses the computation of micro-chunks into a CUDA kernel, and its memory cost depends on the GPU SRAM size and is independent of the sequence length.", "strength_and_weaknesses": "Strength:\n1. The empirical results are promising. Supporting the training with 78848 sequence length is impressive.\n\nWeakness:\n1. This paper lacks sufficient discussion and comparison with the related works. Similar idea of splitting sequence into smaller chunks and incrementally accumulating results has been proposed in (Dao, Tri, et al., 2022). This work investigates the bottleneck of Transformer training from IO perspective. The resultant algorithm fuses the whole Transformer layer into a fused kernel and the memory cost of computing attention map is constant independent of the sequence length. This algorithm has been implemented in multiple libraries including OpenAI triton and Xformers https://github.com/facebookresearch/xformers/blob/91c7c0846455d78b934d6705026b4c7e35b89bb1/xformers/ops.py#L280. We can easily combine it with the Megatron implementation.\n2. Distributed softmax is also not new. This has been used in Megatron for computing output logits: https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/mpu/cross_entropy.py\n\nDao, Tri, et al. \"FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.\" arXiv preprint arXiv:2205.14135 (2022).", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "XXpQZ5ELaY", "reviewer": "Reviewer_L8Ft", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper introduces a novel dimension to partition self-attention layers for distributed training of long-sequence transformers. The proposed technique, called Micro-Query Sequence Parallelism (MQSP), maintains only partial columns of attention map and applies a low-cost distributed softmax. To reuse memory over multiple rows of the attention map, MQSP also utilizies a finer-grained query, Micro-Q, to scale memory overhead linearly instead of quadratically. Evaluation demonstrates MQSP can scale sequence length significantly longer than both ColAISP and Megatron-LM3.\n", "summary_of_the_review": "This paper focuses on improving scalability of multi-head self-attention block with an increasing sequence length. The self-attention block is known to be a scalability bottleneck due to it's quadratic overhead for both memory capacity and computation.  That said, this paper tackles a timely problem.  Evaluation demonstrates the sequence parallelism can handle longer sequences and is more memory space-efficient than both ColAISP (leveraging sequence parallelism in a different way) and Megatron (leveraging tensor parallelism).\n\nWhile the results show some promise, I still have the following questions/concerns about this proposal:\n\n* *Performance Sensitivity* - In Figure 9, there is a tradeoff between reduction in memory usage and token (training) throughput. This makes sense as $m$ increases, attention maps are partitioned into finer-grained chunks to increase the overhead of communication limiting overall throughput. In this particular setup of *2n16g*, the throughput is maintained at a high level before the benefits of memory savings saturate. However, this might not necessarily true depending on the hardware configuration (e.g., # of nodes, # of GPUs, GPU generations, communication bandwidth, etc.).  Would the proposed technique provide robust performance for other system configurations with different compute/communication ratios?  \n* *Comparison against Existing Work* - This work compares mostly against ColAISP, which also leverages sequence parallelism.  I was not aware of this work, and hence do not know how strong a baseline it is. Comparison against Megatron-LM3 alleviates the concerns but I am still wondering how MQSP compares against other existing works, especially leveraging multiple forms of parallelism (e.g., pipeline parallelism together with tensor parallelism) in terms of overall token throughput. \n* *Motivation* - This work could have been better motivated by suggesting some realistic use cases requiring longer sequences. The authors suggest high-resolution image cases as an example, but wouldn't scaling down the input image be a more practical solution, especially for classification and segmentation tasks? What is the practical range of sequence lengths today (and how does it scale) and what kind of applications could potentially benefit from a long sequence with >78K tokens?\n\n\n", "strength_and_weaknesses": "**Strengths** \n\n* The sequence length of transformers is likely to keep increasing, and a scalable solution to the sequence length has values.\n* The proposed solution effectively increases the sequence length over 70K on 32 A100 GPUs.\n\n** Weaknesses **\n\n* It is not clear whether this technique introduces another bottleneck (especially with increasing $m$) for inter-GPU communication and synchronization.\n* The paper lacks comparison against other works leveraging different forms of parallelism (e.g., pipeline parallelism + tensor parallelism).\n* The necessity of scaling sequence length needs to be better motivated.\n ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "Sh3OrhvbLNP", "reviewer": "Reviewer_nVbE", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper proposed MQSP, an optimized sequence parallelism technique for long sequence transformer. Finally, the evaluation result shows that MQSP can scale to longer sequences. In long sequence scenarios, MQSP can significantly reduce memory occupation and have higher throughput compared to Tensor Parallelism and Sequence Parallelism.", "summary_of_the_review": "Tend to reject; Please check my comments above for more information.", "strength_and_weaknesses": "Strengths\n\nWith Micro-Query, MQSP scales up sequence length linearly and shows a pretty good performance improvement. Greatly increases the length of sequences that can be handled in Sequence Parallelism.\nMQSP has a better communication pattern and significantly improved performance compared to previous sequential parallelism techniques.\nThis paper provides a thorough theoretical analysis, while comparing other methods including the most advanced Tensor Parallelism and Sequence Parallelism systems from theory and experiment.\n\n\nWeakness - major concerns\n\nIn the evaluation, only Bert-Large was used for testing. Testing with more models with different sizes can make the results more convincing.\nIn Section 3.1 “MICRO -QUERY SEQUENCE PARALLELISM - Comparison”, a comparison is made with ColAISP only. If a comparison can be made with Tensor Parallelism, it can give the reader more information to choose between Tensor Parallelism and Sequence Parallelism\nIn Section 4.3 “MEMORY FOOTPRINT”, the paper shows the performance of MQSP with the activation checkpoint turned on and off, but does not specify whether the activation checkpoint is turned on in other implementations\nIt can be observed that within a sequence length of 7500, Megatron-LM3 has a clear performance advantage over MQSP, and there needs to be a more prominent expression in the paper to demonstrate why we need long sequences for training, and how long we really need them. One concern is that extra-long sequences can lead to excessive computational cost without actually bringing much gain effect.\n\n\nWeakness - minor concerns\n\nIn Section 4.1 “QUALITY OF CONVERGENCE”, the setting of sequence length for MQSP is missing.\nFor some readers who have not been exposed to much tensor parallelism, the background of this paper is slightly brief. Consider adding more detailed background on Tensor Parallelism.\nSome letter symbols are used in advance in the figure, such as \"m\" in figure one. Maybe highlight the meaning of these letters in the article or caption, it will be easier to understand the article.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "HNh5LeoIM_T", "reviewer": "Reviewer_zrYy", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper proposed a new way to further reduce the memory footage of the self-attention mechanism over long sequences by parallelization on the sequence dimension. The paper proposed to communicate the query tensors to handle the sentence-dimension computation used in self-attention. ", "summary_of_the_review": "I have doubts about the effectiveness of the proposed method. I will raise my score if my concerns are addressed during the rebuttal.", "strength_and_weaknesses": "The paper proposed a more fine-grained method to partition the query dimension to reduce the memory footprint of the self-attention layer. The proposed method is new. However, I have doubts on the actual effectiveness of the method. Specifically,\n\n1. The $O(L^2)$ memory usage for the attention operator can be reduced to $O(L)$ by changing the way it was implemented in GPU without introducing any communication overhead across different GPUs. See the following two references:\n    1. [https://arxiv.org/pdf/2205.14135.pdf](https://arxiv.org/pdf/2205.14135.pdf)\n    2. [https://arxiv.org/pdf/1911.02150.pdf](https://arxiv.org/pdf/1911.02150.pdf)\n    \n    Please compare with these works and show the effectiveness of the proposed method since the biggest benefit of the proposed does not exist any more with the optimized single-GPU kernels, and the proposed method brings much more extra communication.\n    \n2. There is no analytical analysis on the extra communication brought by the newly proposed algorithm, which is the major overhead brought by the proposed method.\n3. Why do you need to evaluate the convergence of the proposed method in 4.1? The method does not change the behavior of the original self-attention algorithm.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject