Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper introduces a novel dimension to partition self-attention layers for distributed training of long-sequence transformers. The proposed technique, called Micro-Query Sequence Parallelism (MQSP), maintains only partial columns of attention map and applies a low-cost distributed softmax. To reuse memory over multiple rows of the attention map, MQSP also utilizies a finer-grained query, Micro-Q, to scale memory overhead linearly instead of quadratically. Evaluation demonstrates MQSP can scale sequence length significantly longer than both ColAISP and Megatron-LM3.

strength_and_weaknesses: **Strengths** 

* The sequence length of transformers is likely to keep increasing, and a scalable solution to the sequence length has values.
* The proposed solution effectively increases the sequence length over 70K on 32 A100 GPUs.

** Weaknesses **

* It is not clear whether this technique introduces another bottleneck (especially with increasing $m$) for inter-GPU communication and synchronization.
* The paper lacks comparison against other works leveraging different forms of parallelism (e.g., pipeline parallelism + tensor parallelism).
* The necessity of scaling sequence length needs to be better motivated.
 
clarity,_quality,_novelty_and_reproducibility: This paper is mostly written clearly and the writing quality is good. As for originality, this is not the first work to propose sequence parallelism demonstrates superior scalability over the prior art leveraging sequence parallelism (ColAISP, in particular).  The idea of distributed softmax seemingly has been explored elsewhere.  


summary_of_the_review: This paper focuses on improving scalability of multi-head self-attention block with an increasing sequence length. The self-attention block is known to be a scalability bottleneck due to it's quadratic overhead for both memory capacity and computation.  That said, this paper tackles a timely problem.  Evaluation demonstrates the sequence parallelism can handle longer sequences and is more memory space-efficient than both ColAISP (leveraging sequence parallelism in a different way) and Megatron (leveraging tensor parallelism).

While the results show some promise, I still have the following questions/concerns about this proposal:

