We sincerely appreciate the reviewer's thorough assessment and insightful comments. Our rebuttal addresses concerns and introduces a significant optimization: chunk-based mini-sequence transformers (MST). This advancement directly addresses the "Potential for Further Optimization" and "Sensitivity to Hyperparameters" concerns, with additional data to support our claims.

1. Further Optimization: Chunk-based Mini-Sequence Transformers

We've developed chunk-based MST to enhance efficiency in memory management and reduce computational overhead. This method:

• Partitions input sequences into fixed-size sub-chunks
• Integrates seamlessly with existing MST design

Key concept: The input sequence $S$ is split into sub-chunks of size $C$, where chunk size $C = S/M$ (M being the number of mini-sequence). This approach effectively addresses challenges in training both long-sequence transformers and small sequences, with no slowdown for the latter.

Hyperparameter selection: We've found that setting chunk_size $C = d$ (d being hidden size) yields optimal performance. Analysis: MST returns I/O complexity as $O(Sd + SI + dI * M)$. $S$ is sequence length, $d$ is hidden size, $I$ is intermediate size, $M$ is the number of mini-sequence. As long as $dI * M<= SI$, IO complexity remains as $O(SI)$. In conclusion, $C = S / M >= d$ is the best selection for the MLP block.

LM-head uses original MST with hyperparameters $M = V/d$. $V$ represents the vocabulary size.

2. Addressing Weaknesses

W1 Limited Model Scope:

To demonstrate the broader applicability of MST, we conducted additional experiments with Qwen, Mistral, and Gemma-2 models:

These results show significant improvements: 12x, 18x, and 24x context length increases for Mistral-7B, Qwen2-7B-Instruct, and Gemma-2-9b, respectively. MST performs exceptionally well for Gemma, likely due to its large vocab size (256k) and MLP intermediate size (14k).

W2 Comparison with Lossy Methods:

We've extended our comparisons to include state-of-the-art lossy quantization methods:

MST outperforms both 8-bit and 4-bit quantization for sequence length enablement (60k vs 28k vs 22k). Moreover, MST combined with 8-bit quantization achieves 110K tokens, a 22x improvement over standard 8-bit training.

W3 Training Time:

We conducted additional experiments with chunk-based MST to address concerns about training time:

Table 2: MST training with two epochs on LongAlpaca-12k

These results demonstrate that MST can train without significant throughput reduction (<5% is considered negligible).

W4 Implementation Complexity:

MST's core idea is conceptually straightforward, primarily targeting MLP and LM-Head blocks. We offer two integration methods:

a) Customized Hugging Face Transformer: Modify existing Hugging Face Transformer implementation.

git clone transformer_mst
pip install transformer_mst

b) Wrapper Mode: One-line wrapper for MST:

from mini-s import mst
model = mst(model)

Addressing Reviewer Questions

Q1 Impact on Model Quality:

Take this evaluation suggestion: Our experiments show that LLAMA3 with MST and a 30K context achieved a 23%- 270% better perplexity than the baseline.

Q2 Sensitivity to Hyperparameters:

The development of chunk-based MST allowed us to optimize hyperparameter selection. Setting chunk_size $C$ = hidden size $d$ provides the best balance between keeping the MLP layer compute-bound and avoiding additional memory movement. Setting $M = V/d$ provides the best memory efficiency for the LM head.

Q3 Applicability to Other Memory Optimization Techniques:

Our comparison with lossy methods shows that MST combines effectively with quantization techniques. MST + 8-bit achieves a 22x improvement in maximum sequence length over standard 8-bit training, while MST + 4-bit pushes this boundary even further to 140K tokens.

Also, MST + activation checkpoint outperforms the gradient checkpointing (also known as activation recomputation) model from 1.5x to 11x for long sequence enable, as shown in the paper and Qwen, Mistral, and Gemma-2 model experiments. As the reviewer claims, "The effectiveness of MST is closely tied to its combination with activation recomputation," so the applicability with gradient checkpointing is obvious.

Q4 Long-Term Impact on LLM Research:

