swiftLLM

SwiftLLM

このプロジェクトはまだ開発中です。一部の機能はまだ実装されていない可能性があり、ドキュメントが不完全である可能性があります。

研究目的に合わせて調整された、小型ながら強力な LLM 推論システム。

わずか 2,000 行のコード (vLLM の 2%) で vLLM と同等のパフォーマンス。

SwiftLLM を選ぶ理由

HuggingFace Transformers、vLLM、LightLLM、DistServe、DeepSpeed-MII など、LLM サービス用のオープン ソース フレームワークが多数あります。なぜ SwiftLLM なのか?

その理由は、これらのフレームワークが研究用ではなく実稼働用に調整されているためです。 100 を超えるモデルのサポート、さまざまなハードウェア サポート、LoRA、量子化、マルチモーダル、プレフィックス キャッシング、ビーム サーチなどの多数の機能が装備されています。実稼働向けのオールインワン ソリューションではありますが、コードベースが大きすぎて複雑すぎて理解や変更ができないため (たとえば、vLLM のコード行数は 100,000 行を超えます)、研究目的で使用するのは困難です。また、彼らの歴史的負担も問題です。

SwiftLLM は、研究目的に合わせた小型ながら強力な LLM 推論システムとして設計されています。 「小型」は研究に不可欠な機能のみを保持することを意味し、「強力」はパフォーマンスに一切の妥協がないことを意味し、最後に「迅速」は理解と変更が容易であることを意味します。基本機能 (以下のリストを参照) をサポートし、vLLM と同等のパフォーマンスを達成できる一方で、SwiftLLM のコードベースは2,000行未満のコード (vLLM の約 2%) であり、Python と OpenAI Triton (書き込み用の DSL) で書かれています。 CUDA カーネル）により、読み取り、変更、デバッグ、テスト、拡張が容易になり、斬新で素晴らしい研究アイデアと簡単に統合できます。

機能一覧

現在、SwiftLLM は次の機能をサポートしています。

• 反復スケジュールと選択的バッチ処理 (Orca で提案)
• PagedAttenton (vLLM で提案、論文)
• LLaMA / LLaMA2 / LLaMA3 モデル (リンク) とそのバリアント
• ピギーバッキングのプリフィルとデコード (SARATHI で提案)
• フラッシュ アテンション (FlashAttendant および FlashAttendant-2 で提案)
• ページド アテンション v2 (フラッシュ デコーディングとも呼ばれ、ここで提案)

また、将来的には次の機能のサポートを追加する予定です。

• テンソル並列処理とパイプライン並列処理

コードベースを小さく保つため、次の機能はサポートされません。研究プロジェクトで使用したい場合は、自分で実装する必要がある場合があります。

• 量子化
• LoRA
• マルチモーダル
• LLaMAのアーキテクチャに従わないモデル
• 貪欲サンプリング以外のサンプリング方法
• NVIDIA GPU 以外のハードウェアもサポート (ただし、OpenAI Triton がサポートしている限り、他のハードウェアへの移行は簡単です)

SwiftLLM は実稼働向けのオールインワン ソリューションではないことに注意してください。これは研究プロジェクトの「基礎」として考えることをお勧めします。また、一部の機能を自分で実装する必要がある場合もあります。研究者の皆様、コードを読み、理解し、修正し、研究のニーズに合わせて拡張することをお勧めします。

建築

SwiftLLM のアーキテクチャは、コントロール プレーンとデータ プレーンの 2 つの主要な部分に分割できます。

簡単に説明すると、コントロールプレーンは「何を計算するか」または「どのようにスケジュールするか」を決定し、データプレーンは「どのように計算するか」または「どのように実装するか」を決定し、具体的な計算を実行します。これらはマスター/ワーカー方式で動作します。コントロール プレーンはマスターのように機能し、高レベルのスケジューリングと調整を実行し、データ プレーンにジョブを送信します。データ プレーンは、低レベルの計算を実行するワーカーのように機能します。

