Paper: BurstAttention: An Efficient Distributed Attention Framework for Extremely Long Sequences (ICLR-2024)

Reviewer comment (Experiments):
The experiments are conducted on a single hardware setup, providing experiments on different configurations would help.

Author response:
Frankly, the current experiments are conducted on a specific hardware setup due to the limited resources we can use for this paper. Inspired by your suggestions, we have recognized the value of evaluating our method across different configurations. 

However, because of some force majeure, we may need time to prepare other hardware configurations (e.g., 4090 and H100). If these hardware configurations are met, our future research plans include conducting experiments on clusters and devices with more extreme lower/higher communication constraints. We believe this will provide us with a deeper understanding of how BurstAttention performs under varied conditions and will help us in tailoring our attention algorithms and parallelism strategies to more specialized scenarios.

Reviewer comment (Reproducibility):
Which PCIe is used in the experiments and what is the bandwidth?

Author response:
Regarding your first question, our experiments utilize PCIe 4.0, and we do not enable CUDA direct access between GPUs for experiments. The reason for this setting can refer to our additional experiments for reviewer Gb5f & J3Ya. 

Concerning your second question, our current experimental results indicate a clear advantage of our approach in scenarios with limited communication capabilities. Since one great advantage of BurstAttention is its low communication overhead and highly overlapping communication and computation operations, the benefits of BurstAttention are less pronounced in environments with high communication bandwidth (e.g., using NVLink and InfiniBand). 

Although enabling NVLink and InfiniBand is expensive and such device conditions are uncommon for LLM inference, we have recognized the potential for further optimization when communication bandwidth is higher. We are exploring various communication scenarios to better understand and optimize the performance of BurstAttention under different communication conditions. Thanks again for your insightful suggestion.

Reviewer comment (Evaluation):
The latency comparison against Megatron is not well analyzed, see questions below.

Author response:
Besides the results of training LLMs, we also show the results of inferring LLMs. In terms of inferring LLMs, BurstAttention still achieves a speed improvement than Megatron V3. Note that although TP (Megatron V3) is more memory efficient than TP (Megatron V1), the all-reduce operation used by TP (Megatron V1) is better optimized than the reduce-scatter and all-gather operations used by TP (Megatron V3). In the actual inference process, TP (Megatron V1) is even slightly faster than TP (Megatron V3).


#### LLaMA-7b Inference First Token Latency (seconds) on 8xA100 GPUs
| Sequence Length                 | 4096      | 8192      | 16384     | 32768     | 65536      | 131072     | 262144     |
|:-----------------------|:----------|:----------|:----------|:----------|:-----------|:-----------|:-----------|
| Ring Self-Attention    | 0.42±0.01 | 0.87±0.01 | 2.00±0.01 | 5.13±0.05 | OOM        | OOM        | OOM        |
| TP (Megatron V1) w/ Flash   | 0.67±0.01 | 1.29±0.01 | 2.58±0.01 | 5.27±0.01 | 11.63±0.02 | 27.54±0.01 | 71.52±0.06 |
| TP (Megatron V3) w/ Flash   | 0.73±0.02 | 1.36±0.01 | 2.68±0.01 | 5.67±0.01 | 12.25±0.01 | 28.73±0.03 | 75.52±0.05 |
| BurstAttention w/o LAO | 0.46±0.01 | 0.88±0.01 | 1.79±0.01 | 3.88±0.01 | 10.78±0.01 | OOM        | OOM        |
| BurstAttention         | 0.44±0.01 | 0.84±0.01 | 1.68±0.01 | 3.27±0.01 | 6.49±0.01  | 16.01±0.01 | 49.32±0.11 |

