添加新模型

添加新模型#

本文档提供了一个关于将 HuggingFace Transformers 模型集成到 vLLM 的高级指南。

备注

添加新模型的复杂性很大程度上取决于模型的架构。如果新模型的架构与vLLM中现有模型的架构相似，那么这个过程就相当直接。然而，对于包含新操作符（例如，新的注意力机制）的模型，这个过程可能会更加复杂。

备注

默认情况下，vLLM 模型不支持多模态输入。要启用多模态支持，请在实现模型后遵循 此指南。

小技巧

如果在将您的模型集成到 vLLM 时遇到问题，请随时在我们的 GitHub 仓库中提出问题。我们将很乐意帮助您！

0. 分叉 vLLM 仓库#

首先通过分叉我们的 GitHub 仓库，然后 从源代码构建 。这使您能够修改代码库并测试您的模型。

小技巧

如果你不想fork仓库并修改vLLM的代码库，请参考下面的“树外模型集成”部分。

1. 带上你的模型代码#

从 HuggingFace Transformers 仓库克隆 PyTorch 模型代码，并将其放入 vllm/model_executor/models 目录中。例如，vLLM 的 OPT 模型 是从 HuggingFace 的 modeling_opt.py 文件改编而来的。

警告

在复制模型代码时，请确保审查并遵守代码的版权和许可条款。

2. 重写 forward 方法#

接下来，您需要按照以下步骤重写模型的 forward() 方法：

  1. 移除任何不必要的代码，例如仅用于训练的代码。

  2. 更改输入参数：

  def forward(
      self,
      input_ids: torch.Tensor,
-     attention_mask: Optional[torch.Tensor] = None,
-     position_ids: Optional[torch.LongTensor] = None,
-     past_key_values: Optional[List[torch.FloatTensor]] = None,
-     inputs_embeds: Optional[torch.FloatTensor] = None,
-     labels: Optional[torch.LongTensor] = None,
-     use_cache: Optional[bool] = None,
-     output_attentions: Optional[bool] = None,
-     output_hidden_states: Optional[bool] = None,
-     return_dict: Optional[bool] = None,
- ) -> Union[Tuple, CausalLMOutputWithPast]:
+     positions: torch.Tensor,
+     kv_caches: List[torch.Tensor],
+     attn_metadata: AttentionMetadata,
+ ) -> Optional[SamplerOutput]:

  1. 更新代码，考虑到 input_ids 和 positions 现在是扁平化的张量。

  2. 根据模型的架构，将注意力操作替换为 PagedAttention、PagedAttentionWithRoPE 或 PagedAttentionWithALiBi。

备注

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

3. (可选) 实现张量并行和量化支持#

如果你的模型太大，无法放入单个GPU中，你可以使用张量并行来管理它。为此，请将模型的线性和嵌入层替换为它们的张量并行版本。对于嵌入层，你可以简单地将 torch.nn.Embedding 替换为 VocabParallelEmbedding。对于输出语言模型头部，你可以使用 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. 注册你的模型#

最后，将您的 *ForCausalLM 类注册到 vllm/model_executor/models/registry.py 中的 _VLLM_MODELS。

6. 树外模型集成#

我们还提供了一种在不修改 vLLM 代码库的情况下集成模型的方法。步骤 2、3、4 仍然是必需的，但您可以跳过步骤 1 和 5。

只需在你的代码中添加以下几行：

from vllm import ModelRegistry
from your_code import YourModelForCausalLM
ModelRegistry.register_model("YourModelForCausalLM", YourModelForCausalLM)

如果你的模型导入了初始化CUDA的模块，考虑改为延迟导入以避免类似 RuntimeError: Cannot re-initialize CUDA in forked subprocess 的错误：

from vllm import ModelRegistry

ModelRegistry.register_model("YourModelForCausalLM", "your_code:YourModelForCausalLM")

重要

如果你的模型是多模态模型，请确保模型类实现了 SupportsMultiModal 接口。更多信息请阅读 这里。

如果你正在使用 vllm serve <args> 运行 API 服务器，你可以用以下代码包装入口点：

from vllm import ModelRegistry
from your_code import YourModelForCausalLM
ModelRegistry.register_model("YourModelForCausalLM", YourModelForCausalLM)
import runpy
runpy.run_module('vllm.entrypoints.openai.api_server', run_name='__main__')

将上述代码保存到一个文件中，并使用 python your_file.py <args> 运行它。