大規模モデルの無限ストリーミング入力推論が 46% 上昇しました。国内オープンソースアクセラレーション「ファミリーバケット」が複数ラウンドの会話の長さの制限を突破

大規模言語モデル (LLM) の世界では、複数ターンの会話を処理することは常に課題でした。 MITのGuangxuan Xiao氏らが最近立ち上げたStreamingLLMは、推論速度と生成効果を犠牲にすることなく、複数回の会話で合計400万トークンのストリーミング入力を実現し、推論速度を22.2倍向上させることができます。

ただし、StreamingLLM はネイティブ PyTorch を使用して実装されており、マルチラウンドの対話推論シナリオのアプリケーションに対する低コスト、低レイテンシ、高スループットの要件を満たすために最適化する余地がまだあります。

Colossal-AI チームはSwiftInferをオープンソース化し、 TensorRT に基づく StreamingLLM を実装しました。これにより、大規模モデルの推論パフォーマンスがさらに 46% 向上し、複数ラウンドの対話推論のための効率的で信頼性の高い実装ソリューションが提供されます。

オープンソースアドレス: https://github.com/hpcaitech/SwiftInfer

StreamingLLM の紹介

大規模言語モデルが記憶できるコンテキストの長さは、ChatGPT などの大規模モデル アプリケーションとユーザー間のインタラクションの品質に直接影響します。

LLM は事前トレーニング中に限られた注意ウィンドウの制限下でのみトレーニングできるため、複数ラウンドの対話シナリオで LLM の生成品質を維持するには、推論システムに高い要件が課せられます。

一般的な KV キャッシュ メカニズムはモデルの計算時間を節約できますが、複数ラウンドのダイアログのシナリオでは、キーと値のキャッシュが大量のメモリを消費し、限られたビデオ メモリでコンテキストを無限に拡張することはできません。

同時に、トレーニングされたモデルは、二次的な微調整を行わないと、トレーニング シーケンスの長さよりも長いテキストにうまく一般化できず、生成結果が悪くなります。

出典: https://arxiv.org/pdf/2309.17453.pdf

この問題を解決するために、StreamingLLM は、アテンション モジュールの Softmax の出力を観察することで、アテンション シンクの現象を発見しました。

注意メカニズムは各トークンに注意値を割り当て、テキストの最初のいくつかのトークンは常に多くの無駄な注意を引くことがわかっています。

スライディング ウィンドウ ベースのアテンション メカニズムを使用する場合、これらのトークンがウィンドウから追い出されると、モデルの生成効果はすぐに崩壊します。しかし、これらのトークンがウィンドウ内に保持されている限り、モデルは安定して高品質のテキストを生成できます。

Dense Attention、Window Attention、Sliding Window w/ Re-computing と比較すると、StreamingLLM の Attention シンクに基づく Attention メカニズムは、計算の複雑さと生成効果の両方において優れたパフォーマンスを発揮します。

モデルを再トレーニングする必要がなく、StreamingLLM は現在主流の大規模言語モデルと直接互換性があり、推論パフォーマンスを向上させることができます。

SwiftInfer: TensorRT に基づく StreamingLLM 実装

StreamingLLM テクノロジーを実際のシナリオにより適切に適用するために、Colossal-AI チームは StreamingLLM メソッドと TensorRT 推論最適化を組み合わせることに成功しました。これにより、元の StreamingLLM の利点をすべて継承しただけでなく、より高い操作効率も実現しました。

さらに、TensorRT-LLM API を使用すると、PyTorch API に近いモデル作成エクスペリエンスも得られます。チームは TensorRT-LLM に基づいて、KV キャッシュ メカニズムと位置オフセットを備えたアテンション モジュールを再実装しました。

下の図に示すように、ウィンドウ サイズが 10 トークンであると仮定すると、生成されたトークンが増加すると (黄色の四角で表されます)、KV キャッシュ内の中間のトークンが排除され、テキストの最初のいくつかのトークン (青色の四角で表されます) は常に保持されます。黄色の四角形の位置が変わるため、注目度を計算するときに位置情報も再注入する必要があります。

StreamingLLM は、モデルがアクセスできるコンテキスト ウィンドウを直接増やすわけではありませんが、ストリーミングのマルチラウンド会話をサポートしながら、モデルの生成効果を確保できることに留意する必要があります。

大規模モデルの無制限入力ストリーム推論速度が 46% 向上

StreamingLLM のオリジナルバージョンは、400 万トークンを超えるストリーミング入力を確実に実装でき、計算量の多いスライディング ウィンドウ アテンション メカニズムよりも 22.2 倍の速度向上を実現します。

Colossal-AI チームによってリリースされた SwiftInfer は推論パフォーマンスをさらに向上させ、推論スループット速度を最大 46% 向上させ、大規模モデルのマルチラウンド対話推論に低コスト、低レイテンシ、高スループットのベスト プラクティスを提供します。 TensorRT-LLM チームも同じ時期に StreamingLLM に対して同様のサポートを提供しました。

Colossal-AI コミュニティニュース

Colossal-AI は現在、GitHub スターを 35,000 個以上獲得しており、グローバル TOP400 にランクインし、ニッチ分野では世界第 1 位を獲得しています。効率的な多次元並列処理と異種メモリにより、大規模な AI モデルのトレーニング/微調整/推論の開発およびアプリケーション コストを削減し、モデル タスクのパフォーマンスを向上させ、GPU 要件を削減できます。主流のオープンソース AI 大規模モデル システム コミュニティとして、Colossal-AI エコシステムは多くの面でアクティブであり、更新され続けています。

Colossal-LLaMA-2-13B オープンソース

Colossal-LLaMA-2-13B モデルは、わずか 250 億トークンのデータと 10,000 元の計算能力しか使用せず、そのパフォーマンスは LLaMA-2 に基づく他の中国製モデルをはるかに上回ります。

中国語のコーパスを使用し、ゼロから事前トレーニングするのに数千万元の費用がかかる可能性のある他の有名なモデルと比較しても、Colossal-LLaMA-2 は同じ規模でも依然として優れたパフォーマンスを発揮します。

バージョン 13B では、より完全なデータ システムを構築することにより、知識コンテンツの習得、自然言語処理タスクの理解、セキュリティ、価値などの問題において質的な向上が達成されました。

巨大なAIクラウドプラットフォーム

Colossal-AI システムの最適化と安価なコンピューティング能力の統合に基づき、Colossal-AI クラウド プラットフォームは最近 AI クラウド ホスト機能をリリースしました。これにより、ユーザーはベアメタルと同様の方法で大規模な AI モデルを開発およびデバッグできます。また、Jupyter Notebook、ssh、サービス ローカル ポート マッピング、Grafana モニタリングなど、さまざまな使用方法を提供し、ユーザーに包括的で便利な開発体験を提供します。

同時に、ColossalAI コード リポジトリと動作環境を含む Docker イメージがユーザー向けに事前に作成されています。ユーザーは、環境やリソースの設定をすることなく、ワンクリックで ColossalAI コード リポジトリ内のコード サンプルを実行できます。

Colossal-AI オープンソース アドレス: https://github.com/hpcaitech/ColossalAI