#### LLama-13b Inference First Token Latency (seconds) on 8xA100 GPUs
| Sequence Length                 | 4096      | 8192      | 16384     | 32768     | 65536      | 131072     | 262144      |
|:-----------------------|:----------|:----------|:----------|:----------|:-----------|:-----------|:------------|
| Ring Self-Attention    | 0.66±0.01 | 1.36±0.01 | 3.08±0.01 | 7.98±0.02 | OOM        | OOM        | OOM         |
| TP (Megatron V1) w/ Flash   | 1.05±0.01 | 2.01±0.01 | 4.03±0.01 | 8.41±0.01 | 18.56±0.02 | 44.39±0.04 | OOM         |
| TP (Megatron V3) w/ Flash   | 1.07±0.01 | 2.09±0.01 | 4.20±0.01 | 8.76±0.01 | 19.06±0.06 | 45.46±0.03 | 119.03±0.04 |
| BurstAttention w/o LAO | 0.72±0.01 | 1.39±0.01 | 2.77±0.05 | 5.99±0.01 | 16.95±0.01 | OOM        | OOM         |
| BurstAttention         | 0.69±0.01 | 1.40±0.05 | 2.57±0.03 | 5.08±0.02 | 9.92±0.01  | 25.91±0.01 | 78.80±0.07  |

Reviewer comment (Experiments):
The argument with sparse attention is vague. There are no experiments.

Author response:
We acknowledge your observation that the sparse attention aspect of our work may appear vague. Since implementing sparse attention based on BurstAttention just requires changing the attention masking matrix, we thus only briefly introduce the compatibility of BurstAttention with sparse attention.

This question is very insightful, especially as it highly relates to the workload imbalance issue. We are still working on addressing the workload imbalance issue associated with causal language modeling, and the imbalance issue is more serious for sparse attention. We are developing more refined strategies to effectively manage this challenge, providing greater clarity and depth in our future iterations. 

Thanks again for your suggestion, and we will add more discussions about the workload imbalance issue and sparse attention in our revision, which will make our work more solid.

Reviewer comment (Experiments):
In the experiment, the objective is causal language modeling, where the system seems to have an unbalanced workload, e.g., the machine hosting the first query will be idle most of the time. However, in Megatron, the workload is balanced because they do not partition sequences. Does the system balance the workload? If not, the system should be 2x slower than Megatron due to the unbalanced workload.

Author response:
In the context of multi-device distributed training, we have overlapped communication and computation. This overlapping is especially pertinent in communication-constrained scenarios, which are our primary focus. In such scenarios, even though BurstAttention has an unbalanced computation workload, all the computation is overlapped in communication, making communication the bottleneck. In contrast, although Megatron has a balanced computation workload, their computation needs to be conducted after communication. Given that BurstAttention has a much lower communication volume compared to Megatron TP, it is more practical in scenarios with extremely long sequence lengths or in communication-constrained environments.

On the other hand, to achieve load balancing, we can adjust the sequence sharding strategy using the attention head dimension. For example, in a 2-GPU scenario, instead of letting the first GPU maintain vectors of [0..B, 0..H, 0..(L/2), 0..D] and letting the second GPU maintain vectors of [0..B, 0..H, (L/2)..L, 0..D], we can let the first GPU maintain [0..B, 0..(H/2), 0..(L/2), 0..D] and [0..B, (H/2)..H, (L/2)..L, 0..D], and let the second GPU maintain [0..B, 0..(H/2), (L/2)..L, 0..D] and [0..B, (H/2)..H, 0..(L/2), 0..D]. Through this cyclic shifting, load balancing across multiple devices in distributed training can be achieved. This will be left for our future work.

Reviewer comment (Novelty):
While integrating two prior methods (RingAttention and FlashAttention) to be more efficient by itself could be a contribution, the paper could elaborate more on the new challenges, if any, and different method designs to establish its new contributions.

Author response:
Our main contribution is to build a unified distributed framework for extremely-long sequence parallelism. Under the unified framework, by comprehensively considering the characteristics of clusters and devices, we propose local attention optimization to reduce memory overhead, global attention optimization to reduce communication overhead, and further overlap of global operations and local operations to minimize latency. Rather than an integration of global attention operations and local attention operations, BurstAttention is a significant reimagining and enhancement of how devices interact and function together in a distributed system to perform attention for long sequences. 

