vLLM 模型加载流程分析

1. 概述

vLLM 的模型加载是一个分层的、可插拔的系统。从高层入口到底层 weight loader,整个过程分为 5 个阶段:

  1. 根据 load_format 选择 ModelLoader
  2. 解析模型架构并实例化 nn.Module
  3. 从磁盘/HuggingFace Hub 加载权重
  4. 后处理权重(量化重打包、attention 初始化)
  5. 返回 model.eval() 模式的模型

2. 加载流程全景

GPUWorker.load_model()                           [gpu_worker.py:335]
  └─ GPUModelRunner.load_model()                 [gpu_model_runner.py:5032]
       └─ get_model(vllm_config)                 [model_loader/__init__.py:133]
            │
            ├─ get_model_loader(load_config)     ─── 选择 Loader 类
            │
            └─ loader.load_model()               ─── 执行实际加载
                 │
                 ├─ initialize_model()           ─── 1. 实例化 nn.Module
                 │    ├─ 解析架构名 (config.json)
                 │    ├─ ModelRegistry 查找
                 │    └─ 调用模型构造函数
                 │
                 ├─ load_weights()               ─── 2. 加载权重
                 │    ├─ 发现/下载权重文件
                 │    ├─ 迭代 (name, tensor) 对
                 │    └─ model.load_weights(iterator) → TP 分片分发
                 │
                 ├─ finalize_layerwise_processing() ─ 3. 在线量化延迟处理
                 │
                 └─ process_weights_after_loading() ─ 4. 权重后处理
                      ├─ quant_method.process_weights_after_loading()
                      └─ 各类 Attention 初始化

3. Model Loader 的选择

3.1 入口

model_loader/init.py:125get_model_loader()

def get_model_loader(load_config: LoadConfig) -> BaseModelLoader:
    return _LOAD_FORMAT_TO_MODEL_LOADER[load_config.load_format](load_config)

_LOAD_FORMAT_TO_MODEL_LOADERL52-69)建立了格式字符串到 Loader 类的映射:

load_format Loader 类 文件
"auto", "hf", "safetensors", "fastsafetensors", "instanttensor", "mistral", "npcache", "pt" DefaultModelLoader default_loader.py
"dummy" DummyModelLoader dummy_loader.py
"bitsandbytes" BitsAndBytesModelLoader bitsandbytes_loader.py
"gguf" GGUFModelLoader gguf_loader.py
"tensorizer" TensorizerLoader tensorizer_loader.py
"sharded_state", "runai_streamer_sharded" ShardedStateLoader sharded_state_loader.py
"runai_streamer" RunaiModelStreamerLoader
"modelexpress" ModelExpressModelLoader

自定义 loader 可通过 @register_model_loader() 装饰器(L72)注册。

3.2 LoadConfig

config/load.py:26 — 控制所有加载参数:

参数 默认值 说明
load_format "auto" 权重格式,自动检测
download_dir HF 缓存目录
safetensors_load_strategy None "lazy" / "eager" / "prefetch" / "torchao"
safetensors_prefetch_num_threads 8 预取线程数
safetensors_prefetch_block_size 16 MiB 预取块大小
model_loader_extra_config loader 特定额外配置
ignore_patterns 下载时跳过的文件模式
use_tqdm_on_load 是否显示进度条

4. 模板方法:load_model()

base_loader.py:43-82BaseModelLoader.load_model() 定义了固定的 5 步模板:

def load_model(self, vllm_config, model_config, prefix=""):
    # Step 1: 确定加载设备
    load_device = self._get_load_device(vllm_config)

    # Step 2: 设置默认 dtype
    with set_default_torch_dtype(model_config.dtype):
        # Step 3: 实例化模型架构(构造 nn.Module,此时权重未初始化)
        model = initialize_model(
            vllm_config=vllm_config,
            model_config=model_config,
            prefix=prefix,
        )

    # Step 4: 加载权重
    self.load_weights(model, model_config)

    # Step 5: 在线量化延迟处理(若有)
    if uses_online_quantization:
        finalize_layerwise_processing(model, ...)

    # Step 6: 权重后处理(量化重打包、attention 初始化)
    process_weights_after_loading(model)

    return model.eval()

每个子类只需实现 download_model()load_weights() 两个抽象方法。BaseModelLoader 本身通过模板方法控制整体顺序。


5. 模型架构解析与实例化

5.1 initialize_model()

