LightSeq: Sequence Level Parallelism for Distributed Training of Long Context Transformers

We thank the reviewer for the insightful comments! We would like to answer your questions and address your concerns in this response:

Q1: Your approach for DistAttn seems to me very similar to what FlashAttention2 does to parallelize on a single GPU. Can you please elaborate a bit more on the algorithmic differences between the two?

A1: In the distributed setting, unbalanced computation and the cross-GPU communication overhead become two challenges that do not exist in the single-GPU case. Therefore, we propose and implement a load-balancing scheduling and an asynchronous communication schema to overcome this. DistAttn without these two optimizations can be understood as a distributed version of FlashAttention2. We extended the FlashAttention2 kernel to support multi-GPU computation by sharding the outer loop onto different devices and communicating the keys and values upon need.

Q2: LightSeq is especially advantageous when the sequences are long and there are few heads or the number of heads is not well divisible by the number of GPUs. Would it be possible/make sense to use tensor-parallelism (in addition to the sequence-parallelism) in the DistAttn to retain the advantages of Megatron-LM and also scale well in scenarios with shorter sequences and many heads?

A2: This is a great suggestion that we do recommend in practice to replace (instead of combining at the same time) TP with DistAttn at long sequences as DistAttn requires less communication and supports scaling beyond the number of heads.

Q3: The 33H case seems rather constructed and unlikely in practice (why would you choose such a number specifically?). It is a valid limitation of parallelization over the heads but should not be that big a problem in actual use, reducing the realistic speedup to ~1.5x (which is still substantial)

A3: Yes, we designed the 33-head model to illustrate the advantages of keeping the distributed training design agnostic to the model configuration. We picked the number 33 for this illustration which may sound unusual. However, we would like to note that it is a common case in practice to have a non-power-of-2 number of heads as we discussed in the second paragraph of the introduction. For example, GPT-2-XL has 25 attention heads, GPT-2 has 12 attention heads, Llama-33B and its fine-tuned versions (e.g., Tulu-30B) have 52 attention heads, Whisper-large has 20 attention heads, and Falcon-7B has 71 attention heads. We would also note that our design can facilitate future model architectures as we discussed in the second paragraph of the introduction:

Moreover, many works have shown that the future Transformer architecture design may have even fewer attention heads. For example, Bian et al. (2021) demonstrate that Transformers with a single head outperforms its multi-head counterparts, representing a challenging scenario for solutions like Megatron-LM.

Q4: One issue is the assumption of causal learning tasks and thus the limitation of the approach to uni-directional Transformer architectures. The approach does not translate to bi-directional models, such as BERT, where long sequences are becoming a problem just as well.

A4: We thank the reviewer for pointing this out. We focus on the optimization for the GPT-based model (specifically, the Llama models) as it’s the most popular model architecture in recent days. The evaluation of bi-directional models is out-of-scope in this paper. However, we respectfully disagree that our approach is therefore only limited to uni-directional Transformer architectures. The DistAttn can be easily applied to bi-directional models by simply disabling the load-balancing scheduler. The load-balancing optimization is proposed to address the bubble caused by causal modeling. But there’s no bubble when handling bi-directional attention, in which case the vanilla design of DistAttn w/o load-balancing scheduling will be good enough.

Dear reviewer xpDg,

We are really encouraged by your helpful and positive feedback, and sincerely appreciate your efforts. We hope our answer, updated experiments and the corresponding manuscript makes the paper clearly and stronger. Please let us know if you have any further comments, and we are more than happy to address them.

Best,

The authors

Q1: Overstatement on proposing sequence parallelism, “why sequence parallelism is proposed in this work, instead of enhanced by this work?”

A1: We didn’t claim we are the first on sequence parallelism. In the related work, we have clearly stated that [1] is among the first to propose sequence parallelism. We also discussed [2] as another previous work.

[1] Shenggui Li, Fuzhao Xue, Yongbin Li, and Yang You. Sequence parallelism: Making 4D parallelism possible. [2] 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.

Q2: Overstatement on flash attention “By default, FlashAttention uses recompute during backward pass. How does this become a contribution of this work?”

A2: We did not claim it is our contribution. Our contribution is a new checkpointing strategy that utilizes this recomputation to avoid another recomputation at the HuggingFace gradient checkpointing level, not the recomputation inside the FlashAttention kernel.

Q3: The optimal way to shift compute: “Besides shifting… kv8.”

A3: Thanks for this important and interesting question. The criteria of optimality is to evenly assign workload to different workers. Assume there are P workers, then the number of units of work is (1+2+...+P) = (1+P) * P / 2. As long as the scheduling finishes in ceil[(1+P) / 2] time steps, it is optimal in terms of compute balancing. Using 8 workers as examples, if the scheduling finishes in ceil[(1+8)/2] = 5 time steps, it is optimal. While there are certainly other schedules to achieve this, we provide one of these.

Q4: How does the DistAtten & load balancing work for MQA models?