Although we mentioned in the paper that RingAttention and FlashAttention can be seen as the degradation of BurstAttention in specific scenarios, this cannot simply be seen as BurstAttention being a combination of global RingAttention and local FlashAttention. In fact, since RingAttention will first compute all local attention scores and then conduct softmax operations across machines to obtain global results, this attention mode makes it difficult to cooperate with the online softmax operations of FlashAttention. This means that combining RingAttention and FlashAttention is non-trivial. Even if RingAttention and FlashAttention are combined, this simple combination cannot fully utilize the cluster and device characteristics for performance optimization.

In summary, the unified framework and the derived optimization strategies are central to our contribution, enabling us to specifically optimize the ring-collective operations, fully overlap computation and communication, and process much longer sequences more efficiently than existing baselines. We will add more discussion of this in our revision.

Reviewer comment (Novelty):
The local attention optimization does not seem to bring any new insights than using FlashAttention in a distributed system.

Author response:
Our main contribution is to build a unified distributed framework for extremely-long sequence parallelism. Under the unified framework, by comprehensively considering the characteristics of clusters and devices, we propose local attention optimization to reduce memory overhead, global attention optimization to reduce communication overhead, and further overlap of global operations and local operations to minimize latency. Rather than an integration of global attention operations and local attention operations, BurstAttention is a significant reimagining and enhancement of how devices interact and function together in a distributed system to perform attention for long sequences. 

Although we mentioned in the paper that RingAttention and FlashAttention can be seen as the degradation of BurstAttention in specific scenarios, this cannot simply be seen as BurstAttention being a combination of global RingAttention and local FlashAttention. In fact, since RingAttention will first compute all local attention scores and then conduct softmax operations across machines to obtain global results, this attention mode makes it difficult to cooperate with the online softmax operations of FlashAttention. This means that combining RingAttention and FlashAttention is non-trivial. Even if RingAttention and FlashAttention are combined, this simple combination cannot fully utilize the cluster and device characteristics for performance optimization.

In summary, the unified framework and the derived optimization strategies are central to our contribution, enabling us to specifically optimize the ring-collective operations, fully overlap computation and communication, and process much longer sequences more efficiently than existing baselines. We will add more discussion of this in our revision.

Reviewer comment (Novelty):
What is the new contribution over FlashAttention? Or simply integrating it into the BurstAttention framework?

Author response:
Our main contribution is to build a unified distributed framework for extremely-long sequence parallelism. Under the unified framework, by comprehensively considering the characteristics of clusters and devices, we propose local attention optimization to reduce memory overhead, global attention optimization to reduce communication overhead, and further overlap of global operations and local operations to minimize latency. Rather than an integration of global attention operations and local attention operations, BurstAttention is a significant reimagining and enhancement of how devices interact and function together in a distributed system to perform attention for long sequences. 

Although we mentioned in the paper that RingAttention and FlashAttention can be seen as the degradation of BurstAttention in specific scenarios, this cannot simply be seen as BurstAttention being a combination of global RingAttention and local FlashAttention. In fact, since RingAttention will first compute all local attention scores and then conduct softmax operations across machines to obtain global results, this attention mode makes it difficult to cooperate with the online softmax operations of FlashAttention. This means that combining RingAttention and FlashAttention is non-trivial. Even if RingAttention and FlashAttention are combined, this simple combination cannot fully utilize the cluster and device characteristics for performance optimization.

In summary, the unified framework and the derived optimization strategies are central to our contribution, enabling us to specifically optimize the ring-collective operations, fully overlap computation and communication, and process much longer sequences more efficiently than existing baselines. We will add more discussion of this in our revision.

