Thank you for your constructive feedback on our submission. We appreciate your recognition of the strengths of our work and the areas where improvements can be made. We have carefully considered your comments and would like to address your concerns and questions as follows.

1. For Data/Pipeline Parallelism Baselines:

We acknowledge the absence of comparisons with common parallelism approaches like data parallelism and pipeline parallelism in our current studies. Our decision to omit these from the initial set of experiments is due to inherent limitations with sequence lengths. From our pilot experiments, directly adopting data parallelism and pipeline parallelism can only handle those sequences shorter than 8192.

We understand the importance of these popular approaches in the field and are planning to conduct further experiments. We aim to explore the applicability of sequence parallelism in various scenarios and compare the differences between different parallelism strategies. This will be a significant aspect of our ongoing research. In our additional experiments for reviewer Gb5f & J3Ya, we have explored combining ZeRO (one optimization for data parallelism) and BurstAttention, and this combination has shown some promising results to trade off runtime and memory efficiency.

2. Experiments on More Hardware Configurations:

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.

3. Questions Regarding PCIe and Communication Bandwidth:

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.

We hope these responses address your concerns and questions. We are committed to continuously improving our work and value your input as it guides our future research directions.

For training, the results of the LLaMA-7b model training experiments indicate that BurstAttention has higher training throughput compared to the tensor parallelism of Megatron V3 with FlashAttention, especially as the batch size increases. This points to BurstAttention's capability to process larger batches more efficiently, which is crucial for training LLMs effectively.

While BurstAttention exhibits impressive performance in various aspects, it's important to acknowledge certain limitations, as highlighted by experimental results. Specifically, it has been observed that in certain scenarios, Megatron V3's tensor parallelism demonstrates a more efficient memory footprint compared to BurstAttention. It's noteworthy that when integrated with the ZeRO (Zero Redundancy Optimizer) technique, BurstAttention not only achieves comparable efficiency but can also surpass Megatron in terms of throughput. This integration effectively leverages the strengths of both technologies, optimizing memory usage while enhancing overall processing speed.

2. For Peer-to-Peer (P2P) Communication settings

Firstly, we don't know that the reviewer is referring to the P2P communication operation such as send-recv or the P2P transport in the PCIe network, which means CUDA direct access between GPUs. We will presume that the reviewer is referring to CUDA direct access between GPUs before we continue the following discussion. However, if the reviewer's focus is indeed on the communication operations such as send-recv, then it should be noted the ring-collective operations are implemented using send-recv operation.

We acknowledge your observation regarding our experiments' absence of peer-to-peer (P2P) communication. Since we focus on proposing a distributed attention framework that solves extremely long sequences, it is hard for communication across machines to utilize P2P features. Moreover, it is also difficult for the environments of inferring LLMs (often using NVIDIA 4090 or L40 for inference) to meet P2P conditions. Hence, in our experiments, we turn off P2P communication as it is not a common setting in many communication-restricted scenarios (e.g., NVIDIA 3090/4090, NVIDIA A100/V100 PCIe, other non-NVIDIA devices, and the scenarios mentioned above).

In our latest experiments, we have optimized the original BurstAttention framework by overlapping the computation and communication phases. As a result, we can observe an acceleration in performance, evident in both multi-node and single-node experiments. Intriguingly, this optimization appears to be particularly efficacious in the variant of BurstAttention without FlashAttention, which is inherently more computation-intensive. This observation suggests a significant interplay between computational density and the effectiveness of our communication-computation overlapping strategy, opening avenues for further exploration in the optimization of distributed attention mechanisms.

3. For the Workload Balancing of Causal Language Modeling:

Regarding your query about workload balancing, we have not optimized the causal language modeling process in BurstAttention or other experiments. In our current implementation, we adhere to the conventional attention-masking approach without skipping computations. Hence, for both BurstAttention and Megatron V1/V3, the workload is balanced, and the experimental comparison is fair.

In fact, adopting the conventional attention-masking approach for parallel computing brings unnecessary computational overheads. For future work, we aim to address the workload balancing issue of causal language modeling by exploring more complex communication topologies beyond simple ring communication. A potential strategy we are considering involves repartitioning the sequence dimension across those dimensions unrelated to computation to ensure an even distribution of the attention mask across devices.

4. For the Argument with Sparse Attention:

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.

LLaMA-7b Inference First Token Latency (seconds) on 8xA100 GPUs

LLama-13b Inference First Token Latency (seconds) on 8xA100 GPUs

1.3 The Results of Scaling Experiments

Some reviewers suggest that we should 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.

Part 1: Performance on 1x and 2x GPU Configurations

Part 2: Performance on 4x and 8x GPU Configurations

