实现基础模型

本指南将引导您完成实现基础 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,
) -> 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：用于多头和分组查询注意力机制的查询、键和值投影的并行线性层。当键/值头数量少于世界大小时，此类会正确复制键/值头。此类处理权重矩阵的权重加载和复制。

注意，上述所有线性层都接收 linear_method 作为输入。vLLM 将根据不同的量化方案设置此参数以支持权重量化。

4. 实现权重加载逻辑¶

现在您需要在 *ForCausalLM 类中实现 load_weights 方法。此方法应从 HuggingFace 的检查点文件中加载权重，并将其分配给模型中对应的层。特别地，对于 MergedColumnParallelLinear 和 QKVParallelLinear 层，如果原始模型具有分离的权重矩阵，您需要分别加载不同的部分。

5. 注册你的模型¶

请参阅此页面获取有关如何注册新模型以便由 vLLM 使用的说明。

常见问题¶

如何支持具有交错滑动窗口的模型？¶

对于具有交错滑动窗口的模型（例如 google/gemma-2-2b-it 和 mistralai/Ministral-8B-Instruct-2410），调度程序将把模型视为全注意力模型，即所有 token 的 kv 缓存不会被丢弃。这是为了确保前缀缓存对这些模型起作用。滑动窗口仅作为注意力核计算的一个参数出现。

要支持具有交错滑动窗口的模型，我们需要注意以下细节

确保模型的 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。作为参考，请查看此行。

完成这两步，交错滑动窗口应该就能在该模型上工作了。