Paper: LightSeq: Sequence Level Parallelism for Distributed Training of Long Context Transformers (ICLR-2024)

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

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

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

Reviewer comment (Experiments):
While the authors argue that pipeline parallelism is limited by its unbalanced memory consumption, it would be good to back this up with experimental results. A comparison to TeraPipe [1], which employs token-level pipelining, would also be interesting.

Author response:
**Q6:** While the authors argue that pipeline parallelism is limited by its unbalanced memory consumption, it would be good to back this up with experimental results.

**A6:** We added a discussion of the unbalanced memory consumption in the related work section. 
Pipeline parallelism keeps a high memory pressure to the first stage when applying interleaved pipeline parallelism to minimize the computation bubble. We back this up by experiments in Table 7 in the updated paper. When disabling the interleaved pipeline scheduling, the computation bubble increases, causing inefficient computation. We refer to Korthikanti et al.[1] for more details.

### Table 7: The memory consumption of Megatron-LM when training Llama-2H with tensor parallelism (degree=2) and pipeline parallelism (degree=8) on 16xA100 40GB GPUs at the sequence length of 128K. The memory consumption is highly uneven across pipeline stages.

|        | Worker 1 | Worker 2 | Worker 3 | Worker 4 | Worker 5 | Worker 6 | Worker 7 | Worker 8 |
|--------|----------|----------|----------|----------|----------|----------|----------|----------|
| node 1 | 31.5GB   | 31.4GB   | 28.7GB   | 28.7GB   | 26.0GB   | 26.0GB   | 24.6GB   | 24.6GB   |
| node 2 | 21.8GB   | 21.8GB   | 20.5GB   | 20.5GB   | 17.9GB   | 17.8GB   | 32.0GB   | 32.1GB   |

[1] Korthikanti, Vijay Anand, et al. "Reducing activation recomputation in large transformer models." Proceedings of Machine Learning and Systems 5 (2023).

Reviewer comment (Experiments):
The 33H case seems rather constructed and unlikely in practice (why would you chose such a number specifically?). It is a valid limitation of parallelization over the heads but should not be that big a problem in actual use, reducing the realistic speedup to ~1.5x (which is still substantial). Also, the title of the paragraph does not really fit the content as it never gives any support why one would want to use an arbitrary numbers of heads.

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

**A3:** Yes, we designed the 33-head model to illustrate the advantages of keeping the distributed training design agnostic to the model configuration. We picked the number 33 for this illustration which may sound unusual. However, we would like to note that it is a common case in practice to have a non-power-of-2 number of heads as we discussed in the second paragraph of the introduction. For example, GPT-2-XL has 25 attention heads, GPT-2 has 12 attention heads, Llama-33B and its fine-tuned versions (e.g., Tulu-30B) have 52 attention heads, Whisper-large has 20 attention heads, and Falcon-7B has 71 attention heads. We would also note that our design can facilitate future model architectures as we discussed in the second paragraph of the introduction: 
> Moreover, many works have shown that the future Transformer architecture design may have even fewer attention heads. For example, Bian et al. (2021) demonstrate that Transformers with a single head outperforms its multi-head counterparts, representing a challenging scenario for solutions like Megatron-LM.

Reviewer comment (Reproducibility):
Quite a lot of detail on the algorithmic implementation of DistAttention, e.g. the communication scheme, is lacking. One has to dig through the source code to get a glimpse of what is going on.

Author response:
**Q7:** Quite a lot of detail on the algorithmic implementation of DistAttention, e.g. the communication scheme, is lacking. 

**A7:**  We thank the reviewer for pointing this out. We added detailed explanations of the computation and communication scheme in Appendix C and updated Figure 7 for clearer illustration.

Reviewer comment (Writing):
The difference between the different models in the model setup is not described clearly enough; GQA should be cited and explained a bit better, the 33H model is only clear from context, and it is not clear how exactly the 16H-2H models are scaled "properly".

Author response:
**Q8:** The difference between the different models in the model setup is not described clearly enough, GQA should be cited and explained a bit better, the 33H model is only clear from context, and it is not clear how exactly the 16H-2H models are scaled "properly"

**A8:** We have included more explanation about the GQA (and cited) and Llama-33H in the experiment setup section in the updated version. For 16H-2H models, we keep the number of attention heads the same in the entire model. For instance, Llama-7B has 32 attention heads per layer and 32 layers, Llama-16H has 16 attention heads per layer and thus 64 layers in total, according to [1]. We have also updated this in the experiment section. 

