Mini-Sequence Transformers: Optimizing Intermediate Memory for Long Sequences Training
Authors: Cheng Luo
Jiawei Zhao
Zhuoming Chen
Beidi Chen
Anima Anandkumar
Venue: NeurIPS 2024 poster (2024)
Keywords: Long-Context
Foundation Models
Systems for ML
LLM Training
GPUs
Memory-efficient Training
Primary area: infrastructure
TL;DR: We propose Mini Sequence to reduce intermediate memory overhead for long sequence training, with 12X longer than the standard implementation of LLaMA3-8b training on a single A100 device.

Abstract: We introduce Mini-Sequence Transformer (MsT), a simple and effective methodology for highly efficient and accurate LLM training with extremely long sequences. MsT partitions input sequences and iteratively processes mini-sequences to reduce intermediate memory usage. Integrated with activation recomputation, it enables significant memory savings in both forward and backward passes. In experiments with the Llama3-8B model, with MsT, we measure no degradation in throughput or convergence even with 12x longer sequences than standard implementations. MsT is fully general, implementation-agnostic, and requires minimal code changes to integrate with existing LLM training frameworks. Integrated with the huggingface library, MsT successfully extends the maximum context length of Qwen, Mistral, and Gemma-2 by 12-24x.

MINI-SEQUENCE TRANSFORMER: Optimizing
Intermediate Memory for Long Sequences Training
Cheng Luo California Institute of Technology chengluo@caltechh edu
Jiawei Zhao Meta FAIR jwzhao@meta.cm com
Zhuoming Chen
Carnegie Mellon University
chuominc@andree. cmu edu
Beidi Chen
Anima Anandkumar
Carnegie Mellon University beidic@andrew. cmu.
California Institute of Technology animaccaltech. edu
cmu.edu edu
Abstract
We introduce MINI-SEQUENCE TRANSFORMER (MST), simple and effective methodology for highly efficient and accurate LLM training with extremely long se- quences. MST partitions input sequences and iteratively processes mini-sequences to reduce intermediate memory usage. Integrated with activation recomputation, it enables significant memory savings in both forward and backward passes. experiments with the Llama3 8B model, with MsT, we measure no degradation in throughput or convergence even with 12x longer sequences than standard imple- mentations. MST is fully general, implementation-agnesse code changes to integrate with existing LLM training frameworks. Integrated with the huggingface library, MST successfully extends the maximum context length of Qwen, Mistral, and Gemma-2 by 12-24x.
and requires minimal
Introduction
The development Transformer [56] has been a remarkable journey, with each iteration pushing the boundaries of what is possible regarding model size. performance, and efficiency. One of the critical challenges in this journey has been managing the memory requirements of these models, particularly during training. As Transformers have significantly grown in size(10) and complexity 44 the memory demand has increased exponentially, necessitating innovative solutions to optimize memory usage while maintaining performance.
significant milestone in this journey was the introduction of multi-query attention [50] This technique dramatically reduced the size of the KV-cache during inference, which uses multiple query heads but single key and value heads. The idea was first adopted in the large-scale training of PaLM [12], then adopted and empirically tested in LLaMA [55]. As the field progressed, multi- query attention evolved into grouped query attention (GQA) (2). which relaxes the single key and value head restriction to multiple heads, and each head is coupled with group of queries. It significantly improves the quality and is adopted by Llama2-70B
55 and Mistral-7B 24
To further improve model quality, Llama3 (3) introduced a tokenizer with a vocabulary of 128K tokens, enabling more efficient language encoding than Llama2's 32K vocabulary. Additionally, Llama3 increased its MLP intermediate size from 1 to 14k. These changes reflect trend toward more extensive vocabulary and intermediate sizes for better quality. Meanwhile, Llama3 maintains its hidden size of 4k for inference efficiency. This trend is also reflected in the Microsoft development of Phi-3 II] compared with Phi-2 (23]
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Loss
MsT(our work)
84k(12x)
Output
80
Activation Recompute
LM head:
Intermediate (5)
Mini-Seq.
PyTorch
70
60k(12x)
60
Mini-Out.1
Mini-Dut.M
Intermediate
Max Sequence Size|K)
50
45k/7x)
MLP/LM head (S/M)
MLP/LMM head (S/M)
Norm
40
Mini-Seq. 1 (S/M)
Mini-Seq. M (5/M)
30
Attention
20
14K(3K)
Norm
MLP/LM head inputs
5k(1x)
7k(bx)
Tensor (5)
Llama2 (C) Max Sequence Size Different Models, with no Throughput and Convergence Degradation
Ulama3
(a) Conventional Transformer
Proposed Mini-Sequenner Transformer Details
Architecture
Figure I: (a) Standard Transformer architecture. MIP's and L.--Hedd's activation sequence length is annotated with S. (b) Head block, which splits the input sequence S into M mini-sequences with sequence length S/M, where M = on this figure. (c) Max sequence size for training Llama2/Llama' on A100-80GB GPU, with no degradation of throughput or convergence using our approach.
TRANSFORMER is used to replace MLP blocks and LM-
These advancements have also brought about new memory challenges, particularly in the intermediate value of linear layers of multilayer perception (MLP) and language modeling head (LM-Head). The substantial increase in intermediate variables, which can be nearly ten times larger than the input variables, has severely limited the network's ability expand sequence length and batch size. This limitation has made difficult to train large models without restricting sequence length to 8K or relying on gradient accumulation distributed systems expand batch size.
Recognizing these challenges, we introduce MINI-SEQUENCE TRANSFORMER (MsT), simple and effective methodology for enabling highly efficient and highly accurate LLM training with extremely long sequence lengths by reducing intermediate memory overhead. MsT introduces per-layer mini-sequence where the input partitions work for each MLP and LM-Head block. MsT partitions individual samples along the sequence dimension and iteratively processes each mini-sequence sequence combining all mini-sequence results to recover full-sequence outputs for these blocks. Our work also adopts activation recomputation g We find no degradation in throughput or convergence as shown in Figure Tc).
Our Approach:
with sequences up 12x compared to standard implementation of Llama3-8B,
To summarize, make the following contributions advance the long mengseeuenne training:
MsT trains 12 - 24x longer sequence lengths than existing systems on single A100 GPU with no degradation throughput and convergence of training.
Fully general and implementation agnostic: MsT supports most parameter-efficient training it works independently with attention layers.
Support for large-scale distributed training: MST works together with DeepSpeed-Ulysses 21 to support linear scaling sequence length by the number of GPUs.
Easy-to-use and portable, requiring minimal code changes to the existing training frameworks like Huggingface 222] The details can referred to Appendix
In subsequent sections, we provide background and related work, a detailed discussion of MINI- SEQUENCE TRANSFORMER (MST) design, Hardware-efficient analysis, experimental evalua- tion, and comparison with existing work. This work is open-source under an MIT license on https://github.c http://hhhbbbbooo/////c//////i//

---

