跳到内容

与 HuggingFace 集成

本文档描述了 vLLM 如何与 HuggingFace 库集成。我们将逐步解释运行 vllm serve 时其内部发生了什么。

假设我们想通过运行 vllm serve Qwen/Qwen2-7B 来部署流行的 QWen 模型。

  1. model 参数是 Qwen/Qwen2-7B。vLLM 通过检查相应的配置文件 config.json 来确定该模型是否存在。有关实现,请参见此 代码片段。在此过程中

    • 如果 model 参数对应一个现有的本地路径,vLLM 将直接从该路径加载配置文件。
    • 如果 model 参数是包含用户名和模型名称的 HuggingFace 模型 ID,vLLM 将首先尝试从 HuggingFace 本地缓存中使用配置文件,其中 model 参数用作模型名称,--revision 参数用作修订版本。有关 HuggingFace 缓存工作原理的更多信息,请参见其网站
    • 如果 model 参数是 HuggingFace 模型 ID 但未在缓存中找到,vLLM 将从 HuggingFace 模型中心下载配置文件。有关实现,请参考 此函数。输入参数包括用作模型名称的 model 参数、用作修订版本的 --revision 参数,以及用作访问模型中心令牌的环境变量 HF_TOKEN。在我们的例子中,vLLM 将下载 config.json 文件。
  2. 确认模型存在后,vLLM 将加载其配置文件并将其转换为字典。有关实现,请参见此 代码片段

  3. 接下来,vLLM 会检查配置字典中的 model_type 字段,以生成要使用的配置对象。vLLM 直接支持某些 model_type 值;有关列表,请参见此处。如果 model_type 不在此列表中,vLLM 将使用 AutoConfig.from_pretrained 加载配置类,其中 model--revision--trust_remote_code 作为参数。请注意

    • HuggingFace 也有自己的逻辑来确定要使用的配置类。它将再次使用 model_type 字段在 transformers 库中搜索类名;有关支持的模型列表,请参见此处。如果未找到 model_type,HuggingFace 将使用配置 JSON 文件中的 auto_map 字段来确定类名。具体来说,它是 auto_map 下的 AutoConfig 字段。有关示例,请参见DeepSeek
    • auto_map 下的 AutoConfig 字段指向模型仓库中的一个模块路径。为了创建配置类,HuggingFace 将导入该模块并使用 from_pretrained 方法加载配置类。这通常可能导致任意代码执行,因此仅当启用 --trust_remote_code 时才执行。
  4. 随后,vLLM 会对配置对象应用一些历史补丁。这些主要与 RoPE 配置相关;有关实现,请参见此处

  5. 最后,vLLM 可以找到我们想要初始化的模型类。vLLM 使用配置对象中的 architectures 字段来确定要初始化的模型类,因为它在其注册表中维护了从架构名称到模型类的映射。如果在注册表中找不到架构名称,则表示 vLLM 不支持此模型架构。对于 Qwen/Qwen2-7Barchitectures 字段是 ["Qwen2ForCausalLM"],它对应于vLLM 的代码中的 Qwen2ForCausalLM 类。此类将根据各种配置进行初始化。

除此之外,vLLM 还依赖 HuggingFace 来完成两项任务。

  1. 分词器:vLLM 使用 HuggingFace 的分词器对输入文本进行分词。分词器使用 AutoTokenizer.from_pretrained 加载,其中 model 参数作为模型名称,--revision 参数作为修订版本。也可以通过在 vllm serve 命令中指定 --tokenizer 参数来使用其他模型的分词器。其他相关参数是 --tokenizer-revision--tokenizer-mode。请查阅 HuggingFace 的文档以了解这些参数的含义。这部分逻辑可以在 get_tokenizer 函数中找到。获得分词器后,值得注意的是,vLLM 将在 get_cached_tokenizer 中缓存分词器的一些开销大的属性。

  2. 模型权重:vLLM 使用 model 参数作为模型名称,--revision 参数作为修订版本,从 HuggingFace 模型中心下载模型权重。vLLM 提供了 --load-format 参数来控制从模型中心下载哪些文件。默认情况下,它将尝试以 safetensors 格式加载权重,如果 safetensors 格式不可用,则回退到 PyTorch bin 格式。我们还可以传递 --load-format dummy 来跳过权重下载。

    • 建议使用 safetensors 格式,因为它在分布式推理中加载效率高,并且可以避免任意代码执行。有关 safetensors 格式的更多信息,请参见文档。这部分逻辑可以在此处找到。请注意

至此,vLLM 与 HuggingFace 的集成已完成。

综上所述,vLLM 从 HuggingFace 模型中心或本地目录读取配置文件 config.json、分词器和模型权重。它使用 vLLM、HuggingFace transformers 中的配置类,或者从模型的仓库中加载配置类。