[1] Multi-head or single-head? an empirical comparison for transformer training.

Reviewer comment (Presentation):
In Section 4.3, when discussing the effect of load balancing, a diagram illustrating the work assigned to each time step would help visualizing the difference between the schedules. It would also help to remind the reader again that the balanced schedule completes the 36 work units in 5 time steps, thus resulting in a maximum speedup of 7.2, instead of having the reader backtrack to understand where the 7.2 came from.

Author response:
**Q9:** In Section 4.3 when discussing the effect of load balancing, a diagram illustrating the work assigned to each time step would help visualize the difference between the schedules.

**A9:** We updated the illustration in Appendix C to also contain time-step information, making it easier for readers to understand! Specifically, Figure 6 shows DistAttn design before applying load-balancing and Figure 7 shows DistAttn after applying load-balancing. The communication scheduling is also reflected in both figures by comparing the tensors each workder holds at two consecutive time steps.

Reviewer comment (Presentation):
Section 4.4 seems a bit odd, as it is titled "Discussion" but contains a mix of future work, introducing a very brief comparison to DeepSpeed Ulysses which is never mentioned before or after, and reiterating on the drawbacks of pipeline parallelism without giving much new information that had not already been stated in Section 4.2.

Author response:
**Q10:** Section 4.4 seems a bit odd, it is titled "Discussion" but contains a mix of future work, introducing a very brief comparison to DeepSpeed Ulysses which is never mentioned before or after, and reiterating the drawbacks of pipeline parallelism without giving much new information that had not already been stated in Section 4.2.

**A10:** Thanks for pointing this out! We apologize for the confusion. We have reorganized the paper to make the content fit the flow better. Specifically, we moved the discussion of comparison with DeepSpeed-Ulysses to the experiment section and the discussion of the pipeline parallelism to the related work section, which fits better the paper flow.

Reviewer comment (Novelty):
The paper states "we propose partitioning solely the input tokens (i.e., sequence parallelism)." Why is sequence parallelism proposed in this work, instead of being enhanced by this work?

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

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

[1] Shenggui Li, Fuzhao Xue, Yongbin Li, and Yang You. Sequence parallelism: Making 4D parallelism possible.
[2] Vijay Anand Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. Proceedings of Machine Learning and Systems, 5, 2023.

Reviewer comment (Novelty):
In Figure 4, by default, FlashAttention uses recompute during the backward pass. How does this become a contribution of this work?

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

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

Reviewer comment (Experiments):
In Figure 1 (right), besides shifting worker 6,7,8's compute on (kv1, kv2, kv3) to worker 1,2,3, can we also shift worker 6,7,8's compute on (kv6, kv7, kv8)?

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

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

Reviewer comment (Reproducibility):
How does the DistAttn & load balancing work for the attention operation with sliding local windows?

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

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

For local sliding windows, the workload is naturally (near) balanced, regardless of single directional or bidirectional attention. Thus, simply disregarding the attention logic to non-local workers suffices. For instance, in exact attention, worker 7 needs to compute attention to all other workers. If the sliding window has a number of tokens equal to that of one worker, then worker 7 only needs to attend to itself and tokens in worker 6. In other words, it only needs to fetch key and value from worker 6, and compute attention. In terms of implementation change, it just needs to change the end condition of the for loop (from looping worker 1 - worker 7 to looping only from worker 6 - worker 7).

Global windows can be interpreted in two ways. We’re not sure which one corresponds to the reviewer’s question so we answer for both. The first interpretation is one special design of the sparse attention (e.g., global attention in Longformer), where there are a certain number of global tokens that all later tokens need to attend to, which are used to capture the global information. To adapt DistAttn to this, one easy way is to keep a replica of all the global tokens in each worker, which is simple and practical as otherwise, the global tokens will need to be all-gathered at each time step. One can also split the global tokens evenly onto all workers and use all-gather upon computation to further reduce the memory requirement. The second interpretation is the bidirectional modeling in the full attention situation (e.g., BERT). In this situation, the users can simply disable the load-balancing schema and use the vanilla DistAttn as illustrated in Figure 6 of the updated manuscript. We introduce the load-balancing only to handle the bubble issue caused by the causal modeling, so there’s no need to apply it if it’s bidirectional attention as there is no computation bubble in this case.