Reviewer comment (Experiments):
The paper only evaluates two decoder-only models at inference time, and a single small model at training time.

Author response:
Thank you for your insights. Indeed, our current hardware resources are limited, which constrains our ability to conduct experiments on larger models. So far, we have managed to train the LLaMA-70b model using a 32xA100 GPU cluster, with enable checkpointing activation. We plan to provide detailed results of the LLaMA-70b experiments in the future as our resources allow.

Reviewer comment (Experiments):
The scaling law along the sequence dimension is present but a scaling law along the model capacity dimension is also important given that the dimensionality of the Q/K/V will also increase with model size and therefore the communication and memory overhead as well.

Author response:
We acknowledge the importance of evaluating the scaling law along the model capacity dimension. In this paper, our analysis primarily focuses on the O(N^2) attention complexity and its optimization through BurstAttention. To clearly evaluate the effectiveness and efficiency of different methods to handle extremely long attention, we choose a fixed model size for our experiments, which can bring an optimization perspective that is independent of the trade-offs introduced by parallel methods and optimizations.

We understand that experiments on larger models would require additional optimization strategies, and the performance bottlenecks may vary based on different hardware configurations. Therefore, simply comparing the performance on larger models and clusters without these considerations may not yield clear insights. We agree this suggestion is meaningful and plan to extend our analysis to larger models in future work, taking into account these complex optimization and device factors.

Reviewer comment (Reproducibility):
For the tensor parallelism/TP+Flash Attention baselines, is the implementation being used the Megatron-LM implementation which parallelizes each head across devices (and therefore reduces the memory footprint/communication)?

Author response:
Yes, we use the tensor parallelism implementation of Megatron for our experiments, which parallelizes attention heads across devices to reduce the memory footprint. However, even with head-wise parallelization, Megatron V1's tensor parallelism still exhibits substantial memory usage. Notably, the communication volume in Megatron V1's attention mechanism remains significant, particularly the all-reduce operation to obtain attention outputs. Hence, the communication and memory overheads are not notably reduced by parallelizing along the attention head dimension. More importantly, the increase of memory overhead brought by parameter growth is linear, while the rise brought by sequence length growth is quadratic, i.e., the core memory overhead is to store attention intermediate results rather than storing parameters in processing extremely long sequences, which cannot be solved by only parallelizing attention heads.

In supplementary experiments, we compare our method with the tensor parallelism approach in Megatron V3, which employs all-gather and reduce-scatter operations instead of all-reduce operations in Megatron V1. For Megatron V3's approach, tensor parallelism's communication volume remains akin to that of Megatron V1, but the memory footprint is significantly reduced. These supplementary results can refer to the responses to Reviewer Gb5f and Reviewer J3Ya.

Reviewer comment (Experiments):
For the time-to-first-token experiment is the sequence length equivalent to the prompt input length e.g. the model generates the first token conditioned on 4096/8192/.../262144 tokens?

Author response:
Regarding your question on the first token latency, you are correct. This metric indeed refers to the time taken for the first inference step, specifically the encoding time. It is an essential measure since the computation for long sequence attention primarily occurs during the encoding phase in inference. In real-time AI services such as ChatGPT, the system's responsiveness significantly impacts the user experience, and these applications usually output results in a streaming manner to improve responsiveness. Since the first token's latency is the longest, reducing the first token's latency is thus a crucial performance metric, as it directly influences the perceived responsiveness and efficiency of the model in these streaming scenarios. [1] and [2] also explain why the first token latency is an important metric for LLM serving, which can provide more details.
   
[1] https://www.databricks.com/blog/llm-inference-performance-engineering-best-practices  
[2] https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/appnotes/transformers-neuronx/generative-llm-inference-with-neuron.html

