DefaultModelLoader: Gewichte vom HF Hub auf die GPU bringen

源码版本v0.25.1

Verantwortung

Der Model Loader (Modellladeprogramm) ist in vLLMs Startphase die Komponente, die Modellgewichte von der Festplatte oder dem Hugging Face Hub in den GPU-Speicher lädt. DefaultModelLoader ist die am häufigsten verwendete Implementierung und verarbeitet verschiedene Festplattenformate wie safetensors / .pt / Mistral-eigenes Format / npcache. Sie fügt mehrteilige Gewichte zu einem (name, tensor)-Iterator zusammen und übergibt ihn an die eigene load_weights-Methode des Modells, die die Tensoren nacheinander in nn.Module einträgt. Auf Basis von BaseModelLoader verkettet es die vier Schritte "Herunterladen, Filtern, Parsen, Iterieren" und kümmert sich weder um Tensor-Sharding noch um Quantisierung – diese beiden Aspekte werden von Schichten wie ColumnParallelLinear und von QuantizationConfig übernommen.

Ein typischer Ladevorgang läuft groß so: __init__ validiert die model_loader_extra_config aus der LoadConfig(default_loader.py:74-126) und lässt nur die Schlüssel enable_multithread_load, num_threads, enable_weights_track zu. download_model läuft über _prepare_weights, das anhand von load_format die allow_patterns auswählt, lokale Verzeichnisse direkt per glob durchsucht und Remote-Pfade über download_weights_from_hf abwickelt(default_loader.py:194-207). load_weights ist der eigentliche Arbeitseinstieg: Zuerst erfolgt die EP-Filterinitialisierung(default_loader.py:425-426), dann wird der Iterator get_all_weights an model.load_weights übergeben und abschließend mit track_weights_loading geprüft, dass keine Parameter fehlen(default_loader.py:444-445).

Entwurfsmotivation

Warum ist der Loader in eine Loader- und eine Source-Schicht aufgeteilt?

• Mehrere Gewichtsquellen: Manche Modelle haben neben dem Hauptkörper secondary_weights (z. B. bei multimodalen Modellen ein separater Checkpoint für den Visual Encoder). Das Source-Dataclass sorgt dafür, dass jede Quelle ihre eigenen prefix, fall_back_to_pt und allow_patterns_overrides mitbringt(default_loader.py:49-69), und get_all_weights fügt sie zu einem Iterator zusammen(default_loader.py:321-340).
• Formatverzweigung: load_format entscheidet über den weiteren Pfad. Bei auto wird zuerst nach consolidated*.safetensors gesucht und bei Treffer auf mistral umgeschaltet(default_loader.py:151-163). safetensors / fastsafetensors / instanttensor erzwingen ausschließlich .safetensors, während pt / npcache den alten Pfad nehmen.
• Multithread-Beschleunigung: enable_multithread_load schaltet auf multi_thread_safetensors_weights_iterator oder multi_thread_pt_weights_iterator um, standardmäßig mit 8 Threads(default_loader.py:278-308). Der Multithread-Modus unterstützt jedoch nur die Standard-Lazy-Strategie und ist mit anderen safetensors_load_strategy-Strategien inkompatibel(default_loader.py:116-126).
• EP-Gewichtsfilterung: Bei MoE + Expert Parallelism (Expertenparallelität) werden nur die Experten geladen, für die dieser Rank zuständig ist. _init_ep_weight_filter berechnet die local_expert_ids(default_loader.py:351-412), woraufhin der safetensors-Iterator nicht-lokale Expert-Tensoren überspringt und so Festplatten-IO einspart.
• safetensors-Index-Deduplizierung: Bei Modellen wie Mistral-7B-Instruct-v0.3 liegen gleichzeitig sharded- und consolidated-safetensors vor. Ein direktes glob würde beide lesen und Konflikte verursachen. filter_duplicate_safetensors_files nutzt model.safetensors.index.json zum Filtern(default_loader.py:217-235).
• Lade-Tracking: Für nicht-quantisierte Modelle gilt standardmäßig enable_weights_track=True. Nach dem Laden wird mit named_parameters abgeglichen, um fehlende Gewichte zu erkennen und zu verhindern, dass stillschweigend mit zufälligen Gewichten inferiert wird(default_loader.py:434-445).

Schlüsseldateien

