DefaultModelLoader:从 HF Hub 把权重搬到显卡

源码版本v0.25.1

职责

模型加载器 (model loader) 是 vLLM 启动流程里把模型权重从硬盘或 Hugging Face Hub 搬到 GPU 显存的那一环。DefaultModelLoader 是其中最常用的一种实现,负责处理 safetensors / .pt / Mistral 自家格式 / npcache 等多种磁盘格式,把多分片权重拼成一个 (name, tensor) 迭代器,再交给模型自己的 load_weights 方法逐个填进 nn.Module。它在 BaseModelLoader 之上把"下载、过滤、解析、迭代"四件事串起来,既不关心张量怎么分片也不关心怎么量化——这两件事由 ColumnParallelLinear 之类的层和 QuantizationConfig 接管。

具体一次加载大致是:__init__ 校验 LoadConfig 里的 model_loader_extra_config(default_loader.py:74-126),只允许 enable_multithread_load、num_threads、enable_weights_track 三个键。download_model 走的是 _prepare_weights,该函数根据 load_format 选 allow_patterns,本地目录直接 glob,远端则走 download_weights_from_hf(default_loader.py:194-207)。load_weights 才是真正干活的入口,它先做 EP 过滤初始化(default_loader.py:425-426),把 get_all_weights 迭代器交给 model.load_weights,再用 track_weights_loading 校验没漏的参数(default_loader.py:444-445)。

设计动机

为什么把加载器拆成 Loader + Source 两层?

• 多源权重:某些模型主体之外还有 secondary_weights(比如多模态 visual encoder 单独一份 checkpoint),Source dataclass 让每个源都带自己的 prefix、fall_back_to_pt、allow_patterns_overrides(default_loader.py:49-69),get_all_weights 把它们拼成一个迭代器(default_loader.py:321-340)。
• 格式分叉:load_format 一项决定后续走哪条路,auto 会先探测 consolidated*.safetensors 自动切到 mistral(default_loader.py:151-163),safetensors / fastsafetensors / instanttensor 强制只用 .safetensors,pt / npcache 走老路径。
• 多线程加速:enable_multithread_load 切到 multi_thread_safetensors_weights_iterator 或 multi_thread_pt_weights_iterator,默认 8 线程(default_loader.py:278-308);但多线程模式只支持默认 lazy 策略,跟 safetensors_load_strategy 别的策略互斥(default_loader.py:116-126)。
• EP 权重过滤:MoE + 专家并行时只加载本 rank 负责的 expert,_init_ep_weight_filter 算出 local_expert_ids(default_loader.py:351-412),后续 safetensors 迭代器据此跳过非本地 expert 张量,省磁盘 IO。
• safetensors 索引去重:像 Mistral-7B-Instruct-v0.3 同时存在 sharded 和 consolidated 两份 safetensors,直接 glob 会全读导致冲突,filter_duplicate_safetensors_files 用 model.safetensors.index.json 过滤(default_loader.py:217-235)。
• 加载追踪:非量化模型默认 enable_weights_track=True,加载完比对 named_parameters 找出漏掉的权重,避免静默地用随机权重跑推理(default_loader.py:434-445)。

关键文件

• DefaultModelLoader class:43-47 — 类定义和 docstring,DEFAULT_NUM_THREADS = 8。
• Source dataclass:49-69 — Source 内嵌数据类,描述一个权重来源。
• _prepare_weights:128-242 — 解析 load_format、下载、glob、去重,返回 (hf_folder, hf_weights_files, use_safetensors)。
• _get_weights_iterator:244-319 — 根据 load_format 和多线程开关挑出 weights iterator,最后给张量名加 prefix。
• get_all_weights:321-340 — 把 primary source 和 secondary_weights 拼成单个生成器。
• _init_ep_weight_filter:351-412 — EP 模式下计算 local_expert_ids 以跳过非本地 expert 张量。
• load_weights:414-445 — 加载入口,torchao 切策略 → EP 过滤 → model.load_weights → 追踪校验。
• track_weights_loading:447-475 — 比对 named_parameters 与 loaded_weights,报出缺失或意外的权重。
• weight_utils.py:1-1 — 提供 safetensors_weights_iterator、download_weights_from_hf、filter_duplicate_safetensors_files 等底层工具。
• base_loader.py:1-1 — BaseModelLoader 抽象基类,定义 load_weights 接口。

数据流

调用方(通常是 worker 的 load_model)先 __init__ 一个 DefaultModelLoader,然后调 load_weights(model, model_config)。函数体里第一步根据 quantization == "torchao" 调整 safetensors 加载策略,接着 _init_ep_weight_filter 算 EP 过滤集合,然后 get_all_weights 把每个 Source 喂给 _get_weights_iterator 拼成生成器,最后 model.load_weights 自己消费:

@instrument(span_name="Load weights")
def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
    if model_config.quantization == "torchao":
        quant_config = get_quant_config(model_config, self.load_config)
        if (
            hasattr(quant_config, "is_checkpoint_torchao_serialized")
            and quant_config.is_checkpoint_torchao_serialized
            and torchao_version_at_least("0.15.0")
        ):
            self.load_config.safetensors_load_strategy = "torchao"

    self._init_ep_weight_filter(model_config)

    loaded_weights = model.load_weights(self.get_all_weights(model_config, model))

    self.counter_after_loading_weights = time.perf_counter()
    logger.info_once(
        "Loading weights took %.2f seconds",
        self.counter_after_loading_weights - self.counter_before_loading_weights,
    )
    # We only enable strict check for non-quantized models
    # that have loaded weights tracking by default.
    default_enable_weights_track = (
        model_config.quantization is None and loaded_weights is not None
    )

(default_loader.py:414-438)注意 model.load_weights(...) 才是真正把张量拷进 nn.Module 的地方,vLLM 的每个模型实现都自定义这个方法,在那里根据张量名匹配对应的 weight_loader(例如 ColumnParallelLinear.weight_loader_v2),分片和量化都在那一层完成。详见 张量并行线性层 和 量化层加载。

边界与失败

• load_format 未知:任何不在 hf / safetensors / fastsafetensors / instanttensor / mistral / pt / npcache 列表里的值直接 raise ValueError(default_loader.py:183-184)。
• 找不到权重文件:glob 不到任何文件时抛 Cannot find any model weights(default_loader.py:237-240)。
• 多线程 + safetensors_load_strategy 不兼容:enable_multithread_load=True 只支持默认 lazy 策略,组合其他策略会 raise ValueError(default_loader.py:118-126)。
• extra_config 类型错误:model_loader_extra_config 不是 dict、键不在白名单、num_threads 非正整数都会立即报错(default_loader.py:79-110)。
• EPLB 时跳过 EP 过滤:enable_eplb 打开时冗余专家槽位会指向别的 rank 的逻辑 expert,过滤会导致这些槽位漏填,所以整段直接 return(default_loader.py:370-375)。
• 量化模型跳过追踪校验:量化方法带 process_weights_after_loading 或 uses_meta_device 时,checkpoint 里可能没有对应的 scale,这些参数从追踪集合里剔除(default_loader.py:453-462)。

小结

DefaultModelLoader 把"找文件、下文件、迭代张量"三件事做到位,具体怎么分片 / 量化 / 反量化都丢给具体层和量化方法。它和 张量并行线性层 的衔接点是 model.load_weights 内部对每个张量名查 weight_loader,跟 量化层加载 的衔接点是 quant_config.get_quant_method 给每个层装上 LinearMethodBase,加载器本身不感知。

对照官方资料:vLLM 文档 · README。