A4: DIstAttn distributes the sequence dimension, and the system does not change whether it is a single head (MQA) or multi-heads (MHA).

Q5: Is Figure 3 measured in a distributed setting or on a single GPU?

A5: We thank the reviewer for this question. It is measured in an A100 40GB. We have updated the manuscript to reflect this.

Q6: How does DistAttn and load balance works for sliding local windows or global windows

A6: We thank the reviewer for this important comment. The focus of current work is on exact attention, which is also described in the related work section. Incorporating LightSeq with a sparse attention pattern is a natural next step. In this answer, we give an initial idea on how to extend LightSeq with local and global windows.

For local sliding windows, the workload is naturally (near) balanced, regardless of single directional or bidirectional attention. Thus, simply disregarding the attention logic to non-local workers suffices. For instance, in exact attention, worker 7 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 7 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, it just 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 windows can be interpreted in two ways. We’re not sure which one corresponds to the reviewer’s question so we answer for both. The first interpretation is one special design of the sparse attention (e.g., global attention in Longformer), where 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 DistAttn to this, one easy way is to keep a replica of all the global tokens in each worker, which is simple and practical as otherwise, the global tokens will need to be all-gathered at each time step. One can also split the global tokens evenly onto all workers and use all-gather upon computation to further reduce the memory requirement. The second interpretation is the bidirectional modeling in the full attention situation (e.g., BERT). In this situation, the users can simply disable the load-balancing schema and use the vanilla DistAttn as illustrated in Figure 6 of the updated manuscript. We introduce the load-balancing only to handle the bubble issue caused by the causal modeling, so there’s no need to apply it if it’s bidirectional attention as there is no computation bubble in this case.

Dear reviewer AsB4,

We would like to appreciate again on your precious time and helpful comments. We hope our response and updated paper has adequately addressed your concerns. And we are more than happy to address any of your further comments.

Best,

The authors

We thank the reviewer for the insightful and helpful feedback! We would like to address your questions in the below response.

Q1: The concept of load balancing and overlapping communication and computation in parallel systems is not novel on their own even if the details here might be.

A1: We agree with the reviewers that load balancing and overlapping communication has always been important topics in parallel systems. However, we want to point out LightSeq is the first work that discusses how to design and implement these two techniques in the space of sequence parallelism.

Q2: Details of the specific load balancing technique could be explained more clearly.

A2: We thank the reviewer for raising this clarification question. We added a section in the appendix (Appendix C) to illustrate the load-balancing design. Table 6 shows the computation and communication pattern before balancing, while Table 7 shows the pattern after load-balancing. We will also answer your question in the next paragraph and please feel free to let us know if there is still anything unclear to you.

Each circle is a unit of computation. Circles in the same color mean that they are computed in the same time step. For instance, the rightmost and bottommost circle means that at time step 1 (t1), worker 8 is executing attn(q8, k8, v8). Similarly, green color denotes computations that happen at the second time step (t2). At t2, worker 1 is executing attn(q8, k1, v1). We have updated the figure and the caption to communicate the idea better.

Q3: The relationship between Figure 1 and Figure 2 is unclear - “I found it … on a given timestep”?

A3: Following up on the previous answer, at the first step (blue color), worker 7 is computing the circle corresponding to q7 and kv7, i.e. attn(q7, k7, v7).

Yes, we would like to provide a formula here. (Before rescheduling) the computation a worker w at time t is $attn(q_w, k_{(w-t+1)}, v_{(w-t+1)})$, i.e. in a backward manner - worker 7 computes attn(q7, k7, v7) on time 1, and attn(q7, k6, v6) on time 2. After rescheduling, if a worker w helps other workers (e.g. worker 1, 2, 3 in Figure 1), it is computing $attn(q_{(w+P-t+1)}, k_w, v_w)$, i.e. also in a backward manner - worker 1 computes attn(q8, k1, v1) on time 2, attn(q7, k2, v2) on time 3. For workers who do not help others, the computation remains the same (but the maximal time step gets reduced because of load balancing).

For communication, the worker pre-fetches the values needed for the next time step. Before rescheduling, a worker w at time t is computing $attn(q_w, k_{(w-t+1)}, v_{(w-t+1)})$. Thus, it will prefetch $k_{(w-t+1)}, v_{(w-t+1)}$ at time (t-1). In other words, at time t, worker w prefetches $k_{(w-t)}$ and $v_{(w-t)}$. As an example, worker 7 computes attn(q7,k6, v6) in time 2, so it prefetches k6 and v6 in time 1. With rescheduling, the logic is very similar. If a worker w helps other workers, it will compute $attn(q_{(w+P-t+1)}, k_w, v_w)$ at time t, and thus prefetches $q_{(w+P-t+1)}$ at time t-1. Additionally, it sends the output o back when it is ready. For workers who do not help others, the prefetching logic for key and value is the same.

Q4: The configuration of Table 1: “On page 5 it is mentioned that… and supplementary”.

