跳至内容

基础模型

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

1. 导入您的模型代码

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

警告

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

2. 让你的代码兼容vLLM

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

初始化代码

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

  • 运行时支持：vLLM的注意力算子会以其完整名称注册到模型状态中。每个注意力算子必须拥有唯一的层名前缀以避免冲突。
  • 非均匀量化支持：量化后的检查点可以选择性地对某些层进行量化，同时保持其他层为全精度。通过在初始化时提供prefix，vLLM可以将当前层的prefix与量化配置进行匹配，以确定该层是否应以量化模式初始化。

初始化代码应如下所示：

Code 
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。对于输出语言模型头部，可以使用ParallelLMHead。至于线性层，我们提供以下并行化选项：

  • ReplicatedLinear: 在多个GPU上复制输入和权重。不节省内存。
  • RowParallelLinear: 输入张量沿隐藏维度进行分区。权重矩阵沿行（输入维度）进行分区。在矩阵乘法后执行all-reduce操作以归约结果。通常用于第二个FFN层和注意力层的输出线性变换。
  • ColumnParallelLinear: 输入张量会被复制。权重矩阵沿列维度（输出维度）进行分区。结果会沿列维度分区。通常用于原始Transformer中第一个FFN层和注意力层的分离QKV变换。
  • MergedColumnParallelLinear: 合并多个ColumnParallelLinear算子的列并行线性层。通常用于具有加权激活函数(如SiLU)的第一个前馈神经网络层。该类处理多个权重矩阵的分片权重加载逻辑。
  • 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参数。具体实现可参考这一行。

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

优云智算