Reviewer comment (Novelty):
The novelty of the paper is unclear. It seems to me that it is essentially a straightforward combination of three existing ideas: RingAttention (Li et al., 2022); online softmax (Milakov & Gimelshein, 2018); and FlashAttention (Dao et al., 2022). Further, the ring-based approach to attention seems essentially to be a variant on classic 1D distributed matrix-matrix multiplication algorithms or Cannon's algorithm with one dimension of the mesh trivial (see, e.g., Golub's "Matrix Computations" book, or the Cannon or Fox algorithm). One could imagine pipelining the distributed attention computation (it is unclear to me whether this paper does this), but then it still seems very similar to SUMMA (Van De Geijn & Watts, 1997).

Author response:
### 1. **For the Novelty**

Our main contribution is to build a unified distributed framework for extremely-long sequence. Under the unified framework, by comprehensively considering the characteristics of clusters and devices, we propose local attention optimization to reduce memory overhead, global attention optimization to reduce communication overhead, and further overlap of global operations and local operations to minimize latency. Rather than an integration of global attention operations and local attention operations, BurstAttention is a significant reimagining and enhancement of how devices interact and function together in a distributed system to perform attention for long sequences. 

Although we mentioned in the paper that RingAttention and FlashAttention can be seen as the degradation of BurstAttention in specific scenarios, this cannot simply be seen as BurstAttention being a combination of global RingAttention and local FlashAttention. In fact, since RingAttention will first compute all local attention scores and then conduct softmax operations across machines to obtain global results, this attention mode makes it difficult to cooperate with the online softmax operations of FlashAttention. This means that combining RingAttention and FlashAttention is non-trivial. Even if RingAttention and FlashAttention are combined, this simple combination cannot fully utilize the cluster and device characteristics for performance optimization.

In summary, the unified framework and the derived optimization strategies are central to our contribution, enabling us to specifically optimize the ring-collective operations, fully overlap computation and communication, and process much longer sequences more efficiently than existing baselines. We will add more discussion of this in our revision.

Reviewer comment (Theory):
The analysis section seems quite perfunctory. It adequately describes the characteristics of the particular implementation, but does not address optimality. E.g., there are lower bounds for distributed matrix-matrix multiplication (via pebbling and in communication-avoiding algorithms) that could be used. Similarly, FlashAttention includes a lower-bound analysis that might be leveraged here.

Author response:
### 3. **For the Analysis Section**

Thank you for your insightful suggestion. In our paper, since we mainly focus on the practical performance of BurstAttention, we thus briefly give the communication and memory complexities of different distributed methods and then compare the end-to-end results of these methods. We highly agree that contextualizing our analysis with lower bounds or comparative benchmarks would provide a clearer understanding of BurstAttention's efficiency relative to theoretical optima.

In recent days, we have added more detailed experiments based on reviewer suggestions. In the future, we will further enrich our analysis to include a more detailed examination of the efficiency of BurstAttention, providing insights from a theoretical perspective as well as comparisons with existing benchmarks. These more detailed analyses will be included in our further revisions.


We are grateful for the opportunity to address these points and hope that our responses clarify the choices made in our research. We remain open to further suggestions and are committed to continual improvement of our work.

Reviewer comment (Evaluation):
The experimental evaluation does not include any error bars or measures of variance. Especially since this includes cross-device communication, and BurstAttention introduces additional communication on top of existing methods (e.g., tensor or data parallelism) which may interfere, these are important to include.

Author response:
#### For Error Bars

Thank you for your insightful suggestion. Our experiments in the paper are the average of multiple experiments, yet we neglect to give error bars as the results of our multiple experiments are very close. We have followed your suggestion to add error bars for both existing experiments and new experiments. Thanks again for your suggestion.

Reviewer comment (Experiments):
The evaluation is conducted only up to 8 GPUs which are interconnected via PCIe. This is not especially large scale for training LLMs (although could be reasonable for fine-tuning), and it is not clear how the method will scale.

Author response:
#### For Cluster Scale

