基础模型¶

本指南将引导你完成实现基础 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.model_executor.layers.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 模块内添加一个 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 | None,
    positions: torch.Tensor,
    intermediate_tensors: IntermediateTensors | None = None,
    inputs_embeds: torch.Tensor | None = None,
) -> torch.Tensor:
    ...

注意

目前，vLLM 支持基本的注意力机制及其带有旋转位置嵌入（RoPE）的变体。如果你的模型采用不同的注意力机制，则需要在 vLLM 中实现一个新的注意力层。

作为参考，请查看我们的 Llama 实现。vLLM 已经支持了大量的模型。建议找到一个与你的模型相似的模型，并将其适配到你模型的架构中。查看 vllm/model_executor/models 获取更多示例。

3. （可选）实现张量并行和量化支持¶

如果你的模型太大无法放入单个 GPU，可以使用张量并行来管理它。为此，请用它们的张量并行版本替换模型的线性层和嵌入层。对于嵌入层，只需用 VocabParallelEmbedding 替换 torch.nn.Embedding。对于输出 LM head，可以使用 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 使用的说明，请参见此页面。

常见问题¶

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

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

确保模型的 config.json 包含 layer_types。
在建模代码中，解析每一层的正确滑动窗口值，并将其传递给注意力层的 per_layer_sliding_window 参数。作为参考，请检查这一行。

完成这两个步骤后，交错滑动窗口应能与该模型配合工作。

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

我们考虑 3 种不同的场景：

使用 Mamba 层（Mamba-1 或 Mamba-2）但不使用注意力层的模型。
将 Mamba 层（Mamba-1 或 Mamba-2）与注意力层结合的模型。
将类 Mamba 机制（例如线性注意力、ShortConv）与注意力层结合的模型。

对于情况 (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）类。该模型还应添加到 vllm/model_executor/models/config.py 中的 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 中的 MambaAttentionBackendEnum。最后，如果想要支持 torch compile 和 CUDA 图，必须将对类 Mamba 层的调用包装在自定义 op 中并进行注册。请参阅 vllm/model_executor/models/minimax_text_01.py 或 vllm/model_executor/layers/mamba/short_conv.py 中对 direct_register_custom_op 的调用示例。新的自定义 op 然后应该添加到 vllm/config/compilation.py 中的 _attention_ops 列表，以确保分段 CUDA 图按预期工作。