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

TL;DR

Mini-Sequence Transformer (MsT) is a novel method that enables efficient training of large language models with extremely long sequences by partitioning input and reducing memory usage without sacrificing performance.

Contribution

MsT introduces a simple, general approach to extend sequence length in LLM training, significantly reducing memory requirements with minimal code modifications.

Findings

MsT achieves 12x longer sequences without throughput loss.

MsT extends context length of several models by 12-24x.

Memory savings enable training with longer sequences.

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.