We acknowledge your observation regarding the limited scale of the cluster used in our experiments. We agree that evaluating our BurstAttention framework on a larger-scale cluster would be beneficial for further validating its scalability and efficiency with longer sequences. The initial experiments were conducted on a modestly sized cluster to demonstrate the feasibility and efficacy of our approach.

Our expanded experiments on this 32-GPU setup, with the training length extended to 128k, have demonstrated notable improvements in speed over the tensor parallel approach, especially under these specific network and hardware configurations. Even models with hundreds of billions of parameters like GPT-3 rarely use more than 8 GPUs for inference, so this 32-GPU experiments are only conducted on training. These results underscore the robustness and scalability of our method, affirming its effectiveness in a real-world, high-performance computing environment.

Reviewer comment (Novelty):
Please clarify the novelty of the paper (see above).

Author response:
### 1. **For the Novelty**

Our main contribution is to build a unified distributed framework for extremely-long sequence. Under the unified framework, by comprehensively considering the characteristics of clusters and devices, we propose local attention optimization to reduce memory overhead, global attention optimization to reduce communication overhead, and further overlap of global operations and local operations to minimize latency. Rather than an integration of global attention operations and local attention operations, BurstAttention is a significant reimagining and enhancement of how devices interact and function together in a distributed system to perform attention for long sequences. 

Although we mentioned in the paper that RingAttention and FlashAttention can be seen as the degradation of BurstAttention in specific scenarios, this cannot simply be seen as BurstAttention being a combination of global RingAttention and local FlashAttention. In fact, since RingAttention will first compute all local attention scores and then conduct softmax operations across machines to obtain global results, this attention mode makes it difficult to cooperate with the online softmax operations of FlashAttention. This means that combining RingAttention and FlashAttention is non-trivial. Even if RingAttention and FlashAttention are combined, this simple combination cannot fully utilize the cluster and device characteristics for performance optimization.

In summary, the unified framework and the derived optimization strategies are central to our contribution, enabling us to specifically optimize the ring-collective operations, fully overlap computation and communication, and process much longer sequences more efficiently than existing baselines. We will add more discussion of this in our revision.

Reviewer comment (Theory):
Can you contextualize the analysis in Section 4 with lower bounds or otherwise provide an indication of how far BurstAttention is from optimal?

Author response:
### 3. **For the Analysis Section**

Thank you for your insightful suggestion. In our paper, since we mainly focus on the practical performance of BurstAttention, we thus briefly give the communication and memory complexities of different distributed methods and then compare the end-to-end results of these methods. We highly agree that contextualizing our analysis with lower bounds or comparative benchmarks would provide a clearer understanding of BurstAttention's efficiency relative to theoretical optima.

In recent days, we have added more detailed experiments based on reviewer suggestions. In the future, we will further enrich our analysis to include a more detailed examination of the efficiency of BurstAttention, providing insights from a theoretical perspective as well as comparisons with existing benchmarks. These more detailed analyses will be included in our further revisions.


We are grateful for the opportunity to address these points and hope that our responses clarify the choices made in our research. We remain open to further suggestions and are committed to continual improvement of our work.

Reviewer comment (Evaluation):
Please add error bars (or similar) to all experimental results.

Author response:
#### For Error Bars

Thank you for your insightful suggestion. Our experiments in the paper are the average of multiple experiments, yet we neglect to give error bars as the results of our multiple experiments are very close. We have followed your suggestion to add error bars for both existing experiments and new experiments. Thanks again for your suggestion.

Reviewer comment (Evaluation):
Please expand the experimental analysis; see above for suggestions.

Author response:
#### The Detailed Analysis of Experiments Results
We sincerely appreciate your valuable suggestion regarding the need for more detailed experimental analysis. We understand that fine-grained analysis can offer deeper insights into the practical bottlenecks of BurstAttention.

