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