model_loader/utils.py:40-96 — 构造模型 nn.Module

def initialize_model(vllm_config, model_config, prefix=""):
    # Step 1: 解析模型类
    model_cls = get_model_architecture(model_config)

    # Step 2: 配置量化
    configure_quant_config(model_cls, vllm_config)

    # Step 3: 调用构造函数
    # 新式模型接受 vllm_config + prefix 参数
    with set_current_vllm_config(vllm_config):
        model = model_cls(vllm_config=vllm_config, prefix=prefix)

    return model

新式模型(大多数)的构造函数签名:__init__(self, vllm_config, prefix)。旧式模型则通过 **kwargs 做向后兼容。

5.2 架构解析

model_loader/utils.py:179-184_get_model_architecture()

def _get_model_architecture(model_config):
    architectures = model_config.hf_config.architectures  # config.json 中的字段
    return model_config.registry.resolve_model_cls(architectures, model_config)

architectures 来自 HuggingFace config.json"architectures" 字段——例如 ["LlamaForCausalLM"]

5.3 ModelRegistry

models/registry.py:1353ModelRegistry 是一个全局单例字典,将架构名(如 "LlamaForCausalLM")映射到 _LazyRegisteredModel 对象。每个条目包含 module_nameclass_name

注册模型的大字典(registry.py:683)组合自 5 个子字典:

字典 行号 内容
_TEXT_GENERATION_MODELS L71 纯文本生成模型(Llama、Qwen、Mistral 等)
_EMBEDDING_MODELS L224 Embedding 模型
_MULTIMODAL_MODELS L344 多模态模型(LLaVA、Qwen2-VL 等)
_SPECULATIVE_DECODING_MODELS L590 推测解码的 draft 模型
_TRANSFORMERS_BACKEND_MODELS L648 通过 Transformers 后端加载的模型

resolve_model_cls()registry.py:1201)的解析逻辑:

1. 若 model_impl == "transformers" → 优先通过 Transformers 后端解析
2. 若架构名不在注册表中且 model_impl == "auto" → 回退到 Transformers 解析
3. 遍历所有 architectures,通过 _normalize_arch() 做名称归一化
4. 调用 _LazyRegisteredModel.load_model_cls() → importlib.import_module()
5. 全部失败 → _raise_for_unsupported() 抛出清晰错误

5.4 模型实现

models/ — 每个架构一个文件:

llama.py         qwen2.py        mistral.py      mixtral.py
gemma.py         phi.py          falcon.py       deepseek_v2.py
qwen2_vl.py      llava.py        internvl.py     ...

每个文件包含一个或多个模型类(继承自相应的 vLLM 基类),并实现 load_weights() 方法。


6. 权重加载

6.1 DefaultModelLoader.load_weights()

default_loader.py:382 — 主流 loader 的权重加载:

def load_weights(self, model, model_config):
    # Step 1: torchao 量化检测
    if quant_config == "torchao":
        # 自动设置 safetensors 加载策略

    # Step 2: Expert Parallelism 权重过滤(EP 时跳过非本地 expert)
    _init_ep_weight_filter(model_config, model)

    # Step 3: 获取权重并加载
    loaded_weights = model.load_weights(self.get_all_weights(model_config, model))

    # Step 4: 严格模式检查(默认开启)
    # 验证每个 model parameter 都被加载了

6.2 权重文件发现

default_loader.py:97_prepare_weights()

_prepare_weights(model_config):
  ├─ 检查模型路径是否本地
  ├─ 若非本地: 检查 ModelScope → maybe_download_from_modelscope()
  ├─ auto 检测 Mistral 格式 (consolidated*.safetensors)
  ├─ 根据 load_format 设置 allow_patterns:
  │   ├─ "hf":     ["*.safetensors", "*.bin"]
  │   ├─ "safetensors" / "fastsafetensors" / "instanttensor": ["*.safetensors"]
  │   ├─ "mistral": ["consolidated*.safetensors"]
  │   ├─ "pt":      ["*.pt"]
  │   └─ "npcache": ["*.bin"]
  ├─ 若非本地: download_weights_from_hf() → HF Hub snapshot_download()
  ├─ 下载/读取 safetensors index 文件,过滤重复文件
  └─ 排除训练专用文件: optimizer.bin, scheduler.pt, ...

weight_utils.py:504download_weights_from_hf() 使用 HfFileSystem 列出文件、优化 allow_patterns,然后调用 snapshot_download()