However, in the context of multi-device distributed training, we have overlapped communication and computation. This overlapping is especially pertinent in communication-constrained scenarios, which are our primary focus. In such scenarios, computation are often subsumed within communication times, rendering traditional metrics like sustained flop percentage less applicable. Furthermore, our algorithm includes non-matmul FLOPs, which further complicates the use of flop-based metrics for comprehensive performance assessment.

Given these considerations, we opted for end-to-end results as our primary performance indicator. This metric directly reflects the extent of optimization achieved by BurstAttention.

Nevertheless, we understand that fine-grained would help to explain where the bottleneck is and how BurstAttention sloved it. Moving forward, we plan to conduct additional experiments on a variety of cluster configurations and attempt to pinpoint the optimal scenarios for BurstAttention's deployment. This will involve investigating how different input shapes, GPU performances, and communication bandwidths impact the effectiveness of BurstAttention compared to other algorithms. In the course of these experiments, we will endeavor to include more detailed profiling of BurstAttention to provide a clearer understanding of its performance characteristics.

Reviewer comment (Experiments):
How does BurstAttention scale to larger numbers of GPUs or on different networks?

Author response:
### 2.3 The Results of Scaling Experiments

Some reviewers suggest us to provide some experimental results of scaling GPU number, sequence length, and batch size. We list the results of scaling experiments in the tables below.

#### LLama-13b Inference First Token Latency (seconds) with Scaling the GPU Number and Sequence Length.

| Setting                   | 1xGPUs SeqLen 16384   | 1xGPUs SeqLen 32768   | 1xGPUs SeqLen 65536   | 2xGPUs SeqLen 16384   | 2xGPUs SeqLen 32768   | 2xGPUs SeqLen 65536   | 4xGPUs SeqLen 16384   | 4xGPUs SeqLen 32768   | 4xGPUs SeqLen 65536   | 8xGPUs SeqLen 16384   | 8xGPUs SeqLen 32768   | 8xGPUs SeqLen 65536   |
|:-------------------------|:----------------------|:----------------------|:----------------------|:----------------------|:----------------------|:----------------------|:----------------------|:----------------------|:----------------------|:----------------------|:----------------------|:----------------------|
| TP(Megatron V3) w/ Flash | 3.70±0.01             | 10.90±0.01            | 35.52±0.01            | 4.01±0.01             | 9.85±0.01             | 26.45±0.01            | 3.37±0.01             | 7.52±0.01             | 18.36±0.01            | 4.13±0.01             | 8.88±0.01             | 19.08±0.01            |
| BurstAttention           | 3.86±0.01             | 11.23±0.01            | 36.35±0.02            | 3.52±0.01             | 8.93±0.01             | 24.83±0.01            | 2.66±0.01             | 6.06±0.01             | 15.77±0.01            | 2.73±0.01             | 5.77±0.01             | 13.05±0.01            |

Reviewer comment (Experiments):
How does BurstAttention perform when combined with other standard parallelism techniques for training, e.g., data-parallelism/FSDP, pipelining, etc.? Does its additional communication cause interference?

Author response:
#### For the Combination of BurstAttention and Data Parallelism/FSDP

Since our framework is for both training and inference, and data/pipeline parallelism methods are seldomly adopted for inference, we mainly focus on comparing our framework with existing sequence and tensor parallelism methods. In our new experiments, we combine BurstAttention and ZeRO together to show the effect of combing BurstAttention and Data Parallelism/FSDP, and these results can refer to 2.1. 

Note that, combing data/pipeline methods and tensor/sequence parallelism methods only introduces additional communication costs, and has no side effects on our analysis on attention with extremely long sequences.

Reviewer comment (Experiments):
In Section 5.3, why fix the batch size to 1 for training? This seems quite small.

Author response:
#### For the Choice of Batch Size and Sequence Length

