添加新模型

添加新模型#

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

备注

添加新模型的复杂程度在很大程度上取决于模型的架构。如果模型与 vLLM 中的现有模型具有相似的架构，则该过程相当简单。但是，对于包含新运算符（例如，新的注意力机制）的模型，该过程可能会稍微复杂一些。

备注

默认情况下，vLLM 模型不支持多模态输入。要启用多模态支持，请在在此处实现模型后，按照 本指南 进行操作。

小技巧

如果你在将模型集成到 vLLM 时遇到问题，请随时在我们的 GitHub 存储库中打开一个问题。我们很乐意帮助你！

0. 分叉 vLLM 存储库#

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

小技巧

如果你不想分叉存储库并修改 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. 更新代码，考虑到：code:input_ids 和 positions 现在是扁平化的张量。

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

备注

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

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

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

最后，将你的 *ForCausalLM 类注册到 vllm/model_executor/models/__init__.py 中的 _MODELS。

6. 树外模型集成#

我们还提供了一种无需修改 vLLM 代码库即可集成模型的方法。步骤 2、3、4 仍然需要，但你可以跳过步骤 1 和 5。

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

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

如果你使用 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> 运行它。

目录