We've demonstrated that MST works seamlessly with context parallelism (also known as sequence parallelism). Our experiments have achieved a remarkable 480k sequence length by combining context parallelism and MST, showcasing the potential for training with extremely large sequences. As we know, Llama 3.1 incorporates context parallelism to introduce additional sequences of 128k into the model architecture with 16 A100 GPUs. Our research demonstrates that MST integrates seamlessly with context parallelism, enabling significantly longer sequence processing than Llama 3.1 without throughput and performance downgrade.

This official comment focuses on the reply to limitations.

Limitations 1: Performance Degradation for Short Sequences

We acknowledge this limitation, and our further optimization of chunk-based MST effectively solves this problem with negligible slowdown by carefully tuning Hyperparameters of $M$ and chunk size $S$

Limitations 2: Dependence on Activation Recomputation

We also recognize that the Dependence on Activation Recomputation is one of the critical limitations. Our key is that the additional computational overhead can be compensated by memory saving where large batch training is enabled to improve performance, as illustrated in section 4.2, Faster Long Sequence Training with MINI-SEQUENCE TRANSFORMER (MST). Also, we have proved that this combination would not affect the total training performance by whole training experiments:

Table 2: MST training with two epochs on LongAlpaca-12k

Limitations 3 Distributed setting:

With our optimization for MST, we successfully enabled the training of LLAMA3 with 480k context on 8 A100, a significantly longer sequence length than 128k of LLAMA 3.1 can achieve on 16 A100s. The limited evaluation of distributed settings is partly due to our limited budget to extend a more significant scale over 8 GPUs. Therefore, we would open-source our code and welcome the LLM community to try our work on large-scale distributed systems.

The maximum sequence length of Llama3-8B runs on the distributed setting.

Limitations 4: Further Optimization

The chunk-based MST implementation provides a way to answer the limitations of Performance Degradation for Short Sequences and Potential for Further Optimization.

Again, we thank the reviewer for the valuable feedback that helped us improve our work. We really learned a lot from the suggestions of "Potential for Further Optimization" and "Sensitivity to Hyperparameters."

Thank the authors for their detailed feedback. I have raised the rating from 5 to 6.

We thank the reviewer for the thoughtful feedback. We have conducted new experiments to address concerns and demonstrate the effectiveness of Mini-Sequence Transformer (MST).

1. Addressing Weakness

Weakness: However, the model's training loss can not effectively demonstrate the method's performance on long sequences with around 1000 steps.

To address this concern, we have conducted new experiments on the LongAlpaca dataset:

• Trained Llama3-8B model for two epochs (around 10000 steps)
• Evaluated the trained model using perplexity
• Split the dataset for 90% training dataset and 10% evaluation dataset
• Carefully select Hyparameter M (the number of mini-sequences) to avoid throughput downgrade (M=2 for 8k, M=4 for 16k, M=8 for 30k)

Table 1: MST training with one epoch on LongAlpaca-16k

Table 2: MST training with two epochs on LongAlpaca-12k

These results show that MST enables training with much longer contexts (up to 30k) and improves perplexity compared to the 8k baseline (23% for LongAlpaca-16k and 270% for LongAlpaca-12k).

Also, we highly valued the reviews' suggestion that:

Question 2.2: Is there another way to prove its effectiveness, for example, evaluation on the long-context benchmark? Therefore, we evaluate the long-context datasets with perplexity to demonstrate MST as an effective way to enable long-context training.

2. Addressing Questions

Question 1: For the training performance (Sec 4.2 Table 6), would that be possible for the author to explain the reason MsT can not significantly improve the TFLOPS?

Question 1.1: If the long-sequence input makes the training memory-bounded, then the reduction of each intermediate feature should be able to speed up the I/O.

We clarify that long sequences make training computation-bound, not memory-bound. Moreover, MST doesn't reduce intermediate features but reuses memory space. That means all $M$ mini-sequence intermediate data are still generated, albeit in the same memory space. Therefore, MST is trading IO complexity with memory overhead, and it is worth for MLP and LM-head as they are computation-bound. We provide complexity analysis to support our claims:

