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.

我们提出了闪电注意力（Lightning Attention），这是第一个在固定的内存消耗下保持不同序列长度的训练速度恒定的线性注意力实现。

不同长度，匀速进行：高效语言建模与闪电注意力

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

We present TransNormerLLM, the first linear attention-based Large Language
Model (LLM) that outperforms conventional softmax attention-based models in
terms of both accuracy and efficiency. TransNormerLLM evolves from the previous
linear attention architecture TransNormer by making advanced modifications that
include positional embedding, linear attention acceleration, gating mechanism,
tensor normalization, inference acceleration and stabilization. Specifically,
we use LRPE together with an exponential decay to avoid attention dilution
issues while allowing the model to retain global interactions between tokens.
Additionally, we propose Lightning Attention, a cutting-edge technique that
accelerates linear attention by more than twice in runtime and reduces memory
usage by a remarkable four times. To further enhance the performance of
TransNormer, we leverage a gating mechanism to smooth training and a new tensor
normalization scheme to accelerate the model, resulting in an impressive
acceleration of over 20%. Furthermore, we have developed a robust inference
algorithm that ensures numerical stability and consistent inference speed,
regardless of the sequence length, showcasing superior efficiency during both
training and inference stages. Scalability is at the heart of our model's
design, enabling seamless deployment on large-scale clusters and facilitating
expansion to even more extensive models, all while maintaining outstanding
performance metrics. Rigorous validation of our model design is achieved
through a series of comprehensive experiments on our self-collected corpus,
boasting a size exceeding 6TB and containing over 2 trillion tokens. To ensure
data quality and relevance, we implement a new self-cleaning strategy to filter
our collected data. Our pre-trained models will be released to foster community
advancements in efficient LLMs.

我们提出了 TransNormerLLM，这是第一个基于线性注意力的大型语言模型（LLM），在准确性和效率方面均超过了传统的基于 softmax 注意力的模型。