LLaMA-7b Training Throughput (Token/s) with Scaling the Batch Size

The scaling experiment results proved that BurstAttention can be more efficient in both training and inferencing.

For inferencing, BurstAttention improves its performance as the number of GPUs and sequence length increase. This suggests that BurstAttention scales effectively across different hardware configurations and can handle larger sequences efficiently. Also, it demonstrates lower first token latency across various sequence lengths compared to TP (Megatron V1 and V3) with FlashAttention and Ring Self-Attention.

[To be continued]

Thank you for your thoughtful and detailed review of our submission. We appreciate the time you took to analyze our work and the insightful questions you raised. Here are our responses to your questions.

1. For Comparison with Megatron V3

Thanks for your insightful suggestion, we have conducted 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.

Furthermore, we have expanded 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.

1.1 The Results of Training LLMs

The results below indicate that Megatron V3's tensor parallelism shows better memory efficiency compared to BurstAttention ( since BurstAttention does not shard model parameters ), while BurstAttention has shorter training runtime and lower latency. Moreover, the longer the sequence length is, the more time-efficient BurstAttention is.

Furthermore, combining BurstAttention with ZeRO optimization (ZeRO has very limited benefits over TP since TP itself shards model parameters) brings significant improvements in memory efficiency. Although BurstAttention+ZeRO brings little additional communication overheads, BurstAttention+ZeRO still achieves memory efficiency comparable to Megatron V3 and demonstrates superior speed in both multi-node setups than Megatron V3. This suggests that BurstAttention, with its current optimizations, offers a more efficient solution in terms of speed, even when faced with a memory-efficient competitor like Megatron V3.

LLaMA-7b Training Runtime (seconds) on 8xA100 GPUs

LLaMA-7b Training Memory Usage (GB) on 8xA100 GPUs

LLaMA-7b Training Runtime (seconds) on 32xA100 GPUs

LLaMA-7b Training Memory Usage (GB) on 32xA100 GPUs

1.2 The Results of Inferring LLMs

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).

[To be continued]

Thank you for your insightful review and valuable questions regarding our paper. We appreciate your recognition of our work's motivation and effectiveness in addressing the challenges of processing extremely long sequences in distributed settings. We would like to address your questions and concerns as follows.

1. The Significance of Reducing the First Token Latency for Inference:

Although the first token latency might represent a small fraction of the overall inference process, it is critically important for real-time AI services such as ChatGPT. In such applications, 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 user-facing 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

2. The Contributions of BurstAttention over RingAttention and FlashAttention:

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.

Thanks again for your suggestion. We hope these responses will clarify the aspects highlighted in your review. We value your feedback and think these suggestions can significantly help us improve our work and its presentation.

Thank you for your insightful feedback on our submission. We greatly appreciate your recognition of the thorough analysis and detailed algorithmic contributions in our paper. Below, we address each of your concerns and questions to clarify our approach and its implications.

1. For the Details of Tensor Parallelism Implementation:

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.

2. For the First Token Latency:

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

3. For the Scaling Law Along the Model Capacity Dimension:

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.

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.

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.

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.

2.1 The Results of Training LLMs

The results below indicate that while Megatron V3's tensor parallelism shows better memory efficiency compared to BurstAttention, since BurstAttention does not do model parameters' sharding, BurstAttention has shorter training runtime and better time efficiency. Moreover, the longer the sequence length, the more time-efficient BurstAttention is.

Furthermore, combining BurstAttention with ZeRO optimization (ZeRO has very limited benefits over TP since TP itself shards model parameters) brings significant improvements in memory efficiency. Although BurstAttention+ZeRO brings little additional communication overheads, BurstAttention+ZeRO still achieves memory efficiency comparable to Megatron V3 and demonstrates superior speed in both multi-node and single-node setups than Megatron V3. This suggests that BurstAttention, with its current optimizations, offers a more efficient solution in terms of speed, even when faced with a memory-efficient competitor like Megatron V3.

[To be continued]

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.

LLama-7b Training Throughput (Token/s) with Scaling the Batch Size

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.

LLama-7b Training Runtime (seconds) on 8xA100 GPUs

LLama-7b Training Memory Usage (GB) on 8xA100 GPUs

LLama-7b Training Runtime (seconds) on 32xA100 GPUs

LLama-7b Training Memory Usage (GB) on 32xA100 GPUs

2.2 The Results of Inferring LLMs

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

LLama-13b Inference First Token Latency (seconds) on 8xA100 GPUs

[To be continued]

Thank you for your insightful comments and questions regarding our manuscript. Your feedback is invaluable in enhancing the quality and comprehensiveness of our research. We have carefully considered each point you raised and provide the following responses.

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.

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.

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.

[To be continued]