Reviewer comment (Reproducibility):
How does the DistAttn & load balancing work for global windows (i.e., each token attends to all tokens)?

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

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

For local sliding windows, the workload is naturally (near) balanced, regardless of single directional or bidirectional attention. Thus, simply disregarding the attention logic to non-local workers suffices. For instance, in exact attention, worker 7 needs to compute attention to all other workers. If the sliding window has a number of tokens equal to that of one worker, then worker 7 only needs to attend to itself and tokens in worker 6. In other words, it only needs to fetch key and value from worker 6, and compute attention. In terms of implementation change, it just needs to change the end condition of the for loop (from looping worker 1 - worker 7 to looping only from worker 6 - worker 7).

Global windows can be interpreted in two ways. We’re not sure which one corresponds to the reviewer’s question so we answer for both. The first interpretation is one special design of the sparse attention (e.g., global attention in Longformer), where there are a certain number of global tokens that all later tokens need to attend to, which are used to capture the global information. To adapt DistAttn to this, one easy way is to keep a replica of all the global tokens in each worker, which is simple and practical as otherwise, the global tokens will need to be all-gathered at each time step. One can also split the global tokens evenly onto all workers and use all-gather upon computation to further reduce the memory requirement. The second interpretation is the bidirectional modeling in the full attention situation (e.g., BERT). In this situation, the users can simply disable the load-balancing schema and use the vanilla DistAttn as illustrated in Figure 6 of the updated manuscript. We introduce the load-balancing only to handle the bubble issue caused by the causal modeling, so there’s no need to apply it if it’s bidirectional attention as there is no computation bubble in this case.

Reviewer comment (Reproducibility):
How does the DistAttn & load balancing work for MQA models?

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

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

Reviewer comment (Experiments):
Regarding Figure 3, is it measured in a distributed setting or on a single GPU?

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

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

Reviewer comment (Reproducibility):
Does this work use FSDP for cross-node and LIGHTSEQ for intra-node communication?

Author response:
**Q7:**  Intertwined with Other Work and the usage of FSDP, “LightSeq … FSDP”

**A7:** We thank the reviewer for asking about the relationship between LightSeq and FSDP. LightSeq is completely orthogonal to FSDP, and by default can be used by itself to reduce the activation volume in long context training. In the paper, we use FSDP (cross-node) so that the system consumes similar memory for the model parameters, for a fair comparison with Megatron-LM. This is intended for a rigorous analysis. The author believes that this combination is not a weakness of the method, rather a strength that LightSeq can be flexibly combined with other memory reduction techniques. We have updated the method and experiment section to avoid confusion.

It is possible to design a more fine-grained hybrid parallelism strategy on LightSeq sequence parallelism and FSDP. It is highly dependent on the model size. For smaller models, it may be better to do FSDP cross nodes because the volume is small. For larger models, it may be preferable to FSDP only for intra-node.

Regarding the question on “whether LightSeq can scale up to more nodes without leveraging FSDP”, our answer is positive. In fact, FSDP only helps reduce the model memory, and actually negatively affects the system scalability. In explanation, the communication volume is the sum of that in FSDP and that in LightSeq. In our main experiment, we have demonstrated scalability with the combination of LightSeq and FSDP. Using LightSeq alone will give faster training speed due to less communication when the model weights can fit in a single device memory.

Reviewer comment (Experiments):
How about using only LIGHTSEQ for cross-node and intra-node communication?

Author response:
**Q7:**  Intertwined with Other Work and the usage of FSDP, “LightSeq … FSDP”

**A7:** We thank the reviewer for asking about the relationship between LightSeq and FSDP. LightSeq is completely orthogonal to FSDP, and by default can be used by itself to reduce the activation volume in long context training. In the paper, we use FSDP (cross-node) so that the system consumes similar memory for the model parameters, for a fair comparison with Megatron-LM. This is intended for a rigorous analysis. The author believes that this combination is not a weakness of the method, rather a strength that LightSeq can be flexibly combined with other memory reduction techniques. We have updated the method and experiment section to avoid confusion.

It is possible to design a more fine-grained hybrid parallelism strategy on LightSeq sequence parallelism and FSDP. It is highly dependent on the model size. For smaller models, it may be better to do FSDP cross nodes because the volume is small. For larger models, it may be preferable to FSDP only for intra-node.

