[論文レビュー] TransNormerLLM: A Faster and Better Large Language Model with Improved TransNormer

Zhen Qin, Li Dong+9 more|arXiv (Cornell University)|Jul 27, 2023

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 兆トークンを超える大規模で自作収集したコーパスを用いたモデルの妥当性を検証し、強力な汎化性とスケーラビリティを示すこと。

提案手法

• トランスフォーマーの DiagAttention を線形アテンションに置き換えることで、線形計算量を達成し、グローバルなトークン相互作用を向上させる。
• アテンションの希釈を軽減しつつ長距離依存関係を保持するため、指数関数的減衰を伴う学習済み相対位置エンコーディング（LRPE）を導入。
• Lightning Attention を開発。これは入出力（I/O）に配慮した新技術であり、トレーニング中、線形アテンションの処理を 2 倍以上高速化し、メモリ使用量を 4 分の 1 に削減する。
• 訓練の安定化のためゲーティング機構を適用し、推論の高速化を図る新しいテンソル正規化方式を導入。これにより、20% 以上の速度向上を達成。
• さまざまなシーケンス長にかかわらず数値安定性と一定の推論速度を保証する堅牢な推論アルゴリズムを設計。
• 最大 8 方向並列化をサポートする効率的なモデル並列スキームを実装。これにより、大規模クラスタ上で動作し、175B パラメータのモデルにスケーリング可能。

実験結果

リサーチクエスチョン

• RQ1 線形アテンションベースの LLM は、標準的な Transformer モデルと比較して、より高い精度とより優れた効率性を達成できるか？
• RQ2 線形アテンション機構におけるアテンションの希釈は、どのように軽減可能か？ また、長距離シーケンスモデリング能力を保持したまま？
• RQ3 Lightning Attention のような新規加速技術は、大規模 LLM のトレーニング時間とメモリ消費量をどの程度削減できるか？
• RQ4 統一された推論アルゴリズムは、さまざまなシーケンス長にわたり、一貫した速度と数値安定性を維持できるか？
• RQ5 大規模データセットでトレーニングされ、大規模 GPU クラスタに分散された線形アテンションベースの LLM のスケーラビリティの限界はどこか？

主な発見

• TransNormerLLM-7B は、同一条件で 1 GPU 当たり 4081.0 tokens/sec を達成し、Transformer-7B の 3362.7 tokens/sec/GPU を上回った。
• 8 方向のモデル並列化を適用した場合、TransNormerLLM-7B の GPU メモリ使用量は 24.1GB にまで低下し、1× 設定と比較して 62.3% 減少した。
• 175B パラメータのモデルにおいて、同じコンテキスト長で、TransNormerLLM は Transformer ベースライン比で相対的なトレーニング速度を 1.35× に向上させた。
• 堅牢な推論アルゴリズムのおかげで、シーケンス長にかかわらず一貫した推論速度を維持しており、予測可能なパフォーマンスを実現した。
• Lightning Attention は、4 倍のメモリ削減と 2 倍以上のトレーニング加速を実現し、実世界のトレーニングシナリオにおいて顕著な効率性向上を示した。
• 6TB、2 兆トークンの自作収集コーパス上で、同サイズのモデルと比較して競争力のあるパフォーマンスを達成し、最先端の結果を同等または上回った。