基本模型¶

本指南將引導您完成實現基本vLLM模型的步驟。

1. 引入您的模型程式碼¶

首先，從源倉庫克隆PyTorch模型程式碼。例如，vLLM的 OPT模型改編自HuggingFace的modeling_opt.py檔案。

警告

請務必審查並遵守原始程式碼的版權和許可條款！

2. 使您的程式碼與vLLM相容¶

為確保與vLLM相容，您的模型必須滿足以下要求

初始化程式碼¶

模型中的所有vLLM模組必須在其建構函式中包含一個prefix引數。此prefix通常是模組在模型狀態字典中的完整名稱，對於以下方面至關重要：

執行時支援：vLLM的注意力運算子透過其完整名稱在模型的狀態中註冊。每個注意力運算子必須具有唯一的prefix作為其層名稱，以避免衝突。
非均勻量化支援：量化檢查點可以選擇性地量化某些層，同時保持其他層為全精度。透過在初始化期間提供prefix，vLLM可以匹配當前層的prefix與量化配置，以確定該層是否應以量化模式初始化。

初始化程式碼應如下所示

程式碼

from torch import nn
from vllm.config import VllmConfig
from vllm.attention import Attention

class MyAttention(nn.Module):
    def __init__(self, vllm_config: VllmConfig, prefix: str):
        super().__init__()
        self.attn = Attention(prefix=f"{prefix}.attn")

class MyDecoderLayer(nn.Module):
    def __init__(self, vllm_config: VllmConfig, prefix: str):
        super().__init__()
        self.self_attn = MyAttention(prefix=f"{prefix}.self_attn")

class MyModel(nn.Module):
    def __init__(self, vllm_config: VllmConfig, prefix: str):
        super().__init__()
        self.layers = nn.ModuleList(
            [MyDecoderLayer(vllm_config, prefix=f"{prefix}.layers.{i}") for i in range(vllm_config.model_config.hf_config.num_hidden_layers)]
        )

class MyModelForCausalLM(nn.Module):
    def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        self.model = MyModel(vllm_config, prefix=f"{prefix}.model")

計算程式碼¶

在MyModel模組內新增一個get_input_embeddings方法，該方法返回給定input_ids的文字嵌入。這等同於直接呼叫文字嵌入層，但在MyModel用於複合多模態模型時提供了一個統一介面。

class MyModel(nn.Module):
        ...

    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        ...

重寫模型的forward方法，移除任何不必要的程式碼，例如訓練特定程式碼。修改輸入引數，將input_ids和positions視為具有單個批大小維度，不帶最大序列長度維度的扁平張量。

def forward(
    self,
    input_ids: torch.Tensor,
    positions: torch.Tensor,
    intermediate_tensors: Optional[IntermediateTensors] = None,
    inputs_embeds: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    ...

注意

目前，vLLM支援基本的注意力機制及其帶有旋轉位置嵌入的變體。如果您的模型採用不同的注意力機制，您需要在vLLM中實現一個新的注意力層。

作為參考，請檢視我們的 Llama實現。vLLM已經支援大量模型。建議找到一個與您的模型相似的模型，並根據您的模型架構進行調整。請檢視 vllm/model_executor/models以獲取更多示例。

3. (可選) 實現張量並行和量化支援¶

如果您的模型太大而無法放入單個GPU，您可以使用張量並行來管理它。為此，請用模型的張量並行版本替換線性層和嵌入層。對於嵌入層，您只需將torch.nn.Embedding替換為VocabParallelEmbedding。對於輸出LM頭部，您可以使用ParallelLMHead。當涉及到線性層時，我們提供以下選項來並行化它們：

ReplicatedLinear：在多個GPU之間複製輸入和權重。不節省記憶體。
RowParallelLinear：輸入張量沿隱藏維度進行分割槽。權重矩陣沿行（輸入維度）進行分割槽。矩陣乘法後執行all-reduce操作以減少結果。通常用於第二個FFN層和注意力層的輸出線性變換。
ColumnParallelLinear：輸入張量被複制。權重矩陣沿列（輸出維度）進行分割槽。結果沿列維度進行分割槽。通常用於第一個FFN層和原始Transformer中注意力層的分離QKV變換。
MergedColumnParallelLinear：列並行線性層，它合併多個ColumnParallelLinear運算子。通常用於帶有加權啟用函式（例如，SiLU）的第一個FFN層。此類處理多個權重矩陣的分片權重載入邏輯。
QKVParallelLinear：用於多頭和分組查詢注意力機制的查詢、鍵和值投影的並行線性層。當鍵/值頭數小於世界大小（world size）時，此類會正確複製鍵/值頭。此類處理權重矩陣的權重載入和複製。

請注意，上述所有線性層都將linear_method作為輸入。vLLM將根據不同的量化方案設定此引數以支援權重量化。

4. 實現權重載入邏輯¶

您現在需要在您的*ForCausalLM類中實現load_weights方法。此方法應從HuggingFace的檢查點檔案載入權重，並將其分配給模型中的相應層。具體來說，對於MergedColumnParallelLinear和QKVParallelLinear層，如果原始模型具有分離的權重矩陣，您需要單獨載入不同的部分。

5. 註冊您的模型¶

有關如何註冊新模型以供vLLM使用的說明，請參閱此頁面。

常見問題¶

如何支援帶有交錯滑動視窗的模型？¶

對於帶有交錯滑動視窗的模型（例如google/gemma-2-2b-it和mistralai/Ministral-8B-Instruct-2410），排程器將把模型視為全注意力模型，即所有token的kv-cache都不會被丟棄。這是為了確保字首快取（prefix caching）在這些模型中工作。滑動視窗僅作為注意力核計算的引數出現。

為了支援帶有交錯滑動視窗的模型，我們需要注意以下細節：

確保模型的config.json包含sliding_window_pattern。然後vLLM會將self.hf_text_config.interleaved_sliding_window設定為self.hf_text_config.sliding_window的值，並從self.hf_text_config中刪除sliding_window。然後模型將被視為全注意力模型。
在建模程式碼中，解析每層的正確滑動視窗值，並將其傳遞給注意力層的per_layer_sliding_window引數。作為參考，請檢視這一行。

透過以上兩步，交錯滑動視窗應該能在模型中工作。