实现基本模型

实现基本模型#

本指南将引导您完成实现基本 vLLM 模型的步骤。

1. 引入您的模型代码#

首先，从源仓库克隆 PyTorch 模型代码。例如，vLLM 的 OPT 模型是从 HuggingFace 的 modeling_opt.py 文件改编而来。

警告

请务必查看并遵守原始代码的版权和许可条款！

2. 使您的代码与 vLLM 兼容#

为了确保与 vLLM 兼容，您的模型必须满足以下要求

初始化代码#

模型中的所有 vLLM 模块都必须在其构造函数中包含 prefix 参数。此 prefix 通常是模型状态字典中模块的完整名称，对于以下方面至关重要：

运行时支持：vLLM 的注意力运算符通过其完整名称在模型状态中注册。每个注意力运算符必须具有唯一的前缀作为其层名称，以避免冲突。
非均匀量化支持：量化检查点可以选择性地量化某些层，同时保持其他层为全精度。通过在初始化期间提供 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,
    kv_caches: List[torch.Tensor],
    attn_metadata: AttentionMetadata,
) -> 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-cache。这是为了确保前缀缓存适用于这些模型。滑动窗口仅作为注意力内核计算的参数出现。

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

确保此行代码将 has_interleaved_attention 评估为 True 对于此模型，并将 self.hf_text_config.interleaved_sliding_window 设置为模型可以理解的交错滑动窗口格式。然后，self.hf_text_config.sliding_window 将被删除，并且该模型将被视为全注意力模型。
在建模代码中，为每一层解析正确的滑动窗口值，并将其传递给注意力层的 per_layer_sliding_window 参数。有关参考，请查看此行代码。

通过这两个步骤，交错滑动窗口应该可以与该模型一起工作。