DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training
Authors: Dacheng Li
Rulin Shao
Anze Xie
Eric P. Xing
Xuezhe Ma
Ion Stoica
Joseph E. Gonzalez
Hao Zhang
Venue: COLM (2024)
Keywords: Distributed system, memory efficient exact attention, long context, large language model
TL;DR: We develop a distributed extension to FlashAttention, specifically optimized for long context causal LLMs training.

Abstract: FlashAttention effectively reduces the quadratic peak memory usage to linear in training transformer-based large language models (LLMs) on a single GPU. In this paper, we introduce DistFlashAttention, a distributed memory-efficient attention mechanism optimized for long-context LLMs training. We propose three key techniques: token-level workload balancing, overlapping key-value communication, and a rematerialization-aware gradient checkpointing algorithm. We evaluate DistFlashAttention on Llama-7B and variants with sequence lengths from 32K to 512K. DistFlashAttention achieves 8x longer sequences, 4.45 - 5.64x speedup compared to Ring Self-Attention, 2-8x longer sequences, 1.24- 2.01x speedup compared to Megatron-LM with FlashAttention. It achieves 1.67x and 1.26-1.88x speedup compared to recent Ring Attention and DeepSpeed-Ulysses. Codes are available at https://github.com/RulinShao/LightSeq.

Published as a conference paper at COLM 2024
DISTFLASHATTN: Distributed Memory-efficiet Attention for Long-context LLMs Training
Dacheng Li dacheng1/70eerkeeleyyy edu
Rulin Shao "M
rulinsêcs. washington edu
Anze Xie
Eric P. Xing
Xuezhe Ma* "
alxie@ucsd.co edu
ericxingecs cmu edu
xuezhema@ist edu
Ion Stoica h istoica@borkeley.coom edu
Joseph E. Gonzalez" j jegonzal@berkeley. edu
Hao Zhang'
haozhangducsd.edd
UC Berkeley
University of Washington
UCSD
CMU
USC
Abstract
FlashAttention (Dao, 2023) effectively reduces the quadratic peak memory usage to linear in training transformer-basee large language models (LLMs) a single GPU. In this paper, introduce DISTFLASHATTN, distributed memory-efficient attention mechanism optimized for long-contex! LLMs training. We propose three key techniques: token-level workload balancing, overlapping key-value communication, and rematerializationaaaaa gra- dient checkpointing algorithm. We evaluate DISTELASHATTN on Llama-7B and variants with sequence lengths from 32K to 512K. DISTFLASHATIN achieves 8x longer sequences, 4.45 5.64x speedup compared Ring Self-Attention, 2 8x longer sequences, 1.24 - 2.01x speedup compared Megatron-LM, with FlashAttentinn. achieves 1.67x and 1.26 1.88x speedup compared to recent Ring Attention and DeepSpeed-Ulysaeee Codes are available at https: //gitpuu wow//ullinhaa/oiggiaiia
1 Introduction
Large language models (LLMs) capable of processing long context have enabled many novel applications, such generating complete codebase (Osika, 2023) and chatting with long documents (Li et al., 2023). Yet, training these LLMs with long sequences significantly increases activation memory footprints, posing new challenges.
Contemporary approaches to manage the high memory demands of long-context LLMs training involve either reducing activation memory a single device partitioning and distributing the sequences across multiple devices. Memory--eficiceee 2022; Dao, 2023; Rabe & Staats, 2021) represents the former, which reduces the peak memory usage of attention operations on a single device. Despite their effectiveness, the absence of distributed extension limits their application to sequence lengths that a single device can accommodate, Naively combining it with existing tensor or pipeline parallelisms (Shoeybi et al., 2019)) leads to excessive communication (S C) and cannot scale with sequence length (S 4). On the other hand, sequence parallelism systems, Ring Self-A Alff-Attention (Li et al., 2021) and Ring Attention (Liu et al., 2023), distribute the activations of a long sequence across multiple devices, but they lack support for memory-efficien attentions (e.g,, FlashAttention) or scheduling optimizations, making them inefficient in training long sequences (6 4.4).
attention (Dao et al.,
* Authors contributed equally.

---

Published as a conference paper at COLM 2024
Finish in 8 time steps
Finish in 5 timee steps
worker
wailing
waiting
waiting
wailing
waiting
wailing
waiting
22t Ai 12tx de, ti lts An, 21 26 Ao
watting
worker 2 los Ray Va lks, As, val waiting
waiting
waiting
wainling
wailiing
waiting
Kit Ray the lits Aos mnd lits stalgh Ro,, Pha
waitiing
worker
lis Re, Maa Ra. Ats the kas It.
waiting
waiting
waiting
waiting
waiting.
Kis Ray the Rtx Aay stal Rit fu.. 2t 95 Ra.
waiting
worker
Re. ks, v1 Rt. Au, my Rt. ly. th 21. ln,
waiting
waitingg
wailinng
waiting
Rts Re, the 12s. lop, stal Hit dog, 2tg) 2t. dn,
waitiing
worker
ln., ka. th lts ha, the lt. ka, vu &a, eul gts
waitingg
waitting
waiting
Ris kas tha lota Aas stal lts Ras stal 2s. ar yts, Ru.
worker lio. Re. th los Re, the lo. k4. salts. Ks. At. Ro. Ie. lo. "t
Wastinng
wailing
Ria Re. the lit. Ass les
sts, my Me. Ro,,
worker 7 kn. ky, im. Re, the ltt, ks, ex sa,er, Ru, euler, lu, ler, ly. My Wh. My wailing
Rrs Rr, te let. Ke, and 2t. Ae, gr. Rig, Mh Mhr ko, u2
worker 8 lta ls, th Wes ky, thy les ke, th Rules, Rs, 9h ls, ks. Uss ha. Ma Aoa Uss Ra.
Kiss Rt., the kes Koy mn les Nu. Phe Mts Rs. Mts Rus
(a) Ring Scheduling (Unbalanced)
(b) I Load- Balanced Scheduling (Ours)
Figure 1: Per-worker workload at different time steps in (a) ring scheduling (Li et al., 2021) and (b) the proposed load-balanced scheduling in an 8-worker scenario. The causal attention introduces a quadratic work dependency on the prefix of each token, where workers assigned earlier tokens remain idle while waiting for workers with later tokens.
The idle fraction of the ring scheduling is ri- 2p asymptotically 1 when scaling to more number of workers. The idle fraction of the proposed load-balanced scheduling is 1 2p when P is even and 0 when P is odd, asymplotically 0 when scaling to larger number of workers.
This paper introduces DISTFLASHATTN to extend the advantages of FlashAttention (Dao, 2023) to the distributed setting while maintaining high GPU utilization and low commu- nication overhead. DISTFLASHATTY efficiently distributes token chunks across multiple devices, while maintaining the IO-aware benefits of memory-efficient attention. We identify three key challenges in achieving high GPU utilization on distributed FlashAttention design for long-context L.MSs and propose three optimizations to addgress them.
The first challenge is the teven-level workload imbalance caused by causal language model- ing, As shown in Figure 1 (a), the causal attention introduces a quadratic work dependency on the prefix of each token. This leads to workers assigned earlier tokens to remain idle while waiting for workers with later tokens to complete, lowering the GPU utilization almost by half. We address this challenge by introducing a load-balancing schedule that routes the extra attention computation of later tokens to those idle workers (S 3.2). This optimization yields twice throughput of the unbalanced version as shown in Figure 5.
The second challenge is the prohibitive communication overhead. When tokens are dis- tributed to multiple machines, these machines need to communicate key-value tensors and softmax statistics to jointly compute the global attention. The communication volume is nontrivial, leading to large communication overhead, which grows with the context length. By leveraging the attention dependencies, we propose a scheduling technique that overlaps communication and computation by pre-fetching tensors. This successfully hides commu- nication overhead inside the computation time, resulting in 1.32x end- end-to-end speedup (Figure 5) compared to a non-overlapping version.
The third challenge is the high computation overhead due to the re-computation in gradient checkpointing (Chen et al., 2016). Gradient checkpointing effectively trades computation for memory by selectively storing intermediate activations (e.g., the inputs of every layer) and recomputing on-the-fly during the backward pass. It has become a standard technique in the training of long-context LLMs to accommodate the prohibitive activation memory (Zheng et al., 2023). However, the recomputation of the FlashAttention causes a high computation overhead in long sequences where the attention dominates the computation time. In $ 3.3, we show the recomputation of FlashAttention is unnecessary for its backward pass and propose a novel gradient checkpointing strategy to avoid it. Our new strategy results in an 1.31x speedup (S 4.6) without introducing any numerical difference.
Our main contributions are:
1. We develop DISTFLASHATIN, distributed, memory-efficient, exact attention mech- anism with sequence parallelism. We propose new optimization techniques to bal- ance the causal computation workloads and overlap computation and computation
2

