基本模型¶

本指南将引导您完成实现 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 调用示例。