# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import contextlib import dataclasses import glob import json import os import re import time from collections.abc import Callable, Generator, Iterable, Sequence from pathlib import Path from typing import cast import huggingface_hub import torch from torch import nn from vllm.config.load import LoadConfig from vllm.logger import init_logger from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase from vllm.model_executor.model_loader.weight_utils import ( download_weights_from_hf, filter_files_not_needed_for_inference, maybe_download_from_modelscope, multi_thread_safetensors_weights_iterator, pt_weights_iterator, safetensors_weights_iterator, ) from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.utils.torch_utils import set_default_torch_dtype from vllm_omni.diffusion.config import set_current_diffusion_config from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.distributed.hsdp import ( HSDPInferenceConfig, HSDPShardContext, apply_hsdp_to_model, finalize_hsdp_root, prepare_hsdp_shard_context, shard_hsdp_module, ) from vllm_omni.diffusion.lora.manager import LoRABackend from vllm_omni.diffusion.model_loader.checkpoint_adapters import ( get_checkpoint_adapter, ) from vllm_omni.diffusion.model_loader.host_weight_loader import HWRLoaderMixin, _HWRCommitError from vllm_omni.diffusion.model_loader.host_weight_plan import ( HostWeightPlan, TensorBinding, build_checkpoint_binding_plan, build_checkpoint_mmap_plan, has_online_quantization, ) from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline from vllm_omni.diffusion.offloader.component_utils import encoder_component_type from vllm_omni.diffusion.offloader.config import ( DIT_COMPONENT, TEXT_ENCODER_COMPONENT, OffloadStrategy, resolve_offload, ) from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery, PipelineModules from vllm_omni.diffusion.offloader.offload_plan import get_offload_plan from vllm_omni.diffusion.registry import initialize_model from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific from vllm_omni.quantization.component_config import is_prequantized_config from vllm_omni.transformers_utils.repo_utils import hf_api # download_gguf was removed from upstream vLLM (commit 6635279d8). # Inlined from the last upstream version before the GGUF plugin migration. def download_gguf( repo_id: str, quant_type: str, cache_dir: str | None = None, revision: str | None = None, ignore_patterns: str | list[str] | None = None, ) -> str: allow_patterns = [ f"*-{quant_type}.gguf", f"*-{quant_type}-*.gguf", f"*/*-{quant_type}.gguf", f"*/*-{quant_type}-*.gguf", ] folder = download_weights_from_hf( model_name_or_path=repo_id, cache_dir=cache_dir, allow_patterns=allow_patterns, revision=revision, ignore_patterns=ignore_patterns, ) local_files: list[str] = [] for pattern in allow_patterns: glob_pattern = os.path.join(folder, pattern) local_files.extend(glob.glob(glob_pattern)) if not local_files: raise ValueError(f"Downloaded GGUF files not found in {folder} for quant_type {quant_type}") local_files.sort(key=lambda x: (x.count("-"), x)) return local_files[0] 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))] DIFFUSION_MODEL_WEIGHTS_INDEX = "diffusion_pytorch_model.safetensors.index.json" TRANSFORMER_WEIGHTS_INDEX = "model.safetensors.index.json" DIFFUSION_MODEL_BIN_WEIGHTS_INDEX = "diffusion_pytorch_model.bin.index.json" SAFETENSORS_INDEX_FILES = [DIFFUSION_MODEL_WEIGHTS_INDEX, TRANSFORMER_WEIGHTS_INDEX] PT_INDEX_FILES = [DIFFUSION_MODEL_BIN_WEIGHTS_INDEX] SHARDED_SAFETENSORS_PATTERN = re.compile(r"^(?P.+)-\d+-of-(?P\d+)\.safetensors$") def _validate_unindexed_safetensors_layout(weight_files: Sequence[str]) -> None: """Reject ambiguous shard families when no index is available.""" shard_counts: dict[str, set[int]] = {} for weight_file in weight_files: match = SHARDED_SAFETENSORS_PATTERN.fullmatch(os.path.basename(weight_file)) if match is not None: shard_counts.setdefault(match.group("family"), set()).add(int(match.group("count"))) conflicts = {family: sorted(counts) for family, counts in shard_counts.items() if len(counts) > 1} if conflicts: raise ValueError( "Ambiguous unindexed safetensors checkpoint with conflicting shard totals: " f"{conflicts}. Refusing to load potentially stale checkpoint shards." ) 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__}" ) def _is_distributed_rank_zero() -> bool: return ( not torch.distributed.is_available() or not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 ) @dataclasses.dataclass(frozen=True) class _HSDPLoadRoot: name: str module: nn.Module context: HSDPShardContext @dataclasses.dataclass(frozen=True) class _HSDPModuleLoadGroup: name: str module: nn.Module context: HSDPShardContext @dataclasses.dataclass(frozen=True) class _PreShardedHSDPLoadPlan: roots: list[_HSDPLoadRoot] groups: list[_HSDPModuleLoadGroup] bindings: dict[str, TensorBinding] binding_names: frozenset[str] class DiffusersPipelineLoader(HWRLoaderMixin): """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): self.load_config = load_config self.od_config = od_config self.quant_config = od_config.quantization_config self.parallel_config = od_config.parallel_config self.host_weight_plan: HostWeightPlan | None = None self._hwr_state: dict[str, object] | None = None self._last_load_request: dict[str, object] | None = None self._force_canonical_load = False def take_host_weight_plan(self) -> HostWeightPlan | None: """Transfer the loader-produced plan to the offload backend.""" plan = self.host_weight_plan self.host_weight_plan = None return plan @staticmethod def _repo_relative_path(subfolder: str | None, filename: str) -> str: if subfolder is None: return filename prefix = f"{subfolder.rstrip('/')}/" return filename if filename.startswith(prefix) else f"{prefix}{filename.lstrip('/')}" def _resolve_weight_index( self, model_name_or_path: Path | str, subfolder: str | None, revision: str | None, index_files: Sequence[str], ) -> list[str] | None: """Resolve an index and return its authoritative shard manifest.""" is_local = os.path.isdir(model_name_or_path) index_paths: list[tuple[str, Path]] = [] for index_file in index_files: repo_index_path = self._repo_relative_path(subfolder, index_file) if is_local: index_path = Path(model_name_or_path) / repo_index_path if index_path.is_file(): index_paths.append((index_file, index_path)) continue try: index_path = hf_api().hf_hub_download( repo_id=str(model_name_or_path), filename=repo_index_path, cache_dir=self.load_config.download_dir, revision=revision, local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, ) except huggingface_hub.errors.EntryNotFoundError: continue index_paths.append((index_file, Path(index_path))) if len(index_paths) > 1: raise ValueError( f"Multiple index files found in {model_name_or_path} with subfolder {subfolder}: " f"{[index_file for index_file, _ in index_paths]}" ) if not index_paths: return None index_file, index_path = index_paths[0] with open(index_path) as f: index = json.load(f) weight_map = index.get("weight_map") if isinstance(index, dict) else None if not isinstance(weight_map, dict) or not weight_map: raise ValueError(f"Weight index {index_file} must contain a non-empty `weight_map`") if not all(isinstance(filename, str) and filename for filename in weight_map.values()): raise ValueError(f"Weight index {index_file} contains an invalid shard filename") return sorted(set(weight_map.values())) def _prepare_weights( self, model_name_or_path: Path | str, subfolder: str | None, revision: str | None, fall_back_to_pt: bool, allow_patterns_overrides: list[str] | None, ) -> tuple[Path | 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 indexed_weight_files = None if allow_patterns_overrides is None: for index_files in (SAFETENSORS_INDEX_FILES, PT_INDEX_FILES): indexed_weight_files = self._resolve_weight_index( model_name_or_path, subfolder, revision, index_files, ) if indexed_weight_files is not None: break # 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 and indexed_weight_files is not None: hf_folder: Path | str = download_weights_from_hf_specific( model_name_or_path=str(model_name_or_path), cache_dir=self.load_config.download_dir, allow_patterns=[self._repo_relative_path(subfolder, filename) for filename in indexed_weight_files], revision=revision, ignore_patterns=self.load_config.ignore_patterns, require_all=True, ) elif 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) if indexed_weight_files is not None: hf_weights_files = [os.path.join(hf_folder, filename) for filename in indexed_weight_files] missing_files = [filename for filename in hf_weights_files if not os.path.isfile(filename)] if missing_files: raise FileNotFoundError(f"Weight files referenced in index but missing: {missing_files}") use_safetensors = any(filename.endswith(".safetensors") for filename in hf_weights_files) else: hf_weights_files = [] for pattern in allow_patterns: hf_weights_files += glob.glob(os.path.join(hf_folder, pattern)) if hf_weights_files: # 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: _validate_unindexed_safetensors_layout(hf_weights_files) 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, ) quant_config = self._get_source_quant_config(source) use_multithread = ( use_safetensors and getattr(self.od_config, "enable_multithread_weight_load", False) and self.load_config.safetensors_load_strategy != "torchao" ) use_torchao = ( not use_safetensors and quant_config is not None and hasattr(quant_config, "get_name") and quant_config.get_name() == "torchao" and getattr(quant_config, "is_checkpoint_torchao_serialized", False) ) if use_multithread: num_threads = getattr(self.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, ) elif use_torchao: sorted_hf_weights_files = sorted(hf_weights_files, key=_natural_sort_key) weights_iterator = pt_weights_iterator( sorted_hf_weights_files, self.load_config.use_tqdm_on_load, self.load_config.pt_load_map_location, ) 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: quant_config = self.quant_config resolve = getattr(quant_config, "resolve", None) if resolve is not None: return 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, sources: Sequence["ComponentSource"] | None = None, ) -> Generator[tuple[str, torch.Tensor], None, None]: if sources is None: sources = self._get_weight_sources(model) for source in sources: yield from self._get_weights_iterator(source, model=model) @staticmethod def _stream_online_quant_weights_to_cpu( model: nn.Module, weights: Iterable[tuple[str, torch.Tensor]], ) -> Generator[tuple[str, torch.Tensor], None, None]: """Offload each online-quantized layer as soon as it is complete. Upstream vLLM's online layerwise loader materializes and quantizes a layer synchronously while consuming ``weights``. Generator execution resumes after that weight has been consumed, which gives us a safe point to move completed layers to CPU before the next layer is loaded. This bounds accelerator residency during CPU-offloaded model startup instead of retaining the entire quantized model until loading ends. """ from vllm.model_executor.model_loader.reload.layerwise import ( get_layerwise_info, ) pending = { module for module in model.modules() if getattr(getattr(module, "quant_method", None), "uses_meta_device", False) and get_layerwise_info(module).can_load() } offloaded = 0 def offload_completed() -> None: nonlocal offloaded for module in tuple(pending): if get_layerwise_info(module).can_load(): continue module.to("cpu") pending.remove(module) offloaded += 1 for weight in weights: # This runs after the consumer has handled the previous yield. offload_completed() yield weight offload_completed() # Quantization workspaces and the old accelerator-side parameter # storages are now reusable cache blocks. Release them before the # remaining (unquantized) model tensors are copied to CPU; otherwise # the cached quantization footprint and final offload overlap in the # process-level startup peak. if offloaded: torch.accelerator.empty_cache() logger.info( "Stream-offloaded %d online-quantized layers to CPU during weight loading", offloaded, ) 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. A parameter with an initialized checkpoint default can explicitly set ``is_checkpoint_optional``. It is still loaded when present on disk. """ all_parameter_names = { name for name, param in model.named_parameters() if not getattr(param, "is_checkpoint_optional", False) } 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 _maybe_fuse_distilled_lora(self, model: nn.Module) -> None: """Fuse distilled LoRA weights into the model before sharding or quantization.""" if self.od_config is None: return lora_backend = getattr(self.od_config, "lora_backend", None) if lora_backend != LoRABackend.DISTILL and lora_backend != "distill": return if getattr(model, "lora_is_fused", False): return lora_path = getattr(self.od_config, "lora_path", None) if not lora_path: return if isinstance(lora_path, list) and len(lora_path) == 1: lora_path = lora_path[0] if hasattr(model, "load_lora_weights"): lora_scale = getattr(self.od_config, "lora_scale", 1.0) if lora_scale > 1.0: logger.warning("lora_scale > 1.0 may not take any effect when using distilled LoRA backend.") logger.info("Fusing distilled LoRA weights from %s into %s", lora_path, model.__class__.__name__) model.load_lora_weights(lora_path) setattr(model, "lora_is_fused", True) else: logger.warning("Pipeline %s does not support loading distilled LoRA weights.", model.__class__.__name__) def _broadcast_model_weights( self, model: nn.Module, target_device: torch.device, src_rank: int = 0, bucket_cap_bytes: int = 256 * 1024 * 1024, ) -> None: """Broadcast model parameters and buffers from src_rank to all other ranks using coalesced bucketing.""" if not torch.distributed.is_initialized() or torch.distributed.get_world_size() <= 1: return rank = torch.distributed.get_rank() world_size = torch.distributed.get_world_size() backend = torch.distributed.get_backend() is_cpu_backend = backend == "gloo" logger.info( "Worker %d: %s full model weights via %s across %d ranks", rank, "Broadcasting" if rank == src_rank else "Receiving", backend, world_size, ) t0 = time.perf_counter() all_tensors = [p.data for _, p in model.named_parameters() if p.numel() > 0] + [ b.data for _, b in model.named_buffers() if b.numel() > 0 ] if not all_tensors: return count = len(all_tensors) total_bytes = sum(t.numel() * t.element_size() for t in all_tensors) # Group tensors into buckets by dtype and maximum byte size buckets: list[list[torch.Tensor]] = [] curr_bucket: list[torch.Tensor] = [] curr_bytes = 0 for tensor in all_tensors: t_bytes = tensor.numel() * tensor.element_size() if curr_bucket and (curr_bucket[0].dtype != tensor.dtype or curr_bytes + t_bytes > bucket_cap_bytes): buckets.append(curr_bucket) curr_bucket = [] curr_bytes = 0 curr_bucket.append(tensor) curr_bytes += t_bytes if curr_bucket: buckets.append(curr_bucket) for bucket in buckets: tot_numel = sum(t.numel() for t in bucket) dtype = bucket[0].dtype if is_cpu_backend: if rank == src_rank: flat_cpu = torch.cat([t.detach().to("cpu").reshape(-1) for t in bucket]) else: flat_cpu = torch.empty(tot_numel, dtype=dtype) torch.distributed.broadcast(flat_cpu, src=src_rank) if rank != src_rank: offset = 0 for t in bucket: n = t.numel() t.copy_(flat_cpu[offset : offset + n].view_as(t)) offset += n del flat_cpu else: if rank == src_rank: flat_cpu = torch.cat([t.detach().to("cpu").reshape(-1) for t in bucket]) dev_flat = flat_cpu.to(target_device, non_blocking=False) del flat_cpu torch.distributed.broadcast(dev_flat, src=src_rank) del dev_flat else: dev_flat = torch.empty(tot_numel, dtype=dtype, device=target_device) torch.distributed.broadcast(dev_flat, src=src_rank) flat_cpu = dev_flat.to("cpu", non_blocking=False) del dev_flat offset = 0 for t in bucket: n = t.numel() t.copy_(flat_cpu[offset : offset + n].view_as(t), non_blocking=False) offset += n del flat_cpu if not is_cpu_backend and target_device.type != "cpu": from vllm_omni.platforms import current_omni_platform current_omni_platform.synchronize() torch.distributed.barrier() elapsed = time.perf_counter() - t0 gb = total_bytes / 1e9 logger.info( "Worker %d: Shared weight broadcast complete (%d tensors, %d buckets, %.2f GB) in %.2fs (%.2f GB/s)", rank, count, len(buckets), gb, elapsed, gb / max(elapsed, 1e-6), ) def load_model( self, 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.""" self.host_weight_plan = None self._hwr_state = None self._last_load_request = { "load_device": load_device, "load_format": load_format, "custom_pipeline_name": custom_pipeline_name, "device": device, } if load_format is None: load_format = "default" # CPU offload + quantization: for offline-quantized models (e.g., AutoRound MXFP8, # compressed-tensors), 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 offline_quant_direct_cpu = False if load_device == "cpu" and self.quant_config is not None and device is not None: quant_cfg = self.quant_config is_offline = getattr(quant_cfg, "data_type", None) == "mx_fp" or getattr( quant_cfg, "is_checkpoint_quantized", False ) or is_prequantized_config(quant_cfg) 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: offline_quant_direct_cpu = True logger.info("Offline-quantized model with CPU offload, loading weights directly on CPU") target_device = torch.device(load_device) with set_default_torch_dtype(self.od_config.dtype): if self.parallel_config.use_hsdp: model = self._load_model_with_hsdp( target_device=device, load_format=load_format, custom_pipeline_name=custom_pipeline_name, offload_after_quant=offload_after_quant, ) else: # The model is headed back to host memory right after online # quantization, so over-wide NPU-unquantizable fallback weights # load straight into host memory instead of round-tripping # through the accelerator (~24 GiB startup peak on MiniMax H3). from vllm_omni.quantization.int8_config import load_unquantizable_fallback_on_cpu fallback_ctx = load_unquantizable_fallback_on_cpu() if offload_after_quant else contextlib.nullcontext() with fallback_ctx: model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False) resolved_offload = resolve_offload(self.od_config) distributed_offload = resolved_offload.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE dit_distributed_offload = distributed_offload and resolved_offload.offloads(DIT_COMPONENT) dit_uses_allgather = dit_distributed_offload and resolved_offload.uses_allgather(DIT_COMPONENT) tensor_parallel_size = int(getattr(self.parallel_config, "tensor_parallel_size", 1)) use_hsdp = bool(getattr(self.parallel_config, "use_hsdp", False)) data_parallel_size = int(getattr(self.parallel_config, "data_parallel_size", 1)) sequence_parallel_size = int(getattr(self.parallel_config, "sequence_parallel_size", 1)) dlo_group_size = data_parallel_size if data_parallel_size > 1 else sequence_parallel_size modules = ModuleDiscovery.discover(model) plan = get_offload_plan(model) selected_encoders = [ encoder for name, encoder in zip(modules.encoder_names, modules.encoders) if resolved_offload.offloads(TEXT_ENCODER_COMPONENT) and encoder_component_type(name, plan) == TEXT_ENCODER_COMPONENT ] allgather_modules: list[nn.Module] = [] if dlo_group_size > 1: if dit_uses_allgather: allgather_modules.extend(modules.dits) if ( distributed_offload and resolved_offload.offloads(TEXT_ENCODER_COMPONENT) and resolved_offload.uses_allgather(TEXT_ENCODER_COMPONENT) ): allgather_modules.extend(selected_encoders) allgather_online_quant = any(self._has_online_quant(module) for module in allgather_modules) if allgather_online_quant: unsupported_methods = { method for module in allgather_modules for method in self._unsupported_dlo_allgather_online_quant_methods(module) } if unsupported_methods: raise ValueError( "DLO+AllGather supports online quantization only for " "per-tensor FP8, INT8, and MXFP8 linears " "(host-loaded unquantized fallback layers are also allowed); " f"unsupported online methods: {', '.join(sorted(unsupported_methods))}. " "Use rank-local transfer for the affected component or " "disable online quantization." ) logger.info( "Validated online methods (per-tensor FP8, INT8, MXFP8) for every component using DLO+AllGather" ) plan_result = None weight_sources = self._get_weight_sources(model) hwr_state = None if not self._force_canonical_load: try: hwr_state = self._resolve_hwr( model, modules, dist_offload=dit_distributed_offload, use_allgather=dit_uses_allgather, load_format=load_format, sources=weight_sources, ) except _HWRCommitError: from vllm_omni.host_weight_runtime import RuntimeMode mode = RuntimeMode(getattr(self.od_config, "host_weight_runtime_mode", "disabled")) if mode is not RuntimeMode.PREFERRED: raise logger.warning( "HWR restore commit failed; discarding the model and retrying a fresh canonical load", exc_info=True, ) del model return self.load_fresh_canonical_model() self._hwr_state = hwr_state hwr_active = hwr_state is not None if hwr_active and hwr_state is not None: self.host_weight_plan = cast(HostWeightPlan | None, hwr_state.get("plan")) if dit_distributed_offload and not hwr_active and not self._force_canonical_load: lora_backend = getattr(self.od_config, "lora_backend", None) has_distilled_lora = lora_backend in (LoRABackend.DISTILL, "distill") and bool( getattr(self.od_config, "lora_path", None) ) plan_result = build_checkpoint_mmap_plan( model, dit_modules=tuple(zip(modules.dit_names, modules.dits)), sources=weight_sources, model_path=str(getattr(self.od_config, "model", "")) or None, tensor_parallel_size=tensor_parallel_size, use_hsdp=use_hsdp, online_quantization=any(self._has_online_quant(dit) for dit in modules.dits), has_distilled_lora=has_distilled_lora, ) self.host_weight_plan = plan_result.plan host_weight_plan = self.host_weight_plan if host_weight_plan is not None: logger.info( "DLO host-weight plan active (%s, %s): skipping ordinary materialization for %s", "AllGather" if dit_uses_allgather and dlo_group_size > 1 else "rank-local", host_weight_plan.backing_kind, sorted(host_weight_plan.planned_source_prefixes) or "legacy DiT sources", ) ordinary_sources = tuple( source for source in weight_sources if source.prefix not in host_weight_plan.planned_source_prefixes ) if ordinary_sources: logger.info( "Loading %d component weight source(s) outside the DLO host-weight plan", len(ordinary_sources), ) self.load_weights( model, sources=ordinary_sources, planned_weights=host_weight_plan.bindings, ) else: if dit_distributed_offload and plan_result is not None: logger.info( "DLO direct checkpoint mmap unavailable; using ordinary loader: %s", plan_result.fallback_reason, ) logger.debug("Loading weights on %s ...", load_device) if offload_after_quant: marked = self._request_offload_after_quant(model) if marked: logger.info( "Online quantization will return each of %d layers to CPU as it is quantized", marked, ) if load_format == "diffusers": cast(DiffusersAdapterPipeline, model).load_weights() else: if offload_after_quant: self.load_weights(model, stream_online_quant_to_cpu=True) else: self.load_weights(model) self._maybe_fuse_distilled_lora(model) # Post-load transforms (e.g. the Marlin repack for # compressed-tensors) only run on the accelerator, so an # offline-quantized CPU-offload load round-trips each # quantized module through the device and back to the host. process_device = ( device if offline_quant_direct_cpu and device is not None else target_device ) self._process_weights_after_loading(model, process_device) # A warm final-layout hit has already completed all # byte-changing work through the restorer. Shared runtime # finalization happens once at the end for both cold and warm # paths; the warm path never re-enters the ordinary # materialization/finalization pipeline. if offload_after_quant: model.to("cpu") logger.info("Quantization complete, offloaded model back to CPU") try: self._apply_skip_softmax_calibration(model) model = model.eval() if self._hwr_state is not None: warm_snapshot = self._hwr_state.get("warm_snapshot") if warm_snapshot is not None: self._assert_final_layout_tensors_unchanged(model, cast(dict[str, tuple[int, str]], warm_snapshot)) self._publish_hwr_after_load(model, ModuleDiscovery.discover(model), self._hwr_state) except Exception: hwr_plan = self._hwr_state.get("plan") if self._hwr_state is not None else None if isinstance(hwr_plan, HostWeightPlan): carrier = hwr_plan.lease_carrier if carrier is not None: carrier.close() from vllm_omni.host_weight_runtime import RuntimeMode mode = RuntimeMode(getattr(self.od_config, "host_weight_runtime_mode", "disabled")) if mode is RuntimeMode.PREFERRED: logger.warning( "HWR warm finalization failed; discarding the model and retrying a fresh canonical load", exc_info=True, ) del model return self.load_fresh_canonical_model() raise self._log_w4a8_fallback_load_summaries(model) self._attach_offload_startup_state(model) return model @staticmethod def _log_w4a8_fallback_load_summaries(model: nn.Module) -> None: """Ask discovered DiTs to report W4A8 state at the common load exit. Each DiT derives readiness from its processed layers, so a deferred weight plan cannot be mistaken for the ordinary or HSDP load path. """ components = ModuleDiscovery.discover(model) for component_name, dit in zip(components.dit_names, components.dits): candidate = getattr(dit, "_log_w4a8_fallback_load_summary", None) if callable(candidate): reporter = cast(Callable[[str], None], candidate) reporter(component_name) @staticmethod def _request_offload_after_quant(model: nn.Module) -> int: """Ask online-quant layers to return to host memory once quantized. The weights only visit the accelerator so the quant kernels can run on them; ``load_model`` sends the model back to the host afterwards either way. Without this the whole transformer accumulates on device until that final move, which for MiniMax H3 is a ~43 GiB peak that no longer fits beside a resident TP-sharded text encoder — even though layer-wise offload means none of it is supposed to be resident at inference time. Only quant methods that advertise ``supports_offload_after_quant`` are asked, since the implementation has to know when a layer is finished. Deferring materialization to the ``meta`` device does not imply that. """ marked = 0 for module in model.modules(): quant_method = getattr(module, "quant_method", None) if quant_method is None or not getattr(quant_method, "supports_offload_after_quant", False): continue quant_method.enable_offload_after_quant() marked += 1 return marked @staticmethod def _has_online_quant(model: nn.Module) -> bool: """Whether any layer uses an online-quant method that defers weight materialization onto the ``meta`` device (upstream vLLM ``uses_meta_device=True``, e.g. online FP8).""" return has_online_quantization(model) @staticmethod def _unsupported_dlo_allgather_online_quant_methods(model: nn.Module) -> tuple[str, ...]: """Return unsupported online-quant methods for DLO AllGather. Per-tensor online FP8, online INT8, and online MXFP8 are safe after the ordinary loader has finalized their weight and scale parameters. DLO shards those runtime tensors by dtype and reconstructs their recorded shapes and strides before the kernel consumes them. They all keep plain transportable 1-byte dtypes over ordinary strided views: - online INT8: int8 weight plus fp32 scale, either contiguous (NPU, pre-transposed (K, N)) or a transposed view (CUDA, stride (1, K)); - online MXFP8: fp8 weight plus e8m0 block scale, either contiguous (NPU: (K, N) weight with (K_groups/2, N, 2) scale) or with the scale stored as a transposed view (vLLM kernel, .t() over a contiguous (K/32, N) buffer). Both shape families are already covered by the physical-order packing that online FP8 requires. Other online methods may create different scale, packing, or aliasing layouts (e.g. dual-scale fp4 pairs, swizzled or NZ hardware formats) and remain fail-closed until validated. ``UnquantizedHostLinearMethod`` is also allowed: it backs layers too wide for npu_quant_matmul, loads their weights straight into host memory, and its runtime layout is a plain contiguous bf16 weight — identical to the ordinary unquantized path DLO already shards. """ from vllm.model_executor.layers.quantization.online.fp8 import ( Fp8PerTensorOnlineLinearMethod, ) from vllm_omni.quantization.int8_config import ( Int8OnlineLinearMethod, NPUInt8OnlineLinearMethod, UnquantizedHostLinearMethod, ) try: from vllm_omni.quantization.mxfp8_config import ( NPUMxfp8OnlineLinearMethod, VllmMxfp8OnlineLinearMethod, ) mxfp8_online_methods: tuple[type, ...] = ( NPUMxfp8OnlineLinearMethod, VllmMxfp8OnlineLinearMethod, ) except ImportError: # MXFP8 requires a vLLM build with MXFP8 kernel support; treat it # as absent when the module cannot be imported. mxfp8_online_methods = () allowed_online_methods: tuple[type, ...] = ( Fp8PerTensorOnlineLinearMethod, Int8OnlineLinearMethod, NPUInt8OnlineLinearMethod, UnquantizedHostLinearMethod, *mxfp8_online_methods, ) unsupported: set[str] = set() for module in model.modules(): quant_method = getattr(module, "quant_method", None) if not getattr(quant_method, "uses_meta_device", False): continue if not isinstance(quant_method, allowed_online_methods): unsupported.add(type(quant_method).__name__) return tuple(sorted(unsupported)) def _apply_skip_softmax_calibration(self, model: nn.Module) -> None: from vllm_omni.diffusion.attention.backends.trtllm_calibration import ( apply_skip_softmax_calibration, ) cfg = getattr(self.od_config, "diffusion_attention_config", None) apply_skip_softmax_calibration(cfg, model) 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). """ # Newer upstream vLLM online-quant methods (uses_meta_device=True) create # weights on the ``meta`` device and materialize them just-in-time as each # layer's weights finish loading (via the layerwise online-process loader). # Any "straggler" layers whose weights were not fully materialized during # load (padded / partially-loaded layers) remain on ``meta``. Upstream's # base_loader calls finalize_layerwise_processing() to materialize them; # the diffusion loader must mirror that, otherwise the module.to() below # raises "Cannot copy out of meta tensor; no data!". This whole meta-device # handling is gated on online quant actually being in use, so that the # proven code path for everything else (in particular FSDP/HSDP-sharded # params, whose per-parameter .data cannot be cross-device reassigned) is # left untouched. Import lazily so older vLLM (no meta-device quant) is # unaffected. has_online_quant = self._has_online_quant(model) if has_online_quant: from vllm.model_executor.model_loader.reload.layerwise import ( finalize_layerwise_processing, ) # model_config is only dereferenced by finalize for vLLM Attention / # MLAAttention layers; diffusion DiT models use their own attention and # have none, so passing None is safe here. finalize_layerwise_processing(model, model_config=None) for _, module in model.named_modules(): quant_method = getattr(module, "quant_method", None) if quant_method is None or not isinstance(quant_method, QuantizeMethodBase): continue # Layers finished during loading would only be staged onto the target # device for a process call that immediately returns. That round trip # is wasted work in general, and undoes the point of the offload for # layers that already went back to the host. if getattr(module, "_already_called_process_weights_after_loading", False): continue if has_online_quant: # Online quant may leave straggler params on the ``meta`` device. # Move only real (non-meta) params onto the target device for # processing and restore them afterward, mirroring upstream vLLM's # device_loading_context — a blanket module.to(target_device) would # raise NotImplementedError on meta params. Online quant initializes # on the accelerator, so params are normally already on the target # device and this loop is a no-op move; the point is to skip meta. original_devices: dict[str, torch.device] = {} for name, param in module.named_parameters(): if param.device.type != "meta" and param.device != target_device: original_devices[name] = param.device param.data = param.data.to(target_device) quant_method.process_weights_after_loading(module) # Restore pre-existing params to their original device; leave any # newly created (e.g. quantized) params on the target device. for name, param in module.named_parameters(): if name in original_devices: param.data = param.data.to(original_devices[name]) else: # No meta params possible here. Preserve the original FSDP/HSDP-aware # whole-module move (module.to()), which correctly handles sharded # DTensor params that per-parameter .data reassignment cannot. 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, *, stream_online_quant_to_cpu: bool = False, sources: Sequence["ComponentSource"] | None = None, planned_weights: Iterable[str] = (), ) -> None: weights_to_load = self._get_expected_parameter_names(model) weights = self.get_all_weights(model) if sources is None else self.get_all_weights(model, sources=sources) if stream_online_quant_to_cpu: weights = self._stream_online_quant_weights_to_cpu(model, weights) loaded_weights = model.load_weights(weights) if loaded_weights is not None: loaded_weights = set(loaded_weights).union(planned_weights) 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 # Offline formats can require scales that older online loaders # were allowed to synthesize. Do not apply that legacy tolerance # to an explicitly required checkpoint tensor. required_missing = { name for name, param in model.named_parameters() if name in weights_not_loaded and getattr(param, "is_checkpoint_required", False) } if required_missing: raise ValueError(f"Required weights were not initialized from checkpoint: {required_missing}") # NOTE: if the model is quantized, ignore not_loaded check for scale # weights. ModelOpt FP8 carries a per-tensor `weight_scale` and a # static activation `input_scale`, which the quant method may # fold/track differently than plain parameters. weights_scale_not_loaded = { name for name in weights_not_loaded if name.endswith(("weight_scale", "input_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, 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", # 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. """ if self.quant_config is None: raise ValueError( "The quantization config is None, and the following weights " f"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 _init_from_load_format( self, load_format: str, target_device: torch.device, custom_pipeline_name: str | type[nn.Module] | None = None, is_hsdp: bool = False, ) -> nn.Module: """Initialize the model from a specified load format.""" 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. model_cls = _resolve_custom_pipeline_cls(custom_pipeline_name) with set_current_diffusion_config(self.od_config): model = model_cls(od_config=self.od_config) # HSDP normally defers GPU placement to apply_hsdp_to_model to keep peak # load-time memory on CPU. Online quantization (e.g. fp8) runs CUDA-only # kernels inside load_weights via the layerwise loader, so when a quant # config is set we initialize on the accelerator like the non-HSDP path; # apply_hsdp_to_model shards GPU-resident params equally well. hsdp_defer_to_cpu = is_hsdp and self.quant_config is None if not hsdp_defer_to_cpu and target_device.type != "cpu": model.to(target_device) else: hsdp_defer_to_cpu = is_hsdp and self.quant_config is None device_ctx = contextlib.nullcontext() if hsdp_defer_to_cpu else target_device with device_ctx: if load_format == "default": model = initialize_model(self.od_config) elif load_format == "diffusers": model = DiffusersAdapterPipeline(od_config=self.od_config, device=target_device) else: raise ValueError(f"Unknown load_format: {load_format}") return model def _load_model_with_pre_sharded_hsdp( self, model: nn.Module, discovered_modules: PipelineModules, hsdp_config: HSDPInferenceConfig, target_device: torch.device, ) -> nn.Module: """Shard meta DiT parameters before reading their rank-local checkpoint slices.""" from torch.distributed.checkpoint import HuggingFaceStorageReader from torch.distributed.checkpoint import load as dcp_load if self.counter_before_loading_weights == 0.0: self.counter_before_loading_weights = time.perf_counter() load_plan = self._prepare_pre_sharded_hsdp_load( model, discovered_modules, hsdp_config, target_device, ) self._check_no_unplanned_meta_tensors(model, load_plan.binding_names) self._validate_pre_sharded_hsdp_bindings(load_plan) self._process_pre_sharded_hsdp_groups_on_meta(load_plan) for group in load_plan.groups: logger.debug("Applying meta-first HSDP shard to %s", group.name) shard_hsdp_module(group.module, group.context) for root in load_plan.roots: logger.debug("Finalizing meta-first HSDP root %s", root.name) finalize_hsdp_root(root.module, root.context) self._materialize_pre_sharded_hsdp_state(model, load_plan, target_device) targets = self._get_pre_sharded_hsdp_targets(model, load_plan) checkpoint_states, local_payload_bytes = self._prepare_pre_sharded_checkpoint_states( load_plan, targets, ) thread_count = int(getattr(self.od_config, "num_weight_load_threads", 4)) if _is_distributed_rank_zero(): directory_label = "directory" if len(checkpoint_states) == 1 else "directories" logger.info( "Pre-sharded HSDP loading %.2f GiB of rank-local tensor slices directly into " "FSDP-owned storage from %d HF safetensors %s with %d reader thread(s) per rank", local_payload_bytes / 1024**3, len(checkpoint_states), directory_label, thread_count, ) for checkpoint_dir in sorted(checkpoint_states, key=str): reader = HuggingFaceStorageReader(str(checkpoint_dir), thread_count=thread_count) # Each worker has the full checkpoint available and each DTensor # already describes its rank-local destination slice. Avoid DCP's # coordinator collectives, which can leave rank 0 with less GPU # memory headroom than the other HSDP workers. dcp_load(checkpoint_states[checkpoint_dir], storage_reader=reader, no_dist=True) remaining_meta = self._pre_sharded_meta_target_names(model, load_plan) if remaining_meta: raise ValueError( "Pre-sharded HSDP checkpoint loading left tensors on meta; " f"first missing tensors: {remaining_meta[:16]}" ) self._finalize_pre_sharded_hsdp_loaded_weights(load_plan) 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, ) self._move_non_hsdp_modules(discovered_modules, target_device) return model @staticmethod def _validate_pre_sharded_hsdp_bindings(load_plan: _PreShardedHSDPLoadPlan) -> None: unsupported = [name for name, binding in load_plan.bindings.items() if binding.transform is not None] if unsupported: raise ValueError( "Pre-sharded HSDP requires checkpoint tensors in runtime layout; " f"tensor transforms are unsupported: {unsupported[:16]}" ) non_safetensors = sorted( { binding.file_path for binding in load_plan.bindings.values() if not binding.file_path.endswith(".safetensors") } ) if non_safetensors: raise ValueError( f"Pre-sharded HSDP requires existing HF safetensors files; unsupported files: {non_safetensors[:8]}" ) @staticmethod def _get_pre_sharded_hsdp_targets( model: nn.Module, load_plan: _PreShardedHSDPLoadPlan, ) -> dict[str, torch.Tensor]: targets: dict[str, torch.Tensor] = { **dict(model.named_buffers()), **dict(model.named_parameters()), } missing_targets = sorted(load_plan.binding_names - targets.keys()) if missing_targets: raise ValueError( f"Pre-sharded HSDP changed checkpoint target names; first missing targets: {missing_targets[:16]}" ) return targets @staticmethod def _materialize_pre_sharded_hsdp_state( model: nn.Module, load_plan: _PreShardedHSDPLoadPlan, target_device: torch.device, ) -> None: """Materialize the installed FSDP rank-local state before checkpoint I/O. ``fully_shard`` supports meta initialization through ``Module.to_empty``: its registered ``_apply`` hooks update the internal sharded storage to point at the newly allocated rank-local tensors. Modules are materialized in the same HSDP-group order and child-first traversal used by ordinary ``fully_shard`` initialization instead of the forward traversal used by a recursive ``to_empty`` call. PyTorch's swap-on-conversion mode preserves Parameter identity so that post-load kernel objects created before sharding cannot retain stale references to the original meta Parameters. A complete checkpoint plan covers every parameter and persistent buffer in each DiT. Non-persistent buffers are not checkpoint entries, so keep their initialized values across ``to_empty``. """ unplanned_parameters: list[str] = [] saved_buffers: dict[str, torch.Tensor] = {} for root in load_plan.roots: prefix = f"{root.name}." for local_name, _parameter in root.module.named_parameters(): full_name = prefix + local_name if full_name not in load_plan.binding_names: unplanned_parameters.append(full_name) for local_name, buffer in root.module.named_buffers(): full_name = prefix + local_name if full_name in load_plan.binding_names: continue if buffer.device.type == "meta": raise ValueError(f"Pre-sharded HSDP cannot preserve an unplanned meta buffer: {full_name}") saved_buffers[full_name] = buffer.detach().cpu().clone() if unplanned_parameters: raise ValueError( "Pre-sharded HSDP requires checkpoint coverage for every DiT parameter; " f"first unplanned parameters: {unplanned_parameters[:16]}" ) previous_swap_mode = torch.__future__.get_swap_module_params_on_conversion() torch.__future__.set_swap_module_params_on_conversion(True) try: sharded_group_roots = {group.module for group in load_plan.groups} materialized_modules: set[nn.Module] = set() for group in load_plan.groups: for module in DiffusersPipelineLoader._hsdp_managed_modules_post_order( group.module, nested_hsdp_roots=sharded_group_roots - {group.module}, ): if module not in materialized_modules: module.to_empty(device=target_device, recurse=False) materialized_modules.add(module) for root in load_plan.roots: for module in DiffusersPipelineLoader._hsdp_managed_modules_post_order( root.module, nested_hsdp_roots=sharded_group_roots, ): if module not in materialized_modules: module.to_empty(device=target_device, recurse=False) materialized_modules.add(module) finally: torch.__future__.set_swap_module_params_on_conversion(previous_swap_mode) if saved_buffers: materialized_buffers = dict(model.named_buffers()) with torch.no_grad(): for name, value in saved_buffers.items(): target = materialized_buffers[name] target.copy_(value.to(device=target.device, dtype=target.dtype)) @staticmethod def _hsdp_managed_modules_post_order( root: nn.Module, *, nested_hsdp_roots: set[nn.Module], ) -> list[nn.Module]: """Mirror FSDP's child-first managed-module traversal. Nested FSDP roots own their own parameters and are materialized in their corresponding load group, so the enclosing traversal must not descend into them. """ modules: list[nn.Module] = [] visited: set[nn.Module] = set() def visit(module: nn.Module) -> None: if module in visited: return visited.add(module) for child in module.children(): if child not in nested_hsdp_roots: visit(child) modules.append(module) visit(root) return modules @staticmethod def _prepare_pre_sharded_checkpoint_states( load_plan: _PreShardedHSDPLoadPlan, targets: dict[str, torch.Tensor], ) -> tuple[dict[Path, dict[str, torch.Tensor]], int]: from torch.distributed.tensor import DTensor checkpoint_states: dict[Path, dict[str, torch.Tensor]] = {} checkpoint_owners: dict[tuple[Path, str], str] = {} local_payload_bytes = 0 for name, binding in load_plan.bindings.items(): target = targets[name] checkpoint_dir = Path(binding.file_path).parent owner_key = (checkpoint_dir, binding.checkpoint_key) if previous_owner := checkpoint_owners.get(owner_key): raise ValueError( "Pre-sharded HSDP cannot bind one checkpoint tensor to multiple runtime tensors: " f"{binding.checkpoint_key!r} maps to both {previous_owner!r} and {name!r}" ) checkpoint_owners[owner_key] = name if target.device.type == "meta": raise ValueError(f"Pre-sharded HSDP target {name!r} was not materialized") local_target = target.to_local() if isinstance(target, DTensor) else target local_payload_bytes += local_target.numel() * local_target.element_size() checkpoint_states.setdefault(checkpoint_dir, {})[binding.checkpoint_key] = target return checkpoint_states, local_payload_bytes @staticmethod def _pre_sharded_meta_target_names( model: nn.Module, load_plan: _PreShardedHSDPLoadPlan, ) -> list[str]: from torch.distributed.tensor import DTensor targets = { **dict(model.named_buffers()), **dict(model.named_parameters()), } return [ name for name in load_plan.binding_names if (target := targets[name]).device.type == "meta" or (isinstance(target, DTensor) and target.to_local().device.type == "meta") ] @staticmethod def _process_pre_sharded_hsdp_groups_on_meta(load_plan: _PreShardedHSDPLoadPlan) -> None: modules_to_process: list[tuple[nn.Module, QuantizeMethodBase]] = [] for group in load_plan.groups: for _, module in group.module.named_modules(): quant_method = getattr(module, "quant_method", None) if quant_method is None or not isinstance(quant_method, QuantizeMethodBase): continue modules_to_process.append((module, quant_method)) if not modules_to_process: return for root in load_plan.roots: if not getattr(root.module, "_hsdp_pre_sharded_meta_post_load", False): raise ValueError( f"Model {type(root.module).__name__} has not declared its post-load processing safe " "for pre-sharded meta initialization" ) for module, quant_method in modules_to_process: if type(quant_method) is not UnquantizedLinearMethod: raise ValueError( "Pre-sharded HSDP does not support the selected post-load processing; " f"{type(module).__name__} uses {type(quant_method).__name__}" ) for module, quant_method in modules_to_process: quant_method.process_weights_after_loading(module) def _prepare_pre_sharded_hsdp_load( self, model: nn.Module, discovered_modules: PipelineModules, hsdp_config: HSDPInferenceConfig, target_device: torch.device, ) -> _PreShardedHSDPLoadPlan: outer_dit_names, outer_dits = discovered_modules.outermost_dits() if not outer_dits: raise ValueError("No DiT modules discovered for HSDP sharding") sources = self._get_weight_sources(model) plan_result = build_checkpoint_binding_plan( model, dit_modules=tuple(zip(outer_dit_names, outer_dits)), sources=sources, model_path=str(getattr(self.od_config, "model", "")) or None, tensor_parallel_size=int(getattr(self.parallel_config, "tensor_parallel_size", 1)), online_quantization=False, ) if plan_result.plan is None: raise ValueError(f"Pre-sharded HSDP checkpoint loading is incompatible: {plan_result.fallback_reason}") plan = plan_result.plan unsupported_sources = [source.prefix for source in sources if source.prefix not in plan.planned_source_prefixes] if unsupported_sources: raise ValueError( "Pre-sharded HSDP checkpoint loading requires dedicated DiT sources; " f"unsupported source prefixes: {unsupported_sources}" ) roots: list[_HSDPLoadRoot] = [] groups: list[_HSDPModuleLoadGroup] = [] for outer_name, outer_dit in zip(outer_dit_names, outer_dits): conditions = getattr(outer_dit, "_hsdp_shard_conditions", None) if not conditions: raise ValueError(f"Model {type(outer_dit).__name__} has no _hsdp_shard_conditions defined") context = prepare_hsdp_shard_context( outer_dit, hsdp_config, target_device=target_device, ) roots.append(_HSDPLoadRoot(outer_name, outer_dit, context)) for local_name, module in reversed(list(outer_dit.named_modules())): if not any(condition(local_name, module) for condition in conditions): continue if not local_name: raise ValueError("_hsdp_shard_conditions must not select the DiT root") full_name = f"{outer_name}.{local_name}" if local_name else outer_name groups.append(_HSDPModuleLoadGroup(full_name, module, context)) parameters = dict(model.named_parameters()) targets: dict[str, torch.Tensor] = {**dict(model.named_buffers()), **parameters} preserved_params = {param for root in roots for param in root.context.ignored_params or ()} binding_names = frozenset(plan.bindings) self._release_hsdp_checkpoint_targets_to_meta( model, targets, parameters, binding_names, preserved_params, ) return _PreShardedHSDPLoadPlan( roots=roots, groups=groups, bindings=plan.bindings, binding_names=binding_names, ) @staticmethod def _check_no_unplanned_meta_tensors(model: nn.Module, planned_names: frozenset[str]) -> None: unplanned_meta = [ name for name, tensor in (*model.named_parameters(), *model.named_buffers()) if tensor.device.type == "meta" and name not in planned_names ] if unplanned_meta: raise ValueError( "Pre-sharded HSDP initialized tensors on meta that are not covered by the " f"checkpoint plan. First unplanned tensors: {unplanned_meta[:16]}" ) @staticmethod def _release_hsdp_checkpoint_targets_to_meta( model: nn.Module, targets: dict[str, torch.Tensor], parameters: dict[str, nn.Parameter], binding_names: frozenset[str], preserved_params: set[nn.Parameter], ) -> None: released: dict[str, tuple[int, float]] = {} for name in binding_names: target = targets[name] if target in preserved_params or target.device.type == "meta": continue logical_gib = target.numel() * target.element_size() / 1024**3 device_type = target.device.type count, gib = released.get(device_type, (0, 0.0)) released[device_type] = (count + 1, gib + logical_gib) if name in parameters: replacement = DiffusersPipelineLoader._make_parameter_like( target, torch.empty_like(target, device="meta"), ) torch.utils.swap_tensors(target, replacement) else: replacement = torch.empty_like(target, device="meta") torch.utils.swap_tensors(target, replacement) if released: logger.info("Released HSDP checkpoint-covered tensors to meta before pre-sharded load: %s", released) @staticmethod def _make_parameter_like(template: nn.Parameter, tensor: torch.Tensor) -> nn.Parameter: """Create a Parameter on new storage while preserving loader metadata.""" if type(template) is nn.Parameter: replacement = nn.Parameter(tensor, requires_grad=template.requires_grad) replacement.__dict__.update(getattr(template, "__dict__", {})) return replacement replacement = template.__class__.__new__(template.__class__, tensor) replacement.__dict__.update(getattr(template, "__dict__", {})) replacement.requires_grad_(template.requires_grad) return replacement @staticmethod def _finalize_pre_sharded_hsdp_loaded_weights( load_plan: _PreShardedHSDPLoadPlan, ) -> None: """Run model post-load hooks for tensors materialized by pre-sharded loading.""" for root in load_plan.roots: post_load = getattr(root.module, "post_load_weights", None) if callable(post_load): post_load() root.module.eval() validate = getattr(root.module, "validate_loaded_weights", None) if callable(validate): prefix = f"{root.name}." validate({name.removeprefix(prefix) for name in load_plan.binding_names if name.startswith(prefix)}) @staticmethod def _move_non_hsdp_modules( discovered_modules: PipelineModules, target_device: torch.device, ) -> None: 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) def _load_model_with_hsdp( self, target_device: torch.device, load_format: str = "default", custom_pipeline_name: str | type[nn.Module] | None = None, offload_after_quant: bool = False, ) -> 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. """ hsdp_config = HSDPInferenceConfig( enabled=True, hsdp_replicate_size=self.parallel_config.hsdp_replicate_size, hsdp_shard_size=self.parallel_config.hsdp_shard_size, param_dtype=self.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 == "diffusers": raise ValueError("HSDP is not supported with the diffusers adapter load format") strategy = getattr(self.od_config, "hsdp_weight_load_strategy", "full") if strategy == "pre_sharded": if load_format != "default": raise ValueError( f"hsdp_weight_load_strategy={strategy!r} currently supports only diffusion_load_format='default'" ) if self.quant_config is not None: raise ValueError(f"hsdp_weight_load_strategy={strategy!r} does not support quantization yet") if getattr(self.od_config, "lora_path", None): raise ValueError(f"hsdp_weight_load_strategy={strategy!r} does not support LoRA yet") elif strategy != "full": raise ValueError(f"Unknown hsdp_weight_load_strategy: {strategy!r}") if strategy == "pre_sharded": model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True) discovered_modules = ModuleDiscovery.discover(model) return self._load_model_with_pre_sharded_hsdp( model, discovered_modules, hsdp_config, target_device, ) # Same host-fallback bound as the ordinary path: with a quant config the # HSDP path initializes on the accelerator (hsdp_defer_to_cpu=False), so # over-wide NPU-unquantizable fallback weights would otherwise round-trip # through the device before apply_hsdp_to_model shards them. Loading them # straight into host memory keeps the load-time device peak at the # quantizable layers alone; sharding then distributes the fallback like # any other parameter. Broadcast loading excludes online quantization # already, so the context below only ever matters on the ordinary # per-rank branch -- but spanning both is harmless. from vllm_omni.quantization.int8_config import load_unquantizable_fallback_on_cpu fallback_ctx = load_unquantizable_fallback_on_cpu() if offload_after_quant else contextlib.nullcontext() with fallback_ctx: model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True) world_size = 1 rank = 0 if torch.distributed.is_initialized(): world_size = torch.distributed.get_world_size() rank = torch.distributed.get_rank() has_online_quant = self._has_online_quant(model) or ( self.quant_config is not None and not getattr(self.quant_config, "is_checkpoint_quantized", False) ) enable_broadcast = bool(getattr(self.od_config, "enable_broadcast_weight_load", False)) and world_size > 1 if enable_broadcast and has_online_quant: logger.info( "Worker %d: Online quantization detected; falling back to ordinary " "per-rank weight loading for HSDP", rank, ) enable_broadcast = False if enable_broadcast: if rank == 0: self.load_weights(model) self._maybe_fuse_distilled_lora(model) self._broadcast_model_weights(model, target_device=target_device, src_rank=0) if ( rank != 0 and getattr(self.od_config, "lora_backend", None) in (LoRABackend.DISTILL, "distill") and getattr(self.od_config, "lora_path", None) and hasattr(model, "load_lora_weights") ): setattr(model, "lora_is_fused", True) else: self.load_weights(model) self._maybe_fuse_distilled_lora(model) # Quantization methods must finish while parameters are ordinary local # tensors. Some post-load transforms use operations (for example, # torch.unique in ModelOpt NVFP4) that do not support DTensor inputs. self._process_weights_after_loading(model, target_device) # Discover pipeline components (DiT, encoders, VAEs) via # ModuleDiscovery, which consults SupportsComponentDiscovery # when available and falls back to well-known attribute names. # This supports nested pipelines (e.g. LTX2DistilledPipeline # where the transformer lives at "pipe.transformer"). discovered_modules = ModuleDiscovery.discover(model) # Shard only the outermost DiTs. A pipeline may list a DiT and one of its # submodules as separate DiTs (e.g. Cosmos3's transformer and the nested # transformer.language_model) for offload's independent rings; for HSDP an # inner DiT is already covered by its ancestor's _hsdp_shard_conditions, so # sharding it again would double-wrap blocks and require the inner stack to # declare its own conditions. outer_dit_names, outer_dits = discovered_modules.outermost_dits() # Online FP8 quantization (Fp8OnlineLinearMethod) leaves layer weights # as non-contiguous transpose views (qweight.t()) so the Cutlass kernel # gets a column-major B. FSDP2 fully_shard rejects non-contiguous params. # Rewrite affected layers in-place to row-major contiguous storage and # shift the .t() to GEMM-call time. Layers using other quant methods or # already-contiguous weights are left untouched. if self.quant_config is not None: from vllm_omni.diffusion.quantization.hsdp_fp8 import ( prepare_fp8_layers_for_fsdp, ) for trans in outer_dits: prepare_fp8_layers_for_fsdp(trans) if not outer_dits: raise ValueError("No DiT modules discovered for HSDP sharding") # Apply HSDP sharding to each outermost DiT transformer for name, trans in zip(outer_dit_names, outer_dits): logger.debug("Applying HSDP to %s", name) apply_hsdp_to_model(trans, hsdp_config, target_device=target_device) # HSDP only shards transformer modules. All other runtime modules must # be placed on the execution device explicitly after sharding. 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