• Computation complexity: $O(S * d * I)$ for both standard MLP and MST. Let $S$ be the sequence length, $d$ be the hidden dimension, $I$ be the intermediate size, and $V$ be the voice size. Standard MLP returns $O(S * d * I)$ computation. For MST, each mini-sequence computation requires $O(S / M * d * I)$ computation, and MST repeats $M$ times. In the end, MST returns $O(S / M * d * I *M) = O(S * d * I)$ FLOPS. The computation complexity is unchanged for any MST setting.
• I/O complexity: $O(Sd + SI + dI * M)$ for MST, slightly increased. Standard MLP requires $O(Sd + SI + dI)$ HBM access. The $O(Sd)$ represents the movement of input into GPU and output into HBM. The $O(SI)$ represents the movement of the intermediate. The $O(dI)$ represents the movement of weights. Each mini-sequence requires $O(Sd/M + SI/M + dI)$, and MST repeats M times, loading weight for every time. In the end, MST returns $O (((Sd/M + SI/M + dI)* M) = O (Sd + SI + dI* M)$ HBM access. The IO complexity is increasing for MST.

Question 1.2: Since the activation recomputation requires additional computation, it can still achieve better TFLOPS than MsT when using the same batch size.

The key to the TFLOPS problem is that MST adopts activation recomputation during experiments for better memory efficiency.

We have extended training performance (Sec 4.2 Table 6) to clarify this. Here are the key results:

Keys:

• MST independent achieves a 16% speedup over activation recomputation for Llama2-7B-hf (b=1)
• MST independent is only 4% slower than vanilla PyTorch for Llama2-7B-hf (b=1)
• MST + activation recomputation only introduces a 2.4% slowdown compared to activation recomputation for Llama3-8B-hf (b=2)
• Carefully select MST hyparameters $M$ can further improve performance (10% improvement for Llama3-8B-hf MST b=2, it was 2825.12 TFLOPS on the original paper)

Question 2: For the convergence experiments (Sec 4.3), are they trained from the Llama3-8B pre-trained model?

No, we train the Llama3-8B from scratch for the original convergence experiments.

In contrast, the new convergence experiments on the LongAlpaca dataset are trained from the Llama3-8B pre-trained model, known as fine-tuning Llama3-8B.

Question 2-1: the motivation for training LLMs using extremely long sequences is better performance.

We think so, and that motivates us to do all the new experiments.

In conclusion, our new experiments shows MST:

1. Has minimal impact on throughput (<5% loss for step training Llama2-7B and <5% loss for complete training llama3-8b)
2. Enables training with significantly longer contexts (up to 30k)
3. Demonstrate effectiveness on long-context tasks (23%-270% perplexity improve)

Detailed rebuttal on Q1: Training Performance

We appreciate the reviewer's question about why MST doesn't significantly improve TFLOPS. To address this, we offer the following explanation:

1. Trade-off between Memory and Computation:

   • MST reduces memory footprint by processing smaller mini-sequences.
   • Like gradient accumulation, it trades computation time for memory savings.

2. Computational Complexity:

   • MLP and LM-HEAD operators in long sequence training are compute-bound.

   • For sequence length $S$, hidden dimension $d$, and intermediate size $I$:

     • Standard MLP: $O(S * d * I)$ computation

     • MST: $O(S / M * d * I) * M = O(S * d * I)$ computation

   • Computational complexity remains unchanged for MST.

3. I/O Complexity:

   • Standard MLP: $O(Sd + SI + dI)$ HBM access
   • MST: $O(Sd + SI + dI * M)$ HBM access
   • I/O complexity increases slightly for MST due to repeated weight loading.

4. Memory Savings vs. I/O:

   • MST reuses memory space for mini-sequences, reducing HBM memory cost by factor M.
   • All intermediate features are still generated, maintaining I/O complexity for intermediates.
   • Increased I/O for weight loading is offset by the compute-bound nature of long-sequence training.

5. Performance in Different Scenarios:

   • Long-sequence training: Compute-bound, so increased I/O has minimal impact on speed.
   • Short-sequence training: I/O-bound, potentially slowing down TFLOPS when $M$ is large.

6. Optimization:

   • To address potential slowdowns in short-sequence scenarios, we introduced chunk-based MST. The key is to select a small $M$ for a small sequence.
   • This new implementation optimizes the memory-throughput trade-off.

In conclusion, MST's primary benefit is significant memory reduction without substantial TFLOPS loss in long-sequence scenarios. The compute-bound nature of long-sequence training generally outweighs the slight increase in I/O complexity. For short sequences, our new chunk-based MST implementation helps maintain performance.

Training Setting

We train a Llama3-8B MST on the LongAlpaca datasets. The LongAlpaca dataset has two variants, LongAlpaca-12k has 12k text with a maximum of 191k characters, and LongAlpaca-16k has 6.28k text with a maximum of 73.9k characters. The training lasts for one epoch for LongAlpaca-16k and two epochs for LongAlpaca-12k; each epoch requires about 5000 steps. For all implementations, we use the AdamW optimizer. We use a weight decay of 0.001, gradient clipping of 1.0, and a constant learning rate of 1e-4. All batch sizes equal 1, with a gradient accumulation step of 32. The bf16 precision is also deployed.

Thank you for your detailed response!

• Thank you for the extra experiment demonstrating MST's effectiveness on long-context training and the corresponding evaluation for long-context tasks. According to the benchmarking result, MST helps mitigate the GPU capacity limitation for long-context understanding. From my perspective, there is supposed to be no barrier between MST and the current training techniques (e.g., tensor parallel, parameter parallel).
• Thank you for providing a detailed explanation of the model's FLOPS. The computation bound and the memory are related not only to the model's complexity but also to the GPU capacity (See roofline model). I believe this part helps readers to understand the advances brought by MST comprehensively.

Based on the paper's results and rebuttal, I believe this paper provides a promising solution for a critical problem and will boost the community. I will adjust my rating based on the rebuttal and discussion with ACs and other reviewers.

Thank you for your thorough review and valuable feedback. We appreciate your support and address your concerns below.

1. Addressing Weaknesses

1.1 Comparison with lossy methods

We've conducted additional experiments comparing our Mini-Sequence Transformer (MST) approach with quantization techniques:

This table shows the maximum sequence lengths achievable on a single A100 GPU. MST enables significantly longer sequences (60K) than standard and lossy approaches and extends to 140K when combined with 4-bit quantization.

1.2 MLP Algorithm Presentation

We've revised the Mini-Sequence MLP algorithm presentation for clarity, removing some details of how MLP works:

Algorithm: Mini-Sequence MLP

Input: Matrix $X \in \mathbb{R}^{N \times d}$, MLP block

1. Partition matrix X into $M$ blocks $X_1, ..., X_M$ of size $N_m \times d$, where $N_m = N/M$

2. For $i = 1$ to $M$:

   Compute $O_i' = \text{mlp}(X_i)$, where $O_i' \in \mathbb{R}^{N_m \times d}$

3. Concatenate $O = [O_1', ..., O_M'] \in \mathbb{R}^{N \times d}$

4. Return $O$

1.3 Notation Correction

We apologize for the confusion in our notation. Here's a clarified glossary of terms:

• $O'$: Mini-sequence output
• $O$: Final output after concatenation Mini-sequence output
• $I_{up}, I_{gate}$: Intermediate values and outputs of $gate\_{proj}$ and $up\_{proj}$
• $W_{up}, W_{gate}$: Weights of $gate\_{proj}$ and $up\_{proj}$
• $G$: Number of groups used in grouped query attention (GQA)

2. Answering Questions

2.1 Performance across different architectures

We've conducted additional experiments with Mistral, Qwen2, and Google's Gemma-2:

MST provides substantial context length increasing across all tested models:

• 12x for Mistral-7B
• 18x for Qwen2-7B
• 24x for Gemma-2-9b

Interestingly, MST performs best with Gemma-2, which uses the largest vocabulary size (256k) compared to Mistral-7B (32k) and Qwen2 (152k). MST can effectively optimize Gemma-2's large peak memory. A larger $M=64$ also benefits.

2.2 Empirical values of Peak Intermediate sizes

We discussed Peak Intermediate memory usage in Appendix D. Using LLAMA3-8b with 4k training as an example:

• Vanilla: 75GB
• Activation recomputation: 52GB
• MST: 47GB

For vanilla and activation recomputation, peak memory occurs in the LM-head. For MST, peak memory is shared between the LM-head and MLP.

2.3 Activation Recomputation Overview

As requested, we've added a brief explanation of activation recomputation:

Activation recomputation, also known as gradient checkpointing, is a memory-saving technique for training large neural networks. This method trades computation for memory by discarding intermediate activations during the forward pass and recomputing them as needed during the backward pass. In standard training, all activations must be stored to compute gradients, which can lead to significant memory usage for large models or long sequences. Activation recomputation alleviates this by only saving activations at certain checkpoints. The forward pass is partially recomputed during backpropagation to obtain the necessary intermediate values.

2.4 Performance on Smaller Context Lengths

We measured the slowdown on smaller context lengths in Appendix F of table: MLP execution times (in seconds) for various sequence lengths and mini-sequence settings:

As observed, increasing $M$ leads to longer execution times, particularly for shorter sequences (e.g., $M=8, SEQ=1024$ introduces a 2.4x slowdown). However, the impact is minimal for longer sequences (e.g., 80000) as the IO overhead becomes negligible compared to computation overhead .

3. New Implementation: Chunk-based MST

We developed a new implementation of Chunk-based MST to address the concerns about slowdown for shorter sequences. This approach splits the sequence into fixed-size chunks (4096 for LLAMA3)

The core idea is to split the input sequence $S$ into sub-chunks of fixed size $C$, where the number of chunks equals $M = S/C$. This provides an effective solution to the challenges of training with small sequences, as there are no splits (or tiny splits) for short sequences, thus introducing minimal slowdown.

We found that the optimal selection of $C$ and $M$ for best performance is $C = d, M = S / d$, where $d$ is the hidden size. The MLP block uses chunk-based MST to balance speed and memory. The LM-Head uses the original MST for memory saving, and the optimal setting for $M$ of the LM-head is determined by $M = V/d$, 32 for llama3, and 64 for Gemma-2.

We conducted an additional training experiment to show that Chunk-based MST training performance as:

Table 2: MST training with two epochs on LongAlpaca-12k

Details Presentation for Mini-sequence MLP

For clear presentation, we take MistralMLP MLP as an example to present Mini-Sequence MLP notation:

      (mlp): MistralMLP
      (

      (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)

      (up_proj): Linear(in_features=4096, out_features=14336, bias=False)

      (down_proj): Linear(in_features=14336, out_features=4096, bias=False)

      (act_fn): SiLU()

    )

Notation correction

$O′$ is the mini-sequence output.

$O$ is the final output after concat.

$I_{up}$ $I_{gate}$ is the intermediate values and output of $gate_{proj}$ and $up_{proj}$,

$W_{up}$ and $W_{gate}$ are the MLP weights of $gate_{proj}$ and $up_{proj}$.

$W_{down}$ is the MLP weight of $down_{proj}$

$G$ is the number of groups used in grouped query attention (GQA); we add the definition to the paper.

We appreciate the reviewers' suggestions for improving our paper.

Optimization Details: Chunk-based Mini-Sequence Transformers

We've developed chunk-based MST to enhance efficiency in memory management and reduce computational overhead. This method:

• Partitions input sequences into fixed-size sub-chunks
• Integrates seamlessly with existing transformer architectures
• Requires minimal code modifications

Key concept: The input sequence $S$ is split into sub-chunks of size $C$, where chunk size $C = S/M$ (M being the number of mini-sequence). This approach effectively addresses challenges in training both long-sequence transformers and small sequences, with no slowdown for the latter.

Hyperparameter selection: We've found that setting chunk_size $C = d$ (d being hidden size) yields optimal performance.

Analysis: MST maintains the same computation complexity. MST increased I/O complexity as $O(Sd + SI + dI * M)$. $S$ is sequence length, $d$ is hidden size, $I$ is intermediate size, $M$ is the number of mini-sequence. As long as $dI * M<= SI$, IO complexity remains as $O(SI)$. In conclusion, $C = S / M >= d$ is the best selection for MLP block.