2 Background and Related Work
This section briefly overviews the performance characteristics of long sequence transformers on modern hardware(e.g.. GPUs). We also describe some backgrounds of mini-batch training and activation recomputation, which inspire our work.
2.1 Transformer Architecture
Figure T is sketch of the building blocks of typical Transformer architecture [56). It consists of input sequences sent into I. repeated block with attention and MIP, then computed output loss with LM-Head block. The inputs and outputs of each block are typically 3D tensor of size (B, S, d) where B is micro batch size, is sequence length, and is hidden dimension. The intermediate value includes the Q2 K, V tensors of size (B, S, d) within the attention block, the / tensor of size (B, S, I) within the MLP block, and the logits tensor size (B, S, V) of within LM-Head block. Here, I represents the intermediate size of MLP, and V represents the vocabulary size.
2.2 Hardware Performance of Long Sequence Training
Memory Hierarchy. GPUs have memory hierarchy with larger but slower global GPU memory (high bandwidth memory; HBM) and smaller but faster-shared memory (SRAM). Transformers' high memory demand originates from the quadratic complexity of self-attention operations, where the memory needed to store attention scores for each token increases quadratically as the sequence length grows. This dramatic increase memory demand can quickly overwhelm the capacity of the HBM, leading to OOM issues. Flashattention [.5] uses kernel fusion to effectively mitigate the memory overheads associated with the quadratic growth sequence length, and Xformer [41] deploys optimized memory access patterns that achieve linear memory scaling, Our work is partly inspired by memory optimization technologies, where our optimization targets are MLP and LM LM-Head.
Occupancy. GPUs have many threads executed in parallel; threads are grouped into thread blocks, which execute on streaming multiprocessors (SMs). Modern hardware has specialized units like tensor cores on NVIDIA GPU to accelerate mammals. In long sequence training scenarios where the sequence size tends be long (>10k), parallelizing over the sequence dimension usually enables high GPU occupancy.
Performance characteristics. GPU operators be classified as either compute- bound memory- bound, which is determined by the time spent in arithmetic operations and the time spent accessing HBM. Typical self-attention with long sequence, MLP with the long intermediate size compute- bound operator because their core operators are matrix-min mutriplyullppy with a large inner dimension of sequence length. Then, cross-entropy with reduction is memory-bound.
2.3 Mini-Batch Training
Our work inspired by Mini-Batch Training algorithms, also known as gradient accumulation. Mini-batch training algorithms [17 40 can support large batch size by processing the training batch in smaller mini- bathhs, which allows the model to be trained on a subset of the data at a time, accumulating gradients over several mini-batch and only updating the parameter with accumulated gradient. This reduces the memory requirements compared to batch gradient descent [25), which enables training bigger batch sizes than GPU memory constrain. We are inspired by the idea and adapt it to train long sequences instead large batch sizes.
2.4 Activation Recomputation
Activation recomputation 8 also known as gradient checkpointing, is 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. standard training, activations must be stored to compute gradients, which can lead to significant memory usage for large models or long sequences. Activation recomputation is orthogonal with our MST, and we integrate this method for better optimizing intermediate value. We analyze the memory efficiency of activation recomputation and its integration with MsT on Sec 3.2
3

---

MINI-SEQUENCE TRANSFORMER (MsT): Algorithm, Analysis, and Distributed Extensions
We present our MENE-NEQUENCE TRANSFORMER (MsT) mechanism to partition the input sequence into M mini-sequences. We show how to compute the exact transformer block by gradient accu- mulation during the backward pass. Then, we analyze its memory efficiency and IO complexity, showing that our method is memory-efficient and throughput-equalized compared to the standard transformer. Based on the analysis, we found the optimal implementation of MsT by selecting the best hyperparameters. We further show how MST can work on distributed settings by integrating with DeepSpeed [21].
We focus here on the forward pass for case of exposition; Appendix backward.
contains details for the
3.1 Algorithms: Optimizing Intermediate Memory With Mini-Sequence Processing
Our idea arises from the observation of large intermediate values from transformer blocks. Given inputs E Nxd HBM, attention blocks and MLP blocks compute the output O E R and LM-head block computes the output loss E RR. N equals to sequence size S here. We observe that the intermediate values are always larger than the input X and output O, loss, illustrated in Table L Attention has intermediate values Q K, V IR xd which is (1 + 2 d)/G larger than input size, where (1 + 2 x d/G = 1.5) in Llama3 setting. refers to the number of grouped query attention (GQA). MLP has intermediate value Tup. gate R where 2 x T/d in Llama3 setting. LM-Head has logits E RVxd, where V/d 32 in Llama3 setting. The detail setting of Llama3-8B is listed in Appendix C
Table 1: Intermediate value size analysis for transformer blocks
Transformer Blocks
Input/Output Size
Peak Intermediate Value Size
IntermediateIIpput Ratio
Attention
(B, / (B, d) (B, S.d)
(B, 2 x (B,S.d/)) (1+ 2 d/G) xs 1.5 (2 /)/d
MLP
2 (B.S.T)
LM-Head
(B,S,o) d)/1
(B,S.V
V 32
The ratio in Llama3 setting.
As flash attention and group query attention have minimized the intermediate value of attention, we put our focus on the MLP block and LM-Head block. Therefore, our implementation MsT is general enough to work with any attention: self-attention [56), cross-attention [5], causal attention [42). their sparse counterparts
59 48
and their various optimized kernels such different versions of
FlashAttention
15
Our implementation adopts FlashA AesshAttentio22 [14] for the experiments.
Input Partition. We apply the mini-sequence technique to overcome the technical challenge of large intermediate values occupying HBM memory. We describe this in Algorithms [ and 2 which represent MLP blocks and LM-Head from Llama serials. Their MLP block consists of three linear layers and SiLU function
[4). and their LM-Head block consists of one linear layer and CrossEntropyLoss function 44). The corresponding backward implementations can be referred to in Appendix as Algorithm line and Algorithm 2 line 1, then compute the output with respect to those mini- sequences. We get the exact same result as standard implementation by contacting all mini-sequence outputs.
B for more details. The main idea is to partition the input X into mini mini-sequence X;
Gradient Accumulation. One of our goals is to reduce intermediate values for backward passes. The backward pass typically requires the matrices X E Nxd d I Nx xI logits g to compute the gradients with respect to weights. However, by input partition the X E IR N.. xd we can reduce the intermediate value as I E IR R.,,1 logits R... XV by Mx in the backward pass in HBM. With gradient accumulation for all mini-sequences, all gradients are generated in the same way as standard implementation by introducing more memory loading time. However, as MLP is the standard computation-bount operator and LM-Head occupies only small amount of total training time, MsT would not affect the whole training speed with significant reduction in memory overhead.

---

Algorithm Mini-Sequence MLP
Require: Matrices X RNxd, MLP block, Wdown, Rixd Weights of three linear layers W gate, W up E RRdxI W Woown E R
Partition matrices X into M blocks ,i of size N' x d, where Nm = N/M
2: for 1 < S M do
3: Compute 0 = ML W gate: Wapi W down), O, ERN.x
4: end for
5: Contact 0 (O,, ..,, 1 RNdd
6: Return O.
Algorithm 2 Mini-Sequence LM-Head
Require: Matrices X E RR Labels L RR-, Weights W out E RdxV
1: Partition matrices X into M blocks Xi,... Xm of size N' X d, where Nm = N/M
2: Partition labels L into M sub-label, .i,...,Lm of size Nms where Nm. N/M
3: for < S M do
4: Compute logits; = X,Woun out, logits; E NRNVV
5: Compute if (i - 1) Nm S L S i Nm. La L, else L; = 1000
6: Compute loss, crossentropylossllggiss; L)
7: end for
8: Compute loss Σn dM loss:/M
9: Return loss.
3.2 Analysis:
Memory Efficiency of MINI-SEQUENCE TRANSFORMER (MsT)
We analyze the memory efficiency of MST. MST can reduce intermediate value by Mx x while maintaining the same throughput performance.
Theorem 1. Let be the sequence length, Wmem be the weight memory occupation, including weights, gradient, and optimizer. Amem be the activation memory occupation per sequence, Immm be the intermediate memory occupation per sequence. The peak memory of the standard transformer is achieved by M = -mmem= + SS x (J mmmm- - -L x Amem). Note that L x Amem >> I for standard transformer, as Amem lasts for all L layers, but mem only lasts for one layer.
Theorem 2. With OpenAl's activation recomputation (391, the L X Amem could be reduced to SQrt(L) x Amem. Therefore the peak memory is reduced to M = W Wmem + S x (Imem + sgrr(L) x Amem). ). For models with large vocabulary and MLP intermediate, sqr((L) X Amem < Imem-
Theorem MST can reduce intermediate value by M Mx, so the memory occupation becomes M = Wmem + S x (Jmem//N frem/M- - + sqrt(L) timesAmennt For GPU with maximum memory Mmoz. (M - W Wmem)
the maximum sequences length contained by Smas [Tmem / hi tagri Amem.
This sequence
length would be much longer than the standard implementation with Smax
(Mmer W Weem)
[mmm LLXAAmm
3.3 Analysis: 10 Complexity and Memory of MINI-SEQUENCE TRANSFORMER (MST)
We analyze the IO complexity of MsT, compared with consistent compute complexity, which can affect its compute-bound or memory-bound performance characteristics.
Theorem 4. 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 xcc((XW] W gote) * (X,WWpp)
W Wdown with O(SdI) FLOPS and MsT MLP returns O(SdI /M * M) O(SdI) FLOPS. Standard LM-Loss returns loss = wwwssentropylosss..c...........p (XW, W, L) with O(SdV + FLOPS, and MsT LM-Loss returns O(SdV SV)/M M) O(dVV + SV) FLOPS.
Theorem 5. Standard MLP requires O(Sd + SI + df) HBM accesses, while MST 7 requires
O(Sd+ ST + dT HBM accesxes. Standard LM-Head requires @(Sd+ SV +dV) HBM accesses, while MsT 2 requires ((Sd + SV + dV HBM accesses.
For Llama3 values of d (4096), I (14336) and V (128256), SI, Su is many time larger than Sd. For long sequence cases, the compute complexity and 10 complexity are dominated by SI and SV. where MsT is close to standard implementation. However, for small sequence cases where S << d. the compute complexity and IO complexity are dominated by dl and dV while MsT needs di M and dV M. Therefore, MsT would cause throughput downgrades for small sequence lengths.
5

---

3,4 Chunk-based MINI-SEQUENCE TRANSFORMER (MST)
We present an optimized implementation of chunk-based MST designed to mitigate throughput reductions when training with small sequence data. The fundamental approach involves partitioning sequences S into equally sized chunks size (when possible), resulting M = S/C mini- sequences.
Our I0 complexity analysis indicates that the number of mini-sequences M influences the HBM accesses as ((Sd + SI + dTM) and ((Sd + SV + dVM) M).. However, the HBM accesses remain stable at (ST) and ((VV) provided that d & SI and dV M S SV. It means d < S/M.
Therefore, by setting the chunk size to C - S/M 2 d, MST avoids throughput downgrades for small sequences. Intuitively, when the sequence size is smaller than the chunk size, MST does split the input, thereby preventing any performance loss.
We apply chunk-based MsT exclusively MLP blocks by setting constant chunk size C' equal to the hidden dimension, = d. For LM-head blocks, we maintain a constant mini mini-sequence size of M V/d, these blocks contribute minimally the overall training time of transformers.
3.5 Extension: Distributed MINI-SEQEENCE TRANSFORMER (MsT)
We extend MINI-SFQUENCE TRANSFORMER (MST) to the distributed setting: we propose MsT + SP, which can effectively scale the trans- former using sequence parallelism(SP). In SP, the input tensor of each Transformer layer is divided along the sequence dimension, allowing for parallel computation across multiple GPUs. This segmentation, in con- junction with activation recomputation, results in a substantial reduction in activation memory requirements. It is worth noting that proposed approach is orthogonal to most sequence parallelism, such as Megatron- (21), Sequence parallelism [29]. and Ring Attention (3). Here, we take Reepspeed-Ulyseee as an example of how they work together.
Loss
LM head: Mini-Seq
MLP:
Mini-Seq
LM 266), Deepspeed-Ulysee
Norm
F shows the design of extending MsT with DeepSpeed-Ulyeee As with the transformers architecture, the design consists of atten- tion block with DeepSpeed-Ulysses, MLP, and LM-Head with MsT's mini-sequence technology. The design consists of input sequences S partitioned across available devices and mini-sequences. Each attention block Matrices Q K, V communicated through all-to-all collectives before and after the attention computation. The remaining modules of MLP and LM-Head use the sequence parallel and mini-sequence together. As SeeepSppee-Uysseeeeeee MsT is working MLP and LM-Head, it is straightforward to make them work together to scale sequence length.
Attention
Deepspeed
Ulysses
Norm
Tensor (S)
main change is working on attention block and
Figure 2: Distributed
SEQUENE TRANSFORMER.
4 Experiment
We evaluate the impact of using chunk-based MINI-SEQUENCE TRANSFORMER (MST) on Ll6ma3[36), a state-of-the-art model for many NLP tasks. We also evaluate Qwen (6) Mistral 24], and Gemma-2 2 (54) for context length improvements. We validate our claims about scaling sequence length, reporting training time, and memory overhead. Distributed Extension results can be found in appendix E which confirms that the sequence length of MsT can scale linearly with the number GPUs.
Maximun Sequence Length. MST can train Llama3-8B with context length 60k and Llama3-7B with context length 84k on a single A100 GPU, outperforming the standard implementation by 12x Also, it achieves 12 24x than the standard implementation of Qwen, Mistral, and Gemma-2.
Training throughput. MsT maintains the same training throughput compared with standard long-sequence training. Moreover, the throughput can be slightly improved with a large batch size supported by MsT.

---

4.1 Longer Sequence Length with MINI-SEQUENCE TRANSFORMER (MST)
Llama3 and Llama2. We train a Llama3-8B(36) MsT and Llama2 models [3] MST by exploring the sequence length on single A100 GPU with lossless training strategies, such as activation recomputation, fusing the backward operation with the optimizer update [44] and MST. Table 2 compares our maximum sequence and training time to the PyTorrch standard implementation and Huggingface PEFT with activation recomputation. Our implementation trains 4x longer sequence LLAMA-3 compared with activation recomputation and 12x longer sequence compared with standard implementation. Also, our implementation trains 1.8x longer sequence compared with activation recomputation and 12x longer sequence compared with standard implementation.
Table 2: Maximum sequence length of Llama3-8B and Llama2-7B.
LJama3-8B--a Implementation
Maximum Sequence Length (K)
Llama3-888 vanilla
Llama3-883 8B-hf activation recomputation
14
Llama3-8B-h MsT
60
Llama2-7B-hh hf vanilla
Llama2-7B-hn activation recomputation
Ilamm2-7B-h MsT
Qwen, Mistral, and Gemma-2. We've extended our evaluation to include Mistral-7B, Qwen2- 7B, and Gemma-2-99 demonstrating significant increases in maximum sequence length (12x for Mistral-7B. 18x for Qwen2-7B, 24x for Gemma-2-9BB) across these architectures. Among these models, MST provide best sequence extension for Gemma-2 of 24x. The critical observation here that gemma-2 uses the largest vocal size (256k) than Mistral-7B (32k) and (15kk(152k).
Table 3: Maximum sequence length of various models.
Model Implementations
Maximum Sequence Length (K)
Mistral-7B vanilla
Mistral-7B activation recomputation
42
Mistral-7B MST
70
Qwen2-7B 7B vanilla
4
Qwen2- activation recomputation
13
Qwen2-7B MST
74
gemma-2-44 vanilla
1.5
gemma-2-9b activation recomputation
5
gemma-22 2-9b MST
36
Combination with gradient accumulation. Gradient Accumulation has been used during training Llama2 and Llama3, which helps them train larger batch sizes given limited available GPU memory. However, in Gradient Accumulation, instead of updating the model parameters after processing each batch of training data, the gradients are accumulated over multiple batches before updating. This means that the memory usage for gradients would occupy the memory used for activation. Therefore, using gradient accumulation during training would constrain the maximum sequence size.
Table summarizes the maximum sequence length with gradient accumulation. The activation recomputation technology can train up 8K sequences. Then MsT can train up 30k sequence length, which is 4x longer sequence length than activation recomputation, and 21x x longer than vanilla. For Llama2-7B, MsT can also train up to 55k sequence length.
Table 4: Maximum sequence length training with gradient accumulation.
Model Implementation with gradient accumulation Maximum Sequence Length (K)
Llama3-8B-hh vanilla
1.5
Llama3-8B-hf Activation Recomputation
8
Llama3-8B-h hf MST
32
Llama2-7B-hr vanilla
4
Llama2-7B-hf activation recomputation
38
Llama2-7B hf MST
55
Comparison and Combination with Lossy Methods. with quantization methods and the combinations between MsT and quantization on Table 5 All lossy methods are HuggingFace official implementations. This comparison demonstrates MsT's
We've comprehensively compared MsT
7

---

superiority in enabling longer sequences for Llama3 training on a single A 100 GPU. MST alone (60K tokens) outperforms these lossy approaches (4bit 28k). When combined with quantization techniques. MsT achieves even more impressive results: MST + 8-bit reaches 110K tokens (a 22x improvement over standard 8-bit), while MsT + 4-bit pushes the boundary to 1 140K tokens. We did not evaluate the effect of quantization on training loss.
Table 5: Maximum sequence length training with lossy method
Llama3 Implementations
Maximum Sequence Length (K)
8-hit
4-bit
10
MST
MST 8-bit
110
MST 4-bit
140
4.2 Faster Long Sequence Training with MINI-SEQUENCE TRANSFORMER (MsT)
We evaluate the training performance of MsT on Llama3-8B with 8k sequence and Llama2-7B with 4k sequence using a single A100 80G GPU. Table Tableeccomparee the training time per step and TFLOPS achieved by MsT with the vanilla PyTorch implementation and activation recomputation technique.
Table Training performance using MST on single A A100 80G GPU.
Model Implementation
Batch Size
Training Time Per Step (s)
TFLOPS
-J-ma3-8B-ht vanilla
OOM
OOM
Llama3-8B-Ba activation recomputation
2
5.01
3271.42
Llama3-8B-hi MST
2
5.13
3194.90
Llama3-8BB Bf MST
8
19.35
3386.13
Llama2-7B-hr vanilla
1
1.24
3290.88
Llama2-7B-hr activation recomputation
1
1.52
2684.67
Llama2-7B-hf MsT without activation recomputation
1
1.31
3115.03
Llama2-7B-hh activation recomputation
8
8,85
3703.48
Llama2-7B: hf MST
8
9.33
3511.39
Llama2-7B MST
16
17.92
3656.17
For Llama3-8B, the vanilla implementation
of memory (OOM) with batch size of 1.
Activation recomputation allows training with a batch size of 2, achieving 3271.42 TFLOPS and a training time of 5.01 seconds per step. MsT, with the same batch size of 2, achieves comparable 3194.90 TFLOPS with a slightly longer training time of 5.13 seconds per step. However, MsT's memory efficiency allows scaling the batch size resulting an improved 3386.13 TFLOPS and a training time of 19.35 seconds per step.
In the case of Llama2-7B, the vanilla implementation can train with batch size of 1, achieving 3290.88 TFLOPS and a training time 1.24 seconds per step. For the same batch size, MsT without activation recomputation achieves 31 1..03 TFLOPS with training time of 1.31 seconds per step, demonstrating a 16% speedup over activation recomputation (2684.67 TFLOPS) and only 5% slowdown compared vanilla PyTorch. MsT further increases the batch size 16, maintaining similar 3656.17 17 TELOPS with a training time of 17.92 seconds per step.
4.3 Better Models with Longer Sequences
Language Modeling with Long Context. The memory efficiency of MsT allows us to increase the context length of llama by than activation recomputation. Table [able shows that training Llama3-8B with 30K context length achieved a 2.7x improvement in perplexity compared to the 8K baseline. We train Llama3-8B 36) MsT on the LongAlpaca dataset [9] The training lasts for two epochs and 10k steps for demonstration. For all implementation, weight decay 0.001, gradient clipping 1.0, and constant learning rate of le-4. All batch sizes equal 16, with gradient accumulation step of 16. The bf16 precision is also deployed.
use the AdamW optimizer [32]. We use
Table 7: LLAMA3-85 8b with MST, with 4times larger context length compared activation recom- putation.
Llama3-88 8B-B hf Implementation Activation Recomputation
Context length
LongAlpaca- 12k (ppl)
loss
Training Time
8k
9.34
2.23
25.6 hours
MST
Sk
7.41
2.00
26.5 hours
MST
16k
3.53
1.26
62.5 hours
MST
30k
3.45
1.23
233 hours
8

---

Vanilla(BF16)
12x
Vanilla(BF16)
24x
Activation Recomputation
Activation Recomputation
4x
Weight
8x
Weight
A100 80G
Optimizer
A100 80G
Optimizer
Gradient Activation
Gradient
MsT
MsT
1x
Activation
80
160
240
320
400
480
80
160
240
320
400
(a) Llama3-8B with 20k context.
(b) Gemma2-99 9B with 20k context.
Figure 3: Memory consumption of pre-training Llama3-8B and Gemma2-9B models with batch size of on a single A 100 device, with activation recomputation and MST. Note that long-sequent training gradients overlap with activation, gradients are not shown in bars.
5
Ablation Study:
5.1 Memory Optimization of MINI-SEQUENCE TRANSFORMER (MsT)
MsT introduces series of memory optimizations reduce the memory overhead of long- unn-sequence training. To understand the effectiveness of MsT memory optimizations, we perform an ablation study that incrementally turns off these optimizations (mini-sequence, activation recomputation) and measures the memory requirements. We consider three options: vanilla (standard Pytorch with BF16), activation recomputation only, and MsT with activation recomputation.
Figure 3 shows the results. We analyze the peak memory usage of Llama3-8B and Gemma2-9B, with sequence length of 20k. For sequence length 20k of Llama3-8B and Gemma2-9B, only MsT can make the model into A100 GPU, The rest of the memory consumption estimated based on its model architectures and theoretical activation amount. For Llama3, activation recomputation can reduce the memory overhead of activation by 3x, and MsT can further reduce 4x memory overhead based activation recomputation. For Gemma2-99 MsT achieves 24x longer sequence than vanilla and 8x longer sequence than activation recomputation. This improvement from 12x to 24x is due to Gemma2-9B's higher intermediate/input ratio (8 for MIP and 72 for the LM head) compared to Llama3 (7 for MI.P and 32 for the LM head, as shown in Table T Further details on the memory ablation study can be found in Appendix [
5.2 How many mini-sequences needed during training
We observe that increasing M, the number of mini-sequences,, can enhance memory efficiency; however, this enhancement has certain upper limit. Specifically, increasing M can also affect throughput performance. Appendix provides details regarding these limitations and their effects. This observation allows to identify the optimal configuration for memory optimization and achieve the best balance between memory performance, consistent with our analysis Sec 3.2 and 3.3
We found that the best balance for memory and throughput is achieved by the optimal values of C for chunk-based MLP C = d, MM = S/d, where is the hidden size. For the LM-Head, the original MST is employed for memory saving, and the optimal setting for M is determined by M = V/d, specifically 32 for Llama3 and 64 for Gemma-2. This value provides the best memory efficiency.
Limitations and Future Directions
We discuss the limitations and future directions. Related work is also given in Appendix A
Compiling to CUDA. Our current approaches are built on Pytorch implementation. This may constrain performance and low-level memory savings. It can be improved by fused kernel and cuda optimization, which can be our next step.
Combination with memory optimization. Our goal is to increase sequence length while maintaining performance and accuracy. Relaxing these requirements, MsT can be combined with activation offload [Moox W meem)
to extend sequence length as Smar Tmem / M+ Amem.. or with quantization to extend sequence b/16 4bst/8bis (Tmem /M /M4 Lx Amem.
length Smar
(Mmax- Wmam)
This combination can be explored in future research.

---

Acknowledgements
We thank Vast Al for computational resources renting.
A. Anandkumar is supported by the Bren named chair professorship, Schmidt AI 2050 senior fellowship, ONR (MURI grant N00014-18-12624)
References
[1] Marah Abdin, Sam Ade Jacobs, Ammar Ahmad Awan, Jyoti Aneja, Ahmed Awadallah, Hany Awadalla, Nguyen Bach, Amit Bahree, Arash Bakhtiari, Harkirat Behl, al. Phi-3 technical re- port: A highly capable language model locally on your phone. arXiv preprint arXiv:2404.14219 2024.
[2] Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13245. 2023.
[3] Apsod et al. Flashce. http:////ggghaaccccccoooo///aaaa 2023.
[4] Apsod al. optimizer step in backward tutorial. hattps://pytet https///yyooorooooootttorrrraae org/tutorials/
html 2024.
[5] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv: 0rXiv:1409.0473. 2014.
[6] Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang, et al. Qwen technical report. arXiv preprint arXiv:2309. 16609, 2023.
[7] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-doeument transformer.
arXiv preprint / 2020.
[8] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint
2016.
[9] Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya
Jia. Longlora: Efficient fine-tuning of long-context large language models. arXiv preprint arXiv:2309. /2307, 2023.
[10] Zonglei Chen, Minbo Ma, Tianrui Li, Hongjun Wang, and Chongshou Li. Long sequence time-series forecasting with deep learning: survey. Information Fusion, 97:101819, 2023.
[11]
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904. 10n::1904.100099 10509, 2019.
[12] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1- 113, 2023.
[13] Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. arXiv preprint 10XIv:1901.028600 2019.
[14] Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.0869// 2023.
[15] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficiiett exact attention with io-awarensss. Advances in Neural Information Processing Systems, 35:16344-16355, 2022.
[16] Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized llms. Advances in Neural Information Processing Systems, 36, 2024.
[17] Joeri Hermans, Gerasimos Spanakis, and Rico Möckel. Accumulated gradient normalization. In Asian Conference on Machine Learning, pages 439-454. PMLR, 2017.
[18] Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint 210.:2200000088,, 2021.
10

---

[19] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, Hy- oukJoong Lee, Jiquan Ngiam, Quoc V Le. Yonghui Wu, al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32, 2019.
[20] Dongseong Hwang, Weiran Wang, Zhuoyuan Huo, Khe Chai Sim, and Pedro Moreno Mengibar. Transformerfam: Feedback attention is working memory. arXiv preprint arXiv:2404.09173, 2024.
[21] Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Leon Song. Samyam Rajbhandari, and Yuxiong He. Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models. arXiv preprint arXiv:2309. 14509, 2023.
[22] Shashank Mohan Jain. Hugging face. In Introduction to transformers for NLP: With the hugging face library and models to solve problems, pages 51 51-67. Springer, 2022.
[23] Mojan Javaheripi, Sébastien Bubeck, Marah Abdin, Jyoti Aneja, Sebastien Bubeck, Caio
César Teodoro Mendes, Weizhu Chen, Allie Del Giorno, Ronen Eldan, Sivakanth Gopi, et al. Phi-2: The surprising power of small language models. Microsoft Research Blog, 2023.
[24] Albert Q Jiang. Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint wwXiv:231000682. 2023.
[25] Sarit Khirirat, Hamid Reza Feyzmahdavian, and Mikael Johansson. Mini-batch gradient descent: Faster convergence under data sparsity. In 2017 IEEE 56th Annual Conference on Decision and Control (CDC), pages 2880-2887 IEER, 2017.
[26] Vijay Anand Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Ander- sch, Mohammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. Proceedings of Machine Learning and Systems, 5, 2023.
[27] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. Communications of the ACM, 60(6):84-90, 2017.
[28] Wojciech Kryšciñski, Nazneen Rajani, Divyansh Agarwal, Caiming Xiong, and Dragomir Radev. Booksum: A collection of datasets for long-form narrative summarization. arXiv preprint arXiv:2105.08200 2021.
[29] Shenggui Li, Fuzhao Xue, Chaitanya Baranwal, Yongbin Li, and Yang You. Sequence paral- lelism: Long sequence training from system perspective. arXiv preprint arXiv:2105. 2021.
[30] Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near-infinite context. arXiv preprint arXiv:2310.01885 2023.
[31] Xin Liu, Daniel McDuff, Geza Kovacs, Isaac Galatzer-Levy, Jacob Sunshine, Jiening Zhan, Ming-Zher Poh, Shun Liao, Paolo Di Achille, and Shwetak Patel. Large language models are few-shot health learners. arXiv preprint arXiv:2305.15522 2023.
[32] Ila Loshchilov and Frank Hutter. Decoupled weight decay regularization, arXiv preprint 171i:: 2017.
[33] Lefteris Loukas, Manos Fergadiotis, lon Androutsopoulos, and Prodromos Malakasiotis. Edgar- corpus: Billions of tokens make the world go round. arXiv preprint arXiv:2109. /4394, 2021.
[34] Kai Lv, Hang Yan, Qipeng Guo, Haijun Lv, and Xipeng Qiu. Adalomo: Low-memory optimiza- tion with adaptive learning rate. arXiv preprint arXiv:2310.10195 2023.
[35] Mohammad Malek et al. Efficient cross entropy. https: httpp:://gitube giggtit...b... com/ wtt.//om//m///e /ogom//glle///////
EFfiiiinttcrrosseetrrooo
2023.
[36] Meta. Introducing meta laama 3: The most capable openly available Ilm to date. https
//ai mii..meaa.cm wwt.a..tcccooocc ctt.taa..oobooo////oocoo hatpo///tttaaaaaaaaaa1 1lama-3) 3/!
[37] Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, et al. Efficient large-scale language model training on gpu clusters using megatron-Im. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-15, 2021.
11

---

[38]
Tung Nguyen, Johannes Brandstetter, Ashish Kapoor, Jayesh K Gupta. and Aditya Grover. Climax: A foundation model for weather and climate. arXiv preprint arXiv:2301.10343 2023.
[39]
OpenAI.
gradient-checkkonntig
http://gggthuuu..ommeeeeeee cow cynnrrtreetota///////
gradientt-checkooitiin 2018.
[40] Myle Ott, Sergey Edunov, David Grangier, and Michael Auli. Scaling neural machine translation. arXiv preprint arXiv: I 808.00006.777. 2018.
[41] Markus N Rabe and Charles Staats. Self-attention does not need (n2) memory. arXiv preprint arXiv:2112.05682, 2021.
[42] Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. preprint, 2018.
[43] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog. 1(8):9, 2019.
[44] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, (1(140): 2020.
[45] Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He. Zero- infinity: Breaking the gpu memory wall for extreme scale deep learning. In Proceedings of the international conference for high performance computing, networking, storage and analysis, pages 1-14, 2021.
[46] Prajit Ramachandran, Barret Zoph, and Quoe V Le. Searching for activation functions. arXiv 2017.
preprint arXiv:
[47] Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap, Jean- baptiste Alayrac, Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context, arXiv preprint arXivv2403.055330 2024.
[48] Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguisties, 9:53-68, 2021.
[49] Reuven Rubinstein. The cross-entropy method for combinatorial and continuous optimization. Methodology and computing in applied probability, 1: 1:127-190, 1999.
[50] Noam Shazeer. Fast transformer decoding: One write-head all you need. arXiv preprint arXiv:1911.021555 2019.
[51] Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pages 4596-4604. PMLR, 2018.
[52] Ying Sheng, Shiyi Cao. Dacheng Li, Coleman Hooper, Nicholas Lee, Shuo Yang, Christopher Chou, Banghua Zhu, Lianmin Zheng, Kurt Keutzer, et al. S-lora: Serving thousands of concurrent lora adapters. arXiv preprint arXiv:2311.022
2023.
[53] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model parallelism. arXiv preprint
2019.
[54] Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based on gemini research and technology. arXiv preprint 2rXiv:2403.08295 2024.
[55] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprini arXiv:2307.0m (ITi::200.00288 092888, 2023.
[56] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[57] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint 2rXiv:2309.1745.3 2023.
12

---

[58] Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, al. Effective long-context scaling of foundation models. arXiv preprint ArXiv:2309. 2023.
[59] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santi- ago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 13:17371117297, 2020.
[60] Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: arXiv:2403.03507, 2024.
Ilm training by gradient low-rank projection. arXiv preprint
[61] Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, et al. Pytorch fsdp: experiences on scaling fully sharded data parallel. arXiv preprint arXiv:2304.1122 12277, 2023.
[62] Maxim Zvyagin, Alexander Brace, Kyle Hippe, Yuntian Deng, Bin Zhang, Cindy Orozco Bohorquez, Austin Clyde, Bharat Kale, Danilo Rivera, Heng Ma, et al. Genslms: Genome-scale language models reveal sars-cov-2 evolutionary dynamics. The International Journal of High Performance Computing Applications, 37(6):683-705, 2023.
A Related Work
Long Sequences Model. The ability to train large models with long sequences is becoming increas- ingly important across various domains, from generative Al to scientific discovery. In generative AI, tasks such conversational AI, knowledge-rice long document summarization, and video generation demand reasoning over extended contexts in both spatial and temporal dimensions. Multimodal foundation models process speech, images, and waveforms simultaneously, requiring long context reasoning over high-dimensional inputs with lengthy sequences. Similarly, chapter and book-level summarization, estimated to involve tens to hundreds of thousands of words, hold great importance and have demonstrated benefits
in conversational AI and abstractive summarization tasks 7 28 33
from long sequence training (58 47 36
The emergence of ChatGPT and subsequent open-source and commercial large language models has propelled chat applications to the forefront of modern AI, making them more relevant than ever. Efficiently processing long sequences is vital for supporting more extensive conversation histories in these applications [55]. Long sequence capability is equally important for AI in scientific fields, enabling better understanding and advancements in healthcare 31,, climate and weather forecasting (38, and large-scale molecular simulations (62)
Lossy Long Sequence Training. One direction making LLM able to process arbitrarily long sequences efficiently by sacrificing the perception window of the network. Sliding window attention introduced [13] to handle infinitely long sequences as input. However, it disregards information beyond the effective receptive field. Longformer (7] extends this idea, which caches block-by- block basis, which increases the perception window size, so as TransformerFAM (0)) Stream LLM 557 is not constrained by given window but selectively disregards information between the first token and given windows. They struggle to capture long-rang dependencies beyond this fixed range, but the approximation quality also seems to degrade at long sequence lengths. MsT can work directly with them to increase the window size for better quality.
Memory Efficient Training, As the demand for long sequence processing continues to grow across various domains, developing efficient methods for training large models with extended context becomes increasingly essential for advancing the state of the art in AI. Parameter-efffiien fine-tuning (PEFT) methods aim to reduce the memory footprint and computational related to parameters and gradients. Adafactor [5]] achieves sub-linear memory cost by factorizing the second-orde statistics using a row-column outer product. Low-Rank Adaptation (LoRA) (18 reduces the memory footprint of pre-trained models using low-rank adaptors with a low-rank weight adaptor for each layer. Several variants of LoRA have been proposed to enhance its performance.
[52 60). Quantization is another
widely used technique in PEFT to reduce the memory cost of optimizer states [16).
13

---

Other techniques focus on reducing the memory footprint and computational related to activation. Gradient accumulation is method that allows for training larger batch sizes by accumulating gradients over multiple mini-batches before performing an optimization step [40] This approach enables training with larger effective batch sizes while maintaining a smaller physical batch size, reducing activation memory requirements. A similar work of activation offloading [55 moves the checkpointed activations to the CPU asynchronously and prefetches the offloaded activations back from the CPU during backward. There related works in sparse Transformers, mainly focusing on full-attention approximation, such as sparse attention [ [5]] Recent works have also focused on single-GPU memory and compute-efficient attention. A popular example in this category is Flash attention [[5), which leverages known techniques such as tiling and recomputation for computing and memory efficiency. Also, some work put their interest in enss-yy FlashCE B] optimizes cross- entropy by leveraging sparse data structures and CUDA optimizations to enhance speed and memory efficiency. Efficient cross-entropy [35) introduces a memory-efficient variant of cross-entropy loss to reduce activation memory by storing only essential computations. These works are orthogonal to work and can be leveraged accordingly further improve the efficiency of Transfoormrr-bar -baseeee--b
models.
Distributed Training. Distributed training techniques have become essential for training large language models (LLMs) due to their immense computational and memory requirements. By splitting workloads across multiple GPUs, these methods help alleviate memory bottlenecks and enable the training of models that would otherwise be infeasible on a single device. Data parallelism [27] replicates the model on multiple devices, processing different data batches in parallel and synchronizing gradients across GPUs. Tensor parallelism (53) divides individual layers of the model across GPUs, enabling more efficient memory usage for extremely large models. Another prominent method is fully sharded data parallelism (FSDP) [6]], which extends the data parallel approach by sharding both model parameters and optimizer states across devices, thus further reducing memory overhead. Sequence parallelism [26] [2 [29 30) specializes in optimizing memory usage for transformer-baseee models by partitioning sequences across GPUs and reducing activation memory. In addition to these, pipeline parallelism [19] splits the model into segments, with each segment assigned to different GPU, and processes data in a pipeline fashion, improving efficiency by overlapping computation and communication. Hybrid parallelism (37 combines data, tensor, and pipeline parallelism to maximize resource utilization depending on the model's architecture and available hardware.
B Algorithm Details
We describe the full details of MINI-SEQUENCE TRANSFORMER (MST) backward pass. Algorithm 3 shows the MLP backward, and Algorithm shows the LM-Head backward.
Algorithm 3 Mini-Sequence MLP Backward
Require:
Gradients of output VO E Rx Matrices X Rd Weights of three linear layers W Woates Wap xd W down /xd
1: Partition matrices X into M blocks X1...., Xm of size Nn x d, where Nm.. = N/MM
2: Partition matrices VO into M blocks VOi,..., vOm of size Nm x d, where Nm = N/M
3: for < i < M do
4: Compute VX,, = V.LP(VV,)
5: Compute WWoomm VW. up: VW gate+= VMLP AMILPgradient gradient (VO,, X,)
6: end for
7: Concatenate VX VXI, VX m NRd
8: Return VX,, VW gate. VW ups VW
Woomn
We now observe about MsT backward pass that when computing the gradients of MLP and LM-Head, do not need to use full input and intermediates data. Instead, we 1 1/M reduced data with mini-sequence, significantly reducing the intermediate value memory overhead.
The main idea of backward is to accumulate gradients VW generated from each mini-sequence X; as Algorithmmliine line and Algorithm line We get the exact same result as standard implementation by accumulating all mini-sequent gradients.
line
14

---

Algorithm 4 Mini-Sequence LM-Heed Backward
Require: Loss gradients Vloss E RR, Logits R Labels L RN, Weights Wout out E RaxV
1: Partition matrices X into M blocks X1,... Xm of size N' x d, where Nm. = N/M
2: Partition labels L into M sub-labels Li...., Lm of size N. where N
M
3: Activation Recomputation with backward
4: for < i S M do
5: Compute logits; X;Woutt out, logits, RNXXV
6: Compute Vlogits; = crnssittropLossBBaavaadddoo
L.)
7: Compute VX; Vlogits, W VX; E R
8: Compute VWout = X Vlogits;
9: Compute VX, VX; O Vloss
10: Compute VW VW out O Vloss
end for
12: Concatenate VX = -X,,,... VX RR
13: Return VX, VW out-
C Llama2 and Llama3: Model Architecture Comparison
This appendix highlights the key architectural differences between the Llama2-7B and Llama3-8B models implemented by Hugging Face. The main distinction lies in the configuration of the MLP blocks and LM-Head (linear with cross-entropy loss) within the model architecture.
Loss
Kilind
Logit:WW..... Nww
lons:Ak)
MLP
Intermediate
xileed
NN
O:Nxd
Attention
Intermediate
Tensor (5)
Figure 4: Standard Transformer architecture with the highlight of LM-head and MLP.
Figure 4 fillustrates a standard Transformer model's architecture, focusing on the MLP and LM-head components. The intermediate tensors (II, 12) have larger dimensions than the input and output. Also, the logit tensor has much larger vocabulary dimension (V) than the hidden dimension (d).
Figure 5: Model Architecture of HuggingFace Implementation of Llama2-7B
(modal): LlamaModel
(wnbed_tokens): Eshadding (32000, 4096)
(layere): []lmrADecderraayyr(
(0-31): 32 (swlf_attn): [lmasPaabdtiessiinnn
(g.proj): Lipear(ia. fnt_featurrs--400 bias-Palse)
out_featuree=-009 bias-Palse)
(k_praj): [imarr(i_ffatrrre--tte
(v.praj): (inffatree. out_featurea-408t hina-Fulae)
(o.pros): Ltnear(1s, (1._featarer-4606 tet_ bias=False)
(ratary_enb): [1amalooaryytneeddiingg
(xip): LlaxaNIP(
(gate_proj):
nut, ant_fencurax-1100: biax-Falas)
(up_proj): Lisear (in.fostures--006, f0ttfesttres-10000 bias=Baise) (doza_proo): (eto: aat_testares-4008 bins-Palse) (act_fn): 8118()
Laypet_ LlaraRNSlore ((4896,), ape-1e-05)
(post_attearioa_ti (ttttearoaa.1yyyaorrr)))) LlanaRMEHoze epe-1o
15

---

(norn): (amaMM5 apa-1e-05)
(zotary_eab); GlassRotaryEabodingg)
(ln.head): Lineer (ia, out, fet.ttttoroo-200 bias-Falaw)
Llama2-7B Model Architecture: As shown in Figure 5 the Llama2-7B model employs an MLP block configuration with the following characteristics:
MLP Block: The first linear layer projects the input from hidden size of 4096 to intermediate size of I 1008.. The second linear layer projects the intermediate representation from I 11008 to 4096, the hidden size.
LM-Head (Output Projection with Linear Loss): The LM-Head in Llama2-7B consists of linear layer that projects the hidden representation from a size of 4096 to a vocabulary size of 32000. The output of the linear layer is then passed through cross-entropy loss function to compute the training loss.
Figure 6: Model Architecture of HuggingFace Implementation of Llama3-8B.
Fa
(model): [lanaKodel
babee__tkeens:: Exbodding (128256, 4096) (layera): (
(0-31): 82 (lexsDecoderlayer(
(seir_atta): [IaasFlaeadciearioo2
(q.proj): Linear Linear (in_featturea-4066
vuc_featurna-408. bisa-Falze)
Lineaz (ia_teatraa -4096, ouc_featuuee=-124, biea=False)
(v_pzoj):
out_featuree=1123. biasrFalse)
(o_praj):
fut_ bias-Falas)
(rotary.eab).
[laadostarSaaeeeiigg
(x1p): MlaxaNIP(
(gate_proj):
out, out_feacuree=14335 bias=False)
(up-proj): Linear (in_ (in_f anttfffattrax-14338 biaz=Palse)
(dove_proj): [iaeec(ia_teattuses-1133 aut_featurea-4008 biea-False)
(act_fa): 811U()
(inputt layurnors) LlazaRNSTozz ((4036,), apa-1u-05) apon__atentson_llyeeeee
((4096,), apa-le-05)
(norn): (1anat)5lorr((4466) spa=le-05) (rotary_esh): (laaskstarySsbeddingg
(nndhaa)): Lineer www.:iiinnnotnnaaaeeee ost_feattrra-119255 bisa-Falze)
Llama3-8B Model Architecture: As shown in Figure
6 the Llama3-8B model employs MLP
block configuration with the following characteristics:
MLP Block: The first linear layer projects the input from a hidden size of 4096 to a larger intermediate size of 13824. The second linear layer then projects the intermediate representation from a size of 13824 back to the hidden size of 4096.
LM-Head (Linear with Cross-Entropy Loss): In Llama3-8B, the LM-Head also consists of linear layer, but it projects the hidden representation from a size of 4096 to larger vocabulary size of 32000. Like Llama2-7B. the output the linear layer is passed through a cross-entropy loss function to compute the training loss.
The increased intermediate size in the MLP block of Llama3-8B allows the model capture complex patterns and transformations more effectively. Additionally, the larger vocabulary size in the L.M- HEAD of Llama3-8B enables the model generate more diverse and nuanced outputs. It is worth noting that while the dimensions of the MLP block and LM-Head differ between Llama2- 7B and Llama3-8B, the overall structure and functionality of these components remain the same. The cross-entropy loss function in the LM-Head measures the discrepancy between the predicted word probabilities and the target words during training, guiding the model to generate more accurate and contextually relevant outputs. These architectural differences contribute to Llama3-8B 's enhanced performance and capabilities compared to predecessor, Llama2-7B, while maintaining a consistent overall structure. However, the unexpected large intermediate value also comes from the change of MIP blocks and LM-HEAD (linear with cross-entropy loss), which we will discuss in the following section.
16

---

Moreover, this trend is also reflected in the Microsoft development of Phi-3 [ compared with Phi-2 23 Whose vocabulary size increased from 50k to 100K (2x), intermediate size increased from 10k to 16k (1.6x), and hidden size slightly increased from 2560 to 3072 (1.2x).
In conclusion, these models share an obvious trend: the ratio between intermediate and hidden size (also vocabulary and hidden size) is becoming larger.
Foooord
Bockooard
Forward
Bockward
opdiiter
Cptimbeens irsernations
60.00
80.00
Optimtnes
10.99
State
20.09
Actuationn
Grads
$0.00
Paramatars
Figure 7: Memory Visualizaion of training Llama3 8B with 4k sequence length on A100-80G.
D MINI-SEQUENCE TRANSFORMER'S Memory Optimization Detials
We compare SINENNE TRANSFORMER (MsT) with capturing and visualizing memory snapshots. We take vanilla Pytorch training for Llama3-SB with 4k sequence length as an example to show how the memory changes with the timeline.
vanilla. 7 the vanilla example. The model parameters had already been loaded into memory before the training step, so we immediately see a chunk of memory devoted to the weights. For Llama3-8B, its weight would be 15GB. As we start our forward pass, memory is allocated gradually for the activations or the tensors we are saving to be able to compute gradients in the backward pass. Here, the memory allocated for activation is larger than the weight, with around 29GB. Once we start the backward pass, the activations are gradually freed while the memory of the gradients starts building up. As the gradient is equal to the size of weights, which is smaller than activation, we can observe the memory usage drop to around 30 GB. Lastly, as the optimizer kicks in, its state will be lazily initialized, during the optimizer step of the first training loop only. In future loops, the optimizer memory will remain and be updated. The memory for the gradients is then freed accordingly at the end of every training loop when it is called zero grade. Here, the optimizer would take 2x of weights when using Adam with 30GB, and the optimizer intermediate is equal to the size of weight with 15. Therefore, the peak memory usage is during the optimizer step, which equals the sum size of weight, gradient, optimizer state, and optimizer intermediates, which roughly equals to 5x of weight size with 75GB as shown in Table g
we should
the optimizer state memory gradually increase
Table 8: Memory overhead of training Llama3-8B on single A A100 80G GPU.
Llama3-8B-hr
Memory Overhead
Within Peak Memory
Activation
29
0
Weight
15
15
Gradient
15
15
Optimizer
45
45
Total
75
optimizer-in-backwart The first memory optimization discussed here is aerimizer-in-backwaar g. It fuses the optimizer state with backward pass to save the memory the gradient and optimizer intermediate. The memory visualization of fiddmizer-in-acckaaae is shown in Figure g where there
17

---

Foward
Badlward
Forward
Backward
70.20
40 P0
Grader Optanzeer
30.00
Act vativsions
Parameters
Figure 8: Memory Visualizaion of training Llama3 8B with 4k sequence length on A100-80G, The optimizer in the Backward technique is deployed here
no optimizer stage but only forward and backward state. The backward time would become larger as result of fusion. Using this technology, the peak memory would change into the sum of weight, optimizer state, and activations. Although we successfully saved 30GB of memory overhead of gradient and optimizer intermediates, it adds up to 29GB memory overhead of activations, with only 1GB of memory saving. Totally consumes 74GB of memory, as shown in Table g It would be worse if sequence length were increased, introducing more activation into LLM training. Therefore, optimizer-in-ban wark can hardly benefit long sequence training, but it simplifies the training process, so we would include this technique in the following discussion.
Table 9: Memory overhead of training Llama3-8B on single A100 80G GPU. The optimizer in the Backward technique is deployed.
Llama3-8B-hr
Memory Overhead
Within Peak Memory
Activation
29
29
Weight
15
15
Gradient
15
0
Optimizer
30
Total
74
Foword
Boolward
Forward
Bookward
Deal
Intornedion Vilee
13.00
Optimine
00.0
15.00
Advalere
10.09
Figure 9: Memory Visualizaion of training Llama3 8B with 4k sequence length on A100-80G, Activation Recomputation technique is deployed here
Activation Recomputation. Activation Recomputation a powerful technique used in our analysis. It can significantly reduce the activation at the cost of small decrease in the training speed due
18

---

to recomputing parts of the graph during back-propagaan propagation. As shown in figure 9 it successfully reduces the total memory overhead from 74GB to 52GB and reduces activation memory overhead from 29GB to 7GB with 4x memory saving. However, we can easily find many singular points in the graph, which appear as an impulse signal. This impulse signal's 's duration is concise, meaning that it intermediate data that is briefly created in the forward/backward pass and immediately removed from the GPU HBM memory. The most prominent intermediate data is several times the total activation (4-5 times in our data analysis). These intermediate data seriously affect the training performance of long sequences and become the activation bottleneck of the row.
Table 10:
Memory overhead of training Llama3-8B on single A A0 80G GPU. Activation Recompu- tation technique is deployed.
Llama3-8B-i
Memory Overhead
Within Peak Memory
Activation
7
Weight
15
15
Gradient
Optimizer
30
30
Total
52
Forward
Forward
Badkward
Intermad
00.00
49.9
40.00
Optinizor 2300
20.09
20.09
Panmmelars
0.600
Figure 10: Memory Visualizaion of training Llama3 8B with 4k sequence length on A A100-80G. MINI-SRQUENCE TRANSFORMER technique is deployed here
MINI-SEQUENCE TRANSFORMER (MsT) Inspired by the observation from Activation Recompu- tation techniques, we propose MsT to reduce the intermediate value during training. We successfully decrease the memory overhead of activation from 7GB to 4GB, while the intermediate value is significantly reduced. This is because only 1/M intermediate values are used for computing the forward outputs, backward errors, and gradients during both the forward and backward.
Table 11: Memory overhead of training Llama3-8B on single A100 80G GPU. MINI-SEQUENCE TRANSFORMER technique is deployed.
Llama3-8B-hi Memory Overhead
Within Peak Memory
Activation
Weight
15
15
Gradient
15
0
Optimizer
30
30
Total
48
Conclusion. we compare the memory usage over time when training the Llama3-8B model using the standard transformer architecture versus using MsT. which is shown on [1 In Figure Ma), which shows the memory timeline for the standard Llama3-8B training, we can see that the memory usage is dominated by three main components: the model weights (in blue), the optimizer state (in green), and the intermediate memory (highlighted by the red circle). The peak memory usage reaches around 67GB.
19

---

LLAMA3 Standard
LLAMA3 Proposed
Peak Memory: 67GB
Peak Memory: 48GB(-30%)
60
Intermediate
50-
50
Memory Cost(GB)
Memory Cost(GB)
Intermediate
40-
40
30
Optimizer
30
Optimizer
20
10-
Weight
10
Weight
0
Time
Time
(a) Conventional Transformer Memory Timeline
(b) Mini-Sequence Transformer Memory Timeline
Figure 11: Memory Visualization. (a) The memory timeline of training Llama3-8B using standard transformer architecture, Red cycle highlights the intermediate memory (b) The memory timeline of training Llama3-8B using MsT, Red cycle highlights the intermediate memory has been narrowed.
In contrast, Figure b) demonstrates the memory timeline when training Llama3-8B with MsT. The critical difference that the intermediate memory, again highlighted by the red circle, has been significantly reduced or "narrowed" compared to the standard transformer case. As result, the peak memory usage with MsT around 47GB, achieving 30% reduction compared to the standard transformer.
The memory timelines illustrate that MsT effectively reduces the intermediate memory footprint during training, significantly contributing to overall memory consumption. By minimizing the intermediate memory. MsT enables more efficient memory utilization and allows training with longer sequence lengths larger batch sizes while staying within the hardware's s available memory limits.
E Scaling to Extreme Long Sequence on Distributed Setting
We evaluate our distributed extension of Ms'T on Llama3-8B Llama2-7B models and compare against vamilla DeepSpeed-Ulysses's Ulysses's sequence parallelism 2.4.8 GPUs, respectively, for the max sequence length and corresponding training time. The results of this evaluation are shown in Table [2
Table 12: Maximum sequence length of Llama3-8B, running on distributed setting.
Model Implementation
GPU numbers
8
Llama:38b-hi MsT
120
240
480
Llama2-7B-hr MsT
160
320
640
F How many mini-sequences are needed during pre-training
We also provide insights into the performance characteristies and memory usage of the Llama3-8B model when using the MINI-SEQUENCE TRANSFORMER (MsT) approach with different numbers of mini-sequences (M) and sequence lengths. Table shows the execution time for different sequence lengths and mini-sequence settings. As the number of mini-sequences (M) increases, the execution time slightly increases, especially for shorter
20

---

sequences. However, for longer sequences (e.g., 80000), the execution time remains relatively stable across different mini-sequences (M).
Table 13: LM-head time for different sequence lengths and mini-sequence settings
LM-head time
1024
2048
4096
8192
20000
40000
80000
standard
0.01
0.02
0.04
0.09
0.2
0.4
0.85
M=2
0.02
0.04
0.07
0.14
0.31
0.67
I.36
M-4
0.03
0.04
0.07
0.14
0.33
0.67
1.33
M=8
0.04
0.05
0.08
0.14
0.34
0.67
1.34
M=16
0.06
0.07
0.09
0.16
0.34
0.68
1.35
M=32
0.11
0.11
0.14
0.19
0.37
0.69
1.36
Table 14]less the memory usage in gigabytes (GB) for different sequence lengths and mini mini-sequence settings for the LM-Head component. The standard setting without mini-sequences consumes 59.92 GB for sequence length of 80000. By increasing the number of mini-sequences,, memory usage decreases significantly. With M=16, the memory usage reduces to 9.12 GB for the same sequence length, achieving 84.8% reduction in memory consumption. This is further improved with M=32
to 89.8%.
Table 14: LM-head memory usage (in GB) for different sequence lengths and mini-sequence settings
LM-head memory
1024
2048
4096
8192
20000
40000
80000
standard
3.20
3.46
4.94
7.91
16.46
30.95
59.92
mini-seq=2
2.59
3.21
4.46
6.95
14.14
26.31
50,66
mini-seq=4
2.28
2.60
3.24
4.52
8.20
14.44
26.92
mini-seq=8
2.13
2.30
2.63
3.30
5.24
8.51
15.05
mini-seq=16
2.06
2.15
2.33
2.70
4.14
5.54
9.12
mini-seq=32
2.02
2.07
2.18
2.39
3.01
4.06
6.15
Table 15 shows the execution time for different sequence lengths and mini-sequence settings. Like the Head, increasing M leads to slightly longer execution times, particularly for shorter sequences. However, the impact on execution time minimal for longer sequences (e.g., 80000).
Table 15: MLP time seconds) for different sequence lengths and mini sequence settings
MLP time
1024
2048
4096
8192
20000
40000
80000
standard
0.05
0.08
0.16
0.31
0.74
1.52
2.96
M=2
0.05
0.10
0.17
0.32
0.76
1.49
3.05
M=4
0.07
0.11
0.19
0.33
0.79
1.52
2.99
M=8
0.12
0.15
0.22
0.38
0.81
1.58
3.05
demonstrate the memory usage for different sequence lengths and mini-sequence settings. The standard setting consumes 14.72 GB for a sequence length of 80000, while using M=8 mini-sequences reduces the memory usage to 11.66 GB, resulting in 20.8% reduction. Here, we can observe.
For the MLP component, Table 16
The analysis suggests that increasing the number of mini-sequences (M) can significantly reduce memory usage, especially for the LM-Heed component, while having a minimal impact on execution time for longer sequences. Memory savings more pronounced for the LM-Head than for the MI.. is important to note that while MST is highly beneficial for training with extremely long sequences, may lead to performance degradation when applied models with shorter sequences due the overhead introduced by partitioning the input and the additional memory movement required for gradient accumulation. It can be easily observed from Table 13 that using M=32 mini-sequences increases the execution time from 0.01s (standard setting) to 0. Is for a sequence length of 1024, using M=8 mini- sequences increases
causing an 11x performance downgrade. Also, from Table the execution time from 0.05s (standard setting) to 0.12s for sequence length of 1024, causing 2x performance downgrade. The performance reduction is more pronounced for the LM-HHea compared to the MLP. Fortunately, L.Mhead accounts for very little of the transformer's 's running time, which smaller than MI.P and much smaller than attention, so our technology will not affect the overall performance, even if it affects its module performance.
21

---

Table 16: MLP memory usage (in
GB) for different sequence lengths and mini- -eni--e settings
MLP memory
1024
2048
4096
8192
20000
40000
80000
standard
0.93
1.09
1.39
2.11
4.18
7.69
14.72
M=2
1.29
1.36
1.50
2.00
3.76
6.73
12.69
M-4
1.32
1.41
1.61
2.00
3.49
6.21
11.66
M=8
1.33
1.44
1.66
2.11
3.42
6.17
11.66
G
Integrated with existing frameworks
MST's core idea is conceptually straightforward, primarily targeting MLP and LM-Head blocks. We offer two integration methods:
Customized Hugging Face Transformer. This method involves directly modifying the Hugging Face Transformer library to incorporate MST functionality. By customizing the library, users can seamlessly integrate MST into their existing workflows that use Hugging Face Transformers. We made this method open-source on https://ittuu..com// com//vvlccccrrasfooneet
To use the customized Hugging Face Transformer with MST, ML developer didn't change any line but install our customized transformers library with MST:
import transformers
Wrapper Mode. The Wrapper Mode provides a less invasive approach integrating MST. This method involves creating wrapper around existing model implementations, intercepting and mod- ifying the forward and backward passes of the MLP and LM-Head blocks. We made this method open-source on bttips:/gitthub.com wwlccc//mniis To use the Wrapper Mode:
from mini-s import mst model mst (model)
Conclusion. Both integration methods offer flexibility adopting MST for long sequence training. The choice between Customized Hugging Face Transformer and Wrapper Mode depends on the specific requirements of the project, the level of integration desired, and willingness to maintain custom libraries.
For users deeply invested in the Hugging Face ecosystem, the Customized Hugging Face Transformer method may be preferable. This method requires minimal changes to existing codebases which is already integrated with the Hugging Face ecosystem, and allows access to all Hugging Face features and optimizations. For those seeking more flexible solution working with multiple model implementations, the Wrapper Mode could be the better choice to used with customized codebase challenges.
22

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility. transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how answer these questions. For each question in the checklist:
You should answer [Yes]. [No] or [NA]
[NA] means either that the question Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]" ", is perfectly acceptable to answer "[No] "provided proper justification is given (c.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). ). general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found. IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Section
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Section 6
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
23

---

The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification:
Section
33
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced
statement of any theorems. The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Appendix D
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
24

---

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g.. registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: supply meterials
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cr cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
7/
gewss CCodeSoodiieeeeee ssi mollicyy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
25

---

Answer: [Yes]
Justification: Section 4
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Error bars are not reported because it would be too computationally exxpensive.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report I-sigma eгrоr bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of exccution) needed to reproduce the experiments?
Question:
Answer:[Yes)
Justification: A100 GPUs
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
26

---

Question: Does the research conducted in the paper conform, in every respect. with the NeurIPS Code of Ethics Sttngg.rrrrr./aarrr..../oo/////r
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation. generating fake profiles, surveillance), fairness considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
27

---

We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: References
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. cow datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: [NA].
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
28

---

According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based on the requirements of your country institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
29


Reviews:
[{"review_id": "mEKy4D52TL", "reviewer": "Reviewer_uLG7", "summary": "This paper introduced minibatching along the sequence length for inputs to the MLP and LM-head parts of a Transformer-based model. This method does not change the functionality of the transformer but improves the memory requirement when inputs are of long sequence length. \n\nOverall, while this paper leaves to be desired with the language and the writing, the idea is simple and it's impact is well quantified. Therefore, I vote to accept this paper.", "questions": "1. For different architectures outside of LLaMa, do you see similar context length improvements? For example, if this was repeated with Qwen or Mistral, does the context length increase stay the same or does it differ more? \n2. What are the empirical values of the Peak Intermediate sizes for Attention, LM-head, and MLP? Despite the schema provided in Table 1, it would be nice to have empirical measurements of the memory requirements for each part of the architecture?\n3. I believe adding a blurb about activation recomputation would be helpful for the reader who is unaware of this method. \n4. Could you measure the slowdown on smaller context lengths? This is mentioned in the paper but I do not believe you point to any exmpirical result.", "limitations": "No negative social impact and limitations are mentioned.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 1, "contribution": 3, "strengths": "1. The method is incredibly simple. Splitting among the sequence length for the MLP and LM-Head parts of the attention head is a simple and easy to implement idea and can be used very broadly. \n2. The ability to handle longer context sequences is very well shown through the experiments. Showing the ablation across different batch sizes is also interesting and a worthwhile addition. \n3. The background in this paper is extremely thorough and well-done. It provides intuition to even a novice in ML system knowledge.\n4. The ability of this method to synergize with existing methods such as DeepSpeed and gradient accumulation is quite nice.", "weaknesses": "1. While the author mentions that lossy methods such as LongLoRa differ from the lossless method presented in this paper, I still believe that providing some further detail on the amount of context length improvement from lossy methods would help contextualize the results in this paper. Indeed, plotting the context length improvement vs. performance degradation would be useful. \n2. The notation in this paper is chosen extremely poorly. $I$ denotes a constant but $I_i$ denotes a vector for example. \n3. What is the difference between $\\mathbf{O}'$ and $\\mathbf{O}$. \n4. G is never defined.\n5. What is $I_{up},I_{gate},W_{up},W_{gate}$? \n6. I do not understand Algorithm 1 at all. When computing a MLP, why are you multiplying two different weight matrices against the input? Why are you doing a matrix elementwise multiplication or convolution operation in Line 3? You define $I'$ but never use it? You define $O'$ but never use it? In general, this is written very poorly.\n7. There are many typos and informal language throughout this paper. The author should take time to go through this paper and fix them. Some of the language reads as a stream of thought and more structured writing would be helpful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough review and valuable feedback. We appreciate your support and address your concerns below.\n\n# 1. Addressing Weaknesses\n\n# 1.1 Comparison with lossy methods\n\nWe've conducted additional experiments comparing our Mini-Sequence Transformer (MST) approach with quantization techniques:\n\n| Llama3 Implementation | Maximum Sequence Length (K) |\n|:-|-:|\n| 8-bit  | 5  |\n| 8-bit + activation checkpointing | 26 |\n| 4-bit | 10 |\n| 4-bit + activation checkpointing | 28 |\n| MST | 60 |\n| **MST + 8-bit** | **110**  |\n| **MST + 4-bit** | **140** |\n\nThis 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.\n\n# 1.2 MLP Algorithm Presentation\n\nWe've revised the Mini-Sequence MLP algorithm presentation for clarity, removing some details of how MLP works:\n\nAlgorithm: Mini-Sequence MLP\n \nInput: Matrix $X \\in \\mathbb{R}^{N \\times d}$, MLP block\n1. Partition matrix X into $M$ blocks $X_1, ..., X_M$ of size $N_m \\times d$, where $N_m = N/M$\n2. For $i = 1$ to $M$:\n\n   Compute $O_i' = \\text{mlp}(X_i)$, where $O_i' \\in \\mathbb{R}^{N_m \\times d}$\n3. Concatenate $O = [O_1', ..., O_M'] \\in \\mathbb{R}^{N \\times d}$\n4. Return $O$\n\n# 1.3 Notation Correction\n\nWe apologize for the confusion in our notation. Here's a clarified glossary of terms:\n\n- $O'$: Mini-sequence output\n- $O$: Final output after concatenation Mini-sequence output\n- $I_{up}, I_{gate}$: Intermediate values and outputs of $gate\\_{proj}$ and $up\\_{proj}$\n- $W_{up}, W_{gate}$: Weights of $gate\\_{proj}$ and $up\\_{proj}$\n- $G$: Number of groups used in grouped query attention (GQA)\n\n# 2. Answering Questions\n\n# 2.1 Performance across different architectures\n\nWe've conducted additional experiments with Mistral, Qwen2, and Google's Gemma-2:\n\n| Implementation | Maximum Sequence Length (K) |\n|:-|-:|\n| Mistral-7B-v0.1 vanilla | 5 |\n| Mistral-7B-v0.1 activation recomputation | 42 |\n| **Mistral-7B-v0.1 MST** | **70** |\n| Qwen2-7B vanilla | 4 |\n| Qwen2-7B activation recomputation | 13 |\n| **Qwen2-7B MST** | **74** |\n| gemma-2-9b vanilla | 1.5 |\n| gemma-2-9b activation recomputation | 5 |\n| **gemma-2-9b MST** | **36** |\n\nMST provides substantial context length increasing across all tested models:\n- 12x for Mistral-7B\n- 18x for Qwen2-7B\n- 24x for Gemma-2-9b\n\nInterestingly, 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.\n\n# 2.2 Empirical values of Peak Intermediate sizes\n\nWe discussed Peak Intermediate memory usage in Appendix D. Using LLAMA3-8b with 4k training as an example:\n\n- Vanilla: 75GB \n- Activation recomputation: 52GB \n- MST: 47GB \n\nFor vanilla and activation recomputation, peak memory occurs in the LM-head. For MST, peak memory is shared between the LM-head and MLP.\n\n# 2.3 Activation Recomputation Overview\n\nAs requested, we've added a brief explanation of activation recomputation:\n\nActivation 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.\n\n# 2.4 Performance on Smaller Context Lengths\n\nWe 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:\n\n| MLP Sequence | 1024 | 2048 | 4096 | 8192 | 20000 | 40000 | 80000 |\n|:-|-:|-:|-:|-:|-:|-:|--:|\n| standard     | 0.05 | 0.08 | 0.16 | 0.31 | 0.74  | 1.52  | 2.96  |\n| M=2          | 0.05 | 0.10 | 0.17 | 0.32 | 0.76  | 1.49  | 3.05  |\n| M=4          | 0.07 | 0.11 | 0.19 | 0.33 | 0.79  | 1.52  | 2.99  |\n| M=8          | **0.12** | 0.15 | 0.22 | 0.38 | 0.81  | 1.58  | 3.05  |\n\nAs 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 .\n\n# 3. New Implementation: Chunk-based MST\n\nWe 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)\n\nThe 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.\n\nWe 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.\n\nWe conducted an additional training experiment to show that Chunk-based MST training performance as:\n\nTable 2: MST training with two epochs on LongAlpaca-12k \n| Implementation | Context length | LongAlpaca-12k (ppl) | loss | Training Time (hours) |\n|:-|-:|-:|-:|-:|\n| Act. recomputation | 8k | 9.34 | 2.23 | 25.6 |\n| chunk-based MST | 8k | 7.41 | 2.003 | 26.5 |\n| chunk-based MST | 16k | 3.53 | 1.26 | 62.5 |\n| chunk-based MST | 30k | 3.45 | 1.23 | 233 |", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "p4s8WzEyMR", "reviewer": "Reviewer_NrY2", "summary": "The paper targets an LLM-specific but significant challenge: training the LLM model for long-context understanding. The author proposed a Mini-Sequence Transformer method that enables LLM training with extremely long sequences. The method is motivated by the mini-batch, which forwards and backs a chunk of data instead of the whole sequence. Incorporating with gradient accumulation, the procedure can be adopted for long sequence training. The experience compares the performance of the proposed method with vanilla implementation and activation re-computation, but the training throughput improvement is limited. Moreover, the convergence experiment is not convincing and might need additional evaluation.", "questions": "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?\n- 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.\n- Since the activation re-computation requires additional computation, it can still achieve better TFLOPS than MsT when using the same batch size. \n\nFor the convergence experiments (Sec 4.3), are they trained from the Llama3-8B pre-trained model?\n- If this is the case, the training loss only reveals the correctness of the prediction's cross-entropy. It can not effectively show the long-sequence performance, which I believe is the motivation for training LLMs using extremely long sequences.\n- is there another way to prove its effectiveness, for example, evaluation on the long-context benchmark?", "limitations": "- The experiment does not well support the motivation.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Training with extremely long sequences is important but challenging for understanding long contexts. This paper proposes a possible solution to enable full parameter training beyond GPU memory constraints.\n- The method shows its potential for training up to 6K sequence length for Llama3-8B, which is helpful for long-context tasks and possibly reduces the barrier of training LLM.\n- The paper is well-written and easy to follow.", "weaknesses": "Although the idea of this method is promising, the experiments can not adequately address my concerns about its effectiveness compared to the existing method.\nFig.3 shows the training loss curve comparison of MsT and baseline. However, the model's training loss can not effectively demonstrate the method's performance on long sequences with around 1000 steps.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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).\n# 1. Addressing Weakness\n> Weakness: However, the model's training loss can not effectively demonstrate the method's performance on long sequences with around 1000 steps.\n\nTo address this concern, we have conducted new experiments on the LongAlpaca dataset:\n- Trained Llama3-8B model for two epochs (around 10000 steps)\n- Evaluated the trained model using perplexity\n- Split the dataset for 90% training dataset and 10% evaluation dataset\n- 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)\n\nTable 1: MST training with one epoch on LongAlpaca-16k\n| Implementation | Context length | LongAlpaca-16k (ppl) | loss | Training Time (hours) |\n|:-|-:|-:|-:|-:|\n| Act. recomputation | 8k | 290.0 | 5.67 | 6.7 |\n| MST | 8k | 301.9 | 5.71 | 7.0 |\n| MST | 16k | 240.3 | 5.48 | 21.2 |\n| MST | 30k | 222.4 | 5.40 | 61.2 |\n\nTable 2: MST training with two epochs on LongAlpaca-12k \n| Implementation | Context length | LongAlpaca-12k (ppl) | loss | Training Time (hours) |\n|:-|-:|-:|-:|-:|\n| Act. recomputation | 8k | 9.34 | 2.23 | 25.6 |\n| MST | 8k | 7.41 | 2.003 | 26.5 |\n| MST | 16k | 3.53 | 1.26 | 62.5 |\n| MST | 30k | 3.45 | 1.23 | 233 |\n\nThese 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). \n\nAlso, we highly valued the reviews' suggestion that:\n> Question 2.2: Is there another way to prove its effectiveness, for example, evaluation on the long-context benchmark?\nTherefore, we evaluate the long-context datasets with perplexity to demonstrate MST as an effective way to enable long-context training.\n\n# 2. Addressing Questions\n> 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?\n\n> 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.\n\nWe 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:\n- 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.\n- 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.\n\n> Question 1.2: Since the activation recomputation requires additional computation, it can still achieve better TFLOPS than MsT when using the same batch size.\n\nThe key to the TFLOPS problem is that MST adopts activation recomputation during experiments for better memory efficiency.\n\nWe have extended training performance (Sec 4.2 Table 6) to clarify this. Here are the key results:\n\n| Model Implementation | Batch Size | Training Time Per Step (s) | TFLOPS |\n|:-|:-:|-:|-:|\n| Llama3-8B-hf vanilla | 1 | OOM | OOM |\n| Llama3-8B-hf Act. recomputation | 2 | 5.01 | 3271.42 |\n| Llama3-8B-hf MST | 2 | 5.13 | 3194.90 |\n| Llama3-8B-hf MST  | 8 | 19.35 | 3386.13 |\n| Llama2-7B-hf vanilla | 1 | 1.24 | 3290.88 |\n| Llama2-7B-hf Act. recomputation | 1 | 1.52 | 2684.67 |\n| **Llama2-7B-hf MST without Act. recomputation**  | 1 |  **1.31** |  **3115.03** |\n| Llama2-7B-hf Act. recomputation | 8 | 8.85 | 3703.48 |\n| Llama2-7B-hf MST  | 8 | 9.33 | 3511.39 |\n| Llama2-7B-hf MST  | 16 | 17.92 | 3656.17 |\n\nKeys:\n- MST independent achieves a 16% speedup over activation recomputation for Llama2-7B-hf (b=1)\n- MST independent is only 4% slower than vanilla PyTorch for Llama2-7B-hf (b=1)\n- MST + activation recomputation only introduces a 2.4% slowdown compared to activation recomputation for Llama3-8B-hf (b=2)\n- 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)\n\n> Question 2: For the convergence experiments (Sec 4.3), are they trained from the Llama3-8B pre-trained model?\n\nNo, we train the Llama3-8B from scratch for the original convergence experiments. \n\nIn contrast, the new convergence experiments on the LongAlpaca dataset are trained from the Llama3-8B pre-trained model, known as fine-tuning Llama3-8B.\n\n> Question 2-1: the motivation for training LLMs using extremely long sequences is better performance.\n\nWe think so, and that motivates us to do all the new experiments.\n\nIn conclusion, our new experiments shows MST:\n1. Has minimal impact on throughput (<5% loss for step training Llama2-7B and <5% loss for complete training llama3-8b)\n2. Enables training with significantly longer contexts (up to 30k)\n3. Demonstrate effectiveness on long-context tasks (23%-270% perplexity improve)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WKjPiGbyFM", "reviewer": "Reviewer_iczK", "summary": "The paper introduces the MINI-SEQUENCE TRANSFORMER (MST), a technique designed to enhance the efficiency and accuracy of LLM training, particularly when dealing with extremely long sequences. The core concept behind MST is the partitioning of input sequences into smaller \"mini-sequences,\" which are then processed iteratively to alleviate the burden on intermediate memory usage. The authors claim that MST, when combined with activation recomputation, leads to substantial memory savings in both the forward and backward passes of LLM training. The technique's efficacy is demonstrated through experiments with the Llama2-8B model, where MST reportedly allows for training with sequences up to 12x longer than standard implementations without compromising throughput or convergence. The authors emphasize that MST is a general, implementation-agnostic solution that can be seamlessly integrated into existing LLM training frameworks with minimal code modifications.", "questions": "1. Impact on Model Quality: While the paper demonstrates that MST does not degrade model quality for the evaluated tasks, a more extensive evaluation on a broader range of NLP tasks would be valuable to confirm its impact on model performance across different domains.\n\n2. Sensitivity to Hyperparameters: The paper could provide more insights into the sensitivity of MST's performance to its hyperparameters, such as the number of mini-sequences and the partitioning strategy. This would help users understand how to tune MST for optimal results in different scenarios.\n\n3. Applicability to Other Memory Optimization Techniques: The paper could explore the potential of combining MST with other memory optimization techniques, such as quantization or gradient checkpointing. This could lead to further improvements in memory efficiency and training capabilities.\n\n4. Long-Term Impact on LLM Research: The paper could discuss the potential long-term impact of MST on LLM research and development. Enabling training with longer sequences could open new avenues for exploring larger and more capable LLMs, potentially leading to breakthroughs in various NLP applications.", "limitations": "1. Performance Degradation for Short Sequences: As acknowledged in the paper, MST may lead to performance degradation when applied to models with short sequences due to the overhead introduced by partitioning and gradient accumulation. This limits its applicability to scenarios where long sequence training is not the primary concern.\n\n2. Dependence on Activation Recomputation: The effectiveness of MST is closely tied to its combination with activation recomputation. While this combination yields significant memory savings, it also introduces additional computational overhead, potentially impacting overall training time.\n\n3. Limited Evaluation on Distributed Settings: Although the paper discusses the extension of MST to distributed training, the evaluation in this setting is relatively limited. Further experiments on larger-scale distributed systems would provide more insights into its performance and scalability in real-world scenarios.\n\n4. Potential for Further Optimization: The current implementation of MST may have room for further optimization, particularly in terms of minimizing the overhead associated with partitioning and gradient accumulation. Exploring more efficient implementation strategies could enhance its performance and broaden its applicability.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Novel Approach to Memory Optimization: The paper introduces a new method, MST, to address the critical challenge of memory management in LLM training, particularly for long sequences. The approach of partitioning sequences and iterative processing of mini-sequences is innovative and shows promising results in reducing intermediate memory usage.\n\n2. Significant Improvement in Sequence Length: The experimental results demonstrate that MST enables training with sequences up to 12 times longer than standard implementations, which is a substantial advancement. This capability can significantly impact various NLP tasks that require reasoning over extended contexts.\n\n3. Generality and Implementation Agnostic: The paper emphasizes that MST is a general methodology that can be applied to various LLM training frameworks with minimal code changes. This generality enhances the technique's potential for widespread adoption and impact.\nCompatibility with Distributed Training: The paper also discusses the extension of MST to distributed settings, showcasing its potential for large-scale training scenarios. This scalability further strengthens the technique's applicability to real-world LLM training.\n\n4. Thorough Evaluation: The paper provides a comprehensive evaluation of MST, including ablation studies and comparisons with existing techniques. The results convincingly demonstrate the effectiveness of MST in reducing memory overhead and enabling longer sequence training.", "weaknesses": "1. Limited Model Scope: The experimental evaluation primarily focuses on the Llama2-8B model. While the authors claim that MST is general, further evaluation on a wider range of LLM architectures would strengthen the paper's claims and demonstrate the technique's broader applicability.\n\n2. Lack of Comparison with State-of-the-Art: The paper could benefit from a more extensive comparison with other state-of-the-art memory optimization techniques for LLM training. This would provide a clearer understanding of MST's relative performance and advantages.\n\n3. Potential Impact on Training Time: While MST shows promising results in reducing memory usage, its potential impact on overall training time, especially for shorter sequences, needs further investigation. The paper acknowledges this limitation but could provide a more in-depth analysis of the trade-offs between memory savings and training time.\n\n4. Implementation Complexity: Although the authors claim that MST is easy to integrate, the actual implementation details and potential challenges in adapting it to different training frameworks could be elaborated further. Providing more guidance on implementation could facilitate wider adoption.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n# 1. Further Optimization: Chunk-based Mini-Sequence Transformers\n\nWe've developed chunk-based MST to enhance efficiency in memory management and reduce computational overhead. This method:\n- Partitions input sequences into fixed-size sub-chunks\n- Integrates seamlessly with existing MST design\n\nKey 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.\n\nHyperparameter selection: We've found that setting chunk_size $C = d$ (d being hidden size) yields optimal performance. \nAnalysis:  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.\n\nLM-head uses original MST with hyperparameters $M = V/d$. $V$ represents the vocabulary size.\n\n# 2. Addressing Weaknesses\n\n# W1 Limited Model Scope:\nTo demonstrate the broader applicability of MST, we conducted additional experiments with Qwen, Mistral, and Gemma-2 models:\n\n| Implementation | Maximum Sequence Length (K) |\n|--|--|\n| Mistral-7B-v0.1 vanilla | 5 |\n| Mistral-7B-v0.1 activation recomputation | 42 |\n| **Mistral-7B-v0.1 MST** |  **70** |\n| Qwen2-7B-Instruct vanilla | 4 |\n| Qwen2-7B-Instruct activation recomputation | 13 |\n| **Qwen2-7B-Instruct MST** | **74** |\n| Gemma-2-9b vanilla | 1.5 |\n| Gemma-2-9b activation recomputation | 5 |\n| **Gemma-2-9b MST** | **36** |\n\nThese 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).\n\n# W2 Comparison with Lossy Methods:\nWe've extended our comparisons to include state-of-the-art lossy quantization methods:\n\n| Implementation | Maximum Sequence Length (K) |\n|-|-|\n| 8-bit | 5 |\n| 8-bit + activation checkpointing | 26 |\n| 4-bit | 10 |\n| 4-bit + activation checkpointing | 28 |\n| MST | 60 |\n| **MST + 8-bit** | **110** |\n| **MST + 4-bit** | **140** |\n\nMST 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.\n\n# W3 Training Time:\nWe conducted additional experiments with chunk-based MST to address concerns about training time:\n\nTable 2: MST training with two epochs on LongAlpaca-12k \n| Implementation | Context length | LongAlpaca-12k (ppl) | loss | Training Time (hours) |\n|:-|-:|-:|-:|-:|\n| Act. recomputation | 8k | 9.34 | 2.23 | 25.6 |\n| MST | 8k | 7.41 | 2.003 | 26.5 |\n| MST | 16k | 3.53 | 1.26 | 62.5 |\n| MST | 30k | 3.45 | 1.23 | 233 |\n\nThese results demonstrate that MST can train without significant throughput reduction (<5% is considered negligible).\n\n# W4 Implementation Complexity:\nMST's core idea is conceptually straightforward, primarily targeting MLP and LM-Head blocks. We offer two integration methods:\n\na) Customized Hugging Face Transformer:\nModify existing Hugging Face Transformer implementation.\n```\ngit clone transformer_mst\npip install transformer_mst\n```\n\nb) Wrapper Mode:\nOne-line wrapper for MST:\n```python\nfrom mini-s import mst\nmodel = mst(model)\n```\n\n# Addressing Reviewer Questions\n\n# Q1 Impact on Model Quality:\nTake this evaluation suggestion: Our experiments show that LLAMA3 with MST and a 30K context achieved a 23%- 270% better perplexity than the baseline. \n# Q2 Sensitivity to Hyperparameters:\nThe 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.\n# Q3 Applicability to Other Memory Optimization Techniques:\nOur 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.\n\nAlso, 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.\n\n# Q4 Long-Term Impact on LLM Research:\nWe'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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We sincerely thank all reviewers for their thorough and insightful feedback. We appreciate the recognition of our work's potential impact on long-context LLM training. In response to the valuable comments received, we have conducted additional experiments and provided further implementation optimization for Mini-Sequence Transformer (MST):

# Implementation and Hyperparameter Optimization:
We introduced the chunk-based MST, which addresses concerns about performance degradation for small sequences and simplifies hyperparameter tuning. This approach splits the sequence into fixed-size chunks $C$ (4096 for LLAMA3). It can be considered a carefully selected MST in which the number of mini-sequence $M$ dynamically changes with sequence length, using small $M$ for the small sequence to maintain performance and large $M$ for the larger sequence to save memory. 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 and $S$ is the sequence length. We use chunk-based MST for MLP block to balance speed and memory, while we use the original MST for LM-head with $M≈V/d=32$ for memory saving. Here $V$ is the vocabulary size.

# Training Performance and TFLOPS:
We've introduced new experiments demonstrating that the Mini-Sequence Transformer (MST) achieves comparable or better TFLOPS than activation recomputation for various models, including Llama2-7B and Llama3-8B. The chunk-based MST implementation is deployed here and shows minimal throughput reduction (<5%) compared to vanilla implementations.
| Model Implementation | Batch Size | Training Time Per Step (s) | TFLOPS |
|:-|:-:|-:|-:|
| Llama3-8B-hf vanilla | 1 | OOM | OOM |
| Llama3-8B-hf activation recomputation | 2 | 5.01 | 3271.42 |
| **Llama3-8B-hf MST** | 2 | 5.13 | **3194.90** |
| **Llama3-8B-hf MST**  | 8 | 19.35 | **3386.13** |
| Llama2-7B-hf vanilla | 1 | 1.24 | 3290.88 |
| Llama2-7B-hf activation recomputation | 1 | 1.52 | 2684.67 |
| **Llama2-7B-hf MST without activation recomputation**  | 1 | 1.31 | **3115.03** |
| Llama2-7B-hf activation recomputation | 8 | 8.85 | 3703.48 |
| **Llama2-7B-hf MST**  | 8 | 9.33 | **3511.39** |
| **Llama2-7B-hf MST**  | 16 | 17.92 | **3656.17** |
# Convergence and Long-Context Performance:
New experiments on the LongAlpaca dataset demonstrate MST's effectiveness in long-context understanding. Training Llama3-8B with 30K context length achieved a 270% improvement in perplexity compared to the 8K baseline, showcasing MST's ability to leverage longer contexts effectively.

Table 2: MST training with two epochs on LongAlpaca-12k 
| Implementation | Context length | LongAlpaca-12k (ppl) | loss | Training Time (hours) |
|:-|-:|-:|-:|-:|
| Act. recomputation | 8k | 9.34 | 2.23 | 25.6 |
| MST | 8k | 7.41 | 2.003 | 26.5 |
| MST | 16k | 3.53 | 1.26 | 62.5 |
| MST | 30k | 3.45 | 1.23 | 233 |

# Applicability to Different Architectures:
We've extended our evaluation to include Mistral-7B, Qwen2-7B, and Google's Gemma-2-9B, demonstrating significant increases in maximum sequence length (12x-24x) across these architectures.
| Implementation | Maximum Sequence Length (K) |
|:-|-:|
| Mistral-7B-v0.1 vanilla | 5 |
| Mistral-7B-v0.1 activation recomputation | 42 |
| **Mistral-7B-v0.1 MST** | **70** |
| Qwen2-7B-Instruct vanilla | 4 |
| Qwen2-7B-Instruct activation recomputation | 13 |
| **Qwen2-7B-Instruct MST** | **74** |
| gemma-2-9b vanilla | 1.5 |
| gemma-2-9b activation recomputation | 5 |
| **gemma-2-9b MST** | **36** |

# Comparison and Combination with Lossy Methods:
We've comprehensively compared MST with quantization methods and the combinations between MST and quantization. This comparison demonstrates MST's superiority in enabling longer sequences for Llama3 training on a single A100 GPU. MST alone (60K tokens) outperforms these lossy approaches (4bit 28k). When combined with quantization techniques, MST achieves even more impressive results: MST + 8-bit reaches 110K tokens (a 22x improvement over standard 8-bit), while MST + 4-bit pushes the boundary to 140K tokens.

| Llama3 Implementation | Maximum Sequence Length (K) |
|:-|-:|
| 8-bit | 5 |
| 8-bit + activation checkpointing | 26 |
| 4-bit | 10 |
| 4-bit + activation checkpointing | 28 |
| MST | 60 |
| **MST + 8-bit** | **110** |
| **MST + 4-bit** | **140** |


# Integration and Compatibility:
We've clarified MST's integration process, offering a customized Hugging Face Transformer implementation and a simple wrapper mode for easy adoption. MST is compatible with other optimization techniques, such as activation recomputation and sequence parallelism (also known as context parallelism).
# Broader Impact and Future Directions:
We highlight MST's potential for enabling extremely long context training combined with lossless techniques like context parallelism, opening new avenues for LLM research and applications.

As we know, Llama 3.1 incorporates context parallelism (also known as sequence parallelism) to introduce additional sequences of 128k into the model architecture with 16 A100 GPUs. Our research demonstrates that Multi-Sequence Training (MST) integrates seamlessly with context parallelism, enabling significantly longer sequence processing. We believe this combination has the potential to handle extremely large sequence lengths. Our experiments with updated implementation successfully achieved a sequence length of 480k tokens by combining context parallelism and MST on 8 A100 GPUs. This represents a substantial advancement in Large Language Model (LLM) training capabilities.

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

| Model Implementation | 2 GPUs | 4 GPUs | 8 GPUs |
|:-|-:|-:|--:|
| Llama3-8b-hf MST | 120 | 240 | **480** |
| Llama2-7B-hf MST | 160 | 320 | 640 |

Meta-review: This paper provides an architectural intervention to improve long context performance for language models. MST proposes minibatching along the sequence length for inputs to the MLP and LM-head. With this minibatching they can support much longer sequence training on the same hardware without hurting performance (since the formulation is mathematically the same). The reviewers all seem to agree that this is a novel and important contribution. Furthermore the authors provided detailed responses to all the reviews in the rebuttal period and even developed a novel implementation that addressed one of the reviewers concerns. For these reasons I vote to accept this paper.
Decision: Accept (poster)