6.3 权重迭代器

default_loader.py:211_get_weights_iterator() 根据格式返回不同的迭代器:

格式 迭代器 特点
npcache np_cache_weights_iterator() 将 PyTorch 权重转为 numpy cache,加速后续加载
fastsafetensors fastsafetensors_weights_iterator() 使用 fastsafetensors 库并行 GPU 加载
instanttensor instanttensor_weights_iterator() 使用 InstantTensor CUDA Direct I/O
safetensors(多线程) multi_thread_safetensors_weights_iterator() 8 线程并行
safetensors(默认) safetensors_weights_iterator() 支持预取优化
pt(多线程) multi_thread_pt_weights_iterator() 8 线程并行
pt(默认) pt_weights_iterator() 标准 PyTorch 加载

多线程默认 DEFAULT_NUM_THREADS = 8

6.4 safetensors_weights_iterator 的智能预取

weight_utils.py:893 — 针对网络文件系统做了专门优化:

safetensors_weights_iterator():
  ├─ 自动检测 NFS/Lustre(网络文件系统)
  ├─ 若文件可装入 RAM → 启动 OS page cache 预取
  │   └─ safetensors_prefetch_num_threads 个线程并发读取
  │   └─ 块大小 = safetensors_prefetch_block_size (16 MiB)
  ├─ 加载策略:
  │   ├─ "eager"    → 立即全部读入内存
  │   ├─ "lazy"     → 按需读取(默认)
  │   ├─ "prefetch" → 后台预取
  │   └─ "torchao"  → 跨 shard 重建 tensor subclass state dict
  └─ Expert Parallelism 过滤:
      └─ should_skip_weight() → 跳过非本 rank 的 expert 权重

6.5 模型如何消费权重——两种模式

模式 1: 逐层 load_weights(如 LlamaDecoderLayer)

llama.py:436 — 每一层的 load_weights() 处理属于自己的权重:

class LlamaDecoderLayer(nn.Module):
    def load_weights(self, weights):
        # 定义 fused 参数到分片的映射
        stacked_params_mapping = [
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
        ]

        for name, loaded_weight in weights:
            # 检查是否需要处理 fused 权重
            for param_name, weight_name, shard_id in stacked_params_mapping:
                if param_name in name:
                    # 将 fused 张量分片后再赋值
                    ...
            # 调用参数上的 weight_loader
            param = dict(self.named_parameters())[name]
            weight_loader = getattr(param, "weight_loader", default_weight_loader)
            weight_loader(param, loaded_weight)

模式 2: AutoWeightsLoader(如 LlamaForCausalLM 整体)

models/utils.py:115AutoWeightsLoader 是自动化的递归 loader:

# LlamaForCausalLM.load_weights() [llama.py:584]
class LlamaForCausalLM:
    def load_weights(self, weights):
        loader = AutoWeightsLoader(self, skip_prefixes=["lm_head."])
        return loader.load_weights(weights)

AutoWeightsLoader._load_module()utils.py:259)的递归逻辑:

_load_module(module, weights):
  ├─ 按首个分隔符 "." 对权重名分组
  ├─ 若子模块有 load_weights 方法 → delegate (模式 1)
  ├─ 若子模块有参数 → 调用 param.weight_loader(param, tensor)
  ├─ 若为子模块 → 递归进入
  └─ 检查 all_params_loaded → 返回 loaded weight names set

6.6 Weight Loader 函数

weight_utils.py — 定义在 tensor 参数上的 loader 属性负责 TP 分片:

Loader 行号 行为
default_weight_loader() L1383 直接 copy,检查 shape 匹配
row_parallel_weight_loader() L1404 对 dim=0 做 TP rank 分片(行并行)
sharded_weight_loader() L1422 创建沿指定轴分片的 loader
composed_weight_loader() L1437 组合 loader 与后处理函数

列并行加载没有独立函数。ColumnParallelLinear 通过自身的 weight_loader() 方法(linear.py:692)处理 TP 分片——在模型构造时将参数对象的 weight_loader 属性绑定为该方法,加载权重时 PyTorch 的 load_state_dict 机制自动调用。

在模型初始化时,每个线性层/embedding 的参数会被打上对应的 weight_loader 属性。例如:

# ColumnParallelLinear.__init__() 中 [linear.py:692]
def weight_loader(self, param, loaded_weight, ...):
    # loaded_weight 沿 dim=0 做 TP rank 分片

