DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training

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. This 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: 1. 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. 2. 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. 3. 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. The 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. Overall, 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. Detailed Comments for Authors ----------------------------- Below is a list if minor typos/presentation comments: * [Pg. 2, line -3] "$5.64\times$ compared"-- "speedup" is missing. * [Pg. 4, line 2] "which" -> "that" * [Pg. 5, 2 lines above eq (3)] The final implementation of DistFlashAttention is Algorithm 2?

Rating

6

Confidence

3

Ethics flag

1

Reasons to accept

* The paper considers an important problem * The paper identifies three issues with an obvious distributed implementation of FlashAttention and then proposes fixes for them. * 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 * How were the three bottlenecks identified? Are these the three most important bottlenecks? If so, how was this determined? * Also the first two bottlenecks are not specific to FlashAttention but to Attention in general? Either way, this should be clarified. * 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) * 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). * 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. As 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? 2. Please clarify which bottlenecks are FlashAttention specific and which ones would apply to any Attention implementation. 3. Please clarify which of the proposed ideas appear in other works (like RingAttention) and which ones are new to this paper. 4. What are the main idea(s) in Algorithm 2? 5. What are the max theoretical gains that one can hope from the ideas presented in this paper?

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.

Rating

5

Confidence

4

Ethics flag

1

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. 1. 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. 2. 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. 3. 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? 4. The proposed method only tests on 7B model sizes. It is unclear whether the method will work for smaller or larger model sizes. 5. 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.

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.

Rating

6

Confidence

3

Ethics flag

1

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. - 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]. - Two primary techniques appear to be straightforward and may lack sufficient novelty. - 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. - 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. [1] Striped Attention: Faster Ring Attention for Causal Transformers: https://arxiv.org/abs/2311.09431 [2] The Case for Co-Designing Model Architectures with Hardware https://arxiv.org/abs/2401.14489 [3] Multi-head or single-head? an empirical comparison for transformer training. arXiv preprint arXiv:2106.09650, 2021.

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. The key contributions are: 1. 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. 2. introduces new methods to balance workload across processors, overlap communication with computation to hide latency, and optimize gradient checkpointing to avoid unnecessary recomputation. 3. Extensively evaluated across various setups and model configurations, demonstrating significant improvements in speed and efficiency.

Rating

7

Confidence

4

Ethics flag

1

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. 2. 8x longer sequences and up to 5.6x speedups are achieved with the proposed techniques 3. 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. 2. 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. 3. More discussions on how the proposed method can benefit CPU inference.

Ethics concerns details

no ethics concerns

The 1.2X speedup from gradient checkpointing

Thanks for your responses. Could y'all please give more details on how the $1.2\times$ speedup for gradient checkpoint was calculated?