DefaultModelLoader: moving weights from HF Hub onto the GPU

源码版本v0.25.1

Responsibilities

The model loader is the stage in vLLM's startup flow that moves model weights from disk or Hugging Face Hub into GPU memory. DefaultModelLoader is the most commonly used implementation; it handles multiple on-disk formats — safetensors / .pt / Mistral's native format / npcache — stitches sharded weights into a (name, tensor) iterator, and hands it to the model's own load_weights method to populate nn.Module one tensor at a time. Built on top of BaseModelLoader, it strings together four jobs — download, filter, parse, iterate — without caring about how tensors are sharded or quantized. Those two concerns are owned by layers like ColumnParallelLinear and by QuantizationConfig.

A typical load goes roughly: __init__ validates model_loader_extra_config inside LoadConfig (default_loader.py:74-126), allowing only enable_multithread_load, num_threads, enable_weights_track. download_model delegates to _prepare_weights, which picks allow_patterns based on load_format, globs the local directory directly, or calls download_weights_from_hf for remote sources (default_loader.py:194-207). load_weights is the entry that actually does the work: it initializes the EP filter (default_loader.py:425-426), feeds the get_all_weights iterator into model.load_weights, then uses track_weights_loading to verify no parameters were missed (default_loader.py:444-445).

Design motivation

Why split the loader into a Loader + Source two-layer design?

• Multi-source weights: Some models carry secondary_weights outside the main body (e.g. a multimodal visual encoder with its own checkpoint). The Source dataclass lets each source carry its own prefix, fall_back_to_pt, allow_patterns_overrides (default_loader.py:49-69), and get_all_weights stitches them into a single iterator (default_loader.py:321-340).
• Format fork: A single load_format field decides which path to take; auto probes for consolidated*.safetensors and switches to mistral (default_loader.py:151-163); safetensors / fastsafetensors / instanttensor force .safetensors only, while pt / npcache take the legacy path.
• Multi-threaded acceleration: enable_multithread_load switches to multi_thread_safetensors_weights_iterator or multi_thread_pt_weights_iterator, default 8 threads (default_loader.py:278-308); but the multi-threaded mode only supports the default lazy strategy and is mutually exclusive with other safetensors_load_strategy options (default_loader.py:116-126).
• EP weight filtering: For MoE + expert parallelism, only the experts owned by this rank are loaded. _init_ep_weight_filter computes local_expert_ids (default_loader.py:351-412), and the safetensors iterator later skips non-local expert tensors, saving disk IO.
• Safetensors index dedup: Models like Mistral-7B-Instruct-v0.3 ship both sharded and consolidated safetensors; a naive glob would read both and conflict. filter_duplicate_safetensors_files uses model.safetensors.index.json to filter (default_loader.py:217-235).
• Load tracking: Non-quantized models default to enable_weights_track=True; after loading, it diffs against named_parameters to flag missing weights, so inference never silently runs with random weights (default_loader.py:434-445).

Key files

• DefaultModelLoader class:43-47 — class definition and docstring, DEFAULT_NUM_THREADS = 8.
• Source dataclass:49-69 — Source inner dataclass describing one weight source.
• _prepare_weights:128-242 — parses load_format, downloads, globs, dedups, returns (hf_folder, hf_weights_files, use_safetensors).
• _get_weights_iterator:244-319 — picks the weights iterator based on load_format and the multi-thread switch, then prefixes tensor names.
• get_all_weights:321-340 — stitches the primary source and secondary_weights into a single generator.
• _init_ep_weight_filter:351-412 — computes local_expert_ids under EP mode to skip non-local expert tensors.
• load_weights:414-445 — load entry: torchao strategy switch → EP filter → model.load_weights → tracking validation.
• track_weights_loading:447-475 — diffs named_parameters against loaded_weights, reporting missing or unexpected weights.
• weight_utils.py:1-1 — provides low-level utilities like safetensors_weights_iterator, download_weights_from_hf, filter_duplicate_safetensors_files.
• base_loader.py:1-1 — BaseModelLoader abstract base class defining the load_weights interface.

Data flow

The caller (typically the worker's load_model) first constructs a DefaultModelLoader via __init__, then calls load_weights(model, model_config). Inside the function body, the first step adjusts the safetensors load strategy when quantization == "torchao", then _init_ep_weight_filter computes the EP filter set, get_all_weights feeds each Source into _get_weights_iterator to build a generator, and finally model.load_weights consumes it:

@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) Note that model.load_weights(...) is where tensors are actually copied into nn.Module. Every vLLM model implementation customizes this method, matching each tensor name to its corresponding weight_loader (e.g. ColumnParallelLinear.weight_loader_v2); sharding and quantization happen at that layer. See Tensor-parallel linear layers and Quantized layer loading for details.

Boundaries and failure modes

• Unknown load_format: Any value outside hf / safetensors / fastsafetensors / instanttensor / mistral / pt / npcache directly raises ValueError (default_loader.py:183-184).
• Weight files not found: When the glob matches nothing, it raises Cannot find any model weights (default_loader.py:237-240).
• Multi-threading + safetensors_load_strategy incompatibility: enable_multithread_load=True only supports the default lazy strategy; combining it with other strategies raises ValueError (default_loader.py:118-126).
• extra_config type errors: model_loader_extra_config not a dict, keys outside the whitelist, or num_threads not a positive integer all error immediately (default_loader.py:79-110).
• Skip EP filter under EPLB: When enable_eplb is on, redundant expert slots point to another rank's logical expert; filtering would leave these slots unfilled, so the entire block returns early (default_loader.py:370-375).
• Quantized models skip tracking validation: When the quantization method has process_weights_after_loading or uses_meta_device, the checkpoint may not contain the corresponding scales, so these parameters are excluded from the tracking set (default_loader.py:453-462).

Summary

DefaultModelLoader handles three things well — find files, download files, iterate tensors — and leaves sharding / quantization / dequantization to the layers and quantization methods. Its seam with Tensor-parallel linear layers is that model.load_weights looks up the weight_loader for each tensor name internally; its seam with Quantized layer loading is that quant_config.get_quant_method installs a LinearMethodBase onto each layer, which the loader itself is unaware of.

See official docs: vLLM docs · README