For training LLMs, the total tokens of a batch are required to be between 2M and 4M, otherwise, the model performance may be degraded. This means that the longer the training sequence length is, the smaller the batch size is. Due to this, several GPUs may need to compute one example. For example, using 2048 GPUs to train 128-layer GPT-3, the sequence length is 4096, the batch size is 1024, data parallelism is 16, pipeline parallelism is 32, and tensor parallelism is 4. In this scenario, the optimal setup is to divide a batch into 64 micro-batches with a micro-batch size of 1. In this case, four GPUs under the tensor parallelism group are inevitably required to calculate one piece of data together. In view of this, in our paper, we use batch size = 1 for our experiments. 

To make our experimental results more sufficient, we follow the reviewers' suggestions and add more experiments with scaling batch and sequence sizes. These results are shown in 2.3.

Reviewer comment (Experiments):
In Section 5.4, why is a different sequence length used for scaling than is presented in the earlier results? This makes it hard to contextualize the scaling.

Author response:
#### For the Choice of Batch Size and Sequence Length

For training LLMs, the total tokens of a batch are required to be between 2M and 4M, otherwise, the model performance may be degraded. This means that the longer the training sequence length is, the smaller the batch size is. Due to this, several GPUs may need to compute one example. For example, using 2048 GPUs to train 128-layer GPT-3, the sequence length is 4096, the batch size is 1024, data parallelism is 16, pipeline parallelism is 32, and tensor parallelism is 4. In this scenario, the optimal setup is to divide a batch into 64 micro-batches with a micro-batch size of 1. In this case, four GPUs under the tensor parallelism group are inevitably required to calculate one piece of data together. In view of this, in our paper, we use batch size = 1 for our experiments. 

To make our experimental results more sufficient, we follow the reviewers' suggestions and add more experiments with scaling batch and sequence sizes. These results are shown in 2.3.

Reviewer comment (Novelty):
In light of the authors' response addressing some of my concerns, I have raised my score. However, I still remain concerned about the novelty and depth of the experiments.

Author response:
### 1. **For the Novelty**

Our main contribution is to build a unified distributed framework for extremely-long sequence. Under the unified framework, by comprehensively considering the characteristics of clusters and devices, we propose local attention optimization to reduce memory overhead, global attention optimization to reduce communication overhead, and further overlap of global operations and local operations to minimize latency. Rather than an integration of global attention operations and local attention operations, BurstAttention is a significant reimagining and enhancement of how devices interact and function together in a distributed system to perform attention for long sequences. 

Although we mentioned in the paper that RingAttention and FlashAttention can be seen as the degradation of BurstAttention in specific scenarios, this cannot simply be seen as BurstAttention being a combination of global RingAttention and local FlashAttention. In fact, since RingAttention will first compute all local attention scores and then conduct softmax operations across machines to obtain global results, this attention mode makes it difficult to cooperate with the online softmax operations of FlashAttention. This means that combining RingAttention and FlashAttention is non-trivial. Even if RingAttention and FlashAttention are combined, this simple combination cannot fully utilize the cluster and device characteristics for performance optimization.

In summary, the unified framework and the derived optimization strategies are central to our contribution, enabling us to specifically optimize the ring-collective operations, fully overlap computation and communication, and process much longer sequences more efficiently than existing baselines. We will add more discussion of this in our revision.

### 2. **For More Experiments and Analysis**

Thanks for your insightful suggestion, we have added some new experiments to make our claim more clear. And we also take the suggestion of Reviewer J3Ya and conduct new experiments for comparison with Megatron V3. For the new experiments, we replace the original tensor parallelism with the implementation of Megatron V3, an implementation of tensor parallelism that considers the features of sequence parallelism. These results can refer to 2.1. 
 
Furthermore, we have expand our training experiments on 32 A100 GPUs. All these experiments are conducted on a cluster interconnected by a RoCE (RDMA over Converged Ethernet) network with an approximate bandwidth of 200 Gb/s, and each machine of the cluster has eight A100 GPUs connected via PCIe 4.0 with peer-to-peer (P2P) communication disabled. The results of the new experiments are shown in the tables below.