7. 权重后处理

7.1 process_weights_after_loading()

model_loader/utils.py:99 — 所有权重加载完成后统一执行:

process_weights_after_loading(model):
  遍历所有 module:
    ├─ 若 module 有 quant_method (QuantizeMethodBase 实例):
    │   └─ quant_method.process_weights_after_loading(module)
    │       (量化后处理:量化参数重打包、scales 计算等)
    │
    ├─ CPU offload 处理:
    │   └─ device_loading_context() → 临时移到 GPU 做处理 → 再移回 CPU
    │
    └─ Attention 模块后处理:
        ├─ Attention.process_weights_after_loading()
        ├─ MLAAttention.process_weights_after_loading()
        └─ MMEncoderAttention.process_weights_after_loading()

7.2 在线量化的延迟处理

model_loader/reload/layerwise.py:216finalize_layerwise_processing()

在线量化(online quantization)场景下,部分层在初始化时位于 meta 设备。等到所有权重加载完成后,再逐层 materialize

  • 处理在初始加载期间被推迟的层
  • 从 meta 设备分配实际显存
  • 执行量化权重计算

8. 量化配置

8.1 量化检测

weight_utils.py:263get_quant_config() 从多处读取量化配置:

get_quant_config():
  ├─ HF model config 的 quantization_config 字段
  ├─ HF model config 的 compression_config 字段 (compressed-tensors)
  ├─ model_config 中的 hf_overrides
  ├─ model_config.quantization_config (在线量化)
  ├─ Inflight BNB 量化
  └─ 模型目录中的 quantize_config.json

8.2 量化方法注册

layers/quantization/init.py:109get_quantization_config() 映射量化方法名到配置类:

支持的量化方法包括:awqfp8fbgemm_fp8modeloptbitsandbytesggufgptq_marlinawq_marlincompressed-tensorsexperts_int8torchaomxfp4 等。

8.3 各 Loader 的量化处理

Loader 量化方式
BitsAndBytesModelLoader 预量化 4-bit/8-bit + inflight 量化,按 TP rank 分片后再量化
GGUFModelLoader gguf_quant_weights_iterator() 处理 GGUF 量化格式
DefaultModelLoader 通过 process_weights_after_loading() 执行后量化
在线量化 权重初始在 meta 设备,在 finalize_layerwise_processing() 中逐层 materialize

9. 特殊 Loader

9.1 DummyModelLoader

dummy_loader.py — 用于显存 profiling 阶段(见初始化文档):

class DummyModelLoader:
    def download_model(self):
        pass  # 无需下载

    def load_weights(self, model, model_config):
        initialize_dummy_weights(model, model_config)
        # → 为每个参数赋随机值 [-1e-3, 1e-3],按参数名确定性种子

weight_utils.py:1450initialize_dummy_weights() 以确定性的方式给出随机权重,使得同一模型每次的 profiling 结果一致。

9.2 ShardedStateLoader

sharded_state_loader.py — 用于预分片的 TP checkpoint:

  • 每个 rank 只读取自己的 shard 文件:model-rank-{rank}-part-{part}.safetensors
  • 使用 safe_open 或 Run:ai streamer 直接读取 tensor
  • 按 key 直接 copy 到模型参数,跳过 model.load_weights() 的匹配逻辑

9.3 TensorizerLoader

tensorizer_loader.py — 使用 CoreWeave 的 tensorizer 库做快速序列化模型加载:

  • 支持 vLLM-tensorized 模型(直接完整 state dict 加载)
  • 也支持标准 HF 模型(通过 tensorizer 加速)

10. 以 Llama 为例的完整加载追踪