A4: Table 1 is using the second cluster setup (2 DGX boxes). We found in the previous We have updated the manuscript to indicate that the default setup is the 2 DGX boxes, and thanks again for the helpful feedback.

To test the performance of inter-node training, we mainly report results on DGX boxes with Infiniband. This is because the current implementation of the P2P component is not yet highly optimized (as we brought up in the discussion section),). We are actively working on the engineering side to improve this. To be more specific, in higher bandwidth inter-node training settings, our communication overlapping component results in a very noticeable speedup, and we thus use the 2 DGX boxes as the default testbed. We have also updated the discussion of the manuscript to make our implementation roadmap clearer.

Q5: Speedup against FlashAttention: “What is the speedup.. to that system too”.

A5: LightSeq and Flash attention are orthogonal work, one is distributed and the other is on a single GPU. Conceptually, LightSeq reduces to Flash attention when P=1. When scaling to P=8, we have observed 7.5x speedup compared to flash attention (Figure 5 Left), a near-linear scaling.

Q6: No accuracy report for longer sequences: “The longer … demonstrating this”.

A6: In this work, we focus on the system/infrastructure side of long-context training. However, the gain of longer sequences training has been validated in several previous works, e.g. in [1], [2].

[1] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

[2] Longformer: The Long-Document Transformer

Dear Reviewer ZAxH,

We are deeply encouraged by your helpful and positive feedback, and sincerely appreciate your efforts. We hope our answer and correspondingly updated manuscript makes the paper clearly and stronger. Please let us know if you have any further comments, and we are more than happy to address them.

Best,

The authors

We thank the reviewer for the insightful and helpful feedback! We would like to address your questions in the below response.

Q1: How does the proposed work compare, both qualitatively and quantitatively, to this 4D sequence parallelism work?

A1: We discussed 4D sequence parallelism in the related work section:

“[1] is among the first to parallelize along the sequence dimension. However, it is not optimized for the computational pattern of causal language modeling and is incompatible with memory-efficient attention, which is crucial to long-context LLM training. ”

In other words, the 4D sequence parallelism work is the first to propose this concept but it’s not optimized for modern LLM training and, thus is unable to support very long sequences efficiently. Quantitatively, LightSeq supports at least 8x longer sequences than this work, and achieves 4.45x - 5.64x speedup. We have added this experiment results and further analysis in the updated paper (appendix: Comparison with Ring Self-Attention). We provided the main experiment numbers below.

Maximal sequence length on Llama-7B on the DGX (A100-80GB) cluster

Per iteration time comparison with RSA (seconds) on the DGX clusters

Q2: The parallelism configurations are unclear. What is the batch size used for training? If the batch size is high enough, wouldn't a combination of tensor/data/pipeline parallelism provide sufficient parallelism to achieve high utilization of GPUs, even without utilizing sequence parallelism? In Table 3, the authors used a batch size of 1, which does not seem to be a reasonable configuration for training. Can you elaborate on this choice?

A2: When combined with tensor/data/pipeline parallelism, LightSeq sequence parallelism can be thought of as a better parallelism approach replacing TP to support long context. To see this, we note the LLM training memory utilization can be decomposed into two parts. The first part is related to the model parameters, i.e., memory used to store parameters, gradients, and optimizer states. The second part is related to the activation. It is worth noting that the memory for activation dominates when increasing the context length and the batch size only affects the memory for activation. Given this premise, we know that, firstly, data parallelism wouldn’t reduce either part, so we could consider it as orthogonal and only discuss the DP=1 case. Secondly, in Table 2, we show that the combination of model parallelism and pipeline parallelism handles shorter sequence lengths in a variety of scenarios, even when batch_size=1. This indicates that our sequence parallelism can support a very long context, as it would avoid TP+PP running out of memory.. In addition, Table 2 shows that using TP alone can effectively support the same length as LightSeq, but it requires more communication and takes longer than LightSeq as shown in Table 1. Increasing the batch size has the same effect on TP and LightSeq sequence parallelism (LightSeq is used with FSDP in experiments). Therefore, LightSeq should be a better choice than TP in supporting long context lengths.

Q3: Regarding the explanation of pipeline parallelism, the authors attribute the low performance to “high memory pressure in the first stage.” Couldn't this be due to poor partitioning between stages?

A3: We believe optimizing pipeline partitioning may not solve the problem. We follow the standard practice in Megatron-LM to evenly partition the layers onto different devices and utilize micro-batches to reduce the bubble size. This partitioning strategy prefers to have an even workload across stages. If we partition it into uneven stages to balance the memory, it will cause unbalanced computation (stragglers), leading to inefficiency.

[1] Shenggui Li, Fuzhao Xue, Yongbin Li, and Yang You. Sequence parallelism: Making 4d parallelism possible. arXiv preprint arXiv:2105.13120, 2021.

Dear reviewer Fg9E,

Once again we sincerely thank your precious time and helpful comments. We hope our response and updated paper has adequately addressed your concerns. And we are more than happy to address any of your further comments.

Best,

The authors