基本模型¶

本指南將引導您完成實現 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.layer 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 模組內新增一個 embed_input_ids 方法，該方法根據 input_ids 返回文字嵌入。這相當於直接呼叫文字嵌入層，但提供了一個統一的介面，以防 MyModel 在複合多模態模型中使用。

class MyModel(nn.Module):
        ...

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

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

def forward(
    self,
    input_ids: torch.Tensor,
    positions: torch.Tensor,
    intermediate_tensors: IntermediateTensors | None = None,
    inputs_embeds: torch.Tensor | None = 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：輸入張量被複制。權重矩陣沿列（輸出維度）分割槽。結果沿列維度分割槽。通常用於 Transformer 原始論文中的 FFN 的第一個層和單獨的 QKV 變換。
MergedColumnParallelLinear：合併多個 ColumnParallelLinear 運算元的列並行線性層。通常用於帶加權啟用函式（例如 SiLU）的 FFN 的第一個層。此類處理多個權重矩陣的分片權重載入邏輯。
QKVParallelLinear：多頭注意力和分組查詢注意力機制的查詢、鍵和值投影的並行線性層。當鍵/值頭的數量小於世界大小時，此類會正確複製鍵/值頭。此類處理權重矩陣的權重載入和複製。

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

4. 實現權重載入邏輯¶

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

5. 註冊您的模型¶

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

常見問題¶

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

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

確保模型的 config.json 包含 layer_types。
在建模程式碼中，解析每個層的正確滑動視窗值，並將其傳遞給注意力層的 per_layer_sliding_window 引數。作為參考，請檢視此行。

透過這兩個步驟，交錯滑動視窗應該可以與模型一起使用。

如何支援使用 Mamba 的模型？¶

我們考慮 3 種不同的場景

使用 Mamba 層（Mamba-1 或 Mamba-2）但不使用注意力層的模型。
結合使用 Mamba 層（Mamba-1 或 Mamba-2）和注意力層的模型。
結合使用 Mamba 類機制（例如，線性注意力、短卷積）和注意力層的模型。

對於場景 (1)，我們建議參考 MambaForCausalLM (用於 Mamba-1) 或 Mamba2ForCausalLM (用於 Mamba-2) 作為參考。模型應繼承協議 IsAttentionFree，並實現類方法 get_mamba_state_dtype_from_config 和 get_mamba_state_shape_from_config 來從配置計算狀態形狀和資料型別。對於 Mamba 層本身，請使用 MambaMixer (用於 Mamba-1) 或 MambaMixer2 (用於 Mamba-2) 類。模型還應新增到 MODELS_CONFIG_MAP 字典中，以確保執行時預設值得到最佳化。

對於場景 (2)，我們建議參考 JambaForCausalLM (一個結合了 Mamba-1 和注意力層的模型的示例) 或 BambaForCausalLM (一個結合了 Mamba-2 和注意力層的模型的示例) 作為參考。這些模型應遵循場景 (1) 的相同說明，但它們應繼承協議 IsHybrid（而不是 IsAttentionFree），並且無需將它們新增到 MODELS_CONFIG_MAP（它們的執行時預設值將從協議推斷）。

對於場景 (3)，我們建議參考 MiniMaxText01ForCausalLM 或 Lfm2ForCausalLM 作為參考，它們分別使用了自定義的“Mamba 類”層 MiniMaxText01LinearAttention 和 ShortConv。請遵循場景 (2) 的相同指南來實現這些模型。我們使用“Mamba 類”來指代具有原地更新狀態的層，而不是像注意力 KV 快取那樣被附加的層。要實現新的自定義 Mamba 類層，應該繼承自 MambaBase 並實現 get_state_dtype、get_state_shape 方法以在執行時計算資料型別和狀態形狀，以及 mamba_type 和 get_attn_backend。還需要實現“注意力元資料”類，該類處理所有層通用的元資料。請參閱 LinearAttentionMetadata 或 ShortConvAttentionMetadata 作為這些的示例。還值得注意的是，在新增新的 Mamba 後端時，我們應該更新 registry.py 中的 MAMBA_TYPE_TO_BACKEND_MAP 和 MambaAttentionBackendEnum。最後，如果希望支援 torch compile 和 CUDA graphs，則需要將 Mamba 類層呼叫包裝在自定義 op 中並進行註冊。請參閱 vllm/model_executor/models/minimax_text_01.py 或 vllm/model_executor/layers/mamba/short_conv.py 中的 direct_register_custom_op 呼叫示例。