Various Lengths, Constant Speed: Efficient Language Modeling with Lightning Attention

TL;DR

Lightning Attention introduces a novel linear attention method that maintains constant training speed across various sequence lengths by splitting attention calculations and optimizing GPU utilization, enabling efficient and accurate language modeling.

Contribution

The paper presents Lightning Attention, a new linear attention approach that overcomes cumsum issues and achieves fixed-speed training across sequence lengths, along with a tailored architecture TNL.

Findings

Lightning Attention maintains constant training speed for different sequence lengths.

TNL outperforms other models in efficiency and matches state-of-the-art performance.

The source code is publicly available for reproducibility.

Abstract

We present Lightning Attention, the first linear attention implementation that maintains a constant training speed for various sequence lengths under fixed memory consumption. Due to the issue with cumulative summation operations (cumsum), previous linear attention implementations cannot achieve their theoretical advantage in a casual setting. However, this issue can be effectively solved by utilizing different attention calculation strategies to compute the different parts of attention. Specifically, we split the attention calculation into intra-blocks and inter-blocks and use conventional attention computation for intra-blocks and linear attention kernel tricks for inter-blocks. This eliminates the need for cumsum in the linear attention calculation. Furthermore, a tiling technique is adopted through both forward and backward procedures to take full advantage of the GPU hardware. To enhance accuracy while preserving efficacy, we introduce TransNormerLLM (TNL), a new architecture that is tailored to our lightning attention. We conduct rigorous testing on standard and self-collected datasets with varying model sizes and sequence lengths. TNL is notably more efficient than other language models. In addition, benchmark results indicate that TNL performs on par with state-of-the-art LLMs utilizing conventional transformer structures. The source code is released at github.com/OpenNLPLab/TransnormerLLM.