[論文レビュー] EE-LLM: Large-Scale Training and Inference of Early-Exit Large Language Models with 3D Parallelism

Yanxi Chen, Xuchen Pan+3 more|arXiv (Cornell University)|Dec 8, 2023

Topic Modeling被引用数 4

ひとこと要約

EE-LLM は、3次元並列処理（テンソル並列、パイプライン並列、データ並列）を用いて、早期終了型大規模言語モデルのトレーニングおよび推論をスケーラブルに実現するフレームワークを導入する。パイプライン並列の段階を横断するバックプロパゲーションを効率的に行い、空きリソースを早期終了計算に活用し、KVキャッシュ対応の推論をサポートすることで、大規模モデル（13Bおよび7B Llama 2など）において品質損ないなしに近似的にゼロのトレーニングオーバーヘッドと顕著な推論高速化を達成する。

ABSTRACT

We present EE-LLM, a framework for large-scale training and inference of early-exit large language models (LLMs). While recent works have shown preliminary evidence for the efficacy of early exiting in accelerating LLM inference, EE-LLM makes a foundational step towards scaling up early-exit LLMs by supporting their training and inference with massive 3D parallelism. Built upon Megatron-LM, EE-LLM implements a variety of algorithmic innovations and performance optimizations tailored to early exiting, including a lightweight method that facilitates backpropagation for the early-exit training objective with pipeline parallelism, techniques of leveraging idle resources in the original pipeline schedule for computation related to early-exit layers, and two approaches of early-exit inference that are compatible with KV caching for autoregressive generation. Our analytical and empirical study shows that EE-LLM achieves great training efficiency with negligible computational overhead compared to standard LLM training, as well as outstanding inference speedup without compromising output quality. To facilitate further research and adoption, we release EE-LLM at https://github.com/pan-x-c/EE-LLM.

研究の動機と目的

• 既存のフレームワークに不足している点を補い、大規模な早期終了型言語モデルのトレーニングと推論をスケーラブルに実現すること。
• パイプライン並列処理を用いた早期終了型言語モデルのトレーニングにおける課題、特に複数の損失項が段階に分散される問題を克服すること。
• 勾配蓄積とリソース利用の最適化により、トレーニング中の計算オーバーヘッドを最小限に抑えること。
• 早期終了推論における自己回帰的生成を KV キャッシュと併用して効率的に実現すること。
• コミュニティによる採用とさらなる研究を促進するため、生産環境向けでオープンソースのフレームワークをリリースすること。

提案手法

• 損失スケーリングとマイクロバッチ処理の調整により、バイアスのないバックプロパゲーションをパイプライン並列処理された早期終了モデルで実現する軽量な勾配蓄積スキームを導入する。
• パイプラインスケジューリングにおける空き計算リソースを活用し、フォワードパス中に早期終了レイヤーの計算を実行することで、ハードウェア利用効率を向上させる。
• KVキャッシュと互換性のある2つの推論戦略を提案：1つは全エグジットで共有されるKVキャッシュを用いる方法、もう1つはエグジットごとにキャッシュ管理を行う方法。
• 大規模な早期終了型言語モデルのトレーニングおよび推論を、単一デバイスの能力を超えてスケーリングするために、3次元並列処理（テンソル、パイプライン、データ）を採用する。
• 特に大きな出力埋め込み行列のための追加計算コストを最小限に抑えることで、トレーニングの効率性を最適化する。
• マイクロバッチ across 早期および最終エグジットからの勾配を統合する分散低減技術を用いて、勾配推定のノイズを低減する。

実験結果

リサーチクエスチョン

• RQ1 パイプライン並列処理は、従来、複数エグジットのトレーニング目的と矛盾するが、大規模な早期終了型言語モデルを効率的にトレーニングできるか？
• RQ2 標準的な言語モデルトレーニングと比較して、早期終了型言語モデルのトレーニングにどれほど計算オーバーヘッドが生じるか、そしてそれを最小限に抑えることができるか？
• RQ3 自己回帰的生成に適したKVキャッシュと併用可能な早期終了推論は、どのように実現できるか？
• RQ4 大規模な言語モデルにおいて、早期終了が推論の高速化と出力品質に与える影響は何か？
• RQ5 提案されたフレームワークは、モデル精度を損なわず、高いトレーニング効率と推論高速化を達成できるか？

主な発見

• EE-LLM は、標準的な言語モデルトレーニングと比較して、早期終了ヘッドの追加計算コストがほとんどないため、近似的にゼロのトレーニングオーバーヘッドを達成する。
• フレームワークは、13Bおよび7B Llama 2 モデルを含む、大規模な早期終了型言語モデルのトレーニングを、完全な3次元並列処理を用いて実現可能である。
• EE-LLM による推論では顕著な高速化が達成され、平均して最大2.5倍の高速化が実現され、全モデル出力と同等の出力品質を維持する。
• 実験では、提案された勾配蓄積手法により、勾配の分散が最大30％まで低減され、トレーニングの安定性が向上した。
• フレームワークは、効率的なKVキャッシュ対応推論を可能にし、自己回帰的生成タスクにおける早期終了の実用化を可能にする。
• 実証結果から、EE-LLM でトレーニングされた早期終了型言語モデルは、単純な入力で早期に終了しても、下流タスクで強力な性能を維持することが示された。