跳到內容

Speculative Decoding

警告

請注意,vLLM 中的推測性解碼尚未最佳化,通常不會為所有提示資料集或採樣引數帶來顯著的 token 間延遲降低。最佳化工作正在進行中,您可以在此處關注: Issue #4630

警告

目前,vLLM 中的推測性解碼與流水線並行不相容。

本文件展示瞭如何將 Speculative Decoding 與 vLLM 一起使用。推測性解碼是一種提高記憶體密集型 LLM 推理中 token 間延遲的技術。

使用草稿模型進行推測

以下程式碼將 vLLM 配置為離線模式,使用草稿模型進行推測性解碼,一次推測 5 個 token。

警告

在 vllm v0.10.0 中,不支援使用草稿模型進行推測性解碼。如果您使用以下程式碼,將會得到一個 NotImplementedError

程式碼
from vllm import LLM, SamplingParams

prompts = [
    "The future of AI is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

llm = LLM(
    model="facebook/opt-6.7b",
    tensor_parallel_size=1,
    speculative_config={
        "model": "facebook/opt-125m",
        "num_speculative_tokens": 5,
    },
)
outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

要在線上模式下執行相同的操作,請啟動伺服器

vllm serve facebook/opt-6.7b \
    --host 0.0.0.0 \
    --port 8000 \
    --seed 42 \
    -tp 1 \
    --gpu_memory_utilization 0.8 \
    --speculative_config '{"model": "facebook/opt-125m", "num_speculative_tokens": 5}'

警告

注意:請使用 --speculative_config 設定所有與推測性解碼相關的配置。之前單獨指定模型(--speculative_model)並新增相關引數(例如 --num_speculative_tokens)的方法已被棄用。

然後使用客戶端

程式碼
from openai import OpenAI

# Modify OpenAI's API key and API base to use vLLM's API server.
openai_api_key = "EMPTY"
openai_api_base = "https://:8000/v1"

client = OpenAI(
    # defaults to os.environ.get("OPENAI_API_KEY")
    api_key=openai_api_key,
    base_url=openai_api_base,
)

models = client.models.list()
model = models.data[0].id

# Completion API
stream = False
completion = client.completions.create(
    model=model,
    prompt="The future of AI is",
    echo=False,
    n=1,
    stream=stream,
)

print("Completion results:")
if stream:
    for c in completion:
        print(c)
else:
    print(completion)

透過匹配提示中的 n-grams 進行推測

以下程式碼將 vLLM 配置為推測性解碼,其中建議透過匹配提示中的 n-grams 來生成。有關更多資訊,請閱讀 此執行緒

程式碼
from vllm import LLM, SamplingParams

prompts = [
    "The future of AI is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

llm = LLM(
    model="facebook/opt-6.7b",
    tensor_parallel_size=1,
    speculative_config={
        "method": "ngram",
        "num_speculative_tokens": 5,
        "prompt_lookup_max": 4,
    },
)
outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

使用 Suffix Decoding 進行推測

以下程式碼將 vLLM 配置為推測性解碼,其中建議使用 Suffix Decoding (技術報告) 生成。

與 n-gram 類似,Suffix Decoding 可以透過匹配最後一個 n 個生成 token 的模式來生成草稿 token。與 n-gram 不同的是,Suffix Decoding (1) 可以匹配提示和先前的生成內容,(2) 使用頻率計數來建議最可能的續寫,以及 (3) 為每個請求在每次迭代中推測自適應數量的 token 以獲得更好的接受率。

Suffix Decoding 在重複性高的任務中可以實現更好的效能,例如程式碼編輯、代理迴圈(如自我反思、自我一致性)和 RL 回滾。

安裝 Arctic Inference

Suffix Decoding 需要 Arctic Inference。您可以使用 pip install arctic-inference 安裝它。

Suffix Decoding 推測 token

Suffix Decoding 將為每個請求在每個解碼步驟推測動態數量的 token,因此 num_speculative_tokens 配置指定的是推測 token 的最大數量。建議使用較高的數字,例如 1632(預設值)。

程式碼
from vllm import LLM, SamplingParams

prompts = [
    "The future of AI is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

llm = LLM(
    model="facebook/opt-6.7b",
    tensor_parallel_size=1,
    speculative_config={
        "method": "suffix",
        "num_speculative_tokens": 32,
    },
)
outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

使用 MLP 推斷器進行推測

以下程式碼將 vLLM 配置為推測性解碼,其中建議透過草稿模型生成,這些模型將草稿預測條件化在上下文向量和取樣 token 上。有關更多資訊,請參閱 此部落格此技術報告

程式碼
from vllm import LLM, SamplingParams

prompts = [
    "The future of AI is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

llm = LLM(
    model="meta-llama/Meta-Llama-3.1-70B-Instruct",
    tensor_parallel_size=4,
    speculative_config={
        "model": "ibm-ai-platform/llama3-70b-accelerator",
        "draft_tensor_parallel_size": 1,
    },
)
outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

請注意,這些推測模型目前需要獨立執行,不能與張量並行一起使用,儘管主模型可以使用張量並行執行(參見上面的示例)。由於推測模型相對較小,我們仍然看到了顯著的速度提升。然而,此限制將在將來的版本中修復。

HF hub 上提供了各種此類推測模型

使用基於 EAGLE 的草稿模型進行推測

以下程式碼將 vLLM 配置為推測性解碼,其中建議由基於 EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency) 的草稿模型生成。有關離線模式的更詳細示例,包括如何提取請求級別的接受率,可以在 此處 找到。

程式碼
from vllm import LLM, SamplingParams

prompts = [
    "The future of AI is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

llm = LLM(
    model="meta-llama/Meta-Llama-3-8B-Instruct",
    tensor_parallel_size=4,
    speculative_config={
        "model": "yuhuili/EAGLE-LLaMA3-Instruct-8B",
        "draft_tensor_parallel_size": 1,
        "num_speculative_tokens": 2,
        "method": "eagle",
    },
)

outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

使用基於 EAGLE 的草稿模型時需要考慮的幾個重要事項

  1. HF 上的 EAGLE 模型倉庫 中提供的 EAGLE 草稿模型應該可以直接被 vLLM 載入和使用,經過 Pull Request #12304。如果您使用的是 vllm 早於 Pull Request #12304 的版本,請使用 指令碼 轉換推測模型,並在 speculative_config 中指定 "model": "path/to/modified/eagle/model"。如果在使用最新版本的 vLLM 時仍然出現權重載入問題,請留下評論或提出問題。

  2. 基於 EAGLE 的草稿模型需要獨立執行,不能使用張量並行(即在 speculative_config 中將 draft_tensor_parallel_size 設定為 1),儘管主模型可以使用張量並行執行(參見上面的示例)。

  3. 當使用基於 EAGLE 的推斷器與 vLLM 時,觀察到的加速比低於參考實現 此處 報告的加速比。此問題正在調查中,並在此處跟蹤: Issue #9565

  4. 當使用基於 EAGLE-3 的草稿模型時,選項 "method" 必須設定為 "eagle3"。也就是說,在 speculative_config 中指定 "method": "eagle3"

Hugging Face hub 上提供了各種 EAGLE 草稿模型

基礎模型 EAGLE on Hugging Face # EAGLE 引數
Vicuna-7B-v1.3 yuhuili/EAGLE-Vicuna-7B-v1.3 0.24B
Vicuna-13B-v1.3 yuhuili/EAGLE-Vicuna-13B-v1.3 0.37B
Vicuna-33B-v1.3 yuhuili/EAGLE-Vicuna-33B-v1.3 0.56B
LLaMA2-Chat 7B yuhuili/EAGLE-llama2-chat-7B 0.24B
LLaMA2-Chat 13B yuhuili/EAGLE-llama2-chat-13B 0.37B
LLaMA2-Chat 70B yuhuili/EAGLE-llama2-chat-70B 0.99B
Mixtral-8x7B-Instruct-v0.1 yuhuili/EAGLE-mixtral-instruct-8x7B 0.28B
LLaMA3-Instruct 8B yuhuili/EAGLE-LLaMA3-Instruct-8B 0.25B
LLaMA3-Instruct 70B yuhuili/EAGLE-LLaMA3-Instruct-70B 0.99B
Qwen2-7B-Instruct yuhuili/EAGLE-Qwen2-7B-Instruct 0.26B
Qwen2-72B-Instruct yuhuili/EAGLE-Qwen2-72B-Instruct 1.05B

Speculative Decoding 的無損保證

在 vLLM 中,推測性解碼旨在提高推理效率,同時保持準確性。本節將推測性解碼的無損保證分為三個關鍵領域:

  1. 理論無損 - 推測性解碼取樣在理論上是無損的,直到硬體數字精度的限制。浮點誤差可能會導致輸出分佈產生細微變化,正如在 Accelerating Large Language Model Decoding with Speculative Sampling 中所討論的。

  2. 演算法無損 - vLLM 的推測性解碼實現經過演算法驗證是無損的。關鍵的驗證測試包括:

    • 拒絕取樣器收斂性:確保 vLLM 的拒絕取樣器的樣本與目標分佈一致。 檢視測試程式碼
    • 貪婪取樣相等性:確認使用推測性解碼的貪婪取樣與不使用推測性解碼的貪婪取樣相匹配。這驗證了 vLLM 的推測性解碼框架,當與 vLLM 前向傳遞和 vLLM 拒絕取樣器整合時,提供了無損保證。 tests/spec_decode/e2e 中的幾乎所有測試都使用 此斷言實現 來驗證此屬性。
  3. vLLM Logprob 穩定性 - vLLM 目前不保證 token 對數機率(logprobs)的穩定性。這可能導致同一請求在不同執行中產生不同的輸出。有關更多詳細資訊,請參閱 FAQ 中的“vLLM 中提示的輸出是否可能在不同執行中有所不同?”部分。

雖然 vLLM 努力確保推測性解碼的無損性,但由於以下因素,在有推測性解碼和沒有推測性解碼的情況下生成的輸出可能存在差異:

  • 浮點精度:硬體數值精度的差異可能導致輸出分佈略有不同。
  • 批處理大小和數值穩定性:批處理大小的變化可能導致 logprobs 和輸出機率的變化,這可能是由於批處理操作中的非確定性行為或數值不穩定性。

有關緩解策略,請參閱 FAQ 中的“vLLM 中提示的輸出是否可能在不同執行中有所不同?”條目。

vLLM 貢獻者資源