From acf66202333ff837c058a0957c778ca14a1460da Mon Sep 17 00:00:00 2001 From: Anjie Hou <149605198+specture724@users.noreply.github.com> Date: Sat, 26 Sep 2026 17:45:20 +0800 Subject: [PATCH] [Refactor][Diffusion] Read the resolved offload policy outside the config layer (#7327) Signed-off-by: specture724 Co-authored-by: Claude Opus 5 Co-authored-by: Hongsheng Liu --- .../models/magi2/test_pipeline_magi2.py | 40 ++++++++++----- .../test_minimax_h3_dlo_lifecycle.py | 25 ++++++--- .../offloader/test_legacy_flag_readers.py | 51 +++++++++++++++++++ .../test_cosmos3_mixed_precision.py | 34 ++++++++----- vllm_omni/diffusion/diffusion_engine.py | 11 ++-- vllm_omni/diffusion/lora/manager.py | 8 ++- .../diffusion/models/bagel/pipeline_bagel.py | 8 ++- .../models/cosmos3/transformer_cosmos3.py | 6 ++- .../pipeline_diffusers_adapter.py | 9 +++- .../models/dreamzero/pipeline_dreamzero.py | 8 +-- .../diffusion/models/flux2/pipeline_flux2.py | 8 ++- .../diffusion/models/lance/pipeline_lance.py | 10 +++- .../models/lingbot_world/pipeline.py | 9 +++- .../diffusion/models/ltx2/ltx2_components.py | 13 +++-- .../diffusion/models/magi2/pipeline_magi2.py | 34 +++++++------ .../diffusion/models/minimax_h3/fasth3.py | 12 ++--- .../models/minimax_h3/pipeline_minimax_h3.py | 14 ++--- .../models/sana_video/pipeline_sana_video.py | 14 ++--- .../models/sana_wm/pipeline_sana_wm.py | 15 +++--- .../models/wan2_2/pipeline_wan2_2_s2v.py | 16 ++++-- vllm_omni/diffusion/offloader/config.py | 13 +++++ .../worker/diffusion_model_runner.py | 18 +++---- 22 files changed, 265 insertions(+), 111 deletions(-) create mode 100644 tests/diffusion/offloader/test_legacy_flag_readers.py diff --git a/tests/diffusion/models/magi2/test_pipeline_magi2.py b/tests/diffusion/models/magi2/test_pipeline_magi2.py index 4e6057457..280caff97 100644 --- a/tests/diffusion/models/magi2/test_pipeline_magi2.py +++ b/tests/diffusion/models/magi2/test_pipeline_magi2.py @@ -78,10 +78,15 @@ class _ParallelStub: @dataclass class _TopologyStub: parallel_config: _ParallelStub = field(default_factory=_ParallelStub) + diffusion_offload_config: object = None enable_cpu_offload: bool = False enable_layerwise_offload: bool = False enable_distributed_layerwise_offload: bool = False dlo_use_allgather: bool = True + dlo_resident_layers: int = 0 + dlo_host_registration_limit_gib: float = 0.0 + host_weight_runtime_mode: str = "disabled" + pin_cpu_memory: bool = True quantization_config: object = None cache_backend: str = "none" custom_pipeline_args: dict[str, object] = field(default_factory=dict) @@ -383,26 +388,33 @@ def test_native_topology_rejects_nondivisible_tensor_parallelism(): def test_native_topology_accepts_single_device_layerwise_with_cpu_staging(): - config = _topology_config( - enable_layerwise_offload=True, - additional_config={}, + _validate_native_topology( + _topology_config( + enable_layerwise_offload=True, + additional_config={}, + ) + ) + _validate_native_topology( + _topology_config( + enable_cpu_offload=True, + enable_layerwise_offload=True, + additional_config={}, + ) ) - _validate_native_topology(config) - config.enable_cpu_offload = True - _validate_native_topology(config) - - config.enable_layerwise_offload = False - with pytest.raises(ValueError, match="Combine --enable-cpu-offload"): - _validate_native_topology(config) + with pytest.raises(ValueError, match="already stages its auxiliary components"): + _validate_native_topology( + _topology_config( + enable_cpu_offload=True, + additional_config={}, + ) + ) def test_native_topology_requires_dlo_rank_local_mode(): - config = _topology_config(enable_distributed_layerwise_offload=True) with pytest.raises(ValueError, match="dlo-no-use-allgather"): - _validate_native_topology(config) - config.dlo_use_allgather = False - _validate_native_topology(config) + _validate_native_topology(_topology_config(enable_distributed_layerwise_offload=True)) + _validate_native_topology(_topology_config(enable_distributed_layerwise_offload=True, dlo_use_allgather=False)) def test_native_topology_accepts_four_device_hsdp_cfg_and_vae_layouts(): diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py index faa6b9a8d..0c09e65c2 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py @@ -13,6 +13,21 @@ class _OffloadAbort(BaseException): pass +def _legacy_distributed_offload_config() -> SimpleNamespace: + return SimpleNamespace( + diffusion_offload_config=None, + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=True, + dlo_use_allgather=True, + dlo_resident_layers=0, + dlo_host_registration_limit_gib=0.0, + host_weight_runtime_mode="disabled", + pin_cpu_memory=True, + parallel_config=SimpleNamespace(data_parallel_size=1), + ) + + def test_no_allgather_example_selects_distributed_backend(): from examples.offline_inference.minimax_h3.dlo_lifecycle import engine_kwargs from vllm_omni.diffusion.offloader.base import OffloadConfig @@ -100,10 +115,7 @@ def test_manual_component_failure_forces_retained_cache_release(mocker): pipeline = object.__new__(MiniMaxH3Pipeline) torch.nn.Module.__init__(pipeline) - pipeline.od_config = mocker.Mock() - pipeline.od_config.diffusion_offload_config = None - pipeline.od_config.enable_layerwise_offload = False - pipeline.od_config.enable_distributed_layerwise_offload = True + pipeline.od_config = _legacy_distributed_offload_config() pipeline._model_cpu_offload_modules = [] pipeline._dlo_component_cache = mocker.Mock() component = mocker.Mock() @@ -123,10 +135,7 @@ def test_manual_component_offload_failure_forces_retained_cache_release(mocker): pipeline = object.__new__(MiniMaxH3Pipeline) torch.nn.Module.__init__(pipeline) - pipeline.od_config = mocker.Mock() - pipeline.od_config.diffusion_offload_config = None - pipeline.od_config.enable_layerwise_offload = False - pipeline.od_config.enable_distributed_layerwise_offload = True + pipeline.od_config = _legacy_distributed_offload_config() pipeline._model_cpu_offload_modules = [] pipeline._dlo_component_cache = mocker.Mock() component = mocker.Mock() diff --git a/tests/diffusion/offloader/test_legacy_flag_readers.py b/tests/diffusion/offloader/test_legacy_flag_readers.py new file mode 100644 index 000000000..34619a8ab --- /dev/null +++ b/tests/diffusion/offloader/test_legacy_flag_readers.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Runtime code reads the resolved offload policy, not the legacy booleans.""" + +import re +from pathlib import Path + +import pytest + +import vllm_omni + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + +# `enable_cpu_offload`, `enable_layerwise_offload` and +# `enable_distributed_layerwise_offload` are compatibility aliases that +# `resolve_offload` materializes from the resolved policy. Only the layer that +# declares and resolves them may read them; everything else asks +# `vllm_omni.diffusion.offloader.config` for the strategy or the selected +# components, so the aliases can be removed without sweeping the runtime again. +COMPATIBILITY_LAYER = { + "config/omni_config.py", + "config/stage_config.py", + "engine/arg_utils.py", + "diffusion/data.py", + "diffusion/offloader/base.py", + "diffusion/offloader/config.py", + "quantization/tools/compare_diffusion_trajectory_similarity.py", +} + +READER = re.compile( + r"""\.enable_(?:cpu|layerwise|distributed_layerwise)_offload\b""" + r"""|getattr\(\s*[^,]+,\s*["']enable_(?:cpu|layerwise|distributed_layerwise)_offload["']""" +) + + +def test_no_runtime_module_reads_the_legacy_offload_flags(): + root = Path(vllm_omni.__file__).parent + offenders = [ + f"{path.relative_to(root)}:{number}" + for path in sorted(root.rglob("*.py")) + if str(path.relative_to(root)) not in COMPATIBILITY_LAYER + for number, line in enumerate(path.read_text().splitlines(), start=1) + if READER.search(line) + ] + + assert offenders == [], ( + "Read the resolved policy instead: resolve_offload_strategy(od_config), " + "offload_enabled(od_config), offload_streams_blocks(od_config), or " + f"should_offload_component(od_config, component). Offending reads: {offenders}" + ) diff --git a/tests/diffusion/quantization/test_cosmos3_mixed_precision.py b/tests/diffusion/quantization/test_cosmos3_mixed_precision.py index 026090fe9..0927409e9 100644 --- a/tests/diffusion/quantization/test_cosmos3_mixed_precision.py +++ b/tests/diffusion/quantization/test_cosmos3_mixed_precision.py @@ -406,28 +406,38 @@ def test_runtime_allows_standard_offload_and_rejects_distributed(monkeypatch) -> "get_tensor_model_parallel_world_size", lambda: 1, ) - od_config = SimpleNamespace( - enable_cpu_offload=True, - enable_layerwise_offload=True, - enable_distributed_layerwise_offload=False, - parallel_config=SimpleNamespace(use_hsdp=False), - ) + + def offload_config(**overrides: object) -> SimpleNamespace: + config = SimpleNamespace( + diffusion_offload_config=None, + enable_cpu_offload=False, + enable_layerwise_offload=True, + enable_distributed_layerwise_offload=False, + dlo_use_allgather=True, + dlo_resident_layers=0, + dlo_host_registration_limit_gib=0.0, + host_weight_runtime_mode="disabled", + pin_cpu_memory=True, + max_num_seqs=1, + parallel_config=SimpleNamespace(data_parallel_size=1, use_hsdp=False), + ) + for name, value in overrides.items(): + setattr(config, name, value) + return config + transformer_cosmos3._validate_mixed_precision_runtime( Cosmos3MixedPrecisionConfig(), - od_config, + offload_config(), ) - od_config.enable_distributed_layerwise_offload = True with pytest.raises(ValueError, match="distributed layer-wise offload"): transformer_cosmos3._validate_mixed_precision_runtime( Cosmos3MixedPrecisionConfig(), - od_config, + offload_config(enable_distributed_layerwise_offload=True), ) - od_config.enable_distributed_layerwise_offload = False - od_config.max_num_seqs = 2 with pytest.raises(ValueError, match="one active request"): transformer_cosmos3._validate_mixed_precision_runtime( Cosmos3MixedPrecisionConfig(), - od_config, + offload_config(max_num_seqs=2), ) diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index e2562ad12..57dddc889 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -44,7 +44,11 @@ from vllm_omni.diffusion.io_support import ( supports_audio_output, supports_multimodal_input, ) -from vllm_omni.diffusion.offloader.config import any_selected_component_uses_allgather +from vllm_omni.diffusion.offloader.config import ( + OffloadStrategy, + any_selected_component_uses_allgather, + resolve_offload_strategy, +) from vllm_omni.diffusion.output_formatter import ( format_diffusion_outputs, format_empty_diffusion_outputs, @@ -573,7 +577,8 @@ class DiffusionEngine: if output.media is not None: if output.output is not None: raise ValueError("DiffusionOutput cannot contain both media and legacy output") - media = output.media.to_cpu() if self.od_config.enable_cpu_offload else output.media + model_level = resolve_offload_strategy(self.od_config) is OffloadStrategy.MODEL_LEVEL + media = output.media.to_cpu() if model_level else output.media output_data = media.video.tensor outputs = finalize_diffusion_media(media, sampling_params=request.sampling_params) else: @@ -585,7 +590,7 @@ class DiffusionEngine: # post-processing to avoid device OOM — model weights may still # reside on the device and leave no headroom for intermediates. output_data = output.output - if self.od_config.enable_cpu_offload: + if resolve_offload_strategy(self.od_config) is OffloadStrategy.MODEL_LEVEL: output_data = _move_tensor_tree_to_cpu(output_data) if self.post_process_func is not None: diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 5670bdacd..b1b7aa6f6 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -73,11 +73,17 @@ class DiffusionLoRAManager: self.pipeline = pipeline self.device = device self.dtype = dtype + # Imported here: vllm_omni.diffusion.data imports this package while it + # is still initializing, and the offloader package imports that module. + from vllm_omni.diffusion.offloader.config import OffloadStrategy, resolve_offload_strategy + od_config = getattr(pipeline, "od_config", None) # DLO owns the base-weight lifecycle. Keep request-switchable LoRA # sidecars resident instead of rebuilding DLO host shards per request. self._resident_lora_device = ( - device if getattr(od_config, "enable_distributed_layerwise_offload", False) is True else None + device + if od_config is not None and resolve_offload_strategy(od_config) is OffloadStrategy.DISTRIBUTED_LAYER_WISE + else None ) # Cache supported/expected module suffixes once, before any layer diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 25f31679c..3d0a467f5 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -32,6 +32,10 @@ from vllm_omni.diffusion.distributed.parallel_state import ( from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery +from vllm_omni.diffusion.offloader.config import ( + OffloadStrategy, + resolve_offload_strategy, +) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch @@ -500,7 +504,9 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf # selectively materialized/moved by the offloader. # 3. HSDP: weights should be loaded on CPU first and sharded afterwards, # rather than eagerly placing the full model on one GPU. - if quant_config is None and not (od_config.enable_layerwise_offload or od_config.parallel_config.use_hsdp): + if quant_config is None and not ( + resolve_offload_strategy(od_config) is OffloadStrategy.LAYER_WISE or od_config.parallel_config.use_hsdp + ): self.to(self.device) self.setup_diffusion_pipeline_profiler( enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index 8510e0da9..7241082e8 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -38,6 +38,10 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available from vllm_omni.diffusion.layers.norm import RMSNorm as _VllmRMSNorm +from vllm_omni.diffusion.offloader.config import ( + OffloadStrategy, + resolve_offload_strategy, +) from vllm_omni.platforms import current_omni_platform from vllm_omni.quantization.component_config import ( ComponentQuantizationConfig, @@ -203,7 +207,7 @@ def _validate_mixed_precision_runtime( raise ValueError("Cosmos3 mixed precision currently supports tensor parallel size 1 only") if int(getattr(od_config, "max_num_seqs", 1)) != 1: raise ValueError("Cosmos3 mixed precision currently supports one active request per worker") - if bool(getattr(od_config, "enable_distributed_layerwise_offload", False)): + if resolve_offload_strategy(od_config) is OffloadStrategy.DISTRIBUTED_LAYER_WISE: raise ValueError( "Cosmos3 mixed precision does not support distributed layer-wise offload " "because its direct loader bypasses ModelOpt post-load transformations" diff --git a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py index 85e679715..f3e1a7432 100644 --- a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py +++ b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py @@ -33,6 +33,10 @@ from vllm_omni.diffusion.models.diffusers_adapter.quantization_utils import ( convert_diffusers_quantization_config, ensure_supported_diffusers_quantization, ) +from vllm_omni.diffusion.offloader.config import ( + OffloadStrategy, + resolve_offload_strategy, +) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniPromptType, OmniTextPrompt @@ -148,9 +152,10 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin): self._accept_call_kwargs = set(inspect.signature(self._pipeline.__call__).parameters.keys()) # CPU offloading - if self.od_config.enable_layerwise_offload: + strategy = resolve_offload_strategy(self.od_config) + if strategy is OffloadStrategy.LAYER_WISE: self._pipeline.enable_sequential_cpu_offload() - elif self.od_config.enable_cpu_offload: + elif strategy is OffloadStrategy.MODEL_LEVEL: self._pipeline.enable_model_cpu_offload() # VAE slicing and tiling: try-catch because not all models have VAE diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py index 3ebda6b6e..2b085b91e 100644 --- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py +++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py @@ -54,6 +54,10 @@ from vllm_omni.diffusion.models.dreamzero.utils import ( DEFAULT_SIGMA_SHIFT, ) from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import FlowUniPCMultistepScheduler +from vllm_omni.diffusion.offloader.config import ( + OffloadStrategy, + resolve_offload_strategy, +) from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.experimental.ar_diffusion.capability import ( @@ -439,9 +443,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): else: self.vae = DistributedAutoencoderKLWan() self.vae.init_distributed() - if not ( - getattr(od_config, "enable_cpu_offload", False) or getattr(od_config, "enable_layerwise_offload", False) - ): + if resolve_offload_strategy(od_config) not in (OffloadStrategy.MODEL_LEVEL, OffloadStrategy.LAYER_WISE): self.vae = self.vae.to(device=get_local_device(), dtype=od_config.dtype) self.register_buffer( "vae_latents_mean", diff --git a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py index 139b5275a..0d737d437 100644 --- a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py +++ b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import inspect import json import logging @@ -34,6 +34,10 @@ from vllm_omni.diffusion.models.flux2 import Flux2Transformer2DModel from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery from vllm_omni.diffusion.models.mistral_encoder import MistralEncoderModel from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.offloader.config import ( + OffloadStrategy, + resolve_offload_strategy, +) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch @@ -1160,7 +1164,7 @@ class Flux2Pipeline( text_encoder_quant_config = _resolve_component_quant_config(self.od_config.quantization_config, "text_encoder") transformer_quant_config = _resolve_component_quant_config(self.od_config.quantization_config, "transformer") if ( - self.od_config.enable_cpu_offload + resolve_offload_strategy(self.od_config) is OffloadStrategy.MODEL_LEVEL and text_encoder_quant_config is not None and transformer_quant_config is None ): diff --git a/vllm_omni/diffusion/models/lance/pipeline_lance.py b/vllm_omni/diffusion/models/lance/pipeline_lance.py index f33a35b96..adb8fda62 100644 --- a/vllm_omni/diffusion/models/lance/pipeline_lance.py +++ b/vllm_omni/diffusion/models/lance/pipeline_lance.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """LancePipeline — Lance (ByteDance) packaged for the vLLM-Omni diffusion engine. Lance is BAGEL-lineage (Qwen2-MoT unified AR+diffusion), so the transformer @@ -59,6 +59,10 @@ from vllm_omni.diffusion.models.bagel.pipeline_bagel import ( BagelPipeline, add_special_tokens, ) +from vllm_omni.diffusion.offloader.config import ( + OffloadStrategy, + resolve_offload_strategy, +) from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) @@ -401,7 +405,9 @@ class LancePipeline(BagelPipeline): ) ) - if quant_config is None and not (od_config.enable_layerwise_offload or od_config.parallel_config.use_hsdp): + if quant_config is None and not ( + resolve_offload_strategy(od_config) is OffloadStrategy.LAYER_WISE or od_config.parallel_config.use_hsdp + ): self.to(self.device) self.setup_diffusion_pipeline_profiler( enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler diff --git a/vllm_omni/diffusion/models/lingbot_world/pipeline.py b/vllm_omni/diffusion/models/lingbot_world/pipeline.py index b236e44b5..215e9b2bb 100644 --- a/vllm_omni/diffusion/models/lingbot_world/pipeline.py +++ b/vllm_omni/diffusion/models/lingbot_world/pipeline.py @@ -67,6 +67,10 @@ from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler from vllm_omni.diffusion.models.utils import _load_json from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import load_transformer_config, retrieve_latents +from vllm_omni.diffusion.offloader.config import ( + OffloadStrategy, + resolve_offload_strategy, +) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch @@ -611,8 +615,9 @@ class LingBotWorldCausalDMDPipeline( dtype = getattr(od_config, "dtype", torch.bfloat16) model = od_config.model local_files_only = os.path.exists(model) - managed_component_placement = bool( - getattr(od_config, "enable_cpu_offload", False) or getattr(od_config, "enable_layerwise_offload", False) + managed_component_placement = resolve_offload_strategy(od_config) in ( + OffloadStrategy.MODEL_LEVEL, + OffloadStrategy.LAYER_WISE, ) # Standard components use from_pretrained; the custom transformer uses the loader. diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py index 718756fe4..8b7d79e64 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_components.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_components.py @@ -29,6 +29,10 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import Dis from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders +from vllm_omni.diffusion.offloader.config import ( + OffloadStrategy, + resolve_offload_strategy, +) from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan from vllm_omni.transformers_utils.repo_utils import hf_api @@ -621,11 +625,10 @@ def _load_ltx25_native_diffusion_decoder( def _place_aux_components(pipeline: Any) -> None: parallel_config = getattr(pipeline.od_config, "parallel_config", None) - use_managed_placement = bool( - getattr(pipeline.od_config, "enable_cpu_offload", False) - or getattr(pipeline.od_config, "enable_layerwise_offload", False) - or getattr(parallel_config, "use_hsdp", False) - ) + use_managed_placement = resolve_offload_strategy(pipeline.od_config) in ( + OffloadStrategy.MODEL_LEVEL, + OffloadStrategy.LAYER_WISE, + ) or bool(getattr(parallel_config, "use_hsdp", False)) if use_managed_placement: return diff --git a/vllm_omni/diffusion/models/magi2/pipeline_magi2.py b/vllm_omni/diffusion/models/magi2/pipeline_magi2.py index 52b95fa4c..dc6d422a7 100644 --- a/vllm_omni/diffusion/models/magi2/pipeline_magi2.py +++ b/vllm_omni/diffusion/models/magi2/pipeline_magi2.py @@ -44,6 +44,12 @@ from vllm_omni.diffusion.models.interface import ( ) from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.offloader import OffloadPlan, PinnedModuleStager +from vllm_omni.diffusion.offloader.config import ( + DIT_COMPONENT, + OffloadStrategy, + offload_streams_blocks, + resolve_offload, +) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( DiffusionPipelineProfilerMixin, ) @@ -298,18 +304,17 @@ def _validate_native_topology(od_config: OmniDiffusionConfig) -> None: raise ValueError(f"MAGI-2 tensor_parallel_size={tp_size} does not divide: " + ", ".join(invalid_tp_dimensions)) configured_world_size = dp_size * cfg_size * tp_size * sp_size - cpu_offload = bool(od_config.enable_cpu_offload) - layerwise_offload = bool(od_config.enable_layerwise_offload) - distributed_offload = bool(od_config.enable_distributed_layerwise_offload) - if cpu_offload and not layerwise_offload: + resolved_offload = resolve_offload(od_config) + strategy = resolved_offload.strategy + layerwise_offload = strategy is OffloadStrategy.LAYER_WISE + distributed_offload = strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE + if strategy is OffloadStrategy.MODEL_LEVEL: raise ValueError( "MAGI-2 already stages its auxiliary components from CPU, while " "the complete Preview transformer cannot fit on one qualified GPU. " - "Combine --enable-cpu-offload with --enable-layerwise-offload, or " - "use --enable-layerwise-offload alone." + "Use layer offload instead: diffusion_offload_config={'mode': 'layer', " + "'components': ['dit']}, or the legacy --enable-layerwise-offload." ) - if layerwise_offload and distributed_offload: - raise ValueError("MAGI-2 ordinary and distributed layerwise offload are mutually exclusive") if layerwise_offload and configured_world_size != 1: raise ValueError( "MAGI-2 ordinary layerwise offload is a single-worker path; " @@ -330,7 +335,6 @@ def _validate_native_topology(od_config: OmniDiffusionConfig) -> None: f"expected vae_patch_parallel_size=1 or {configured_world_size}, got {vae_pp_size}." ) - dlo_allgather = bool(getattr(od_config, "dlo_use_allgather", True)) if cfg_size > 1 and dp_size > 1: raise ValueError("MAGI-2 CFG parallelism is not yet combined with DLO data parallelism") if distributed_offload and getattr(parallel, "use_hsdp", False): @@ -339,7 +343,7 @@ def _validate_native_topology(od_config: OmniDiffusionConfig) -> None: raise ValueError("MAGI-2 data parallelism currently requires distributed layerwise offload") if dp_size > 1 and tp_size > 1: raise ValueError("MAGI-2 DLO data-parallel replicas currently require tensor_parallel_size=1") - if distributed_offload and dlo_allgather: + if distributed_offload and resolved_offload.uses_allgather(DIT_COMPONENT): if dp_size <= 1: raise ValueError( "MAGI-2 DLO AllGather requires data_parallel_size > 1. SP ranks " @@ -574,9 +578,7 @@ class Magi2Pipeline( ) self._is_output_rank = self._parallel_group.rank == 0 self._offload_aux_after_use = True - self._transformer_is_layerwise_offloaded = bool( - od_config.enable_layerwise_offload or od_config.enable_distributed_layerwise_offload - ) + self._transformer_is_layerwise_offloaded = offload_streams_blocks(od_config) self._transformer_is_hsdp = bool(getattr(od_config.parallel_config, "use_hsdp", False)) self._distributed_video_decode = int(od_config.parallel_config.vae_patch_parallel_size) > 1 @@ -585,8 +587,10 @@ class Magi2Pipeline( from .modeling_magi2 import Magi2PreviewTransformer MAGI2_PREVIEW_CONFIG.validate() - mmap_dlo = bool( - od_config.enable_distributed_layerwise_offload and getattr(od_config, "dlo_use_allgather", True) + resolved_offload = resolve_offload(od_config) + mmap_dlo = ( + resolved_offload.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE + and resolved_offload.uses_allgather(DIT_COMPONENT) ) if mmap_dlo: # AllGather DLO binds checkpoint tensors as mmap views and copies diff --git a/vllm_omni/diffusion/models/minimax_h3/fasth3.py b/vllm_omni/diffusion/models/minimax_h3/fasth3.py index dcc513028..9f04c204e 100644 --- a/vllm_omni/diffusion/models/minimax_h3/fasth3.py +++ b/vllm_omni/diffusion/models/minimax_h3/fasth3.py @@ -42,6 +42,7 @@ import torch from safetensors import safe_open from vllm.logger import init_logger +from vllm_omni.diffusion.offloader.config import offload_enabled, resolve_offload_strategy from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule from vllm_omni.errors import OmniClientError from vllm_omni.platforms import current_omni_platform @@ -608,19 +609,14 @@ class FastH3WeightFusion: """Hold a starting server to the ladder this student was trained on.""" if partition == "ref2va": raise ValueError("FastH3 preview v1 distills T2VA only, so it cannot serve a Ref2VA partition") - offloads = [ - flag - for flag in ("enable_cpu_offload", "enable_layerwise_offload", "enable_distributed_layerwise_offload") - if getattr(od_config, flag, False) - ] - if offloads: + if offload_enabled(od_config): # A host-weight plan installs the transformer without going through # load_weights(), which is where the fusion and its completeness # check live. Serving base H3 weights under a four-step schedule # would otherwise degrade output with nothing to signal it. raise ValueError( - f"FastH3 is fused while the checkpoint streams in, so it cannot be combined with " - f"{sorted(offloads)}. Serve it without offload." + "FastH3 is fused while the checkpoint streams in, so it cannot be combined with " + f"{resolve_offload_strategy(od_config).value} offload. Serve it without offload." ) if self.requires_vsa: backend = _resolve_dit_attention_backend(od_config) diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 98722b932..a48fdc385 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -53,7 +53,9 @@ from vllm_omni.diffusion.offloader.config import ( DIT_COMPONENT, TEXT_ENCODER_COMPONENT, OffloadStrategy, + offload_streams_blocks, resolve_offload, + resolve_offload_strategy, should_offload_component, ) from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery @@ -1051,8 +1053,8 @@ class MiniMaxH3Pipeline( self.text_encoder_group = None self.text_encoder = None self._encoder_modules = [] - legacy_manual_components = getattr(od_config, "diffusion_offload_config", None) is None and bool( - od_config.enable_layerwise_offload or getattr(od_config, "enable_distributed_layerwise_offload", False) + legacy_manual_components = getattr(od_config, "diffusion_offload_config", None) is None and ( + offload_streams_blocks(od_config) ) # Preserve the legacy MiniMax-H3 low-residency path. The compact API # deliberately limits explicit component selection to dit/text_encoder, @@ -1078,7 +1080,7 @@ class MiniMaxH3Pipeline( self._dlo_component_cache = None offloads_text_encoder = should_offload_component(od_config, TEXT_ENCODER_COMPONENT) needs_component_cache = legacy_manual_components or offloads_text_encoder - if getattr(od_config, "enable_distributed_layerwise_offload", False) and needs_component_cache: + if resolve_offload_strategy(od_config) is OffloadStrategy.DISTRIBUTED_LAYER_WISE and needs_component_cache: self._dlo_component_cache = BoundedAllocatorCache(self.device) if legacy_manual_components: _register_dlo_component_cache( @@ -1437,10 +1439,8 @@ class MiniMaxH3Pipeline( if od_config is None: return False if getattr(od_config, "diffusion_offload_config", None) is None: - return bool( - getattr(od_config, "enable_layerwise_offload", False) - or getattr(od_config, "enable_distributed_layerwise_offload", False) - ) + # The compatibility topology stages every component it can. + return offload_streams_blocks(od_config) return component is getattr(self, "text_encoder", None) and should_offload_component( od_config, TEXT_ENCODER_COMPONENT ) diff --git a/vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py b/vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py index 6c2d740d4..028ea0c9c 100644 --- a/vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py +++ b/vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py @@ -48,6 +48,11 @@ from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_p from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.models.utils import _load_json +from vllm_omni.diffusion.offloader.config import ( + OffloadStrategy, + offload_enabled, + resolve_offload_strategy, +) from vllm_omni.diffusion.postprocess import interpolate_video_tensor from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import resolve_video_num_frames @@ -280,18 +285,13 @@ def _validate_cache_offload_parallelism(od_config: OmniDiffusionConfig) -> None: f"Cache backend {cache_backend!r} is not supported by the native SANA-Video pipeline; " "use 'cache_dit' or 'none'." ) - if cache_backend == "cache_dit" and od_config.enable_distributed_layerwise_offload: + if cache_backend == "cache_dit" and resolve_offload_strategy(od_config) is OffloadStrategy.DISTRIBUTED_LAYER_WISE: raise NotImplementedError( "SANA-Video does not support Cache-DiT with distributed layerwise offload; " "per-rank cache skips desynchronize the weight AllGather." ) - offload_enabled = ( - od_config.enable_cpu_offload - or od_config.enable_layerwise_offload - or od_config.enable_distributed_layerwise_offload - ) - if cache_backend == "none" and not offload_enabled: + if cache_backend == "none" and not offload_enabled(od_config): return parallel_config = od_config.parallel_config diff --git a/vllm_omni/diffusion/models/sana_wm/pipeline_sana_wm.py b/vllm_omni/diffusion/models/sana_wm/pipeline_sana_wm.py index 9c2843040..c8381e9cd 100644 --- a/vllm_omni/diffusion/models/sana_wm/pipeline_sana_wm.py +++ b/vllm_omni/diffusion/models/sana_wm/pipeline_sana_wm.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Sana-WM pipeline integration. This module wires the registry-visible surface, release-layout validation, and @@ -50,6 +50,10 @@ from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import ( SanaWmTransformer3DModel, ) from vllm_omni.diffusion.models.schedulers import FlowMatchEulerDiscreteScheduler +from vllm_omni.diffusion.offloader.config import ( + OffloadStrategy, + resolve_offload_strategy, +) from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( DiffusionPipelineProfilerMixin, @@ -389,11 +393,10 @@ class SanaWmPipeline( components to still be on CPU when they take over. """ parallel_config = getattr(self.od_config, "parallel_config", None) - if ( - getattr(self.od_config, "enable_cpu_offload", False) - or getattr(self.od_config, "enable_layerwise_offload", False) - or getattr(parallel_config, "use_hsdp", False) - ): + if resolve_offload_strategy(self.od_config) in ( + OffloadStrategy.MODEL_LEVEL, + OffloadStrategy.LAYER_WISE, + ) or getattr(parallel_config, "use_hsdp", False): return modules = ModuleDiscovery.discover(self) for module in (*modules.encoders, *modules.vaes, *modules.resident_modules): diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py index 3cbd727a7..dc96575b8 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py @@ -46,7 +46,12 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import ( create_s2v_transformer_from_config, ) -from vllm_omni.diffusion.offloader.config import DIT_COMPONENT, selected_offload_components +from vllm_omni.diffusion.offloader.config import ( + DIT_COMPONENT, + OffloadStrategy, + resolve_offload_strategy, + selected_offload_components, +) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.chunked_video import ChunkedVideoMP4Session @@ -659,7 +664,10 @@ class Wan22S2VPipeline( t5_checkpoint = os.path.join(model_path, "models_t5_umt5-xxl-enc-bf16.pth") self.text_encoder = UMT5EncoderModel(_WAN_UMT5_CONFIG) - _cpu_offload = self.od_config.enable_cpu_offload or self.od_config.enable_layerwise_offload + _cpu_offload = resolve_offload_strategy(self.od_config) in ( + OffloadStrategy.MODEL_LEVEL, + OffloadStrategy.LAYER_WISE, + ) self.text_encoder = _load_wan_t5_as_umt5(self.text_encoder, t5_checkpoint, dtype=dtype) if not _cpu_offload: self.text_encoder = self.text_encoder.to(self.device) @@ -1108,7 +1116,9 @@ class Wan22S2VPipeline( def _should_release_dit_before_decode(self) -> bool: if getattr(self.od_config.parallel_config, "use_hsdp", False): return False - return self.od_config.enable_cpu_offload and DIT_COMPONENT in selected_offload_components(self.od_config) + return resolve_offload_strategy(self.od_config) is OffloadStrategy.MODEL_LEVEL and ( + DIT_COMPONENT in selected_offload_components(self.od_config) + ) def forward( self, diff --git a/vllm_omni/diffusion/offloader/config.py b/vllm_omni/diffusion/offloader/config.py index 46f5e95f3..d6044530d 100644 --- a/vllm_omni/diffusion/offloader/config.py +++ b/vllm_omni/diffusion/offloader/config.py @@ -356,6 +356,19 @@ def resolve_offload_strategy(config: Any) -> OffloadStrategy: return resolve_offload(config).strategy +def offload_enabled(config: Any) -> bool: + """Return whether any CPU-offload policy is active.""" + return resolve_offload(config).strategy is not OffloadStrategy.NONE + + +def offload_streams_blocks(config: Any) -> bool: + """Return whether an active policy streams weights block by block.""" + return resolve_offload(config).strategy in ( + OffloadStrategy.LAYER_WISE, + OffloadStrategy.DISTRIBUTED_LAYER_WISE, + ) + + def materialize_legacy_offload_flags(config: Any) -> OffloadStrategy: """Keep existing strategy readers working after resolving the compact API.""" resolved = resolve_offload(config) diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 50e801ae7..188228b15 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -54,7 +54,13 @@ from vllm_omni.diffusion.models.interface import ( supports_step_execution, ) from vllm_omni.diffusion.offloader import enable_offload_backend -from vllm_omni.diffusion.offloader.config import TEXT_ENCODER_COMPONENT, resolve_offload +from vllm_omni.diffusion.offloader.config import ( + TEXT_ENCODER_COMPONENT, + OffloadStrategy, + offload_enabled, + resolve_offload, + resolve_offload_strategy, +) from vllm_omni.diffusion.postprocess.device_reduction import prepare_diffusion_media_for_transport from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION from vllm_omni.diffusion.request import OmniDiffusionRequest @@ -321,13 +327,7 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin): device=self.device, ) - load_device = ( - "cpu" - if self.od_config.enable_cpu_offload - or self.od_config.enable_layerwise_offload - or getattr(self.od_config, "enable_distributed_layerwise_offload", False) - else str(self.device) - ) + load_device = "cpu" if offload_enabled(self.od_config) else str(self.device) def get_memory_context() -> AbstractContextManager[Any]: if memory_pool_context_fn is not None: @@ -783,7 +783,7 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin): # better perf. HSDP2's fully_shard pre-forward hooks need tensor version # counters, which inference tensors do not track. use_hsdp = od_config.parallel_config.use_hsdp - use_distributed_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False) + use_distributed_offload = resolve_offload_strategy(self.od_config) is OffloadStrategy.DISTRIBUTED_LAYER_WISE grad_context = torch.no_grad() if (use_hsdp or use_distributed_offload) else torch.inference_mode() with grad_context: for req in reqs: