DefaultModelLoader: del HF Hub a la GPU

源码版本v0.25.1

Responsabilidades

El cargador de modelos (model loader) es la etapa del flujo de arranque de vLLM que lleva los pesos desde el disco o Hugging Face Hub a la memoria de la GPU. DefaultModelLoader es la implementación más común: se encarga de formatos de disco como safetensors / .pt / el formato propio de Mistral / npcache, concatena los pesos multi-shard en un iterador de (name, tensor) y luego se los pasa al método load_weights del propio modelo para que los rellene dentro del nn.Module. Por encima de BaseModelLoader encadena cuatro tareas —descargar, filtrar, parsear, iterar— sin preocuparse por cómo se parten los tensores ni por cómo se cuantizan: esas dos tareas las asumen capas como ColumnParallelLinear y QuantizationConfig.

Una carga concreta es más o menos así: __init__ valida model_loader_extra_config dentro del LoadConfig (default_loader.py:74-126), admitiendo solo enable_multithread_load, num_threads y enable_weights_track. download_model pasa por _prepare_weights, que según load_format elige allow_patterns, hace glob si es un directorio local o llama a download_weights_from_hf si es remoto (default_loader.py:194-207). load_weights es la entrada que realmente trabaja: primero inicializa el filtro EP (default_loader.py:425-426), entrega el iterador get_all_weights a model.load_weights y por último verifica con track_weights_loading que no falte ningún parámetro (default_loader.py:444-445).

Motivación de diseño

¿Por qué partir el cargador en dos capas Loader + Source?

• Pesos multi-fuente: algunos modelos además del cuerpo principal tienen secondary_weights (por ejemplo, el encoder visual de un modelo multimimodal con checkpoint aparte); el dataclass Source permite que cada fuente lleve su propio prefix, fall_back_to_pt y allow_patterns_overrides (default_loader.py:49-69), y get_all_weights los concatena en un único iterador (default_loader.py:321-340).
• Bifurcación de formato: el campo load_format decide qué rama seguir; auto primero detecta consolidated*.safetensors y conmuta a mistral (default_loader.py:151-163), safetensors / fastsafetensors / instanttensor fuerza usar solo .safetensors, y pt / npcache va por la ruta legacy.
• Aceleración multi-hilo: enable_multithread_load activa multi_thread_safetensors_weights_iterator o multi_thread_pt_weights_iterator, con 8 hilos por defecto (default_loader.py:278-308); pero el modo multi-hilo solo soporta la estrategia lazy por defecto y es incompatible con las demás estrategias de safetensors_load_strategy (default_loader.py:116-126).
• Filtrado de pesos EP: con MoE + expert parallelism se cargan solo los experts de los que se encarga este rank; _init_ep_weight_filter calcula local_expert_ids (default_loader.py:351-412) para que el iterador de safetensors salte los tensores de experts no locales y ahorrar IO de disco.
• Deduplicación del índice safetensors: modelos como Mistral-7B-Instruct-v0.3 tienen a la vez safetensors sharded y consolidated; un glob directo leería ambos y daría conflictos, así que filter_duplicate_safetensors_files los filtra usando model.safetensors.index.json (default_loader.py:217-235).
• Seguimiento de carga: en modelos no cuantizados enable_weights_track=True por defecto; al terminar, compara con named_parameters para detectar pesos faltantes y evitar correr inferencia con pesos aleatorios en silencio (default_loader.py:434-445).

Archivos clave

• DefaultModelLoader class:43-47 — definición de la clase y docstring, DEFAULT_NUM_THREADS = 8.
• Source dataclass:49-69 — dataclass anidado Source, describe una fuente de pesos.
• _prepare_weights:128-242 — parsea load_format, descarga, hace glob, deduplica y devuelve (hf_folder, hf_weights_files, use_safetensors).
• _get_weights_iterator:244-319 — según load_format y el flag multi-hilo elige el iterador de pesos y al final añade el prefix al nombre del tensor.
• get_all_weights:321-340 — concatena la fuente primaria y los secondary_weights en un único generador.
• _init_ep_weight_filter:351-412 — en modo EP calcula local_expert_ids para saltar tensores de experts no locales.
• load_weights:414-445 — entrada de carga: cambia la estrategia de safetensors para torchao → filtro EP → model.load_weights → verificación con track.
• track_weights_loading:447-475 — compara named_parameters con loaded_weights y reporta pesos faltantes o inesperados.
• weight_utils.py:1-1 — utilidades de bajo nivel como safetensors_weights_iterator, download_weights_from_hf, filter_duplicate_safetensors_files.
• base_loader.py:1-1 — clase abstracta BaseModelLoader que define la interfaz load_weights.

Flujo de datos

El llamador (normalmente el load_model del worker) primero instancia con __init__ un DefaultModelLoader y luego llama a load_weights(model, model_config). En el cuerpo de la función, el primer paso ajusta la estrategia de carga de safetensors si quantization == "torchao", después _init_ep_weight_filter calcula el conjunto de filtrado EP, get_all_weights alimenta cada Source a _get_weights_iterator para construir un generador y, por último, model.load_weights lo consume:

@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) Nótese que model.load_weights(...) es lo que de verdad copia los tensores al nn.Module; cada implementación de modelo en vLLM define este método y ahí, según el nombre del tensor, hace match con el weight_loader correspondiente (por ejemplo ColumnParallelLinear.weight_loader_v2). El sharding y la cuantización se hacen en esa capa. Véase capas lineales con tensor parallelism y carga de capas cuantizadas.

Límites y fallos

• load_format desconocido: cualquier valor fuera de hf / safetensors / fastsafetensors / instanttensor / mistral / pt / npcache lanza ValueError directamente (default_loader.py:183-184).
• No se encuentran los pesos: si el glob no devuelve ningún archivo lanza Cannot find any model weights (default_loader.py:237-240).
• Multi-hilo + safetensors_load_strategy incompatibles: enable_multithread_load=True solo soporta la estrategia lazy por defecto; combinarla con otra hace raise ValueError (default_loader.py:118-126).
• Tipo erróneo en extra_config: si model_loader_extra_config no es un dict, la clave no está en la lista permitida o num_threads no es un entero positivo, falla de inmediato (default_loader.py:79-110).
• EPLB salta el filtro EP: con enable_eplb activado, los slots redundantes de experts apuntan al expert lógico de otro rank; filtrarlos haría que esos slots se quedaran vacíos, así que toda la sección hace return directo (default_loader.py:370-375).
• Modelos cuantizados saltan la verificación de tracking: cuando el método de cuantización trae process_weights_after_loading o uses_meta_device, el checkpoint puede no tener las scales correspondientes; esos parámetros se excluyen del conjunto de tracking (default_loader.py:453-462).

Resumen

DefaultModelLoader hace bien las tres cosas —encontrar archivos, descargarlos, iterar tensores— y delega toda la lógica de sharding / cuantización / descuantización a las capas concretas y a los métodos de cuantización. El punto de contacto con capas lineales con tensor parallelism es que model.load_weights busca internamente el weight_loader para cada nombre de tensor; el contacto con carga de capas cuantizadas es que quant_config.get_quant_method instala un LinearMethodBase en cada capa, sin que el cargador itself se entere.

Véase la documentación oficial: vLLM 文档 · README