vLLM 模型加载流程分析
1. 概述
vLLM 的模型加载是一个分层的、可插拔的系统。从高层入口到底层 weight loader,整个过程分为 5 个阶段:
- 根据
load_format选择 ModelLoader - 解析模型架构并实例化
nn.Module - 从磁盘/HuggingFace Hub 加载权重
- 后处理权重(量化重打包、attention 初始化)
- 返回
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:125 — get_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_LOADER(L52-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-82 — BaseModelLoader.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:1353 — ModelRegistry 是一个全局单例字典,将架构名(如 "LlamaForCausalLM")映射到 _LazyRegisteredModel 对象。每个条目包含 module_name 和 class_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:504 — download_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:115 — AutoWeightsLoader 是自动化的递归 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:216 — finalize_layerwise_processing():
在线量化(online quantization)场景下,部分层在初始化时位于 meta 设备。等到所有权重加载完成后,再逐层 materialize:
- 处理在初始加载期间被推迟的层
- 从 meta 设备分配实际显存
- 执行量化权重计算
8. 量化配置
8.1 量化检测
weight_utils.py:263 — get_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:109 — get_quantization_config() 映射量化方法名到配置类:
支持的量化方法包括:awq、fp8、fbgemm_fp8、modelopt、bitsandbytes、gguf、gptq_marlin、awq_marlin、compressed-tensors、experts_int8、torchao、mxfp4 等。
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:1450 — initialize_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 |