Skip to content

使用 vLLM 作為動態批次（Dynamic Batching）加速推理的 API 服務

Last Updated on 2024-01-11 by Clay

介紹

我之前曾寫了一篇介紹 vLLM 加速推理框架的筆記（使用 vLLM 進行大型語言模型（LLM）的高速推理），然而因受篇幅與時間限制，沒來得及探討更細緻的功能。

除了將 vLLM 作為加速 LLM 的推理框架，應用在研究用途上之外，vLLM 還實現了更強大的功能 —— 也就是動態批次（Dynamic Batching，亦稱作 Rolling Batch 或是 Continually Batching，但經常被混用）的推理技術。

這種方法在個人使用研究 LLM 效果時感覺不出差異，然而一旦到了生產佈署場景時就變得非常重要 —— 因為這會讓使用者的體驗非常滑順，在 GPU 記憶體充足的情況下，基本不需要等待就可以即時地獲得 LLM 生成的回饋。

首先，我們平常使用 transformers 套件所做的推理，基本上全部都是靜態批次（Static Batching）的模式。以下是一張在網路上流傳已久的示意圖的重製。

假設現在我們正在生成 4 個不同的句子，在靜態批次的模式下，模型一次只會對 batch 內的句子繼續生成下一個 token —— 然後，我們總共需要等待到 T8 時，才能把所有句子生成的結果全部返回。

示意圖中空白的部份，正是浪費掉的運算時間。

然而，若是我們的 LLM 推理過程中，在每次解碼出下一個 token 時都檢查是否可以加入新的生成請求，我們就能在解碼其他仍然正在生成的句子的同時，將已經解碼完成的計算資源留給動態加入的下一筆生成資料（也就是下圖中的 S5、S6、S7）。

而且若是在產品的角度來看，使用者也可以在自己的資料生成完畢時立刻拿回自己的生成資料，不必等待捆綁在同個 batch 的最長生成序列結束，可謂兩全其美。

在 vLLM 中已經實現了這樣的異步推理引擎，我們可以很方便地直接使用。

使用方式

首先只需要確認安裝好 vllm 套件即可。

pip3 install vllm


之後，則可以直接透過以下指令啟動 vllm api service：

python -m vllm.entrypoints.api_server --model facebook/opt-125m


當然，為了測試我選擇了 facebook/opt-125m 的模型，應該也不用指望這個量級的模型回答得有多棒。啟動完畢後，預設的 port 應為 8000。

我們可以直接在終端機中使用以下指令：

curl -X POST -H "User-Agent: Test Client" -H "Content-Type: application/json" \
     -d '{
        "prompt": "San Francisco is a",
        "n": 4,
        "use_beam_search": true,
        "temperature": 0.0,
        "max_tokens": 16,
        "stream": true
     }' \
     http://localhost:8000/generate


即可得到 streaming 的回傳格式。

如果是習慣使用 Python 的話，則可以嘗試使用以下腳本測試 vLLM API：

curl -X POST -H "User-Agent: Test Client" -H "Content-Type: application/json" \
     -d '{
        "prompt": "San Francisco is a",
        "n": 4,
        "use_beam_search": true,
        "temperature": 0.0,
        "max_tokens": 16,
        "stream": true
     }' \
     http://localhost:8000/generate


並且由於 vLLM API 本身是異步處理的，所以你可以同時發送多個請求，應該會看到每個請求幾乎都會在第一時間返回生成結果，是很好的體驗。

接下來若有機會，我會想再分享如果在 Gradio 這樣的 demo 界面上，該如何完成 vLLM 的 streaming 串接。

References

Read More

Tags:

Leave a Reply