[논문 리뷰] TransNormerLLM: A Faster and Better Large Language Model with Improved TransNormer

Zhen Qin, Li Dong외 9명|arXiv (Cornell University)|2023. 07. 27.

Topic Modeling인용 수 5

한 줄 요약

TransNormerLLM는 선형 어텐션 기반의 대규모 언어 모델로, 지수 감쇠가 적용된 LRPE, 트레이닝을 가속화하는 Lightning Attention, 게이팅 메커니즘, 텐서 정규화, 강력한 추론 알고리즘을 통합함으로써 기존의 Transformer 모델보다 정확도와 효율성 면에서 뛰어난 성능을 발휘한다. 모델은 트레이닝 속도를 2배 이상 빠르게 하고 메모리 사용량을 4배 감소시키며, 385M에서 175B 파라미터 스케일에 걸쳐 높은 성능을 유지한다.

ABSTRACT

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 mechanisms, tensor normalization, and 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 for smooth training and a new tensor normalization scheme to accelerate the model, resulting in an impressive acceleration of over $20\%$. Furthermore, we develop 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. We also implement an efficient model parallel schema for TransNormerLLM, enabling seamless deployment on large-scale clusters and facilitating expansion to even more extensive models, i.e., LLMs with 175B parameters. We validate our model design through a series of ablations and train models with sizes of 385M, 1B, and 7B on our self-collected corpus. Benchmark results demonstrate that our models not only match the performance of state-of-the-art LLMs with Transformer but are also significantly faster. Code is released at: https://github.com/OpenNLPLab/TransnormerLLM.

연구 동기 및 목표

• 기존의 소프트맥스 어텐션 기반 Transformer보다 정확도와 추론 효율성 면에서 뛰어난 대규모 언어 모델을 개발하는 것.
• 기존 선형 어텐션 기법의 한계, 즉 어텐션 분산 현상과 실제 성능 저하 문제를 아키텍처적 및 최적화 기법의 혁신을 통해 해결하는 것.
• 효율적인 모델 병렬 처리와 메모리 최적화를 통해 초대규모 모델(최대 175B 파라미터)의 확장 가능한 트레이닝 및 배포를 가능하게 하는 것.
• 다양한 시퀀스 길이에 걸쳐 수치 안정성과 일관된 추론 속도를 확보하여 실무적 배포 가능성 향상.
• 2조 토큰이 넘는 대규모 자가 수집 코퍼스(6TB, 2조 토큰 이상)를 기반으로 한 모델 검증을 통해 강력한 일반화 능력과 확장성을 입증하는 것.

제안 방법

• 전역 토큰 상호작용을 향상시키기 위해 TransNormer의 DiagAttention를 선형 어텐션으로 대체하여 복잡도를 선형으로 유지.
• 주의 분산 현상을 완화하면서 장거리 의존성을 유지하기 위해 지수 감쇠가 적용된 LRPE(Learned Relative Positional Encoding) 도입.
• 입출력(I/O)에 민감한 기법인 Lightning Attention 개발으로, 선형 어텐션의 트레이닝 속도를 2배 이상 높이고 메모리 사용량을 4배 감소.
• 학습 안정화를 위한 게이팅 메커니즘과 추론 가속화를 위한 새로운 텐서 정규화 기법을 적용해 20% 이상의 속도 향상 달성.
• 다양한 시퀀스 길이에 관계없이 수치 안정성과 일관된 추론 속도를 보장하는 강력한 추론 알고리즘 설계.
• 최대 8방향 병렬 처리를 지원하는 효율적인 모델 병렬 스키마 구현으로 대규모 클러스터에서의 배포와 175B 파라미터 모델의 확장성 확보.

실험 결과

연구 질문

• RQ1 선형 어텐션 기반 LLM이 기존의 표준 Transformer 모델보다 정확도와 효율성 면에서 모두 뛰어나게 할 수 있는가?
• RQ2 장거리 시퀀스 모델링 능력을 유지하면서 선형 어텐션 기법에서 발생하는 어텐션 분산 현상을 어떻게 완화할 수 있는가?
• RQ3 Lightning Attention와 같은 새로운 가속 기법이 초대규모 LLM의 트레이닝 시간과 메모리 소비를 얼마나 줄일 수 있는가?
• RQ4 통합된 추론 알고리즘이 다양한 시퀀스 길이에 걸쳐 일관된 속도와 수치 안정성을 유지할 수 있는가?
• RQ5 대규모 데이터셋에서 트레이닝되고 대규모 GPU 클러스터에 분산 배포된 선형 어텐션 기반 LLM의 확장성 한계는 어디까지인가?

주요 결과

• TransNormerLLM-7B는 동일한 조건에서 1GPU당 4081.0 토큰/초의 속도를 기록하여, Transformer-7B의 3362.7 토큰/초를 초월한다.
• 8방향 모델 병렬 처리를 적용한 TransNormerLLM-7B는 GPU 메모리 사용량을 1× 구성 대비 62.3% 감소시켜 24.1GB로 줄였다.
• 175B 파라미터 모델에서 TransNormerLLM는 동일한 컨텍스트 길이에서 기준 Transformer 모델 대비 상대적 트레이닝 속도 1.35배를 달성했다.
• 강력한 추론 알고리즘 덕분에 다양한 시퀀스 길이에 걸쳐 일관된 추론 속도를 유지하여 예측 가능한 성능을 보였다.
• Lightning Attention는 메모리 사용량을 4배 감소시키고 트레이닝 속도를 2배 이상 높여 실제 트레이닝 환경에서 뚜렷한 효율성 향상을 입증했다.
• 자신이 수집한 6TB, 2조 토큰 이상의 코퍼스에서 경쟁력 있는 성능을 달성했으며, 유사 크기의 모델들에 비해 최신 기술 수준의 성능을 충족하거나 초월했다.