* *Performance Sensitivity* - In Figure 9, there is a tradeoff between reduction in memory usage and token (training) throughput. This makes sense as $m$ increases, attention maps are partitioned into finer-grained chunks to increase the overhead of communication limiting overall throughput. In this particular setup of *2n16g*, the throughput is maintained at a high level before the benefits of memory savings saturate. However, this might not necessarily true depending on the hardware configuration (e.g., # of nodes, # of GPUs, GPU generations, communication bandwidth, etc.).  Would the proposed technique provide robust performance for other system configurations with different compute/communication ratios?  
* *Comparison against Existing Work* - This work compares mostly against ColAISP, which also leverages sequence parallelism.  I was not aware of this work, and hence do not know how strong a baseline it is. Comparison against Megatron-LM3 alleviates the concerns but I am still wondering how MQSP compares against other existing works, especially leveraging multiple forms of parallelism (e.g., pipeline parallelism together with tensor parallelism) in terms of overall token throughput. 
* *Motivation* - This work could have been better motivated by suggesting some realistic use cases requiring longer sequences. The authors suggest high-resolution image cases as an example, but wouldn't scaling down the input image be a more practical solution, especially for classification and segmentation tasks? What is the practical range of sequence lengths today (and how does it scale) and what kind of applications could potentially benefit from a long sequence with >78K tokens?
review_title: null

final_score: rating: 3: reject, not good enough
review_content: summary_of_the_paper: The paper proposed a new way to further reduce the memory footage of the self-attention mechanism over long sequences by parallelization on the sequence dimension. The paper proposed to communicate the query tensors to handle the sentence-dimension computation used in self-attention. 
strength_and_weaknesses: The paper proposed a more fine-grained method to partition the query dimension to reduce the memory footprint of the self-attention layer. The proposed method is new. However, I have doubts on the actual effectiveness of the method. Specifically,

1. The $O(L^2)$ memory usage for the attention operator can be reduced to $O(L)$ by changing the way it was implemented in GPU without introducing any communication overhead across different GPUs. See the following two references:
    1. [https://arxiv.org/pdf/2205.14135.pdf](https://arxiv.org/pdf/2205.14135.pdf)
    2. [https://arxiv.org/pdf/1911.02150.pdf](https://arxiv.org/pdf/1911.02150.pdf)
    
    Please compare with these works and show the effectiveness of the proposed method since the biggest benefit of the proposed does not exist any more with the optimized single-GPU kernels, and the proposed method brings much more extra communication.
    
2. There is no analytical analysis on the extra communication brought by the newly proposed algorithm, which is the major overhead brought by the proposed method.
3. Why do you need to evaluate the convergence of the proposed method in 4.1? The method does not change the behavior of the original self-attention algorithm.
clarity,_quality,_novelty_and_reproducibility: Many figures in the paper is not illustrative. I eventually figure out the following questions by 

1. Figure 1: Which dimension corresponds to queries, and which dimension corresponds to keys?
2. Figure 2: I suppose $Q_i, K_i, V_i$ are all $L/n\times d_k$ sized matrices, but one cannot tell the $d_k$ dimension from the figure.
3. Figure 3: The size of $Q_i, K_i, V_i$ in the figure does not reflect their actual shape (which should be exactly the same!). Also $Q_i$ and $Q^j$ are of the same shape, which should not be the case. 
4. Figure 3 caption: “The red dotted arrow lines represent the aligned row for distributed softmax.” Do you mean the orange dotted arrow? What does the arrow direction mean?  

The paper should be able to be reproduced. The proposed method from the paper is new.
summary_of_the_review: I have doubts about the effectiveness of the proposed method. I will raise my score if my concerns are addressed during the rebuttal.
review_title: null

final_score: rating: 3: reject, not good enough
review_content: summary_of_the_paper: This work studies long sequence problem from the system perspective. The self-attention part of a Transformer requires quadratic memory cost, limiting the applicability of full-attention for thousands of tokens. There are major two lines of directions for tacking with the high complexity due to long sequence: 1) full-attention approximation such as sparse attention, and 2) model parallelism approach. This work focuses on the 2nd line. The existing works such as Megatron and ColAISP reduce memory complexity by n times, where n is the tensor parallel size. This paper proposed to further reduce the memory cost by splitting the attention computation into m micro-steps, and in each micro-step, a micro-context output is computed. In this way, the memory buffer can be reused and the memory cost of attention map is reduced by nm times. The experimental results show that the proposed method, dubbed MQSP, supports training a BERT-large model with 78848 sequence length on 32 A100 GPUs.
strength_and_weaknesses: Strength:
1. The empirical results are promising. Supporting the training with 78848 sequence length is impressive.

Weakness:
1. This paper lacks sufficient discussion and comparison with the related works. Similar idea of splitting sequence into smaller chunks and incrementally accumulating results has been proposed in (Dao, Tri, et al., 2022). This work investigates the bottleneck of Transformer training from IO perspective. The resultant algorithm fuses the whole Transformer layer into a fused kernel and the memory cost of computing attention map is constant independent of the sequence length. This algorithm has been implemented in multiple libraries including OpenAI triton and Xformers https://github.com/facebookresearch/xformers/blob/91c7c0846455d78b934d6705026b4c7e35b89bb1/xformers/ops.py#L280. We can easily combine it with the Megatron implementation.
2. Distributed softmax is also not new. This has been used in Megatron for computing output logits: https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/mpu/cross_entropy.py

Dao, Tri, et al. "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness." arXiv preprint arXiv:2205.14135 (2022).
clarity,_quality,_novelty_and_reproducibility: Though the experiments look promising, the novelty is limited and the discussion and empirical comparison with the highly related work FlashAttention are not provided (details are provided in Weakness).

The paper is hard to follow. There is little description of Figures 3 and 4. I found the current main text loosely coupled with the figures, giving me some difficulty in comprehension. I encourage the authors to put more efforts in explaining the proposed algorithm in details especially for pages 4 and 5. Also, I suggest that authors need to ensure that all the symbols occurred in the paper have definitions and shapes such as $c_i^j$ and $C_i$.

It is not clear to me where 2BDL comes from for MEGATRON-LM3 in Table 1. By using the sequence parallelism introduced in (Korthikanti, V., et al., 2022), the memory costs for all the intermediate activations including dropout and layernorm can be reduced by n times.

I didn't run the code, so I do not have any data for reproducibility.

Korthikanti, V., Casper, J., Lym, S., McAfee, L., Andersch, M., Shoeybi, M., & Catanzaro, B. (2022). Reducing Activation Recomputation in Large Transformer Models. arXiv preprint arXiv:2205.05198.


summary_of_the_review: The paper studies an important problem since long sequence has a wide variety of applications such as question-answering and summarization. The empirical results are also promising and interesting. However, I am concerned with the novelty. The paper is lack of discussion and empirical comparison with FlashAttention. FlashAttention adopts very similar technique and fuses the computation of micro-chunks into a CUDA kernel, and its memory cost depends on the GPU SRAM size and is independent of the sequence length.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: The paper proposed MQSP, an optimized sequence parallelism technique for long sequence transformer. Finally, the evaluation result shows that MQSP can scale to longer sequences. In long sequence scenarios, MQSP can significantly reduce memory occupation and have higher throughput compared to Tensor Parallelism and Sequence Parallelism.
strength_and_weaknesses: Strengths

With Micro-Query, MQSP scales up sequence length linearly and shows a pretty good performance improvement. Greatly increases the length of sequences that can be handled in Sequence Parallelism.
MQSP has a better communication pattern and significantly improved performance compared to previous sequential parallelism techniques.
This paper provides a thorough theoretical analysis, while comparing other methods including the most advanced Tensor Parallelism and Sequence Parallelism systems from theory and experiment.


Weakness - major concerns

In the evaluation, only Bert-Large was used for testing. Testing with more models with different sizes can make the results more convincing.
In Section 3.1 “MICRO -QUERY SEQUENCE PARALLELISM - Comparison”, a comparison is made with ColAISP only. If a comparison can be made with Tensor Parallelism, it can give the reader more information to choose between Tensor Parallelism and Sequence Parallelism
In Section 4.3 “MEMORY FOOTPRINT”, the paper shows the performance of MQSP with the activation checkpoint turned on and off, but does not specify whether the activation checkpoint is turned on in other implementations
It can be observed that within a sequence length of 7500, Megatron-LM3 has a clear performance advantage over MQSP, and there needs to be a more prominent expression in the paper to demonstrate why we need long sequences for training, and how long we really need them. One concern is that extra-long sequences can lead to excessive computational cost without actually bringing much gain effect.


Weakness - minor concerns

In Section 4.1 “QUALITY OF CONVERGENCE”, the setting of sequence length for MQSP is missing.
For some readers who have not been exposed to much tensor parallelism, the background of this paper is slightly brief. Consider adding more detailed background on Tensor Parallelism.
Some letter symbols are used in advance in the figure, such as "m" in figure one. Maybe highlight the meaning of these letters in the article or caption, it will be easier to understand the article.
clarity,_quality,_novelty_and_reproducibility: OK
summary_of_the_review: Tend to reject; Please check my comments above for more information.
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: All reviewers find the proposed approach interesting. However they raised significant concerns about limited experimental results and comparisons with existing memory efficient attention methods. Further the paper is lacking in convincing applications for considering long sequence lengths of O(78k). Is there a setting of relevance where using such a long sequence length leads to better performance? Overall I believe the paper is currently ready for publication and addressing reviewers concerns can make the paper stronger.

Strengths: Novel approach to reducing attention cost for long sequence lengths ~ 78k.

Weakness:
1. Missing comparisons to existing work on memory efficient attention.
2. Unclear motivation for considering long sequence lengths of O(78k).
justification_for_why_not_higher_score: The paper currently suffers from the following weakness.
Weakness:
1. Missing comparisons to existing work on memory efficient attention.
2. Unclear motivation for considering long sequence lengths of O(78k).
Decision: Reject