---

Published as a conference paper at COLM 2024
to increase GPU utilization and reduce communication overhead for training long- context LLMs. We also propose a remateriaization-aaare gradient checkpointing strategy that eliminates redundant forward recomputation of FlashAttention.
2. We perform comprehensive evaluation of DISTFLASHATIN on LLaMA models, against four strong distributed systems. DISTFLASHATTN supports 8x longer sequences with 5.64X compared to Ring Self-Attention, 2x longer sequences with 1.24 - 2.01x speedup compared to Megatron-LM. DISTFLASHATTN achieves 1.67x and 1.26 - 1.88x speedup compared to Ring Attention and DeepSpeed-Ulyssee
Related work
Memory-efficient attention. Dao et al. (2022) and Lefaudeux et al. (2022) propose to use online normalizer (Milakov & Gimelshein, 2018) to compute the attention in blockwise and memory-efficiet way. It reduces peak memory usage by not materializing large intermediate states, e.g. the attention softmax matrix. In addition, research on sparse attention computes only a sparse subset of the attention score, which also reduces the memory footprints yet may lead inferior performance (Beltagy et al., 2020; Sun et al., 2022; Zaheer et al., 2020). In this work, we limit our scope to exact attention.
Sequence parallelism and ring attention Ring Self-Attentin (Li et al., 2021) is among the first to parallelize Transformers in the sequence dimension. However, its distributed attention design is not optimized for causal language modeling and incompatible with memory-efficient attention, which are crucial for long-context LLM training. Ring Atten- tion (Liu et al., 2023) proposes to compute distributed attention in memory-efficient blockwise pattern. However, it is also not optimized for causal language modeling, leading to 2x extra computation. DISTFLASHATTN optimizes for both memory-efficiet attention and causal language modeling. More recently, DeepSpeed Ulysses (Jacobs et al., 2023) proposes hybrid parallelism strategy. It computes distributed attention in the tensor model parallelism to address these two problems and utilizes sequence parallelism else- where (Shoeybi et al., 2019). We provide head-to-head comparison in Table 4.
Model Parallelism and FSDP Tensor Model parallelism (Korthikanti at. 2023) partitions model parameters and also distributes the activation in parallel LLM training. Pipeline model parallelism (Huang et al., 2019) also partitions the activations. However, it applies high memory pressure to the first pipeline stage. We show in $ 4.3 that this leads to less effective support for long sequences. Thus, we focus on comparing with tensor model parallelism and only consider pipeline parallelism when the number of heads is insufficient for tensor parallelism. Fully sharded pata-parallelism (FSDP) (Zhao et al., 2023; Rajbhandari et al., 2020) distributes optimizer states, gradients, and model parameters onto different devices and gathers them on-the-fly. Our work focuses on reducing the activation memory that dominates in long-context training. Therefore, FSDP is orthogonal to our work.
Gradient checkpointing. Gradient checkpointing (Chen et al., 2016) trades computation for memory by not storing activations for certain layers and recomputing them during the forward pass. Selective checkpointing (Korthikanti et al., 2023) suggests recomputing only the attention module, it requires significant memory but relatively few FLOPs (in contexts of smaller length). Checkmate (Jain et al., 2020) finds optimal checkpointing positions using integer linear programming. However, none of these designs have considered the effects of memory-efficient attention kernels, which perform recomputation within the computational kernel to avoid materializing large tensors. In this paper, we demonstrate that by simply altering the checkpointing positions, we can avoid the recomputation of these kernels without introducing any numerical difference.
3 Method
In this section, we first present a distributed memory-efficient attention mechanism that distributes the computation along the sequence dimension, DISTFLASHATIN (S 3.1) in its vanilla form. We then introduce two novel optimizations in DISTFLASHAIIN: a load- balanced scheduling strategy for causal language modeling to reduce the computation
3

---