get_model(vllm_config)
│
├─ get_model_loader("auto") → DefaultModelLoader
│
└─ DefaultModelLoader.load_model()
     │
     ├─ initialize_model()
     │   ├─ model_config.hf_config.architectures = ["LlamaForCausalLM"]
     │   ├─ ModelRegistry["LlamaForCausalLM"]
     │   │   → module="vllm.model_executor.models.llama"
     │   │   → class="LlamaForCausalLM"
     │   ├─ importlib.import_module("vllm.model_executor.models.llama")
     │   ├─ configure_quant_config() → (无量化,跳过)
     │   └─ LlamaForCausalLM(vllm_config=vllm_config, prefix="")
     │        ├─ __init__() [llama.py:515](https://github.com/vllm-project/vllm/blob/5bb8d2767a2829b56e58c68fa8f380e9e4e2bd3e/vllm/model_executor/models/llama.py#L515)
     │        │   ├─ LlamaModel(vllm_config)     # Transformer backbone
     │        │   │   ├─ VocabParallelEmbedding   # embed_tokens
     │        │   │   └─ LlamaDecoderLayer × N   # decoder layers
     │        │   │       ├─ QKVParallelLinear    # qkv_proj (fused)
     │        │   │       ├─ RowParallelLinear    # o_proj
     │        │   │       ├─ QKVParallelLinear    # gate_up_proj (fused)
     │        │   │       ├─ RowParallelLinear    # down_proj
     │        │   │       └─ Attention            # self_attn
     │        │   └─ ParallelLMHead / LogitsProcessor
     │        └─ 返回模型实例(参数已创建但未加载权重)
     │
     ├─ load_weights(model)
     │   ├─ _prepare_weights() → 发现 model-00001-of-00005.safetensors ...
     │   ├─ safetensors_weights_iterator() → 迭代 (name, tensor) 对
     │   │   ├─ "model.embed_tokens.weight" → torch.Size([128256, 4096])
     │   │   ├─ "model.layers.0.self_attn.q_proj.weight" → ...
     │   │   └─ ... (数百个参数)
     │   │
     │   └─ LlamaForCausalLM.load_weights(iterator)
     │        └─ AutoWeightsLoader.load_weights()
     │            ├─ embed_tokens.weight → embed_tokens(VocabParallelEmbedding)
     │            │   └─ param.weight_loader(param, tensor)
     │            │       └─ VocabParallelEmbedding:
     │            │           TP rank 0: 取 [0:32000] 行
     │            │           TP rank 1: 取 [32000:] 行
     │            │
     │            ├─ model.layers.0.* → LlamaDecoderLayer.load_weights()
     │            │   ├─ self_attn.qkv_proj.weight → 拆分为 q/k/v 分别加载
     │            │   │   ├─ q_proj: ColumnParallelLinear.weight_loader() → dim=0 TP 分片
     │            │   │   ├─ k_proj: ColumnParallelLinear.weight_loader() → dim=0 TP 分片
     │            │   │   └─ v_proj: ColumnParallelLinear.weight_loader() → dim=0 TP 分片
     │            │   ├─ self_attn.o_proj.weight: row_parallel_weight_loader()
     │            │   ├─ mlp.gate_up_proj.weight → 拆分为 gate/up
     │            │   └─ mlp.down_proj.weight: row_parallel_weight_loader
     │            │
     │            └─ lm_head.weight → (与 embed_tokens 共享,跳过)
     │
     ├─ process_weights_after_loading(model)
     │   ├─ (无量化,跳过 quant_method.process)
     │   └─ 各类 Attention 初始化
     │
     └─ return model.eval()

11. 关键文件索引

组件 文件 核心行号
加载入口 init.py L52-69 (registry), L125 (get_model_loader), L133 (get_model)
模板方法基类 base_loader.py L25 (BaseModelLoader), L43 (load_model)
默认 Loader default_loader.py L97 (_prepare_weights), L211 (_get_weights_iterator), L382 (load_weights)
Dummy Loader dummy_loader.py L28 (DummyModelLoader)
BNB Loader bitsandbytes_loader.py L782 (load_weights)
GGUF Loader gguf_loader.py L418 (load_model)
初始化工具 model_loader/utils.py L40 (initialize_model), L99 (process_weights_after_loading), L179 (_get_model_architecture)
权重工具 weight_utils.py L263 (get_quant_config), L504 (download), L893 (safetensors_iterator), L1383 (default_weight_loader), L1450 (dummy)
模型注册表 models/registry.py L71 (_TEXT_GENERATION_MODELS), L683 (_VLLM_MODELS), L1201 (resolve_model_cls), L1353 (ModelRegistry)
AutoWeightsLoader models/utils.py L115 (AutoWeightsLoader), L259 (_load_module)
Llama 示例 models/llama.py L290 (LlamaForCausalLM.init), L436 (DecoderLayer.load_weights), L584 (ForCausalLM.load_weights)
在线量化处理 reload/layerwise.py L216
量化注册 quantization/init.py L109
LoadConfig config/load.py L26

本站访问量

This site uses Just the Docs, a documentation theme for Jekyll.