コントロール プレーンのコードは、 Engine 、 Scheduler 、API サーバー、 TokenizationEngineなどのコンポーネントを含むswiftllm/serverディレクトリにあります。データ プレーンのコードはswiftllm/workerディレクトリにあり、計算グラフ ( swiftllm/worker/model.py内)、モデル内のレイヤーの実装 ( swiftllm/layers内)、および OpenAI Triton カーネル ( 「カーネル」は GPU 上で実行される関数として想像できます ( swiftllm/kernels内)。

例として、おもちゃの API サーバー ( swiftllm/server/api_server.pyにあります) を見てみましょう。

• 起動時に、 EngineConfig使用してEngine作成します。
• その後、エンジンはEngine.initializeによって初期化され、そこでScheduler 、 TokenizationEngine 、およびワーカーのセット (Tensor Parallelism がサポートされていないため、現在は 1 つだけ) が作成されます。次に、ワーカーにprofile_num_blocks実行して GPU ブロックの数を計算するよう命令し、その後エンジンはすべてのワーカーに KV キャッシュと KV スワップを割り当てるよう命令します。
• 最後に、イベント ループはEngine.start_all_event_loopsによってアクティブ化されます。ループの各ステップで、エンジンは計算するリクエストの次のバッチをスケジューラにクエリし、ワーカーにスワップイン/スワップアウトを実行するように命令してから、計算するためにバッチをワーカーに送信します。
• API サーバーはユーザーのリクエストをリッスンし、エンジンと対話してリクエストを実行します。

現在、コントロール プレーン ( Engine ) とデータ プレーン ( LlamaModel ) は同じノード上に存在します。 Tensor Parallelism / Pipeline Parallelism が実装された後、データ プレーンは複数のノードに分散される場合があります。

使用方法

SwiftLLM を使用するには、コントロール プレーンとデータ プレーンの両方を使用する方法、またはデータ プレーンのみを使用する方法の 2 つを提供します。

アイデアが既存のコントロール プレーンにシームレスに統合できるほどシンプルまたは洗練されている場合は、コントロール プレーンとデータ プレーンの両方を使用できます。素晴らしい IDE を実装したい場合は、データ プレーンのみを利用し、新しいコントロール プレーンを自分で実装することもできます。

構築して実行する

まず環境をセットアップしましょう。

• Python 3.9 以上のクリーンな conda 環境から開始することも、既存の環境を使用することもできます。 conda 環境を使用している場合は、忘れずにアクティブ化してください。
• PyTorchをインストールします。ハードウェアに基づいて正しいバージョンを選択するように注意してください。
• pip install packaging使用してpackagingインストールする

そしてインストールが始まります。

• git clone https://github.com/interestingLSY/swiftLLM.git
• cdリポジトリ ( cd swiftLLM ) に移動し、 pip install -r requirements.txtを介して他の依存関係をインストールします。
• PyTorch は OpenAI Triton の安定バージョンをインストールする場合があります。最先端のパフォーマンスを得るためにナイトリー バージョンを使用したいが、潜在的な問題がある場合は、それをアンインストールしてナイトリー バージョンをインストールできます。
• pip install -e . SwiftLLM を環境にインストールします。
• pip install -e csrcを介していくつかの C バインディングをインストールします。

以下にいくつかの例を示します。

• 現在、SwiftLLM は、HuggingFace からのウェイトの自動ダウンロードをサポートしていません。最初に HuggingFace からモデル ウェイトを複製またはダウンロードする必要がある場合があります。 .bin形式と.safetensors形式の両方がサポートされています。モデルの重みが/data/to/weight/に保存されていると仮定します。
• オフラインでのサービス提供の例については、 python3 examples/offline.py --model-path /data/to/weightを試してください。この例では、データ プレーンのみを利用します。コントロール プレーンなしで SwiftLLM を使用する予定がある場合は、これが良い出発点になります。
• Engineを使用するオンライン サービングのサンプルについては、 python3 examples/online.py --model-path /data/to/weightを試してください。これは、コントロール プレーンとデータ プレーンの両方を使用する予定がある場合に最適な例です。
• より複雑な例については、 swiftllm/server/api_server.pyを参照してください。 API サーバーを起動し、オンライン サービス用の vLLM のようなインターフェイスを提供します。

パフォーマンス

SwiftLLM は小さいにもかかわらず (小さいものも愛らしいです!)、パフォーマンスに一切の妥協がありません。私たちはいくつかのシナリオで SwiftLLM を評価し、SwiftLLM が vLLM と比較して同等、またはそれ以上のパフォーマンスを達成できることを実証しました。

単一の転送操作

最初のシナリオは「単一の前方操作」です。モデルに入力のバッチを供給し、1 つの出力トークン (1 つの「前方」操作と同等) を生成させます。これは LLM 推論 (オンラインとオフラインの両方) の基本的な操作であるため、そのパフォーマンスは非常に重要です。

ここでは、FP16 精度で NVIDIA A100 80G PCIE / RTX 4090 GPU を搭載した LLaMA-3 7B モデルを使用します。結果を以下に示します (数値が低いほど優れています)。

SwiftLLM は、同じ設定下で vLLM と同等のパフォーマンス (またはそれを上回るパフォーマンス) を達成できることがわかります。

オンラインサービス

2 番目のシナリオは「オンライン サービング」です。ここでは、API サーバーを起動し、実世界のデータセットからプロンプトをサンプルし、モデルに補完を生成させます。これは、チャットボットやコード補完などの実世界のアプリケーションで LLM が使用されるシナリオです。

ここでは、ShareGPT データセットを使用してプロンプトをサンプリングし、さまざまなラムダでポアソン プロセスを使用して、さまざまなリクエスト到着率をシミュレートします。結果を以下に示します (数値が低いほど優れています)。

A100 80G PCIE では、SwiftLLM が vLLM と同等のパフォーマンスを達成できるのに対し、RTX 4090 では、SwiftLLM が vLLM を大幅に上回っていることがわかります (主に、コントロール プレーンのオーバーヘッドが低いため)。