基本模型¶

本指南将引导您完成实现基本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,
    intermediate_tensors: Optional[IntermediateTensors] = None,
    inputs_embeds: Optional[torch.Tensor] = 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：输入张量被复制。权重矩阵沿列（输出维度）进行分区。结果沿列维度进行分区。通常用于第一个FFN层和原始Transformer中注意力层的分离QKV变换。
MergedColumnParallelLinear：列并行线性层，它合并多个ColumnParallelLinear操作符。通常用于带有加权激活函数（例如，SiLU）的第一个FFN层。此类处理多个权重矩阵的分片权重加载逻辑。
QKVParallelLinear：用于多头和分组查询注意力机制的查询、键和值投影的并行线性层。当键/值头数小于世界大小（world size）时，此类会正确复制键/值头。此类处理权重矩阵的权重加载和复制。

请注意，上述所有线性层都将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都不会被丢弃。这是为了确保前缀缓存（prefix caching）在这些模型中工作。滑动窗口仅作为注意力核计算的参数出现。

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

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

通过以上两步，交错滑动窗口应该能在模型中工作。