diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index b1776f3d5..2f9aadfd8 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -108,8 +108,8 @@ def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeyp def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): class _DummyLoader: - def __init__(self, load_config): - del load_config + def __init__(self, load_config, od_config=None): + del load_config, od_config def load_model(self, **kwargs): del kwargs diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 775220758..bd5af2263 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -379,6 +379,10 @@ class OmniDiffusionConfig: # Compilation enforce_eager: bool = False + # Parallel weight loading (for faster diffusion model startup) + enable_multithread_weight_load: bool = True + num_weight_load_threads: int = 4 + # Enable sleep mode enable_sleep_mode: bool = False diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index c4271c338..78d9ffd03 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -3,6 +3,7 @@ import dataclasses import glob import os +import re import time from collections.abc import Generator, Iterable from pathlib import Path @@ -22,6 +23,7 @@ from vllm.model_executor.model_loader.weight_utils import ( filter_duplicate_safetensors_files, filter_files_not_needed_for_inference, maybe_download_from_modelscope, + multi_thread_safetensors_weights_iterator, safetensors_weights_iterator, ) from vllm.utils.import_utils import resolve_obj_by_qualname @@ -35,6 +37,12 @@ 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" @@ -42,9 +50,6 @@ DIFFUSION_MODEL_WEIGHTS_INDEX = "diffusion_pytorch_model.safetensors.index.json" class DiffusersPipelineLoader: """Model loader that can load diffusers pipeline components from disk.""" - # default number of thread when enable multithread weight loading - DEFAULT_NUM_THREADS = 8 - @dataclasses.dataclass class ComponentSource: """A source for weights.""" @@ -70,18 +75,9 @@ class DiffusersPipelineLoader: counter_before_loading_weights: float = 0.0 counter_after_loading_weights: float = 0.0 - def __init__(self, load_config: LoadConfig): + def __init__(self, load_config: LoadConfig, od_config: OmniDiffusionConfig | None = None): self.load_config = load_config - - # TODO(Isotr0py): Enable multithreaded weight loading - # extra_config = load_config.model_loader_extra_config - # allowed_keys = {"enable_multithread_load", "num_threads"} - # unexpected_keys = set(extra_config.keys()) - allowed_keys - - # if unexpected_keys: - # raise ValueError( - # f"Unexpected extra config keys for load format {load_config.load_format}: {unexpected_keys}" - # ) + self.od_config = od_config def _prepare_weights( self, @@ -175,18 +171,36 @@ class DiffusersPipelineLoader: def _get_weights_iterator(self, source: "ComponentSource") -> Generator[tuple[str, torch.Tensor], None, None]: """Get an iterator for the model weights based on the load format.""" - hf_folder, hf_weights_files, use_safetensors = self._prepare_weights( + _, 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, ) - weights_iterator = safetensors_weights_iterator( - hf_weights_files, - self.load_config.use_tqdm_on_load, - self.load_config.safetensors_load_strategy, + + 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() diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index a16f2538a..263ce3126 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -102,7 +102,7 @@ class DiffusionModelRunner: # Load model within forward context load_config = LoadConfig() - model_loader = DiffusersPipelineLoader(load_config) + model_loader = DiffusersPipelineLoader(load_config, od_config=self.od_config) time_before_load = time.perf_counter() with get_memory_context(): diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 19c8b07ac..4af60b1a2 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -190,6 +190,8 @@ class AsyncOmni(OmniBase): "enforce_eager": kwargs.get("enforce_eager", False), "diffusion_load_format": kwargs.get("diffusion_load_format", "default"), "custom_pipeline_args": kwargs.get("custom_pipeline_args", None), + "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True), + "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4), }, "final_output": True, "final_output_type": "image", diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index c9ed7150f..f99e6c77a 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -272,6 +272,21 @@ class OmniServeCommand(CLISubcommand): help="Enable VAE tiling for memory optimization (useful for mitigating OOM issues).", ) + # Parallel weight loading (faster diffusion startup) + omni_config_group.add_argument( + "--disable-multithread-weight-load", + action="store_false", + dest="enable_multithread_weight_load", + default=True, + help="Disable multi-threaded safetensors loading (default: enabled with 4 threads).", + ) + omni_config_group.add_argument( + "--num-weight-load-threads", + type=int, + default=4, + help="Number of threads for parallel weight loading (default: 4).", + ) + # diffusion model offload parameters omni_config_group.add_argument( "--enable-cpu-offload",