라이트닝 어텐션-2: 대규모 언어 모델에서 무한한 시퀀스 길이를 처리하기 위한 무료 점심

Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models

January 9, 2024

저자: Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun, Yiran Zhong

cs.AI

초록

선형 어텐션은 최근 기존의 소프트맥스 어텐션을 대체할 유망한 메커니즘으로 등장한 효율적인 어텐션 기법이다. 선형 어텐션은 토큰을 선형 계산 복잡도로 처리할 수 있는 능력을 통해 이론적으로는 속도를 희생하지 않고도 무한한 길이의 시퀀스를 다룰 수 있으며, 즉 고정된 메모리 소비로 다양한 시퀀스 길이에 대해 일정한 학습 속도를 유지할 수 있다. 그러나 누적 합산(cumsum) 문제로 인해 현재의 선형 어텐션 알고리즘들은 인과적 설정에서 이론적 이점을 발휘하지 못하고 있다. 본 논문에서는 선형 어텐션이 이론적 계산 이점을 실현할 수 있도록 하는 최초의 선형 어텐션 구현체인 라이트닝 어텐션-2를 제안한다. 이를 위해 타일링(tiling) 개념을 활용하여 선형 어텐션 계산에서 블록 내(intra-block)와 블록 간(inter-block) 요소를 별도로 처리한다. 구체적으로, 블록 내에서는 기존의 어텐션 계산 메커니즘을 사용하고, 블록 간에는 선형 어텐션 커널 트릭을 적용한다. GPU 하드웨어의 장점을 최대한 활용하기 위해 순방향 및 역방향 과정 모두에서 타일링 기법을 채택하였다. 우리는 이 알고리즘을 Triton으로 구현하여 IO 인지적이고 하드웨어 친화적으로 만들었다. 다양한 모델 크기와 시퀀스 길이에 대해 실험을 수행한 결과, 라이트닝 어텐션-2는 입력 시퀀스 길이에 관계없이 일관된 학습 및 추론 속도를 유지하며 다른 어텐션 메커니즘보다 훨씬 빠른 성능을 보였다. 소스 코드는 https://github.com/OpenNLPLab/lightning-attention에서 확인할 수 있다.

English

Linear attention is an efficient attention mechanism that has recently emerged as a promising alternative to conventional softmax attention. With its ability to process tokens in linear computational complexities, linear attention, in theory, can handle sequences of unlimited length without sacrificing speed, i.e., maintaining a constant training speed for various sequence lengths with a fixed memory consumption. However, due to the issue with cumulative summation (cumsum), current linear attention algorithms cannot demonstrate their theoretical advantage in a causal setting. In this paper, we present Lightning Attention-2, the first linear attention implementation that enables linear attention to realize its theoretical computational benefits. To achieve this, we leverage the thought of tiling, separately handling the intra-block and inter-block components in linear attention calculation. Specifically, we utilize the conventional attention computation mechanism for the intra-blocks and apply linear attention kernel tricks for the inter-blocks. A tiling technique is adopted through both forward and backward procedures to take full advantage of the GPU hardware. We implement our algorithm in Triton to make it IO-aware and hardware-friendly. Various experiments are conducted on different model sizes and sequence lengths. Lightning Attention-2 retains consistent training and inference speed regardless of input sequence length and is significantly faster than other attention mechanisms. The source code is available at https://github.com/OpenNLPLab/lightning-attention.