Mini-Sequence Transformer: Optimizing Intermediate Memory for Long Sequences Training (2407.15892v4)

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.

• The paper introduces a novel Mini-Sequence Transformer (MsT) that reduces intermediate memory through input partitioning and activation recomputation.
• It demonstrates that Llama3-8B can be trained on sequences up to 60k tokens while maintaining similar throughput to conventional methods.
• The approach enables scalable long-sequence training with practical benefits for extended context NLP tasks and distributed model setups.

Mini-Sequence Transformer: Optimizing Intermediate Memory for Long Sequences Training

The paper introduces the Mini-Sequence Transformer (MsT), a methodology designed to address the memory constraints typically associated with training Long LLMs on extremely long sequences. The approach leverages both input partitioning into mini-sequences and activation recomputation, allowing for significant memory savings without sacrificing throughput or convergence. Experiments with the Llama3-8B model demonstrate that MsT can handle sequences up to 12 times longer than standard implementations without performance degradation.

Introduction

Transformers have revolutionized NLP, but their memory demands have grown exponentially with model size and complexity, driving the need for optimization techniques. The introduction of multi-query attention (MQA) and grouped query attention (GQA) has previously helped manage inference memory by reducing KV-cache sizes. However, the problem of increased intermediate activation memory in MLP and LM-Head layers persists, especially as these models adopt larger vocabularies and intermediate dimensions.

Methodology

MsT reduces intermediate memory by partitioning input sequences and processing these mini-sequences iteratively. By integrating activation recomputation, significant memory savings are achieved during both forward and backward passes. The technique splits the input sequence dimensionally, working sequentially through smaller subsections while accumulating results to form full-sequence outputs, effectively handling long-sequence processing.

Experimental Results

Key contributions of the paper are:

• Scaling Sequence Length: MsT enabled training Llama3-8B with context lengths up to 60k and Llama2-7B up to 84k on a single A100 GPU, significantly surpassing the capabilities of activation recomputation and standard implementations.
• Training Throughput: Despite managing longer sequences, MsT maintained similar training speeds to conventional methods. For instance, with a batch size of 2, the MsT approach showed minor reductions in throughput but achieved substantial memory savings, enabling larger batch sizes and, consequently, faster training.

Analysis and Distributed Extensions

The work includes a detailed analysis of the memory hierarchy and performance characteristics of Transformer architectures. MsT's adaptations to MLP and LM-Head layers specifically mitigate the memory overhead by reducing intermediate value size. The IO complexity reveals how MsT scales effectively for long sequences by optimizing the use of GPU memory without incurring significant computational overhead.

Furthermore, the integration with DeepSpeed-Ulysses demonstrates the approach's potential in distributed training environments. This integration allows for linear scalability of sequence lengths with a proportional increase in the number of GPUs, enhancing the practical utility of MsT for extensive LLM training.

Implications and Future Directions

The MsT methodology's ability to manage significantly longer sequences without degrading performance has profound implications:

• Practical Applications: Tasks that require extensive context, such as long document summarization and conversational AI, can benefit from the longer context windows MsT enables. This capability can enhance the performance of models in generating coherent and contextually accurate outputs over extended inputs.
• Theoretical Advancements: From a theoretical perspective, MsT provides a new approach to managing long-sequence training, extending the practical limits of sequence length without necessitating architectural changes. The approach is highly generic and can be adapted across various LLM training frameworks.

Future work may focus on:

1. Optimization through CUDA: Implementing and optimizing the MsT methodology directly in CUDA can further reduce intermediate memory overhead and improve computational efficiency.
2. Balancing Performance: As MsT introduces some overhead in short-sequence scenarios, further work is needed to optimize its use dynamically based on sequence length, ensuring it is beneficial across a broader range of tasks.
3. Hybrid Approaches: Combining MsT with other memory-efficient techniques, such as quantization or sparse attention mechanisms, may yield even greater efficiencies.

Conclusion

Mini-Sequence Transformer (MsT) effectively addresses the memory challenges of training extremely long sequences in LLMs. By partitioning input sequences and utilizing activation recomputation, it achieves significant memory savings and allows for maintaining training speeds. This approach supports practical applications requiring long context and contributes theoretically by providing a system-agnostic, scalable solution. Future research may focus on further optimizations and extending the benefits of MsT across more varied training scenarios and models.