Published as a conference paper at COLM 2024
$3,
3
Communication
20
24,04
09.87
09,87
6
7
5 7
4
Stream
Computation
attn(qr, ky, 27)
attn(gr, No, v6) attn(gg, ks, vs) attm(qr, has vs) attn (9:, Ass v2))
Stream
Figure 2: Overlap example the forward pass of worker 7 out of an 8 worker scnerio. For simplicity, "worker p" is denoted as p.
bubble and an asynchronous communication design that overlaps the communication into computation (S 3.2). Finally, we propose new remateriaizaaioonaaan checkpointing strategy (S 3.3) which effectively cuts off the recomputation time in gradient checkpointing when using DISTFLASHATIN in long-context training.
3.1 DISTFLASHÄTTN distributed memory-efficient attention via sequence parallelism
The goal of DISTFLASHÄTIN is twofold: (1) distribute a single sequence into multiple workers so they jointly utilize the memory to support long sequence training; (2) maintain the IO-aware benefits of memory-efficient attention so that training is fast and incurs less memory footprint. In particular, we choose FlashAttention (Dao, 2023) as the paradigm.
Algorithm 1 (Vanilla) DISTFLASHATTI of worker p
Require: qp, Vp 1: Initialize Op = 0", Sp = =mm==mmm 1m,, where 0 = and m [-o0.... o0]
2: Op, Sp = =attn(qpr kp, vp, Op, Sp)
3: for mrlss do
4: r- (modl P)
5: Fetch from remote: worker p K. worker
6:
Op, Sp attn( (ppn(qq,kek k,, Vr, Op, Sp)
7: end for
Return Op-
To distribute the long sequence. DISTFLASHATTN splits the input sequence consisting of N tokens evenly across P workers (e.g. GPUs) along the sequence dimension. Each worker computes and stores the activations of only a subsequence of N/P tokens. Therefore, supports training Px longer with P workers than single-worker FlashAttention.
Formally, let qp, kp, Vp R xd be the query, key and value of the subsequence the p-th worker (p (I. P)), where is the hidden dimension. Considering the most prevalent causal attention in LLMs, worker p computes the attention output op associated with qp
Op Softmax 9p qplk),..,kpI"" [VI, Vpl
(1)
Vd
To maintain the 10-awareness. Naively, each worker could gather all the keys and values associated with other subsequences and then locally computes op by invoking the existing single machine FlashAttention. However, this gathering introduces memory pressure by having to store the full list of keys and values locally, total size of 2Nxd
Fortunately, the block-wise nature of the single-worker FlashAttention only requires one block of keys and values in each iteration of its algorithm, Leveraging this observation,
4

---

Published as a conference paper at COLM 2024
checkpoint
forward
recomputation
backward
other modules
Fwd
Fwd
FFN
FA
FA
FFN
FA
FA
FFN
Bwd
Bwd
Huggingface Checkpointing
Rematerialization-NNe -waar Checkpointing (Ours)
Figure 3: Comparison of HuggingFace gradient checkpointing strategy and our materialization-aware gradient checkpointing strategy. Note that our checkpointing strategy saves an entire flash attention forward per layer in recomputation by simply shifting the checkpoint boundaries without introducing any numerical difference. The checkpointed tensors, i.e, the outputs of (tahtttteninn((A),, are saved not only for the recomputation of subsequent layers but also the backward computation of the preceding FlashAttention.
compute iteratively: in each iteration when r * p, worker p fetches only one k, v, from remote worker It then computes partial attention results based on qp and k,, v, and perform proper rescaling by invoking the single-worker FlashAttention kernel. To perform proper rescaling between iterations, each worker also needs maintain a copy of softmax statistics! Sp E R A
Computing in this iterative manner, each worker also stores the key and value of one subsequence of size ~ factor of memory of the naïvely design. We refer to Dao et al. (2022) for more details of the single-worker FlashAttention. We denote each iteration of the partial attention result and the rescaling as attn qp, kr, Vr, Sp), and present the vanilla DISTFLASHATIN algorithm in Algorithm 1. In Appendix A, we show how to implement the arm kernel from Dao (2023) in pseudo-code.
3.2 Load balanced scheduling with communication and computation overlap
Load-balanced scheduling. In causal attention, each token only attends to its previous tokens, iee, the p-th worker computes attn(qp, k,,,, for all r S p. This introduces workload imbalance between workers: worker with larger p computes more attn(-)
(Figure (a)). Using the scheduling described in $ 3.1, the idle fraction is på-P (- when 2Pt
p P oo), which means roughly half of the workers are idle. To reduce this idle time, we let worker that has finished all its atmm(.) computations (i.e., the "helper") perform attention computation for worker 72 with heavier workload, as shown in Figure (b).
Notably, the "helper" r needs to communicate softmax statistics and the partial attention output to the original worker r2, so that worker r2 can update its local copy of statistics and output correctly (Algorithm 2). This update function denoted as rescale() and updates the partial output and statistics in the same way as how Dao (2023) updates results from two block execution. This scheduling gives an average idle time fraction:
0, P odd IP,
X
(2)
is even
Note that when P even, the idle time is asymptotically 0 to more workers. We provide illustration with 8 workers in Figure 1. While we focus on the exact attention mechanism, also discuss sparse patterns in Appendix D.
Communication and computation overlap. DISTFLASHATIN relies on peer-to-peer (P2P) communication to fetch k,, (or qr in the load-balanced scheduling) from remote workers before computing attn(-). However, these communications can be naturally overlapped. To simplify the equations, we use the unbalanced schedule to describe the intuition, while the
These statistics and in FlashAttention words.

---

Published as a conference paper at COLM 2024
final DISTFLASHÄTTN implementation equipped with both optimizations. Precisely, these two operations are parallelized:
Fetch worker p Krot/eel
worker r +
(3)
Compute: attn(qp, k,, Sp)
Thus, in the next iteration, K,+1, Vr+1 are already stored in the memory of worker p, without blocking the next iteration's computation. ing the attention computation kernel in the main GPU stream, and the P2P communication kernel in another stream, where they can run parallel (Zhao et al., 2023). We demonstrate the overlapped scheduling for worker in the 8-worker scenario in Figure. 2. Empirically, we find this optimization effectively reduces the communication overhead by hiding the communication time inside the computation time (5 4.6).
modern accelerators, this can be done by plac-
3.3 Rematerializaaionnawaae checkpointing strategy
Gradient checkpointing (Chen et al., 2016) is a de-facto way of training cong-context trans- formers. Often, the system uses heuristics to insert gradient checkpoints at the boundary of each Transformer layer (Wolf et al., 2019). However, with the presence of Dao et al. (2022), we find the previous gradient checkpointing strategy causes a redundant recomputation of the FlashAttention forward kernel. Precisely, when computing the gradient of the MLP layer, Wolf et al. (2019) re-computes the forward of the entire Transformer layer including Attenton.. During this process, the PlashAttention backward kernel re-computes the softmax block-wisely again to reduce memory usage. Essentially, this is because FlashAtten- tion does not materialize the intermediate values during the forward, and recomputes it during the backward, regardless the re-computation the outer system level (e.g., the HuggingFace gradient checkpointing (Wolf et al., 2019)).
To tackle this, we propose to insert checkpoints at the output of the FlashAttention kernel, instead of at the Transformer layer boundary. We use each checkpoint not only for the recomputation of its subsequent mod- ules but also for the backward ard computation of its
Wall clock tim
200
Attention
Other
150
preceding FlashAttention module without recompu- tation. Thus we only need to compute the forward of FleshAttention once, effectively avoiding all recom- putations of PlashAttention as shown in Figure 3.
100
74.4
95.8
50
12.910.5
27.123.7
8K
16K
32K
64K
Sequence length
Figure shows that attention dominates in the for- ward pass with in long sequences, which indicates our method saves ~ 0.23 x 32 (i.e., ~ 7) seconds when training 64K sequence example on Llama-7b on single machine. In addition, this saves a communi- cation brought by our DISTFLASHATTH forward in the distributed training scenario. We benchmark the end-to-end speedup brought by this materialization-awarr checkpointing strategy in & 4.6.
Figure Time breakdown of at- tention versus other modules in forward pass, measured with Flash- Attention (Dao, 2023) on single 40GB A100 GPU. (Unit: ms)
4 Experiments
We evaluate ASTFLLSSHATT together with our new checkpointing strategy against al- ternative distributed approaches for long-context LLMs training. Our primary baseline is Megatron-LM (Shoeybi et al., 2019), used in tandem with FlashAttention, which serves as robust baseline extensively adopted within the industry. theoretical analysis on its high communication volume. We also provide a comparison with the previous requence-parallet system (Li et al., 2021). In addition, we include comparison and Ring Attention (Jacobs et al., 2023; Liu et al., 2023). In the ablation study, we delineate the individual contributions of each component of our methodology, specifically load balancing, computation-commniice. overlapping, and rematerializationawaal checkpointing, towards the overall performance enhancement. Code implementation details can be found in Appendix 4.1.
Appendix C, we also show
to recent systems including DeepSpeed-Ulyssee

---

Published as a conference paper at COLM 2024
Table 1: Per iteration wall-clock time of DISTFLASHATIN and Megatron-LM (Korthikanti et al., 2023) (Unit: seconds). Speedup in bold denotes the better of the two systems in the same configuration. Time measured with DGX boxes.
Method
# GPUs
Sequence
Length Total
LLaMA-7B
LLaMA-GQA
LLaMA-33H
Per GPU
Time
speedup
Time
speedup
Time
speedup
1x8
8K
64K
6.81
1.0x
6.60
1.0x
8.37
1.0x
Megatron-LM
16K
128K
20.93
1.0x
20.53
1.0x
25.75
1.0x
1x8
32K
256K
72.75
1.0x
71.93
1.0x
90.21
1.0x
8K
64K
5,98
1.14x
5.61
1.18x
6.08
1.38x
DISTFLASHIATIN
1x8
16K
128K
17.26
1.21x
16.86
1.22x
17.77
1.45x
1x8
32K
256K
58.46
1.24x
57.01
1.26x
59.96
1.50x
2x8
8K
128K
14.26
1.0x
14.21
1.0x
20.63
1.0x
Megatron-LM
2x8
16K
256K
43.44
1.0%
43.20
1.0x
62.78
1.0x
2x8
32K
512K
147.06
1.0x
146.38
1.0x
216.70
1.0x
2x8
8K
128K
12.75
1.12x
9.74
1.46x
13.12
1.57x
DIST
2x8
16K
256K
30.21
1.44x
28.49
1.52x
31.33
2.00x
2x8
32K
512K
106.37
1.38x
102.34
1.43x
107.76
2.01x
Cluster setup. We evaluate our method and the baselines (1) A single A100 DGX box
with 8x80 GB GPUs. These GPUs are connected with NVLink; (2) 2 DGX boxes with the same setting. These two boxes are interconnected by a stable 100 Gbps Infiniband. This is representative setting for cross-node training, where the communication overhead is large. Unless otherwise stated, this is our default setup. (3) Our in-house development cluster with 2x8 A100 40GB GPUs. This cluster has unstable inter-node bandwidth. Due to the limited computational budget, we report some peripheral results on this cluster.
Model setup. We evaluate our system LLaMA-7B and its variants, encompassing four sets of model architectures in total: two with regular attention heads and two with irregular ones. We note both categories are important in real-world applications.
With regular attention heads. (1) multi-head attention (MHA) models: LLaMA-7B with 4096 hidden size and 32 self-attention heads (Touvron et al., 2023); (2) grouped-quers attention(GQA) models: LLaMA-GQA (Ainslie et al., 2023), same as LLaMA-7B but with key-value heads, each shared by queries group. During attention computation, it will first replicate to 32 heads to perform matrix multiplication with the correct shape.
With irregular attention heads. In addition, we benchmark the following variants that have appeared in applications but have not received much attention regarding their system efficiency: (3) models with an irregular (e.g,, non-power-of-twoo number of attention heads?: We intentionally test our systems and baselines on L.LMA -33H, which has the same configu- ration as LLaMA-7B but with 33 normal self-attention heads per layer. (4) models with fewer attention heads According to the recipe in Liu et al. (2021), we designed LLaMA-16H, LLaMA-8H, LLaMA-4H, and LLaMA-2H with 16, 8, 4, and 2 heads, respectively, as a proof of concept for situations when the number of heads is insufficient to further scale up model parallelism with limited resources. We keep the number of attention heads by scaling the number of layers properly* and keep the intermediate FFN layer size the same to make the model sizes still comparable. For example, LLaMA-16H has 16 attention heads per layer, hidden size of 2048, an FFN layer of size 11008, and 64 layers.
4.1 Implementation Details
We build the kernel of DISTFLASHÄTTN, modifying from the Triton kernel of FlashAt- tention2 in 500 lines of codes (LoCs). We implement the load balancing and overlapping
2or example, GPT-2-XL. has attention heads, GPT-2 has attention heads, LLaMA-33B and its fine-tuned versions (e-g.. Tulu-30B) have 52 attention heads, Whisper-large has attention heads, and Falcon-7B has 71 attention heads (Radford et al., 2019; Almazrouei et al., 2023; Ivison et al., 2023). LLu al. (2021) finds fewer attention heads with more layers increase the performance.
For instance, LII J-aB has 32 attention heads and 32 layers, thus I.LMAA16HH has 16 attention heads per layers and 64 layers.
7

---

Published as a conference paper at COLM 2024
scheduling Python and NCCL Pytorch bindings in 1000 LoCs (Paszke et al., 2019; Jeaugey, 2017), and the checkpointing strategy in 600 lines of Pytorch. We use block sizes of 128 and the number of stages to in the kernel for the best performance in evaluate DISTFLASH ATTN using FSDP (inter-node if applicable) so that it consumes similar memory than the Megatron-LNN baseline for fair comparison (Zhao et al., 2023). For fair comparisons, we run all comparisons using the same attention backend. We also add sup- port for Megatron-LM so that comparing with them can produce more insightful analysis: (1) not materializing the causal attention mask, greatly reducing the memory footprint. For instance, without this support, Megatron-LM will run out of memory with LLaMA-7B at sequence length of 16K per GPU. (2) head padding where the attention heads cannot be divided by device number. All results are gathered with Adam optimizer, 10 iterations of warm-up, and averaged over the additional 10 iterations.
cluster. We
4.2 Comparison with Megatron-LM on MHA and GQA models
Multi-head attention (MHA). On the LLaMA-7B model (Table 1), our method achieves 1.24x and 1.44x speedup com- pared to Megatron-LM in single-node and cross-node setting, up to the longest se- quence length we experiment. This is joint result of our overlapping communication technique and our rematerialization-aaarr checkpointing strategy. We analyze how much each factor contributes to this result in the ablation study (S 4.6). Grouped-ques -uery attention (GQA). On GQA model, DISTFI ASHATTN communicates less volume due to the reduction of size of keys and values. On the contrary, the communication of Megatron-LM remains the same because it does not communicate keys and values. Thus, DISTFLASHATTN achieves a higher speedup on LLaMA-GQA model (Table 1).
Table The maximal sequence length Per GPU supported by DISTFLASHATTN and Megatron-LM with tensor parallelism and pipeline parallelism on 16xA100 40GB GPUs.
16H
8H
4H
2H
Megatron TP+DP
512K
256K
128K
64K
Megatron TP+PP
512K
256K
256K
128K
DISTFLASHATIN
512K
512K
512K
512K
4.3 Comparison with Megatron-LM on models with irregular or less number of heads
In support of irregular numbers of heads. Megatron-LM assumes the number of attention heads is divisible by the model parallelism degree. For example, it supports parallelism degrees of 2, 4, 8, 16, and 32 for models with 32 attention heads. However, it needs pad dummy heads when the number of heads is not divisible by the ideal parallelism degree. For example, needs to pad 15 dummy heads to support parallelism degree of 16 for models with 33 attention heads (e.g., 1ImaMA-33H) leading to a substantial computation wastage of 45.5%. As shown in Table 1, we observe a 1.50x and 2.01x speedup (an additional 20% and 45% speedup compared to LLaMA-7B cases, aligned with the theoretical analysis).
In support of less number of heads. When the number of GPUs exceeds the number of attention heads, Megatron-LM allows three possible solutions: (1) Pad dummy heads as in the LLaMA-33H scenario. However, the percentage of dummy heads almost directly translates to the percentage of slowdown in long sequences where attention computation dominates. (2) Use data parallelism for excess GPUs. However, data parallelism does not reduce per sequence memory usage, and thus can not jointly support longer sequences. (3) Use pipeline parallelism. However, the memory usage at each stage of the pipeline is not evenly distributed, limiting the maximal sequence length supported. For instance, in the LLaMA-2H experiment, we find that different stages consume from 18GB to 32GB a 64K sequence length (Section B). In addition, using pipeline parallelism introduces extra fraction of GPU idle time. We demonstrate the effect of using the latter two solutions in Table 2. In 16 A A100 40GB GPUs, DISTFLASHATTH supports 2x and 8x longer sequences.
8

---

Published as a conference paper at COLM 2024
4.4 Comparison with Ring Self-Attention (RSA) and Ring Attention
Ring self-attention (RSA) (Li et al., 2021) communicates tensors in ring fashion. We first report the maximal sequence length of RSA and DISTFLASHATTN in Table 3, and found that DISTFLASHÄTIN supports at least 8x longer sequences than RSA. This is mainly because RSA is not natively compatible with memory-efficient attention. We further measure the iteration time with the maximal sequence length that RSA can support in Table 3, and find that DISTFLASHATTN is 4.45x 5.64x faster than RSA. This speedup includes a 2x improvement from our causal workload balancing optimization and additional gains from the overlapping optimization and extending memory-efficien attention to the distributed setting. Both experiments are conducted with the Llama-7B model and on the DGX cluster.
Table 3: Max sequence length and per iteration time (seconds) compared with RSA.
Node
2 Nodes
Node (32K)
Nodes (64K)
RSA
14.10
30.49
RSA DISTFLASHATTN
32K
64K
DISTFLASHATTN
2.50
6.85
256K
> 512K
Speedup
5.64x
4.45x
Ring Attention (Liu atal., 2023) implements distributed attention in a memory-efficient man- ner. The key difference between Ring Attention and DISTFLASHATIN DISTFLASHATTN has additional optimization of causal workload balancing and better gradient checkpoint strategy. The implementation of Ring Attention uses a different framework from ours (Jax versus PyTorch). To provide a fair comparison, we consider our ablation version in $ 4.6 as a PyTorch implementation of Ring Attention. $4.6 provides detailed analysis. In 8-GPU setting, we observe 1.67x speedup (7.5x versus 4.5x speedup compared to single GPU Attennittt over the design of Ring Attention.
4.5
Comparison with DeepSpeed Ulysses
Table 4: Per iteration wall-clock time (seconds) of DISTFLASHATIN and DeepSpeed Ulysses.
Method
GPUs
Sequence Length Per GPU
Time
Speedup
PGPss
Sequence Length Per GPU
Time
Speedup
Total
Total
Llama-70
Llama-3311
DeepSpeed Ulysses
2x8
16K
256K
37.53
1.0x
2x8
16K
256K
56.63
1.0x
2x8
32K
512K
134.09
1.0x
2x8
32K
512K
202.89
DISTFLASHATTN
2x8
16K
256K
30.21
1.21x
2x8
16K
256K
31.33
1.81x
2x8
32K
512K
106.37
1.26x
2x8
32K
512K
107.76
1888x
DeepSpeed-Ulysses (Jacobs et al., 2023) uses all-to-all primitive to reduce the communication. We evaluate a representative subset of experiments in Table 4 due to computational budget limit. On experiments with regular heads models (Llama-7B), DISTFLASHATTN achieves 1.26 x speedup. On experiments on irregular heads models (Llama-33H), achieves 1.88x speedup. Essentially, DeepSpeed-Ulyssee also paralleize on the attention head dimension, and suffer from the same problems as analyzed in $ 4.3.
4.6 Ablation Study
Effect of Load Balancing We study load balancing on an attention forward pass of LLaMA- 7B model, on 8 A100 40GB GPUs (Figure 5). The backward pass follows similar analysis. With an unbalanced schedule (Figure 1), the total work done is 36, where the total work could be done in 8 units of time is 64. Thus, the expected speedup is 4.5x. In the balanced schedule, the expected speedup is 7.2x. We scale the total sequence length from 4K to 256K. The unbalanced version saturates in 4.5x speedup compared to single GPU FlashAttention while the balanced version saturates at 7.5x speedup. Both of them align with our theoretical analysis and show the effectiveness of the balanced scheduling.

---

Published as a conference paper at COLM 2024
single machine
unbalanced
balanced
No overtap
overtap
no communication
iter time
2.0
1.95
Speedup
1.5
1.65
1.33
1.08
1.0
0.5
0.0
0K
50K
100K
150K
200K
250K
Sequence length
Sequence length
Figure 5: Effect of balanced schedule (left) and the effect of overlapping (right).
Effect of overlapping communication and computation. We study the overlapping com- munication LLaMA-7B and DGX boxes (Figure 5). We find that overlapping greatly reduces the communication overhead. On a global sequence length of 128K, the commu- nication overhead is reduced from 105% to 44%. This overlapping scheme maximizes its functionality when the communication overhead is less than 100%, where all communica- tion can be potentially overlapped. Empirically, we find the system only exhibits 8% and 1% overhead in these cases, a close performance to an ideal system without communication.
Effect of remaaerializaaionnawaaa checkpointing. We show in Table 5 the effects of the proposed rematerialization-waaal gradient checkpointing. Our method achieves 1.16x, 1.24x, and 1.31x speedup at the sequence length of 8K, 16K, and 32K per GPU respectively. The materialiiation-aware checkpointing strategy speeds up more at longer sequence lengths where the attention dominates the computation.
4.7 Partition on the attention heads or sequence dimension
are distributed systems that partition on atten- tion heads. While it allows seamless inte- gration with the FlashAttention kernel, it has certain limitations. These includes: (1) Not being able to utilize the pattern inside the attention module, missing opportunities to reduce communication for causal, and prouped-query attention (See $ C). (2) not flexible to support arbitrary number of at- tention heads, and (3) Importantly, its scala- bility is limited by the number of attention heads (in the scale several to several dozens), while the maximal number of parallelism degree for sequence parallelism is at least several thousands. Given these reasons, we think worth pursuing the sequence parallelism paradigm when distributing the attention module.
Megatron-LM and DeepSpeed-Ulysee
Table 5: Our checkpointing algorithm ("Our ckpt") versus HuggingFace strategy ("HF ckpt") on 8 A100s (batch size 1, Unit: seconds).
Method
Sequence I Length Per GPU
1K
2K
4K
8K
16K
32K
HF ckpt
0.84
1.29
2.64
6.93
21.44
76.38
Our ckpt
0.84
1.36
2.50
5.98
17.26
58.46
Speedup
1.0x
0.94x
1.06x
1.16x
1.24x
1.31x
Conclusion
In this work, we introduce ATTN,,, distributed memory-efficient attention proto- type for long-context transformer training based sequence parallelism. DISTFI. ASHAATN presents novel system optimizations including load balancing for causal language model- ings, overlapped communication with computation in the distributed attention computation, and re-materiallzation-aavee
checkpointing strategy. Experiments evaluate multiple fami- lies of transformer models and on different cluster types, and over four strong distributed system baselines. In particular, DISTFLASHATIN has demonstrated up to 2.01x speedup and scales up to 8x longer sequences, compared to the popular system, Megatron-LM with FlashA Attention.
10

---

Published as a conference paper at COLM 2024
References
Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint ar Xio:22305.13245, 2023.
Ebtesam Almazrouei, Hamza Alobeidli, Abdulaziz Alshamsi, Alessandro Cappelli, Ruxan- dra Cojocaru, Mérouane Debbah, Étienne Goffinet, Daniel Hesslow, Julien Launay, Quentin Malartic, et al. The falcon series of open language models. arxivv preprint #tXii:2311.16867, 2023.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document trans- former. ar Xie preprint or Xiv:2004.005150, 2020.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. Xrviv preprint ar arXic:1604.06174, 2016.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning.
Xrv preprint ar (rXiv:2307008691, 2023.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344-16359, 2022.
Yanping Huang, Youlong Cheng, Ankur Вapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc Le, Yonghui Wu, al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32, 2019.
Hamish Ivison, Yizhong Wang, Valentina Pyatkin, Nathan Lambert, Matthew Peters, Pradeep Dasigi, Joel Jang, David Wadden, Noah A. Smith, Iz Beltagy, and Hannaneh Hajishirzi. Camels in changing climate: Enhancing im adaptation with tulu 2, 2023.
Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Leon Song, Samyam Rajbhandari, and Yuxiong He. Deepspeed ulysses: System optimizations for enabling training extreme long sequence transformer models. Xrviv preprint ar XrXiv:2309.1550, 2023.
Paras Jain, Ajay Jain, Aniruddha Nrusimha, Amir Gholami, Pieter Abbeel, Joseph Gonzalez, Kurt Keutzer, and lon Stoica. Checkmate: Breaking the memory wall with optimal tensor rematerialization. Proceedings of Machine Learning and Systems, 2:497-511, 2020.
Sylvain Jeaugey. Neel 2.0. In GPU Technology Conference (GTC), volume 2, 2017.
Vijay Anand Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. Proceedings of Machine Learning and Systems, 5, 2023.
Benjamin Lefaudeux, Francisco Massa, Diana Liskovich, Wenhan Xiong, Vittorio Caggiano, Sean Naren, Min Xu, Jieru Hu, Marta Tintore, Susan Zhang, Patrick Labatut, and Daniel Haziza. xformers: A modular and hackable transformer modelling library. https: //github. com/faccbookreeearcch wwwpeeearcc//fooorreeee /fomrrss 2022.
Dacheng Li, Rulin Shao, Anze Xie, Ying Sheng, Lianmin Zheng, Joseph E Gonzalez, Ion Stoica, Xuezhe Ma, and Hao Zhang. How long can open-source llms truly promise on context length, 2023.
Shenggui Li, Fuzhao Xue, Yongbin Li, and Yang You. Sequence parallelism: Making 4d parallelism possible. arXiv preprint MxXi=:2105.13120 2021.
Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near-infinite context. aixiv preprint ar arXio22310.01889, 2023.
11

---

Published as a conference paper at COLM 2024
Liyuan Liu, Jialu Liu, and Jiawei Han. Multi-head or single-head? empirical comparison for transformer training. arXio preprint arXio:2106.0966 2021.
Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. ar Xiv preprint Xit::805.00267,, 2018.
Anton Osika. gpt-engineer, 2023. URL htttp::///gghh htm////tttniiiiaaaggggggegeeeeee
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An im- perative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Markus N Rabe and Charles Staats. Self-attention does not need o(n²) memory. arviv preprint arXivv2112.0568// 2021.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 188):9, 2019.
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.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using mödel parallelism. ar preprint ar arXi::1909.008
2019.
Yutao Sun, Li Dong, Barun Patra, Shuming Ma, Shaohan Huang, Alon Benhaim, Vishrav Chaudhary, Xia Song, and Furu Wei. A length-extrapolatable transformer. arXiv preprint arXtt:211..10554, 2022.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXio:2302.13971, 2023.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Huggingface's transformers: State-of-the-art natural language processing. arXiv preprint drXiv.1910.03771, 2019.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Adoances in neural information processing systems, 33: 17283-17297, 2020.
Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, et al. Pytorch fsdp: experiences on scaling fully sharded data parallel. ar preprint ar arXio:2304.11277, 2023.
Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric. P Xing, Hao Zhang, Joseph E. Gonzalez, and lon Stoica. Judging llm-as-a-judge with mt-bench and chatbot arena, 2023.
12

---

Published as a conference paper at COLM 2024
A From FlashAttention to attn(.) in DISTFLASHATIN
In this section, we provide the details of the (-) kernel in DISTFLASHATTN. 3). For conceptual simplicity, we demonstrate it in the most vanilla version, without the actual scheduling (e.g. load balancing and overlapping). We also demonstrate it with the causal language modeling objective. The standalone attention is mainly borrowed from the FlashAttentions paper (Dao, 2023). To make it compatible with DISTFLASHATTN, we mainly revised the several points:
1. Accumulate results statistics o, m and from previous computation, instead of initializing them inside the function.
2. Pass an extra argument "last", which means whether this is the last chunk of attention computation. Only when it is true, we compute the logsumexp L.
At high level, worker p, DISTFLASHATTI first initializes local statistics o, m, 1, Then LISTFLASHATN loops over all its previous workers. In each iteration, it fetches the key and the value from worker and invokes the revised standalone attention statistics. At the end of the iteration, it needs to delete the remote key and value from HBM so that the memory does not accumulate. At the last iteration of the loop, it additionally calculates the logsumexp according to the final m and (triggered by the "last" variable in the algorithm). At the end of the forward pass, worker p has the correct m, L.. The backward pass is similar and conceptually simpler because we do not need to keep track of statistics such as and Instead, we only need use the logsumexp stored in the forward pass.
update local
Algorithm 2 (Balanced) DISTFLASHATTN of worker p
Require: qp, kp. Vp Initialize Op o Sp = 0 Im²,1%, where o 0, 1 and m [- coo7
2: Op, Sp =alin(qp, kp, Vp, Sp)
3: for 1 s do
4: r= (p-11) 1) (mod P)
5:
if t then
Fetch key, value from remote: p k,
6:
7:
Op, sp =att(q 9p, ky, Vr, Sp)
8:
ift # [ and (p+1)
then 12 = (p+ (mod P)
9:
Fetch result from remote: 0
10:
12
11:
Pp, Sp = rescale(Op, Sp, pp sp))
12:
end if
else
14:
if t # L&J then
15:
Fetch query from remote: p 0,, S, = =tin(qq, kp, vp, 00, )
16:
Send result to remote: o, end if
17:
18:
19:
end if
end for
21:
Return Op.
B Memory Consumption for Pipeline Parallelism
In this section, we show the memory consumption of Megatron-LM when training with tensor parallelism and pipeline parallelism. As presented in table 6, memory consumption
13

---

Published as a conference paper at COLM 2024
Algorithm 3 DISTFLASHATTIN Pseudo code (forward pass)
Require:
Matrices Q", K2, VP # in HBM, block sizes Be, Br, rank function standalone fwdq, k, v, o, f, m, causal, last
Divide q into T, PB,
blocks q1,-.., qT, of size B, x each,
kr. 2: and divide k, v in to T = PB. blocks kI,...,, and Vi...., VT., of size B x each.
3: Divide the output
R N xd into T, blocks logsumexp L into T, blocks Ly, of size B, each.
of size B, each, and divide the
4: for s do
5: Load from HBM to on-chip SRAM.
6: Load o, RR (, RB, m; R from HBM on-chip SRAM as 00) c, m (0)
7: for 1 $ j - Te do
8:
if causal and s then
9:
Continue
end if
11:
Load k,, from HBM to on-chip SRAM.
On chip, compute s/") qikb B...
12:
On chip, compute (1) = max(m (1-1) rowmax(s; m: m))) RB (pointwise), C(1) = mm/--mm c ((i-1) + rrossum (p! E
RB, p. = =
13:
On chip, compute of (i) diag( m! m ----mmm))) - (i-1) + vg.
14:
15:
end for
16:
Write o; to HBM as the i-th block of of
On chip, compute o, = diag (T)) -1
(Tc)
17:
18:
if last then
On chip, compute L, = m m((e) (o)((())))
19:
20:
Write L, to HBM as the i-th block of L. end if
22:
end for
Return o, and the logsumexp
end function
RR,m² 24: Initialize OP (0) y xd x p(p) (0) E R m2 (-00) p. $
25: OP, "M, I.P standalone KP, VP, OP, 1p, mD True, p=1)
forl < do
27: Receive K and v from Remote worker into HBM.
28: OP, LP, mP, L standalone fwd(Q", K², V>, OP, l2, mP, False, r=(p-1)
29: Delete K" and V from HBM.
30: end for
31: Return the output OP and the logsumexp I..
14

---

Published as a conference paper at COLM 2024
Table 6: The memory consumption of Megatron-LM when training Llama-2H with tensor parallelism (degree=2) and pipeline parallelism (degree=8) on 16xA100 40GB GPUs at the sequence length of 128K. The memory consumption highly uneven across pipeline stages.
Worker 1
Worker 2
Worker
Worker 4
Worker 5
Worker 6
Worker 7
Worker 8
node 1
31.5GB 21.8GB
31.4GB
28.7GB 20.5GB
28.7GB
26.0GB 17.9GB
26.0GB 17.8GB
24.6GB 32.0GB
24.6GB 32.1GB
node 2
21.8GB
20.5GB
are uneven across different pipeline stages, making scaling through pipeline parallelism hard.
C Communication and memory analysis
Denote the hidden dimension as In DISTFI ATTN, every worker needs to fetch key and value chunks both of size &d d before performing the corresponding chunk-wise computation. Thus, the total communication volume in the P-workers system is 2 x T x P = 2Nd. With the causal language objective, half of the keys and values do not need to be attended, halving the forward communication volume to Nd. In the backward pass, DISTFLASHATTH needs to communicate keys, values, and their gradients, which has 2Nd volume. It adds up 3Nd as the total communication volume for ATTH In Megatron-LM (Korthikanti et al., 2023), each worker needs to perform six all-gather and four reduce-scatter on a d size tensor, thus giving a total communication volume of 10Nd. Considering gradient check-pointing, Megatron-I.M will perform communication in the forward again, giving a total volume of 14Nd. On the other hand, our communication volume remains 3Nd because of the rematerializationaaaan 4.7x communication volume reduction compared with Megatron-LM.
strategy. In conclusion, DISTFLASHATIN achieves
In large model training, we usually utilize techniques such FSDP to also reduce the memory consumed by model weights. In this case, We note that the communication introduced by FSDP is only proportional to the size of model weights, which does not scale up with long sequence length. We show the end-to-end speedup with FSDP in Table 1. For clarity, we also note that AISTFLASHATTN is orthogonal to FSDP and by default can be used by itself. In the situations where the model uses MQA or GQA, DISTELASHATTN further saves the communication volumes by the shared key and values, which we discuss in detail in $ 4.2. However, we also note that this is a theoretical analysis, where the wall-clock time may differ because of factors such as implementations. In the experiment section, we provide wall-clock end-to-end results for comparison.
D Discussion on sparse attention
While this paper focuses on discussing the exact attention mechanism, we also provide possible solutions for sparse patterns and hope it can inspire future works. In particular, we discuss load balancing for local sliding windows and global attention (Beltagy et al., 2020).
Local sliding windows For local sliding windows, the workload is naturally (near) balanced, regardless of single directional or bidirectional attention. Thus, simply disregarding the attention logic non-local workers suffices. For instance, in exact attention, worker needs to compute attention to all other workers. If the sliding window has a number of tokens equal to that of one worker, then worker only needs to attend to itself and tokens in worker 6. In other words, it only needs to fetch key and value from worker 6, and compute attention. In terms of implementation change, the system merely needs to change the end condition of the for loop (from looping worker 1 worker 7 to looping only from worker 6 worker 7).
Global attention In global attention, there are a certain number of global tokens that all later tokens need to attend to, which are used to capture the global information. To adapt FLSSALA to this, one solution is to keep a replica of all the global tokens in each worker, which is simple and practical otherwise, the global tokens will need to be all-
15

---

Published as a conference paper at COLM 2024
gathered at each time step, The other possibility is to also split the global tokens evenly onto all workers and use all-gather upon computation to further reduce the memory requirement.
16


Reviews:
[{"review_id": "2uNHNtFO29", "reviewer": "Reviewer_8Cac", "summary": "The paper titled introduces a distributed variant of FlashAttention, aimed at training large language models (LLMs) with long sequence contexts more efficiently on distributed systems. The paper presents three main innovations: token-level workload balancing, overlapping key-value communication, and rematerialization-aware gradient checkpointing. These techniques collectively enhance the efficiency of memory usage and reduce computational and communication overheads significantly.\n\nThe key contributions are:\n1. allows for up to 8x longer sequence processing and achieves speedups between 1.24x and 5.64x compared to existing models like Ring Self-Attention and Megatron-LM with FlashAttention.\n2. introduces new methods to balance workload across processors, overlap communication with computation to hide latency, and optimize gradient checkpointing to avoid unnecessary recomputation.\n3. Extensively evaluated across various setups and model configurations, demonstrating significant improvements in speed and efficiency.", "questions": "", "limitations": "", "rating": 7, "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": "We thank Reviewer 8Cac for highlighting the innovations of our load-balanced scheduling, overlapping of communication and computation, and rematerialization aware gradient checkpointing strategy. We would love address the reviewer’s concerns as below:\n\nQ1: Memory consumption comparison with baselines.\n\nA1: We thank the reviewer for raising this question. Our method is memory-efficient mainly because we retained the I/O awareness of FlashAttention when extending it to distributed version. Therefore, we note our method is equivalently memory efficient compared with Megatron-LM(with flash attention integration), DeepSpeed-Uylesses(with flash attention integration), and Ring Attention. However, these methods, including ours, are more memory efficient than pipeline parallelism and Ring Self-Attention, because pipeline parallelism introduce imbalance memory consumption leading to a high pressure on the first stage and Ring Self-Attention is not I/O-aware, consuming much more memory due to materialization of the large softmax matrix. The memory efficiency could be empirically reflected by the longest sequence length, as shown in Table 2 and Table 3, that pipeline parallelism and Ring Self-Attention will support a shorter length than our method.\n\nQ2: Results on larger models or even MoEs.\n\nA2: We thank the reviewer for raising this suggestion. We would love to have evaluated our method on larger models and other architectures if we had enough resources. We are trying to find more resources to supplement this experiment and we will update the new results once we have additional compute.\n\nQ3: How the proposed method can benefit CPU inference.\n\nA3: We thank the reviewer for asking. Our method is implemented for GPU computation, as we assume that most practice of long-context training and inference will be conducted on GPUs, while our method can also be extended to CPU use case as it is a general tiling optimization. Despite our experiments mainly show the speedup for training, we note that our method can accelerate both the forward and backward of distributed FlashAttention, leading to a speedup for both training and inference.", "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": "1. It introduces significant innovations such as load-balanced scheduling, overlapping of communication and computation, and a novel rematerialization aware gradient checkpointing strategy. It critical challenges in training large language models (LLMs) with long sequences, notably improving scalability and efficiency on distributed systems. \n2. 8x longer sequences and up to 5.6x speedups are achieved with the proposed techniques\n3. Code and implementations are available in the open source", "reasons_to_reject": "1. Since the proposed techniques are claimed to be more memory efficient. More direct comparisions with Ringattention and other baselines one the memory consumption would be important.\n2. The evaluation mainly focus on llama 7 b models, more discussions on how the proposed methods can scale to larger size of model or even MoE models would be interesting.\n3. More discussions on how the proposed method can benefit CPU inference.", "questions_to_authors": "", "extra_scores": {"ethics_concerns_details": "no ethics concerns"}}, {"review_id": "wQl7SvSKg0", "reviewer": "Reviewer_EkNX", "summary": "DISTFLASHATTN tackles the issue of load imbalance encountered during the training of causal language models when the input sequence is distributed across multiple worker GPUs. Typically, this distribution results in later GPUs undertaking a disproportionate amount of computation, as they must attend to a larger segment of the input sequence. DISTFLASHATTN addresses this by redistributing some of the computational tasks from later workers to earlier ones, which are underutilized in the initial phases. This strategy enhances overall computational efficiency as evaluated in various settings. This work also introduces a checkpointing trick to avoid waste of computation.", "questions": "", "limitations": "", "rating": 6, "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": "We thank the reviewer EkNX for the questions.\n\nQ1: Lacking discussion of the comparison with RingAttention, DeepSpeed Ulysses, and Striped Attention.\n\nA1: \n- RingAttention: We both distribute FlashAttention along sequence dimension, but our system has two additional innovations: (1) Load-balancing optimization for causal modeling (2x less computation)  (2) Rematerialization-aware gradient checkpointing (1.3x end-to-end speedup).\n\n-  DeepSpeed Ulysses: DeepSpeed Ulysses partitions on the attention heads, and cannot support sequence parallelism degree beyond the head numbers. Meanwhile, we show our method is 1.26-1.88x faster than DeepSpeed Ulysses.\n\n- Striped Attention: Striped Attention is a concurrent work which also found the issue with imbalanced workload. Our method differs from this work by changing the scheduling instead of the order of tokens. \n\nQ2: Two primary techniques appear to be straightforward and may need more novelty.\n\nA2: We agree that the techniques of computation-communication overlapping and load-balancing scheduling are not new. However, we think the application of these ideas in our specific context is novel: (1) identifying the workload imbalance problem in casual language modeling and (2) designing the specific schedule of the asynchronous computation-communication overlap in the distributed attention computation.\n\nIn addition, we indeed made more contributions besides these two. We identified the computation overhead when applying gradient checkpointing with FlashAttention. Our rematerialization-aware gradient checkpointing saves 1 forward computation for free, which could lead to speedup in applications of a wide span of applications. \n\nQ3: Justification for adopting unconventional head counts and widening MLP.\n\nA3: We picked the unconventional head counts to highlight the flexibility of our method and can be applied on any model configuration efficiently while other baselines cannot. Regardless of this setting, our method is still faster than other baselines on more popular setup. \n\nQ4: Comparison with DeepSpeed Ulysses on shorter sequence lengths.\n\nA4: On a shorter 8K sequence length, we obtained 1.12x speedup over Megatron-LM while DeepSpeed obtained 1.23x. DistFlashAtten may not be the best choice for shorter context below 8K tokens, where it is hard to overlap the communication overhead inside computation. However, DistFlashAttn performs better than all baselines starting from 16K, which is the main focus of this paper.", "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": "- Causal training of Transformers holds significant importance. The issue addressed is both practical and valuable. Furthermore, the proposed method demonstrates superior performance in the GQA setting (see Section 4.1, Table 1). Given that GQA is utilized in LLaMa3 training, this methodological improvement is notably beneficial.\n\n- Removing the dependency on the number of attention heads facilitates the utilization of a broader range of devices.", "reasons_to_reject": "-  The section on related work is inadequately written. It lacks comprehensive discussions on directly relevant works like Ring attention and Ulysses. Essential concepts and several critical references are missing e.g. [1].  \n\n- Two primary techniques appear to be straightforward and may lack sufficient novelty.\n\n- I find the experimental setup involving irregular attention heads somewhat untenable for several reasons. First, the use of LLaMa33H appears unrealistic. Although the author cites previous works that employed an irregular number of heads, recent findings by [2] suggest that such configurations can be altered without adversely affecting performance. Consequently, I see no justification for adopting these unconventional head counts. Second, the fourth experimental setting seems far-fetched. According to [3], adjustments in this setting should involve increasing the number of layers rather than merely widening the MLP, as using more layers typically reduces training throughput. However, the author deviates from this approach, opting to increase the MLP width instead. Such comparisons may appear simplistic and are unlikely to be applicable in practical scenarios. \n\n- The evaluation lacks comparisons on moderate sequence lengths against stronger baselines. For instance, it is notable that comparisons involving Ulysses at training lengths of 4K and 8K are absent. Considering the relevance and significance of 8K training in current contexts, its omission is a critical oversight.\n\n\n[1] Striped Attention: Faster Ring Attention for Causal Transformers: https://arxiv.org/abs/2311.09431\n\n[2] The Case for Co-Designing Model Architectures with Hardware https://arxiv.org/abs/2401.14489 \n\n[3] Multi-head or single-head? an empirical comparison for transformer training. arXiv preprint arXiv:2106.09650, 2021.", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "S5ydIM8yOV", "reviewer": "Reviewer_pctD", "summary": "This paper proposes a sequence parallal strategy for training long-context LLMs. By carefully optimizing the overall pipeline, the proposed method achieves relatively better performance than previous sequence parallal strategies such as Magatron-LM and Deepspeed Ulysses. Experiments are performed on language modeling tasks with LLAMA 7B.", "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": "We thank reviewer pctD for the helpful questions.\n\nQ1: Memory comparison with other baseline methods.\n\nA1: We thank the reviewer for raising this question. Our method is memory-efficient mainly because we retained the I/O awareness of FlashAttention when extending it to distributed version. Therefore, we note our method is equivalently memory efficient compared with Megatron-LM(With FlashAttention integration), DeepSpeed-Ulysses(With FlashAttention integration),, and Ring Attention. However, these methods, including ours, are more memory efficient than pipeline parallelism and Ring Self-Attention, because pipeline parallelism introduce imbalance memory consumption leading to a high pressure on the first stage and Ring Self-Attention is not I/O-aware, consuming much more memory due to materialization of the large softmax matrix. The memory efficiency could be empirically reflected by the longest sequence length, as shown in Table 2 and Table 3, that pipeline parallelism and Ring Self-Attention will support a shorter length than our method.\n\nQ2: Accuracy verification of the implementation.\n\nA2: We agree with the reviewer that accuracy verification is necessary. In fact, we have verified the numerical difference between our implementation and the pytorch implementation, which was the same test script used in FlashAttention when they developing the kernel. The test function is included in our submitted codes (`test_op()` function in distFA_async_attn.py). Therefore, we think this proof could provide sufficient guarantee leading to the same accuracy. We note that our method only changes the system implementation for higher efficiency, which does not impact the accuracy and convergence.\n\nQ3: Test with smaller or larger models.\n\nA3: We thank the reviewer for raising this suggestion. We are actively finding more resources to supplement this experiment and we will update the new results once we have additional compute.\n\nQ4: Test with more nodes.\n\nA4: We thank the reviewer for providing the suggestion! Unfortunately, as a research lab, we have very limited access to computational resources, so we cannot find more nodes in a short time to supplement this experiment. However, we will keep finding and update the results once we obtain enough nodes to run.", "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": "At first glance, the paper appears to be well-written. The performance is exceptional, and it addresses the limitation of Megatron-LM sequence parallel, where the number of GPUs should be less than the number of attention heads. Moreover, it has outperformed the speed of ring self-attention.", "reasons_to_reject": "There are several limitations of this paper.\n\n1. The author argues that the proposed method is memory-efficient attention. However, there is no memory comparison in this paper. It shows that it can train 2-8x longer sequences than Ring Self-Attention, which may reflect that it consumes smaller memory. However, when compared with other sequence parallal methods, there is no direct comparison. I would also recommend the author to provide theoritical memory comparison against other methods.\n\n2. It is important to verify the accuracy of distributed algorithm implementations, as there is a risk of bugs causing the optimization to get stuck in local minima. Without providing accuracy verification experiments, it is difficult to justify the correctness of the implementation. Therefore, verifying the accuracy of the implementation is necessary to ensure its reliability.\n\n3. The proposed method only verifies on Llama 7B and its variants. For a general method, it is necessary to test it on other structures or tasks. For example, how will this method perform on bidirectional language modeling?\n\n4. The proposed method only tests on 7B model sizes. It is unclear whether the method will work for smaller or larger model sizes.\n\n5. The conducted experiments were limited to a system with only two nodes. It remains uncertain whether the performance of the system will improve or decline with an increase in the number of nodes, as compared to the existing methods. Further research must be conducted to investigate the impact of scaling on the performance of the system, and to ascertain the optimal configuration for the network.", "questions_to_authors": "As above.", "extra_scores": {}}, {"review_id": "EAJsEfXA0L", "reviewer": "Reviewer_XmqB", "summary": "The recently introduced FlashAttention re-packages the computation of attention ($softmax(QK^\\top)\\cdot V$) so that it reduces the number of I/O operations in a _single_ GPU setup. This led to substantial improvement over the earlier implementations of Attention. However, FlashAttention was only designed for a single GPU case.\n\nThis paper considers the problem of extending FlashAttention to work for the case of a cluster of GPUs. Specifically the paper identifies three bottlenecks in this extension:\n\n1. For the case of causal Attention, one would expect to see a factor $2\\times$ speedup over non-causal Attention (since the former only performs half the number of FLOPs over the latter). However, if one is not careful in how `chunks' of FlashAttention computations are distributed then lot of processors would sit 'idle.' The idea in the paper is to re-distribute the computation of the sub-blocks across processors.\n2. The problem with re-distributing 'chunks' of the Attention computation to various processors is that if not done properly, this could lead to a large communication overhead. Thus, to get towards regaining the $2\\times$ improvement mentioned above, this paper proposes a better scheduling of these partial computations to reduce the communication. Further, the paper shows how to parallelize the communication and computation tasks at each processor to gain further improvement.\n3. To handle really long context length, the now standard practice is to do _gradient checkpointing_, which basically trades memory for re-computation. However, for longer context lengths this recomputation time starts to become a bottleneck. This paper observes that in FlashAttention backward pass recomputation is not needed.\n\nThe paper implements the above into a system it calls DistFlashAttention and compares their implementation to existing distributed Attention implementations including Ring (Self)Attention, Megatron-LM and DeepSpeed-Ulysses. It achieves a minimum speedup of $1.26\\times$ and a max of $5.64\\times$ speedup over various implementations. The speedups for Ring Self-Attention and Megatron-LM also comes along with $2-8\\times$ longer sequence lengths. The paper also presents ablation experiments.\n\nOverall, the paper considers an important problem and proposes new system that achieves speedup over existing systems. However, the current writing leaves certain questions unanswered (please see the Question to Authors section for more). Because of these questions, I am rating this a just above threshold of acceptance. However, if the authors are able to address those questions satisfactorily then this paper _could_ be a clear accept for me.\n\nDetailed Comments for Authors\n-----------------------------\n\nBelow is a list if minor typos/presentation comments:\n\n* [Pg. 2, line -3] \"$5.64\\times$ compared\"-- \"speedup\" is missing.\n* [Pg. 4, line 2] \"which\" -> \"that\"\n* [Pg. 5, 2 lines above eq (3)] The final implementation of DistFlashAttention is Algorithm 2?", "questions": "", "limitations": "", "rating": 6, "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": "We thank the reviewer for acknowledging our contributions. We would like to answer the questions below:\n\nQ1: How were the bottlenecks identified? Clarify which bottlenecks are FlashAttention specific and which ones would apply to any Attention implementation.\n\nA1: We identify these bottlenecks by checking the GPU utilization at each step. To measure the inefficiency, we treat the GPU utilization in the single-GPU setting as an upper bound and compare our GPU utilization in the idstributed setting with it. Based on this metric, we identify the 3 bottlenecks and designed ways to overcome them. The bottleneck of the imbalance workload and the communication overhead are specific to sequence parallelism, and the inefficient gradient checkpointing is specific to FlashAttention in both single-GPU and multi-GPU settings.\n\nQ2: Clarify which ideas appear in other works and which ones are new in this paper.\n\nA2: The idea of sequence parallelism was first introduced by Ring Self-Attention. The extension of FlashAttention with sequence parallelism is new (which is studied in Ring Attention, which is a concurrent work—we developed DistFlashAttn before it was published). In addition, the 3 bottlenecks we identified and solved are new, including imbalance workload (solved with our scheduling algorithm), communication overhead (hidden by asynchronous computation-communication overlap), and inefficient gradient chekcpointing (improved by our rematerialization-aware gradient checkpointing).\n\nQ3: What are the main ideas in Algorithm 2?\n\nA3: The main idea in Algorithm 2 is the load-balancing scheduling added to the naive version of DistFlashAttn, as illustrated in Figure 1. It reduces the time GPUs spend on waiting for others to finish their computation. We thank the reviewer for pointing out its lack of context in the main paper. We will add more description in the main paper.\n\nQ4: What is the max theoretical gians that one can hope from the ideas presented in this paper?\n\nA4: Each component in our algorithm can bring orthogonal benefits that can be combined. Therefore, the speedup can be multiplied. In summary, we could achieve a 2.4x speedup (2x brought by load-balancing scheduling and 1.2 brought by new gradient chekcpointing)  and a support of infinite sequence length (brought by sequence parallelism).\n\nDue to word count limit, we skipped some details. We would love to provide more detailed answers in the second phrase of 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": "* The paper considers an important problem\n* The paper identifies three issues with an obvious distributed implementation of FlashAttention and then proposes fixes for them.\n* The experiments shows that the proposed solutions lead to an overall speedup over existing systems.", "reasons_to_reject": "* Some parts of the paper could do with more explanation\n    * How were the three bottlenecks identified? Are these the three most important bottlenecks? If so, how was this determined? \n    * Also the first two bottlenecks are not specific to FlashAttention but to Attention in general? Either way, this should be clarified.\n    * Some parts of proposed system (e.g. distributing 'shards' of the Attention computation) are not new to the paper (e.g. the sharding is proposed in RingAttention as well-- from what I understand)\n* The paper shunts all discussion of its scheduling ideas to the pseudocode in Algorithm 2. The main part of the paper should at the very least state the main idea(s).\n* It would be helpful to put the improvements in the experiments in context if it would be possible to give a theoretical bound for what is the maximum improvement that can be achieved with the three main ideas. For the causal part, it is clear that handling that would give a $2\\times$ improvement but I do not get a sense of what is the max improvement once can expect from the other ideas. Also do the improvements from the different ideas multiply? I.e. is there reason to believe that the various ideas are not in 'conflict' with each other? Also stating the max theoretical speedups for the ablation experiments would be useful as well.\n\nAs mentioned earlier, none of the above weaknesses are inherent and probably can be fixed by revising the writing.", "questions_to_authors": "1. How were the three bottlenecks identified?\n2. Please clarify which bottlenecks are FlashAttention specific and which ones would apply to any Attention implementation.\n3. Please clarify which of the proposed ideas appear in other works (like RingAttention) and which ones are new to this paper.\n4. What are the main idea(s) in Algorithm 2?\n5. What are the max theoretical gains that one can hope from the ideas presented in this paper?", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper proposes three techniques to do distributed attention on multiple devices, building on FlashAttention and Ring Attention: (1) load balancing in the causal case to avoid devices idling (2) overlapping computation and communication (3) gradient checkpointing that avoids redoing the forward pass of Flash Attention.
These three techniques are evaluated empirically, showing 1.3-2.0x speedup against Megatron-LM with Attention and Ring Attention.
+ The three techniques are conceptually simple and easy to understand
+ The empirical evaluation is comprehensive and convincing
+ Good ablation studies to understand the effects of each of the proposed techniques.
- Techniques (1) and (2) are relatively standard in ML systems. However, demonstrating that these bring speedup to this important problem is still valuable.
Some reviewers have concern over the memory efficiency and scaling of the proposed method. The memory efficiency issue has been addressed in the rebuttal, and scaling to larger model is straightforward.
Decision: Accept