DefaultModelLoader : du HF Hub vers la VRAM du GPU

源码版本v0.25.1

Responsabilités

Le chargeur de modèle (model loader) est, dans le démarrage de vLLM, l'étape qui transporte les poids du modèle depuis le disque ou Hugging Face Hub vers la VRAM du GPU. DefaultModelLoader en est l'implémentation la plus courante : il gère plusieurs formats sur disque (safetensors / .pt / format maison Mistral / npcache), assemble les poids multi-shards en un itérateur de (name, tensor), puis laisse la méthode load_weights du modèle les déposer un à un dans le nn.Module. Il s'appuie sur BaseModelLoader pour enchaîner « télécharger, filtrer, parser, itérer » ; il ne se préoccupe ni du sharding des tensors ni de la quantization — ces deux tâches sont laissées aux couches comme ColumnParallelLinear et à QuantizationConfig.

Une passe de chargement typique : __init__ valide model_loader_extra_config du LoadConfig(default_loader.py:74-126) en n'acceptant que les clés enable_multithread_load, num_threads, enable_weights_track. download_model passe par _prepare_weights, qui choisit les allow_patterns selon load_format, glob directement pour un dossier local, ou utilise download_weights_from_hf pour du distant(default_loader.py:194-207). load_weights est l'entrée qui travaille réellement : initialisation du filtre EP(default_loader.py:425-426), passage de l'itérateur get_all_weights à model.load_weights, puis vérification via track_weights_loading qu'aucun paramètre n'a été oublié(default_loader.py:444-445).

Motivation de conception

Pourquoi séparer le loader en deux couches Loader + Source ?

• Poids multi-sources : certains modèles ont, en plus du corps principal, des secondary_weights (par ex. l'encoder visuel d'un modèle multimodal, avec son propre checkpoint) ; le dataclass Source donne à chaque source son prefix, fall_back_to_pt, allow_patterns_overrides(default_loader.py:49-69) ; get_all_weights les concatène en un seul itérateur(default_loader.py:321-340).
• Embranchement par format : load_format décide du chemin à suivre ; auto détecte d'abord consolidated*.safetensors et bascule en mistral(default_loader.py:151-163) ; safetensors / fastsafetensors / instanttensor force l'usage exclusif de .safetensors, pt / npcache emprunte l'ancien chemin.
• Accélération multi-thread : enable_multithread_load bascule vers multi_thread_safetensors_weights_iterator ou multi_thread_pt_weights_iterator, avec 8 threads par défaut(default_loader.py:278-308) ; mais le mode multi-thread ne supporte que la stratégie lazy par défaut, et est incompatible avec les autres safetensors_load_strategy(default_loader.py:116-126).
• Filtrage EP des poids : en MoE + expert parallelism, on ne charge que les experts gérés par ce rank ; _init_ep_weight_filter calcule les local_expert_ids(default_loader.py:351-412) ; l'itérateur safetensors saute ensuite les tensors d'experts non locaux, économisant des I/O disque.
• Déduplication de l'index safetensors : Mistral-7B-Instruct-v0.3 possède à la fois des versions sharded et consolidated des safetensors ; un glob brut lirait les deux et causerait des conflits, filter_duplicate_safetensors_files utilise model.safetensors.index.json pour filtrer(default_loader.py:217-235).
• Tracking du chargement : pour les modèles non quantifiés, enable_weights_track=True par défaut ; après chargement, on compare aux named_parameters pour repérer les poids oubliés, évitant de lancer une inference en silence avec des poids aléatoires(default_loader.py:434-445).

Fichiers clés

• DefaultModelLoader class:43-47 — définition de classe et docstring, DEFAULT_NUM_THREADS = 8.
• Source dataclass:49-69 — dataclass interne Source, décrit une source de poids.
• _prepare_weights:128-242 — parse load_format, télécharge, glob, déduplique, renvoie (hf_folder, hf_weights_files, use_safetensors).
• _get_weights_iterator:244-319 — selon load_format et le toggle multi-thread, choisit le weights iterator, puis préfixe le nom des tensors avec prefix.
• get_all_weights:321-340 — concatène primary source et secondary_weights en un seul générateur.
• _init_ep_weight_filter:351-412 — en EP, calcule les local_expert_ids pour sauter les tensors d'experts non locaux.
• load_weights:414-445 — entrée du chargement : stratégie torchao → filtre EP → model.load_weights → tracking.
• track_weights_loading:447-475 — compare named_parameters et loaded_weights, signale les poids manquants ou inattendus.
• weight_utils.py:1-1 — fournit safetensors_weights_iterator, download_weights_from_hf, filter_duplicate_safetensors_files, etc.
• base_loader.py:1-1 — classe abstraite BaseModelLoader, définit l'interface load_weights.

Flux de données

L'appelant (typiquement le load_model du worker) instancie d'abord un DefaultModelLoader, puis appelle load_weights(model, model_config). Dans le corps, on ajuste d'abord la stratégie safetensors si quantization == "torchao", puis _init_ep_weight_filter calcule le filtre EP, puis get_all_weights nourrit chaque Source à _get_weights_iterator pour produire le générateur, et enfin model.load_weights le consomme :

@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) Notez que c'est model.load_weights(...) qui copie réellement les tensors dans le nn.Module ; chaque implémentation de modèle dans vLLM redéfinit cette méthode, et y associe chaque nom de tensor au bon weight_loader (par ex. ColumnParallelLinear.weight_loader_v2) — sharding et quantization sont faits à ce niveau. Voir couches linéaires parallèles par tensor et chargement des couches quantifiées.

Limites et échecs

• load_format inconnu : toute valeur hors de hf / safetensors / fastsafetensors / instanttensor / mistral / pt / npcache lève ValueError(default_loader.py:183-184).
• Fichiers de poids introuvables : si le glob ne retourne rien, on lève Cannot find any model weights(default_loader.py:237-240).
• Multi-thread + safetensors_load_strategy incompatibles : enable_multithread_load=True ne supporte que la stratégie lazy par défaut ; les combinaisons avec d'autres stratégies lèvent ValueError(default_loader.py:118-126).
• Type d'extra_config erroné : si model_loader_extra_config n'est pas un dict, qu'une clé n'est pas dans la whitelist, ou que num_threads n'est pas un entier positif, on lève immédiatement(default_loader.py:79-110).
• Saut du filtre EP en EPLB : quand enable_eplb est activé, les slots d'experts redondants pointent vers des experts logiques d'autres ranks, et le filtrage les laisserait vides ; on return donc directement(default_loader.py:370-375).
• Modèles quantifiés hors tracking : si la méthode de quantization a process_weights_after_loading ou uses_meta_device, le checkpoint peut ne pas contenir les scales correspondants ; ces paramètres sont retirés de l'ensemble de tracking(default_loader.py:453-462).

Résumé

DefaultModelLoader fait correctement les trois choses « trouver les fichiers, les télécharger, itérer les tensors » ; le sharding / quantization / déquantization proprement dit est laissé aux couches et aux méthodes de quantization. La jointure avec les couches linéaires parallèles par tensor se fait à l'intérieur de model.load_weights, où chaque nom de tensor est associé à un weight_loader ; avec le chargement des couches quantifiées, elle se fait via quant_config.get_quant_method qui équipe chaque couche d'un LinearMethodBase — le loader lui-même reste ignorant.

Voir la documentation officielle : Documentation vLLM · README