# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import dataclasses import glob import os import re import time from collections.abc import Generator, Iterable from pathlib import Path from typing import cast import torch from huggingface_hub import hf_hub_download from torch import nn from vllm.config import ModelConfig from vllm.config.load import LoadConfig from vllm.logger import init_logger from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase from vllm.model_executor.model_loader.weight_utils import ( download_gguf, download_safetensors_index_file_from_hf, download_weights_from_hf, filter_duplicate_safetensors_files, filter_files_not_needed_for_inference, maybe_download_from_modelscope, multi_thread_safetensors_weights_iterator, safetensors_weights_iterator, ) from vllm.transformers_utils.repo_utils import file_exists from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.utils.torch_utils import set_default_torch_dtype from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig from vllm_omni.diffusion.model_loader.checkpoint_adapters import ( get_checkpoint_adapter, ) from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery from vllm_omni.diffusion.registry import initialize_model logger = init_logger(__name__) def _natural_sort_key(filepath: str) -> list: """Natural sort key for filenames with numeric components, e.g. model-00001-of-00005.safetensors -> ['model-', 1, '-of-', 5, '.safetensors'].""" return [int(s) if s.isdigit() else s for s in re.split(r"(\d+)", os.path.basename(filepath))] MODEL_INDEX = "model_index.json" DIFFUSION_MODEL_WEIGHTS_INDEX = "diffusion_pytorch_model.safetensors.index.json" TRANSFORMER_WEIGHTS_INDEX = "model.safetensors.index.json" INDEX_FILES = [DIFFUSION_MODEL_WEIGHTS_INDEX, TRANSFORMER_WEIGHTS_INDEX] def _resolve_custom_pipeline_cls(custom_pipeline_name: str | type | None) -> type: """Resolve a custom pipeline reference to a class. Accepts either a fully qualified name string (resolved via import) or an already-imported class object (returned as-is). """ if custom_pipeline_name is None: raise ValueError("custom_pipeline_name is required for load_format='custom_pipeline'") if isinstance(custom_pipeline_name, str): return resolve_obj_by_qualname(custom_pipeline_name) if isinstance(custom_pipeline_name, type): return custom_pipeline_name raise TypeError( f"custom_pipeline_name must be a qualified name string or a class, got {type(custom_pipeline_name).__name__}" ) class DiffusersPipelineLoader: """Model loader that can load diffusers pipeline components from disk.""" @dataclasses.dataclass class ComponentSource: """A source for weights.""" model_or_path: str """The model ID or path.""" subfolder: str | None """The subfolder inside the model repo.""" revision: str | None """The optional model revision.""" prefix: str = "" """A prefix to prepend to all weights.""" fall_back_to_pt: bool = True """Whether .pt weights can be used.""" allow_patterns_overrides: list[str] | None = None """If defined, weights will load exclusively using these patterns.""" counter_before_loading_weights: float = 0.0 counter_after_loading_weights: float = 0.0 def __init__(self, load_config: LoadConfig, od_config: OmniDiffusionConfig | None = None): self.load_config = load_config self.od_config = od_config def _prepare_weights( self, model_name_or_path: Path, subfolder: str | None, revision: str | None, fall_back_to_pt: bool, allow_patterns_overrides: list[str] | None, ) -> tuple[str, list[str], bool]: """Prepare weights for the model. If the model is not local, it will be downloaded.""" model_name_or_path = maybe_download_from_modelscope(model_name_or_path, revision) or model_name_or_path is_local = os.path.isdir(model_name_or_path) load_format = self.load_config.load_format use_safetensors = False possible_index_files = [ f"{subfolder}/{index_file}" if subfolder is not None else index_file for index_file in INDEX_FILES ] available_index_file = [ f for f in possible_index_files if file_exists(model_name_or_path, f, revision=revision) ] if len(available_index_file) > 1: raise ValueError( f"Multiple index files found in {model_name_or_path} with subfolder {subfolder}: {available_index_file}" ) index_file = available_index_file[0] if available_index_file else "" # only hf is supported currently if load_format == "auto": load_format = "hf" # Some quantized models use .pt files for storing the weights. if load_format == "hf": allow_patterns = ["*.safetensors", "*.bin"] else: raise ValueError(f"Unknown load_format: {load_format}") if fall_back_to_pt: allow_patterns += ["*.pt"] if allow_patterns_overrides is not None: allow_patterns = allow_patterns_overrides if not is_local: hf_folder = download_weights_from_hf( model_name_or_path, self.load_config.download_dir, allow_patterns, revision, subfolder=subfolder, ignore_patterns=self.load_config.ignore_patterns, ) else: hf_folder = model_name_or_path if subfolder is not None: hf_folder = os.path.join(hf_folder, subfolder) hf_weights_files: list[str] = [] for pattern in allow_patterns: hf_weights_files += glob.glob(os.path.join(hf_folder, pattern)) if len(hf_weights_files) > 0: # Decide by actual files rather than pattern name (patterns may include subfolders). use_safetensors = any(f.endswith(".safetensors") for f in hf_weights_files) break if use_safetensors: # For models like Mistral-7B-Instruct-v0.3 # there are both sharded safetensors files and a consolidated # safetensors file. Using both breaks. # Here, we download the `model.safetensors.index.json` and filter # any files not found in the index. if not is_local: download_safetensors_index_file_from_hf( model_name_or_path, index_file, cache_dir=self.load_config.download_dir, subfolder=subfolder, revision=revision, ) hf_weights_files = filter_duplicate_safetensors_files(hf_weights_files, hf_folder, index_file) else: hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files) if len(hf_weights_files) == 0: raise RuntimeError(f"Cannot find any model weights with `{model_name_or_path}`") return hf_folder, hf_weights_files, use_safetensors def _get_weights_iterator( self, source: "ComponentSource", model: nn.Module | None = None, ) -> Generator[tuple[str, torch.Tensor], None, None]: """Get an iterator for the model weights based on the load format.""" _, hf_weights_files, use_safetensors = self._prepare_weights( source.model_or_path, source.subfolder, source.revision, source.fall_back_to_pt, source.allow_patterns_overrides, ) od_config = self.od_config use_multithread = ( use_safetensors and od_config is not None and getattr(od_config, "enable_multithread_weight_load", False) and self.load_config.safetensors_load_strategy != "torchao" ) if use_multithread: num_threads = getattr(od_config, "num_weight_load_threads", 4) # Keep deterministic shard order before passing to vLLM helper. sorted_hf_weights_files = sorted(hf_weights_files, key=_natural_sort_key) weights_iterator = multi_thread_safetensors_weights_iterator( sorted_hf_weights_files, self.load_config.use_tqdm_on_load, max_workers=num_threads, ) else: weights_iterator = safetensors_weights_iterator( hf_weights_files, self.load_config.use_tqdm_on_load, self.load_config.safetensors_load_strategy, ) if self.counter_before_loading_weights == 0.0: self.counter_before_loading_weights = time.perf_counter() # Apply the prefix. prefixed_weights_iterator = ((source.prefix + name, tensor) for (name, tensor) in weights_iterator) if model is not None: checkpoint_adapter = self._get_checkpoint_adapter(model, source, use_safetensors) if checkpoint_adapter is not None: return checkpoint_adapter.adapt(prefixed_weights_iterator) return prefixed_weights_iterator def _get_source_quant_config(self, source: "ComponentSource") -> object | None: if self.od_config is None: return None quant_config = self.od_config.quantization_config if hasattr(quant_config, "resolve"): return quant_config.resolve(source.prefix.rstrip(".")) return quant_config def _get_checkpoint_adapter( self, model: nn.Module, source: "ComponentSource", use_safetensors: bool, ): return get_checkpoint_adapter( model=model, source=source, quant_config=self._get_source_quant_config(source), use_safetensors=use_safetensors, ) def get_all_weights( self, model: nn.Module, ) -> Generator[tuple[str, torch.Tensor], None, None]: sources = self._get_weight_sources(model) for source in sources: yield from self._get_weights_iterator(source, model=model) def _get_weight_sources(self, model: nn.Module) -> tuple["ComponentSource", ...]: return tuple( cast( Iterable[DiffusersPipelineLoader.ComponentSource], getattr(model, "weights_sources", ()), ) ) def _get_expected_parameter_names(self, model: nn.Module) -> set[str]: """Return parameter names that should be covered by strict load checks.""" all_parameter_names = {name for name, _ in model.named_parameters()} sources = self._get_weight_sources(model) # Keep strict behavior if no source metadata exists. if not sources: return all_parameter_names # Empty prefix means "root" source, i.e. entire model should be covered. if any(source.prefix == "" for source in sources): return all_parameter_names source_prefixes = tuple(source.prefix for source in sources if source.prefix) if not source_prefixes: return all_parameter_names return {name for name in all_parameter_names if name.startswith(source_prefixes)} def download_model(self, model_config: ModelConfig) -> None: self._prepare_weights( model_name_or_path=model_config.model, subfolder=None, revision=model_config.revision, fall_back_to_pt=True, allow_patterns_overrides=None, ) def load_model( self, od_config: OmniDiffusionConfig, load_device: str, load_format: str | None = "default", custom_pipeline_name: str | type[nn.Module] | None = None, device: torch.device | None = None, ) -> nn.Module: """Load a model with the given configurations.""" if load_format is None: load_format = "default" self.od_config = od_config # CPU offload + quantization: for offline-quantized models (e.g., AutoRound MXFP8), # weights are already quantized in the checkpoint — load directly on CPU. # For online quantization, load on device so quantization can run on accelerator, # then move back to CPU afterward. offload_after_quant = False if load_device == "cpu" and od_config.quantization_config is not None: quant_cfg = od_config.quantization_config is_offline = getattr(quant_cfg, "data_type", None) == "mx_fp" or getattr( quant_cfg, "is_checkpoint_quantized", False ) if not is_offline: load_device = device.type offload_after_quant = True logger.info( "Online quantization with CPU offload, using %s for weight loading (will offload back to CPU)", load_device, ) else: logger.info("Offline-quantized model with CPU offload, loading weights directly on CPU") target_device = torch.device(load_device) with set_default_torch_dtype(od_config.dtype): if od_config.parallel_config.use_hsdp: model = self._load_model_with_hsdp( od_config, target_device=device, load_format=load_format, custom_pipeline_name=custom_pipeline_name ) else: if load_format == "custom_pipeline": # NOTE: Custom pipelines call HuggingFace `from_pretrained(...).to(device)` # internally. If we construct them under `with target_device:` (CUDA), # safetensors takes a direct-to-GPU fast path that calls `cudaMalloc` # via the driver API and BYPASSES PyTorch's caching allocator. # That makes those bytes invisible to CuMemAllocator, so `sleep()` # cannot offload/unmap them and GPU memory stays pinned. # # Fix: build the custom pipeline on CPU first (no default device # context), then explicitly move it to the target device. The # subsequent `.to(target_device)` issues `torch.empty(..., device=cuda)` # + `copy_`, which goes through the caching allocator and is fully # tracked by CuMemAllocator. from vllm_omni.diffusion.config import set_current_diffusion_config model_cls = _resolve_custom_pipeline_cls(custom_pipeline_name) with set_current_diffusion_config(od_config): model = model_cls(od_config=od_config) if target_device.type != "cpu": model.to(target_device) else: with target_device: if load_format == "default": model = initialize_model(od_config) elif load_format == "diffusers": model = DiffusersAdapterPipeline(od_config=od_config, device=target_device) else: raise ValueError(f"Unknown load_format: {load_format}") logger.debug("Loading weights on %s ...", load_device) if load_format == "diffusers": # DiffusersAdapterPipeline.load_weights() calls # DiffusionPipeline.from_pretrained() internally; it does # not use our native customized pipeline classes. cast(DiffusersAdapterPipeline, model).load_weights() elif self._is_gguf_quantization(od_config): self._load_weights_with_gguf(model, od_config) else: # Quantization does not happen in `load_weights` but after it self.load_weights(model) # Process weights after loading for quantization (e.g., FP8 online quantization) # This is needed for vLLM's quantization methods that need to transform weights self._process_weights_after_loading(model, target_device) if offload_after_quant: model.to("cpu") logger.info("Quantization complete, offloaded model back to CPU") return model.eval() def _process_weights_after_loading(self, model: nn.Module, target_device: torch.device) -> None: """Process weights after loading for quantization methods. This handles vLLM's quantization methods that need to process weights after loading (e.g., FP8 online quantization from BF16/FP16 weights). """ for _, module in model.named_modules(): quant_method = getattr(module, "quant_method", None) if isinstance(quant_method, QuantizeMethodBase): # Move module to target device for processing if needed module_device = next(module.parameters(), None) if module_device is not None: module_device = module_device.device needs_device_move = module_device != target_device if needs_device_move: module.to(target_device) quant_method.process_weights_after_loading(module) if needs_device_move: module.to(module_device) def load_weights(self, model: nn.Module) -> None: weights_to_load = self._get_expected_parameter_names(model) loaded_weights = model.load_weights(self.get_all_weights(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, ) # TODO(Isotr0py): Enable weights loading check after decoupling # all components' weights loading (AutoModel.from_pretrained etc). # We only enable strict check for non-quantized models # that have loaded weights tracking currently. if loaded_weights is not None: weights_not_loaded = weights_to_load - loaded_weights # NOTE: if the model is quantized, ignore not_loaded check for scale weights weights_scale_not_loaded = {name for name in weights_not_loaded if name.endswith("weight_scale")} weights_not_loaded = weights_not_loaded - weights_scale_not_loaded if weights_not_loaded: self._check_unloaded_weights(weights_not_loaded) if weights_scale_not_loaded: logger.warning( f"Following weight_scale weights were not initialized from checkpoint: {weights_scale_not_loaded}" ) @staticmethod def _is_expected_quantized_weight(name: str) -> bool: """Return True if *name* is a quantization-specific parameter. Quantization methods (GPTQ, AWQ, FP8, GGUF, Autoround, etc.) create extra parameters that have no counterpart in an unquantized checkpoint. These are expected to be absent and should not trigger a load error. """ # Weight suffixes that quantization methods register in the model but # are not present in unquantized checkpoints. _QUANTIZED_WEIGHT_SUFFIXES = ( # GPTQ / AWQ / AutoRound – g_idx is optional (not all checkpoints include it) ".g_idx", # FP8 ".weight_scale", ".weight_scale_inv", ".input_scale", # GGUF ".qweight_type", # INT8 (weight_scale already covered above) ) return name.endswith(_QUANTIZED_WEIGHT_SUFFIXES) def _check_unloaded_weights( self, weights_not_loaded: set[str], ) -> None: """Validate unloaded weights, tolerating expected quantization artifacts. For quantized models, weights matching known quant-specific suffixes are logged as a warning. Any *other* missing weight raises ``ValueError`` regardless of quantization. """ od_config = getattr(self, "od_config", None) if od_config is None or od_config.quantization_config is None: raise ValueError(f"Following weights were not initialized from checkpoint: {weights_not_loaded}") expected_missing = {w for w in weights_not_loaded if self._is_expected_quantized_weight(w)} unexpected_missing = weights_not_loaded - expected_missing if expected_missing: logger.warning( "Following weights were not initialized from checkpoint (expected for quantized models): %s", expected_missing, ) if unexpected_missing: raise ValueError(f"Following weights were not initialized from checkpoint: {unexpected_missing}") def _is_gguf_quantization(self, od_config: OmniDiffusionConfig) -> bool: quant_config = od_config.quantization_config if quant_config is None: return False # New API: DiffusionGGUFConfig (QuantizationConfig subclass) from vllm_omni.quantization.gguf_config import DiffusionGGUFConfig if isinstance(quant_config, DiffusionGGUFConfig): if quant_config.gguf_model is None: raise ValueError("GGUF quantization requires gguf_model") return True # Dict-style config: {"method": "gguf", "gguf_model": "..."} if isinstance(quant_config, dict): if quant_config.get("method") == "gguf": if not quant_config.get("gguf_model"): raise ValueError("GGUF quantization requires gguf_model") return True return False # Check by name for any config that reports as "gguf" if hasattr(quant_config, "get_name") and quant_config.get_name() == "gguf": gguf_model = getattr(quant_config, "gguf_model", None) if gguf_model is None: raise ValueError("GGUF quantization requires gguf_model") return True # Fallback: object with gguf_model attribute but no get_name if hasattr(quant_config, "gguf_model") and quant_config.gguf_model: return True return False def _is_transformer_source(self, source: "ComponentSource") -> bool: if source.subfolder == "transformer": return True return source.prefix.startswith("transformer.") def _get_model_loadable_names(self, model: nn.Module) -> set[str]: # Avoid model.state_dict() here because GGUF uses UninitializedParameter # which raises during detach(). Collect names directly. return {name for name, _ in model.named_parameters()} | {name for name, _ in model.named_buffers()} def _resolve_gguf_model_path(self, gguf_model: str, revision: str | None) -> str: if os.path.isfile(gguf_model): return gguf_model # repo_id/filename.gguf if "/" in gguf_model and gguf_model.endswith(".gguf"): repo_id, filename = gguf_model.rsplit("/", 1) return hf_hub_download( repo_id=repo_id, filename=filename, revision=revision, cache_dir=self.load_config.download_dir, ) # repo_id:quant_type if "/" in gguf_model and ":" in gguf_model: repo_id, quant_type = gguf_model.rsplit(":", 1) return download_gguf( repo_id, quant_type, cache_dir=self.load_config.download_dir, revision=revision, ignore_patterns=self.load_config.ignore_patterns, ) raise ValueError( f"Unrecognized GGUF reference: {gguf_model!r} (expected local file, " "/.gguf, or :)" ) def _get_gguf_weights_iterator( self, source: "ComponentSource", model: nn.Module, od_config: OmniDiffusionConfig, ) -> Generator[tuple[str, torch.Tensor], None, None]: quant_config = od_config.quantization_config gguf_model = getattr(quant_config, "gguf_model", None) if gguf_model is None: raise ValueError("GGUF quantization requires quantization_config.gguf_model") gguf_file = self._resolve_gguf_model_path(gguf_model, od_config.revision) adapter = get_gguf_adapter(gguf_file, model, source, od_config) weights_iter = adapter.weights_iterator() return ((source.prefix + name, tensor) for (name, tensor) in weights_iter) def _load_weights_with_gguf(self, model: nn.Module, od_config: OmniDiffusionConfig) -> set[str]: sources = self._get_weight_sources(model) loaded: set[str] = set() loadable_names: set[str] | None = None for source in sources: if self._is_transformer_source(source): loaded |= model.load_weights(self._get_gguf_weights_iterator(source, model, od_config)) # GGUF checkpoints can be transformer-only or partially quantized. # Only fall back to HF if this source still has missing loadable weights. loadable_names = loadable_names or self._get_model_loadable_names(model) has_missing_for_source = any( name.startswith(source.prefix) and name not in loaded for name in loadable_names ) if not has_missing_for_source: continue hf_iter = self._get_weights_iterator(source) hf_iter = ( (name, tensor) for (name, tensor) in hf_iter if name in loadable_names and name not in loaded ) loaded |= model.load_weights(hf_iter) else: loaded |= model.load_weights(self._get_weights_iterator(source)) weights_to_load = self._get_expected_parameter_names(model) weights_not_loaded = weights_to_load - loaded if weights_not_loaded: raise ValueError(f"Following weights were not initialized from checkpoint: {weights_not_loaded}") return loaded def _load_model_with_hsdp( self, od_config: OmniDiffusionConfig, target_device: torch.device, load_format: str = "default", custom_pipeline_name: str | type[nn.Module] | None = None, ) -> nn.Module: """Load model with HSDP sharding for inference. The pipeline contains multiple components (text_encoder, VAE, transformer). Only the transformer is sharded with HSDP. Other components are loaded normally. Approach: Load weights first using model's load_weights (handles QKV fusion etc.), then apply HSDP sharding to redistribute weights across GPUs. """ from vllm_omni.diffusion.distributed.hsdp import apply_hsdp_to_model parallel_config = od_config.parallel_config hsdp_config = HSDPInferenceConfig( enabled=True, hsdp_replicate_size=parallel_config.hsdp_replicate_size, hsdp_shard_size=parallel_config.hsdp_shard_size, param_dtype=od_config.dtype, ) # Initialize model WITHOUT device context (weights start on CPU). # Unlike the non-HSDP path which uses `with target_device:` to create weights # directly on GPU, HSDP needs weights on CPU first so they can be redistributed # across GPUs by apply_hsdp_to_model. The model's load_weights handles weight # mapping (QKV fusion, etc.). if load_format == "default": model = initialize_model(od_config) elif load_format == "custom_pipeline": model_cls = _resolve_custom_pipeline_cls(custom_pipeline_name) from vllm_omni.diffusion.config import set_current_diffusion_config with set_current_diffusion_config(od_config): model = model_cls(od_config=od_config) self.load_weights(model) # Collect all transformers to shard (some models have transformer_2 for MoE) transformers_to_shard = [] transformer = getattr(model, "transformer", None) if transformer is None: raise ValueError("Model has no transformer attribute for HSDP") transformers_to_shard.append(("transformer", transformer)) # Check for transformer_2 (MoE two-stage models like Wan2.2-I2V) transformer_2 = getattr(model, "transformer_2", None) if transformer_2 is not None: transformers_to_shard.append(("transformer_2", transformer_2)) # Apply HSDP sharding to all transformers for name, trans in transformers_to_shard: logger.debug("Applying HSDP to %s", name) apply_hsdp_to_model(trans, hsdp_config) # # HSDP only shards transformer modules. All other runtime modules must # # be placed on the execution device explicitly after sharding. discovered_modules = ModuleDiscovery.discover(model) modules_to_move: list[nn.Module] = [] if discovered_modules.vaes is not None: modules_to_move.extend(discovered_modules.vaes) if discovered_modules.encoders is not None: modules_to_move.extend(discovered_modules.encoders) if discovered_modules.resident_modules is not None: modules_to_move.extend(discovered_modules.resident_modules) for module in modules_to_move: module.to(target_device) return model