• DefaultModelLoader class:43-47 — Klassendefinition und Docstring, DEFAULT_NUM_THREADS = 8.
• Source dataclass:49-69 — Source-innere Datenklasse, beschreibt eine Gewichtsquellen.
• _prepare_weights:128-242 — parst load_format, lädt herunter, glob, dedupliziert, gibt (hf_folder, hf_weights_files, use_safetensors) zurück.
• _get_weights_iterator:244-319 — wählt anhand von load_format und dem Multithread-Schalter den Weights-Iterator aus und fügt dem Tensornamen zuletzt prefix hinzu.
• get_all_weights:321-340 — fasst die primäre Source und secondary_weights zu einem einzigen Generator zusammen.
• _init_ep_weight_filter:351-412 — berechnet im EP-Modus die local_expert_ids, um nicht-lokale Expert-Tensoren zu überspringen.
• load_weights:414-445 — Ladeeinstieg: torchao-Strategie wählen → EP-Filter → model.load_weights → Tracking-Validierung.
• track_weights_loading:447-475 — gleicht named_parameters mit loaded_weights ab und meldet fehlende oder unerwartete Gewichte.
• weight_utils.py:1-1 — stellt Low-Level-Werkzeuge wie safetensors_weights_iterator, download_weights_from_hf, filter_duplicate_safetensors_files bereit.
• base_loader.py:1-1 — BaseModelLoader-abstrakte Basisklasse, definiert die load_weights-Schnittstelle.

Datenfluss

Der Aufrufer (üblicherweise load_model des Workers) instanziiert zuerst einen DefaultModelLoader über __init__ und ruft dann load_weights(model, model_config) auf. In der Funktion wird anhand von quantization == "torchao" zuerst die safetensors-Ladestrategie angepasst. Anschließend berechnet _init_ep_weight_filter die EP-Filtermenge, get_all_weights füttert jede Source in _get_weights_iterator und erzeugt so einen Generator, den model.load_weights schließlich konsumiert:

@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) Beachten Sie, dass erst model.load_weights(...) die Tensoren tatsächlich in nn.Module kopiert. Jede vLLM-Modellimplementierung überschreibt diese Methode selbst und wählt dort anhand des Tensornamens den passenden weight_loader (z. B. ColumnParallelLinear.weight_loader_v2) aus – Sharding und Quantisierung erfolgen erst auf dieser Ebene. Details siehe Tensorparallele lineare Schichten und Quantisierungsschicht laden.

Grenzen und Fehler

• Unbekanntes load_format: Jeder Wert außerhalb der Liste hf / safetensors / fastsafetensors / instanttensor / mistral / pt / npcache führt sofort zu raise ValueError(default_loader.py:183-184).
• Keine Gewichtsdatei gefunden: Wenn glob keine Datei findet, wird Cannot find any model weights geworfen(default_loader.py:237-240).
• Multithread + safetensors_load_strategy inkompatibel: enable_multithread_load=True unterstützt nur die Standard-Lazy-Strategie; Kombinationen mit anderen Strategien lösen raise ValueError aus(default_loader.py:118-126).
• extra_config-Typfehler: Wenn model_loader_extra_config kein dict ist, Schlüssel nicht in der Whitelist liegen oder num_threads keine positive ganze Zahl ist, wird sofort ein Fehler geworfen(default_loader.py:79-110).
• EP-Filter bei EPLB überspringen: Wenn enable_eplb aktiv ist, zeigen redundante Expert-Slots auf logische Experten eines anderen Ranks; eine Filterung würde diese Slots leer lassen, daher wird der gesamte Block direkt per return beendet(default_loader.py:370-375).
• Tracking-Validierung bei quantisierten Modellen überspringen: Wenn die Quantisierungsmethode process_weights_after_loading oder uses_meta_device mitbringt, kann der Checkpoint die entsprechenden Scales nicht enthalten – diese Parameter werden aus der Tracking-Menge entfernt(default_loader.py:453-462).

Zusammenfassung

DefaultModelLoader erledigt die drei Aufgaben "Dateien finden, herunterladen, Tensoren iterieren" ordentlich; wie genau geschhardet, quantisiert oder dequantisiert wird, überlässt es den konkreten Schichten und Quantisierungsmethoden. Die Nahtstelle zu Tensorparallele lineare Schichten ist die weight_loader-Auflösung pro Tensorname innerhalb von model.load_weights; die Nahtstelle zu Quantisierungsschicht laden ist quant_config.get_quant_method, das jeder Schicht ein LinearMethodBase aufsetzt – der Loader selbst bekommt davon nichts mit.

Siehe offizielle Dokumentation: vLLM 文档 · README.