swiftLLM

史威夫特法學碩士

該項目仍在開發中。某些功能可能尚未實現，且文件可能不完整。

一個專為研究目的而定制的小型但功能強大的法學碩士推理系統。

僅 2k 行程式碼（vLLM 的 2%）即可獲得與 vLLM 相當的效能。

為什麼選擇 SwiftLLM

用於 LLM 服務的開源框架有很多，包括 HuggingFace Transformers、vLLM、LightLLM、DistServe 和 DeepSpeed-MII。為什麼選擇 SwiftLLM？

原因是，這些框架是為生產量身定制的，而不是為研究量身定制的。它們配備了眾多功能，例如 100 多個模型支援、各種硬體支援、LoRA、量化、多模態、前綴快取、波束搜尋等。雖然它是一種一體化的生產解決方案，但它們的程式碼庫太大、太複雜，難以理解和修改（例如，vLLM 有 10 萬多行程式碼），因此很難將它們用於研究目的。此外，他們的歷史負擔也是一個問題。

SwiftLLM 被設計為一個小型但功能強大的 LLM 推理系統，專為研究目的而客製化。 「微小」意味著它只保留了研究所必需的功能，「強大」意味著它在性能上沒有妥協，最後「迅速」意味著它易於理解和修改。在支援基本功能（請參閱下面的列表）並且能夠實現與 vLLM 相當的性能的同時，SwiftLLM 的程式碼庫少於2k行程式碼（約 vLLM 的 2%），用 Python 和 OpenAI Triton（一個用於編寫CUDA核心），使其易於閱讀、修改、調試、測試、擴展，並且可以輕鬆地與您新穎而精彩的研究想法整合。

功能列表

目前，SwiftLLM支援以下功能：

• 迭代調度和選擇性批次（Orca 中提出）
• PagedAttenton（vLLM 論文中提出）
• LLaMA / LLaMA2 / LLaMA3 模型（連結）及其變體
• 搭載預填和解碼（SARATHI 中提出）
• Flash Attention（在FlashAttention和FlashAttention-2中提出）
• 分頁注意力 v2（也稱為 Flash-Decoding，此處提出）

我們計劃在未來添加對以下功能的支援：

• 張量並行和管道並行

為了保持程式碼庫較小，我們將不支援以下功能。如果你想在你的研究計畫中使用它們，你可能需要自己實現它們：

• 量化
• 洛拉
• 多式聯運
• 不遵循 LLaMA 架構的模型
• 貪婪採樣以外的採樣方法
• 硬體支援 NVIDIA GPU 之外的其他硬體（但只要 OpenAI Triton 支持，就應該很容易遷移到其他硬體）

請記住，SwiftLLM 並不是一款一體化的生產解決方案。建議將其視為您研究專案的“基礎”，您可能需要自己實現一些功能。我親愛的研究人員，我們鼓勵您閱讀程式碼、理解它、修改它並擴展它以滿足您的研究需求。

建築學

SwiftLLM的架構可以分為兩個主要部分：控制平面和資料平面。

簡單來說，控制面決定“計算什麼”或“如何調度”，而資料面決定“如何計算”或“如何實現”並進行具體計算。它們以主從方式工作：控制平面充當主站，執行高級調度和協調，並將作業發送到資料平面，資料平面充當執行低階計算的工作執行緒。

控制平面的程式碼位於swiftllm/server目錄中，包括Engine 、 Scheduler 、 API 伺服器和TokenizationEngine等元件。資料平面的程式碼位於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和一組（由於不支援張量並行，目前只有一個）工作執行緒。然後它命令worker執行profile_num_blocks來計算GPU區塊的數量，之後引擎命令所有worker分配他們的KV快取和KV交換。
• 最後，透過Engine.start_all_event_loops啟動事件循環。在循環的每個步驟中，引擎向調度程序查詢下一批要計算的請求，命令工作執行緒執行換入/換出，然後將批次傳送給工作執行緒進行計算。
• API 伺服器監聽使用者請求並與引擎互動來滿足這些請求。

目前，控制平面 ( Engine ) 和資料平面 ( LlamaModel ) 駐留在同一節點上。實現張量並行/管道並行後，資料平面可能會分散到多個節點。

如何使用

我們提供兩種使用 SwiftLLM 的方法：同時使用控制平面和資料平面，或僅使用資料平面。

如果您的想法足夠簡單或優雅，可以無縫整合到現有的控制平面中，那麼您可以同時使用控制平面和資料平面。在另一種情況下，如果你想實作一個出色的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 相當甚至更好的效能。

單次轉送操作

第一個場景是“單一前向操作”，我們向模型提供一批輸入，並讓它產生一個輸出令牌（相當於一個“前向”操作）。這是LLM推理（線上和離線）的基本操作，因此其性能至關重要。

這裡我們使用 FP16 精度下配備 NVIDIA A100 80G PCIE / RTX 4090 GPU 的 LLaMA-3 7B 模型。結果如下圖所示（越低越好）：

可以看出，在相同設定下，SwiftLLM 可以達到與 vLLM 相當（甚至優於）的效能。

線上服務

第二個場景是“線上服務”，我們啟動一個 API 伺服器，從真實資料集中取樣提示，並讓模型產生補全。這是 LLM 在聊天機器人或程式碼完成等現實應用程式中使用的場景。

這裡我們使用 ShareGPT 資料集對提示進行取樣，並使用不同 lambda 的泊松流程來模擬不同的請求到達率。結果如下圖所示（越低越好）：

可以看到，在A100 80G PCIE上，SwiftLLM可以達到與vLLM相當的效能，而在RTX 4090上，SwiftLLM顯著優於vLLM（主要是因為我們的控制平面的開銷較低）。