Regarding the question on “whether LightSeq can scale up to more nodes without leveraging FSDP”, our answer is positive. In fact, FSDP only helps reduce the model memory, and actually negatively affects the system scalability. In explanation, the communication volume is the sum of that in FSDP and that in LightSeq. In our main experiment, we have demonstrated scalability with the combination of LightSeq and FSDP. Using LightSeq alone will give faster training speed due to less communication when the model weights can fit in a single device memory.

Reviewer comment (Experiments):
Can LIGHTSEQ scale up to more nodes without leveraging FSDP?

Author response:
**Q7:**  Intertwined with Other Work and the usage of FSDP, “LightSeq … FSDP”

**A7:** We thank the reviewer for asking about the relationship between LightSeq and FSDP. LightSeq is completely orthogonal to FSDP, and by default can be used by itself to reduce the activation volume in long context training. In the paper, we use FSDP (cross-node) so that the system consumes similar memory for the model parameters, for a fair comparison with Megatron-LM. This is intended for a rigorous analysis. The author believes that this combination is not a weakness of the method, rather a strength that LightSeq can be flexibly combined with other memory reduction techniques. We have updated the method and experiment section to avoid confusion.

It is possible to design a more fine-grained hybrid parallelism strategy on LightSeq sequence parallelism and FSDP. It is highly dependent on the model size. For smaller models, it may be better to do FSDP cross nodes because the volume is small. For larger models, it may be preferable to FSDP only for intra-node.

Regarding the question on “whether LightSeq can scale up to more nodes without leveraging FSDP”, our answer is positive. In fact, FSDP only helps reduce the model memory, and actually negatively affects the system scalability. In explanation, the communication volume is the sum of that in FSDP and that in LightSeq. In our main experiment, we have demonstrated scalability with the combination of LightSeq and FSDP. Using LightSeq alone will give faster training speed due to less communication when the model weights can fit in a single device memory.

Reviewer comment (Experiments):
It seems that this work is only evaluated on Llama models.

Author response:
**Q9:** It seems that this work is only evaluated on Llama models.

**A9:** This paper focuses on optimizing decoder-based LLMs with casual masks, which we believe are sufficiently important models and are of the strongest interest to the ML community. We pick Llama for its popularity, and further evaluate LightSeq on several variants of Llama that have different system characteristics. Other autoregressive Transformers (like GPT, MPT, of a similar size) will have similar system performance to Llama.

Reviewer comment (Evaluation):
Why does sequence length at 128K show 1.44x in the overlapped case vs 1.x in the no communication case in Figure 5?

Author response:
**Q10:** Why sequence length at 128K show 1.44x overlapped cased vs 1.x no communication in Figure 5?

**A10:** We would like to ask the reviewer to clarify the question: what does “1.x” mean here. In this figure, the wall clock time is normalized so that the no communication implementation has a ratio 1. And the overlapped scheme has 44% higher wall clock time, while a non-overlapped implementation has 105% higher wall clock time. This shows the effectiveness of our overlapping optimization.

Reviewer comment (Novelty):
This is not the first work to leverage sequence-level parallelism for training transformer models. For example, how does the proposed work compare, both qualitatively and quantitatively, to [this 4D sequence parallelism work](https://arxiv.org/abs/2105.13120)?

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

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

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

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

### Maximal sequence length on Llama-7B on the DGX (A100-80GB) cluster
 |           | 1 Node (8 GPUs) | 2 Nodes (16 GPUs) |
|-----------|-----------------|-------------------|
| RSA       | 32K             | 64K               |
| LightSeq   | > 256K          | > 512K            |

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

|           | 1 Node (32K sequence length) | 2 Nodes (64K sequence length) |
|-----------|--------------|---------------|
| RSA       | 14.10        | 30.49         |
| LightSeq   | 2.50         | 6.85          |
| Speedup   | 5.64x        | 4.45x         |


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

Reviewer comment (Experiments):
The parallelism configurations are unclear. What is the batch size used for training? If the batch size is high enough, wouldn't a combination of tensor/data/pipeline parallelism provide sufficient parallelism to achieve high utilization of GPUs, even without utilizing sequence parallelism?

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

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

Reviewer comment (Experiments):
In Table 3, the authors used a batch size of 1, which does not seem to be a reasonable configuration for training. Can you elaborate on this choice?

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

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

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

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

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