跳到主要内容

添加一个新的模型

这份文档提供了有关将 HuggingFace Transformers 模型集成到 vLLM 的高级指南。

提示

添加新模型的复杂性在很大程度上取决于模型的架构。如果这个模型和vLLM目前支持的模型架构相似，那么添加新模型就相对简单。 然而，对于包含新运算符（例如新的注意力机制）的模型，这就可能会更复杂一些。

信息

当你整合新模型到vLLM遇到问题时，随时在github上向我们提问，我们会很高兴的帮你解答

1. Fork vLLM 仓库​

Fork我们的GitHub仓库，然后从源代码构建它。能够修改代码库并测试您的模型。

2. 带来你的模型代码​

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

注意

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

3. 重写前向传播方法。​

接下来，您需要按照以下步骤重写模型的前向传播方法

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

  • 修改输入参数

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[KVCache],
+    input_metadata: InputMetadata,
+) -> Optional[SamplerOutput]:

  • 更新代码要考虑到目前input_ids和positions现在是扁平的张量

  • 根据模型的架构，将attention operation替换为PagedAttention、PagedAttentionWithRoPE或PagedAttentionWithALiBi。

注意

目前, vLLM 支持基本的multi-head attention 机制和有旋转位置嵌入的变体. 如果你的模型运用不同的 attention 机制, 您需要在vLLM中实现一个新的attention层.

4. （可选）实现张量并行性和量化支持​

如果你的模型太大不适合单个GPU,可以使用张量并行性来管理它，为此，请将模型的线性层和嵌入层替换为tensor-parallel版本。对于嵌入层，您可以简单地将 nn.Embedding 替换为 VocabParallelEmbedding。对于输出LM头，可以使用ParallelLMHead。对于线性层，我们提供以下选项来并行化它们：

  • ReplicatedLinear: 跨多个 GPU 复制输入和权重。不节省内存。

  • RowParallelLinear: 输入张量沿隐藏维度进行分区。权重矩阵沿行（输入维度）划分。矩阵乘法之后执行全约简操作以约简结果。通常用于第二个 FFN 层和注意力层的输出线性变换。

  • ColumnParallelLinear: 输入张量被复制。权重矩阵沿列（输出维度）划分。结果沿列维度分区。通常用于原始 Transformer 中的第一个 FFN 层和注意力层的分离 QKV 变换。

  • MergedColumnParallelLinear: 合并多个 ColumnParallelLinear 运算符的列并行线性。通常用于具有加权激活函数（例如 SiLU）的第一个 FFN 层。该类处理多个权重矩阵的分片权重加载逻辑。

  • QKVParallelLinear: 用于多头和分组查询注意机制的查询、键和值投影的并行线性层。当键/值头的数量小于世界大小时，此类会正确复制键/值头。此类处理权重矩阵的权重加载和复制。

请注意，上面的所有线性层均采用 Linear_method 作为输入。 vLLM会根据不同的量化方案设置该参数，以支持权重量化

5. 实现权重加载逻辑​

您现在需要在 *ForCausalLM 类中实现 load_weights 方法。此方法应该从 HuggingFace 的检查点文件加载权重，并将它们分配给模型中的相应层。具体来说，对于 MergedColumnParallelLinear 和 QKVParallelLinear 图层，如果原始模型具有分离的权重矩阵，则需要分别加载不同部分。

6. 注册你的模型​

最后，将 *ForCausalLM 类包含在 vllm/model_executor/models/init.py 中，并将其注册到 vllm/model_executor/model_loader.py 中的 _MODEL_REGISTRY 中。