diff --git a/docs/design/feature/offloader/README.md b/docs/design/feature/offloader/README.md index 7b767f41f..96cba8f60 100644 --- a/docs/design/feature/offloader/README.md +++ b/docs/design/feature/offloader/README.md @@ -53,6 +53,11 @@ must use distinct hook names and remove only hooks they own. ## Discovery and topology +`resolve_offload_plan(pipeline, config)` is the single topology entry point. +Every backend calls it once when `enable()` starts and then consumes the +resolved components, block stacks, resident heads, and staged-residency flags +it returns. Backends do not re-read model declarations. + Pipeline component discovery prefers `SupportsComponentDiscovery` declarations: - `_dit_modules`; @@ -61,16 +66,26 @@ Pipeline component discovery prefers `SupportsComponentDiscovery` declarations: - `_resident_modules`. Dotted paths are supported. Legacy pipelines may use the fallback scan of -well-known attribute names, but new integrations should declare components -explicitly. +well-known attribute names; it now warns once per pipeline class, and new +integrations must declare components explicitly. -Both layerwise backends first consume pipeline `OffloadPlan` metadata. A -plan's `block_attrs` maps each DiT path to its ordered block containers, while -`encoder_block_attrs` declares streamable encoder stacks. DiTs absent from the -plan fall back to `_layerwise_offload_blocks_attrs` (including the deprecated -singular-name compatibility path). Discovery metadata describes structure -only; the backend remains responsible for transfer, synchronization, and -storage ownership. +Block topology comes from the pipeline `OffloadPlan` first: `block_attrs` maps +each DiT path to its ordered block containers, `encoder_block_attrs` declares +streamable encoder stacks, `on_demand_component_paths` marks pipeline-managed +residency, `resident_dit_paths` marks the DiTs that may hold resident layers, +and `encoder_dlo_weight_replication` marks the encoders whose loader-produced +weights are safe for AllGather. DiTs absent from the plan fall back to +`_layerwise_offload_blocks_attrs` (including the deprecated singular-name +compatibility path). + +The resolver is pure: it moves no tensor, installs no hook, writes no module +attribute, and reads no process group — multi-rank facts come from +`OffloadConfig`. It owns topology validation, so a configuration the model +cannot serve fails before the first component is placed. Explicit component +selection turns every declaration problem into an error, while the +compatibility topology (no component selector) keeps the historical +warn-and-skip behavior. Metadata still describes structure only; the backend +remains responsible for transfer, synchronization, and storage ownership. ## Cross-strategy invariants diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index 05d7d93b5..c24639434 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -52,6 +52,7 @@ from vllm_omni.diffusion.offloader.offload_plan import ( OffloadPlan, get_offload_plan, ) +from vllm_omni.diffusion.offloader.plan_resolver import resolve_offload_plan from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state from vllm_omni.host_weight_runtime import MappedHostRegion from vllm_omni.platforms import current_omni_platform @@ -2813,18 +2814,21 @@ class TestDistributedComponentSelection: def test_encoder_allgather_rejects_stub_rank_before_block_discovery(self): """Every rank must reject an unsafe encoder group, including stub ranks.""" - backend = DistributedLayerwiseOffloadBackend( - OffloadConfig( - strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, - pin_cpu_memory=False, - dp_size=2, - components=frozenset({"text_encoder"}), - dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"}, - ), - torch.device("cpu"), + + class StubEncoderPipeline(nn.Module): + _offload_plan = OffloadPlan(encoder_block_attrs={"text_encoder": ("missing.blocks",)}) + + def __init__(self): + super().__init__() + self.text_encoder = nn.Module() + + config = OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dp_size=2, + components=frozenset({"text_encoder"}), + dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"}, ) - backend.dp_group = object() - stub_plan = OffloadPlan(encoder_block_attrs={"text_encoder": ("missing.blocks",)}) with pytest.raises(ValueError, match="not declared replicated"): - backend._try_layerwise_offload_encoder(nn.Module(), "text_encoder", stub_plan) + resolve_offload_plan(StubEncoderPipeline(), config) diff --git a/tests/diffusion/offloader/test_layerwise_backend.py b/tests/diffusion/offloader/test_layerwise_backend.py index 2dd8804cf..099e93f79 100644 --- a/tests/diffusion/offloader/test_layerwise_backend.py +++ b/tests/diffusion/offloader/test_layerwise_backend.py @@ -488,7 +488,8 @@ class TestLayerwiseComponentSelection: with pytest.raises(ValueError, match="Selected text encoder 'text_encoder_2' requires"): backend.enable(pipeline) - assert not pipeline.text_encoder._omni_layerwise_enabled + # The plan resolver rejects the topology before any encoder is hooked. + assert not getattr(pipeline.text_encoder, "_omni_layerwise_enabled", False) def test_default_selection_preserves_unplanned_auxiliaries(self, patched_offload_runtime): pipeline = _LegacyComponentPipeline() diff --git a/tests/diffusion/offloader/test_plan_resolver.py b/tests/diffusion/offloader/test_plan_resolver.py new file mode 100644 index 000000000..483ce4060 --- /dev/null +++ b/tests/diffusion/offloader/test_plan_resolver.py @@ -0,0 +1,458 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Unit tests for the shared offload plan resolver.""" + +from typing import ClassVar + +import pytest +import torch +from torch import nn + +from tests.diffusion.offloader.helpers import ( + _DummyBlock, + _PlainEncoder, + _SingleBlockModel, + _StagedEncoder, + _StagedVAE, + patch_offload_runtime, +) +from vllm_omni.diffusion.offloader import layerwise_backend as layerwise_backend_module +from vllm_omni.diffusion.offloader import plan_resolver as plan_resolver_module +from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy +from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend +from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan +from vllm_omni.diffusion.offloader.plan_resolver import resolve_offload_plan + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + +CPU = torch.device("cpu") + + +@pytest.fixture +def patched_offload_runtime(monkeypatch): + patch_offload_runtime(monkeypatch, layerwise_backend_module.current_omni_platform) + + +def _config( + components: set[str] | None = None, + strategy: OffloadStrategy = OffloadStrategy.LAYER_WISE, + **kwargs, +) -> OffloadConfig: + return OffloadConfig( + strategy=strategy, + pin_cpu_memory=False, + components=None if components is None else frozenset(components), + **kwargs, + ) + + +def _stack_of(component): + assert len(component.stacks) == 1 + return component.stacks[0] + + +class _WanStylePipeline(nn.Module): + """Streamable DiT plus a plan-declared, replicated text encoder.""" + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] + _offload_plan = OffloadPlan( + encoder_component_types={"text_encoder": "text_encoder"}, + encoder_block_attrs={"text_encoder": ("encoder.block",)}, + encoder_dlo_weight_replication=frozenset({"text_encoder"}), + ) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=4) + self.text_encoder = _PlainEncoder() + self.vae = _StagedVAE() + + +class _StagedComponentPipeline(nn.Module): + """MiniMax-H3 shape: two encoder stacks plus pipeline-staged components.""" + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] + _offload_plan = OffloadPlan( + encoder_component_types={"text_encoder": "text_encoder"}, + encoder_block_attrs={"text_encoder": ("vision.blocks", "text_model.layers")}, + on_demand_component_paths=frozenset({"text_encoder", "vae"}), + resident_dit_paths=frozenset({"transformer"}), + ) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=4) + self.text_encoder = _StagedEncoder() + self.vae = _StagedVAE() + + +class _LegacyScanPipeline(nn.Module): + """No SupportsComponentDiscovery declaration and no OffloadPlan.""" + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=3) + self.text_encoder = _PlainEncoder() + + +class TestProducerPrecedence: + def test_plan_block_attrs_win_over_class_attributes(self): + class TwoStackDiT(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([_DummyBlock() for _ in range(2)]) + self.alt_blocks = nn.ModuleList([_DummyBlock() for _ in range(3)]) + + class Pipeline(nn.Module): + _offload_plan = OffloadPlan(block_attrs={"transformer": ("alt_blocks",)}) + + def __init__(self): + super().__init__() + self.transformer = TwoStackDiT() + + pipeline = Pipeline() + resolved = resolve_offload_plan(pipeline, _config({"dit"})) + + stack = _stack_of(resolved.dits[0]) + assert stack.attrs == ("alt_blocks",) + assert list(stack.blocks) == list(pipeline.transformer.alt_blocks) + + def test_class_attributes_are_used_without_a_plan(self): + pipeline = _LegacyScanPipeline() + resolved = resolve_offload_plan(pipeline, _config({"dit"})) + + assert list(_stack_of(resolved.dits[0]).blocks) == list(pipeline.transformer.blocks) + + def test_declared_component_type_wins_over_leaf_name(self): + class Pipeline(nn.Module): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["mllm"] + _vae_modules: ClassVar[list[str]] = [] + _resident_modules: ClassVar[list[str]] = [] + _offload_plan = OffloadPlan( + encoder_component_types={"mllm": "text_encoder"}, + encoder_block_attrs={"mllm": ("encoder.block",)}, + ) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + self.mllm = _PlainEncoder() + + resolved = resolve_offload_plan(Pipeline(), _config({"text_encoder"})) + + encoder = resolved.encoders[0] + assert encoder.path == "mllm" + assert encoder.selected + assert len(_stack_of(encoder).blocks) == 2 + + def test_component_declaration_wins_over_attribute_scan(self): + class Pipeline(nn.Module): + _dit_modules: ClassVar[list[str]] = ["custom_dit"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = [] + _resident_modules: ClassVar[list[str]] = [] + + def __init__(self): + super().__init__() + self.custom_dit = _SingleBlockModel(num_blocks=2) + # Would be picked up by the deprecated attribute scan. + self.transformer = _SingleBlockModel(num_blocks=2) + + resolved = resolve_offload_plan(Pipeline(), _config({"dit"})) + + assert [component.path for component in resolved.dits] == ["custom_dit"] + + def test_attribute_scan_warns_once_per_pipeline_class(self, monkeypatch): + warnings: list[str] = [] + + class _Recorder: + def warning(self, message, *args): + warnings.append(message % args if args else message) + + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + monkeypatch.setattr(plan_resolver_module, "logger", _Recorder()) + + class UndeclaredPipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + + resolve_offload_plan(UndeclaredPipeline(), _config({"dit"})) + resolve_offload_plan(UndeclaredPipeline(), _config({"dit"})) + + scan_warnings = [message for message in warnings if "SupportsComponentDiscovery" in message] + assert len(scan_warnings) == 1 + assert "deprecated" in scan_warnings[0] + + +class TestResolvedTopology: + def test_encoder_plan_resolves_one_stack_per_declared_path(self): + pipeline = _WanStylePipeline() + resolved = resolve_offload_plan(pipeline, _config({"dit", "text_encoder"})) + + encoder = resolved.encoders[0] + assert encoder.selected and not encoder.on_demand + assert list(_stack_of(encoder).blocks) == list(pipeline.text_encoder.encoder.block) + assert [component.path for component in resolved.vaes] == ["vae"] + + def test_staged_plan_resolves_every_declared_stack(self): + pipeline = _StagedComponentPipeline() + resolved = resolve_offload_plan(pipeline, _config({"text_encoder"})) + + encoder = resolved.encoders[0] + assert encoder.on_demand + assert [len(stack.blocks) for stack in encoder.stacks] == [2, 2] + # An explicit selection never stages a VAE through the legacy plan. + assert not resolved.vaes[0].on_demand + # DiT is unselected, so it keeps no streaming topology. + assert resolved.dits[0].selected is False + assert resolved.dits[0].stacks == () + + def test_omitted_selector_keeps_the_legacy_staged_vae(self): + resolved = resolve_offload_plan(_StagedComponentPipeline(), _config()) + + assert resolved.vaes[0].on_demand + assert resolved.encoders[0].on_demand + assert resolved.dits[0].selected + + def test_on_demand_only_plan_resolves_without_stacks(self): + class Pipeline(nn.Module): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] + _offload_plan = OffloadPlan(on_demand_component_paths=frozenset({"text_encoder", "vae"})) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + self.text_encoder = _StagedEncoder() + self.vae = _StagedVAE() + + resolved = resolve_offload_plan(Pipeline(), _config({"text_encoder"})) + + assert resolved.encoders[0].on_demand + assert resolved.encoders[0].stacks == () + + def test_module_strategy_resolves_selection_without_stacks(self): + resolved = resolve_offload_plan( + _WanStylePipeline(), + _config({"dit", "text_encoder"}, strategy=OffloadStrategy.MODEL_LEVEL), + ) + + assert [component.path for component in resolved.components] == ["transformer", "text_encoder", "vae"] + assert all(component.stacks == () for component in resolved.components) + assert resolved.encoders[0].selected + + def test_module_strategy_leaves_staged_residency_to_the_pipeline(self): + resolved = resolve_offload_plan( + _StagedComponentPipeline(), + _config({"dit", "text_encoder"}, strategy=OffloadStrategy.MODEL_LEVEL), + ) + + assert not any(component.on_demand for component in resolved.components) + + +class TestResidentLayers: + def _dlo_config(self, resident_layers: int, components: set[str] | None = None) -> OffloadConfig: + return _config( + components, + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + dlo_transfers={"dit": "rank-local", "text_encoder": "rank-local"}, + dlo_resident_layers=resident_layers, + ) + + def test_resident_head_splits_the_declared_dit(self): + pipeline = _StagedComponentPipeline() + resolved = resolve_offload_plan(pipeline, self._dlo_config(2, {"dit"})) + + stack = _stack_of(resolved.dits[0]) + assert stack.resident_head == 2 + assert list(stack.resident) == list(pipeline.transformer.blocks)[:2] + assert list(stack.streaming) == list(pipeline.transformer.blocks)[2:] + + def test_resident_count_beyond_the_stack_keeps_every_block(self): + resolved = resolve_offload_plan(_StagedComponentPipeline(), self._dlo_config(9, {"dit"})) + + stack = _stack_of(resolved.dits[0]) + assert stack.resident_head == 4 + assert stack.streaming == () + + def test_resident_head_leaving_one_streaming_block_is_rejected(self): + with pytest.raises(ValueError, match="leaves only one streaming block"): + resolve_offload_plan(_StagedComponentPipeline(), self._dlo_config(3, {"dit"})) + + def test_legacy_single_streaming_block_stays_resident(self): + resolved = resolve_offload_plan(_StagedComponentPipeline(), self._dlo_config(3)) + + stack = _stack_of(resolved.dits[0]) + assert stack.resident_head == 4 + assert stack.streaming == () + + def test_resident_layers_without_a_declared_path_is_rejected(self): + with pytest.raises(ValueError, match="no matching resident_dit_paths"): + resolve_offload_plan(_WanStylePipeline(), self._dlo_config(2, {"dit"})) + + +class TestValidation: + def test_selected_dit_requires_a_discovered_module(self): + pipeline = nn.Module() + pipeline.text_encoder = _PlainEncoder() + + with pytest.raises(ValueError, match="No DiT/transformer modules found"): + resolve_offload_plan(pipeline, _config({"dit"})) + + def test_model_level_selection_requires_both_swap_sides(self): + pipeline = nn.Module() + pipeline.transformer = _SingleBlockModel(num_blocks=2) + + with pytest.raises(ValueError, match="requires an encoder execution stage"): + resolve_offload_plan(pipeline, _config({"dit"}, strategy=OffloadStrategy.MODEL_LEVEL)) + + def test_selected_text_encoder_requires_a_matching_module(self): + pipeline = nn.Module() + pipeline.transformer = _SingleBlockModel(num_blocks=2) + + with pytest.raises(ValueError, match="No text encoder modules found"): + resolve_offload_plan(pipeline, _config({"text_encoder"})) + + def test_selected_encoder_requires_a_streamable_or_staged_plan(self): + pipeline = nn.Module() + pipeline.transformer = _SingleBlockModel(num_blocks=2) + pipeline.text_encoder = _PlainEncoder() + + with pytest.raises(ValueError, match="Selected text encoder 'text_encoder' requires"): + resolve_offload_plan(pipeline, _config({"text_encoder"})) + + def test_staged_component_requires_the_explicit_lifecycle(self): + class Pipeline(nn.Module): + _offload_plan = OffloadPlan(on_demand_component_paths=frozenset({"text_encoder"})) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + self.text_encoder = _PlainEncoder() + + with pytest.raises(ValueError, match="must implement load_to_device"): + resolve_offload_plan(Pipeline(), _config({"text_encoder"})) + + def test_single_block_dit_is_rejected_for_layer_offload(self): + pipeline = nn.Module() + pipeline.transformer = _SingleBlockModel(num_blocks=1) + + with pytest.raises(ValueError, match="requires at least two streamable"): + resolve_offload_plan(pipeline, _config({"dit"})) + + def test_encoder_allgather_requires_declared_replication(self): + config = _config( + {"text_encoder"}, + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + dp_size=2, + dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"}, + ) + + with pytest.raises(ValueError, match="not declared replicated"): + resolve_offload_plan(_StagedComponentPipeline(), config) + + def test_two_dits_cannot_own_the_same_blocks(self): + class Pipeline(nn.Module): + _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = [] + _resident_modules: ClassVar[list[str]] = [] + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + self.transformer_2 = _SingleBlockModel(num_blocks=2) + self.transformer_2.blocks = self.transformer.blocks + + with pytest.raises(ValueError, match="claimed by both"): + resolve_offload_plan(Pipeline(), _config({"dit"})) + + def test_legacy_selection_keeps_unstreamable_components(self): + pipeline = nn.Module() + pipeline.transformer = _SingleBlockModel(num_blocks=1) + pipeline.text_encoder = _PlainEncoder() + + resolved = resolve_offload_plan(pipeline, _config()) + + assert resolved.dits[0].stacks == () + assert resolved.encoders[0].stacks == () + + +class TestPurity: + def _snapshot(self, pipeline: nn.Module): + return { + "devices": [(name, tensor.device) for name, tensor in pipeline.named_parameters()], + "attrs": {name: sorted(vars(module)) for name, module in pipeline.named_modules()}, + } + + def test_resolution_does_not_touch_the_model(self): + pipeline = _StagedComponentPipeline() + before = self._snapshot(pipeline) + + resolve_offload_plan(pipeline, _config({"dit", "text_encoder"})) + + assert self._snapshot(pipeline) == before + assert pipeline.text_encoder.offload_calls == 0 + assert pipeline.text_encoder.to_calls == 0 + assert pipeline.vae.to_calls == 0 + + def test_resolution_is_repeatable(self): + pipeline = _WanStylePipeline() + config = _config({"dit", "text_encoder"}) + + first = resolve_offload_plan(pipeline, config) + second = resolve_offload_plan(pipeline, config) + + assert [component.path for component in first.components] == [component.path for component in second.components] + for left, right in zip(first.components, second.components, strict=True): + assert left.selected == right.selected + assert left.on_demand == right.on_demand + assert [list(stack.blocks) for stack in left.stacks] == [list(stack.blocks) for stack in right.stacks] + + def test_backend_enable_rejects_before_touching_earlier_components(self, patched_offload_runtime): + class Pipeline(nn.Module): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder", "text_encoder_2"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] + _offload_plan = OffloadPlan( + encoder_component_types={ + "text_encoder": "text_encoder", + "text_encoder_2": "text_encoder", + }, + encoder_block_attrs={"text_encoder": ("encoder.block",)}, + ) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + self.text_encoder = _PlainEncoder() + self.text_encoder_2 = _PlainEncoder() + self.vae = _StagedVAE() + + pipeline = Pipeline() + backend = LayerWiseOffloadBackend(_config({"text_encoder"}), CPU) + + with pytest.raises(ValueError, match="Selected text encoder 'text_encoder_2' requires"): + backend.enable(pipeline) + + assert not backend.enabled + assert not getattr(pipeline.text_encoder, "_omni_layerwise_enabled", False) + for block in pipeline.text_encoder.encoder.block: + assert getattr(block, "_hook_registry", None) is None + assert pipeline.vae.to_calls == 0 diff --git a/vllm_omni/diffusion/offloader/__init__.py b/vllm_omni/diffusion/offloader/__init__.py index 0df4039eb..42ba45cc9 100644 --- a/vllm_omni/diffusion/offloader/__init__.py +++ b/vllm_omni/diffusion/offloader/__init__.py @@ -20,6 +20,12 @@ from .distributed_layerwise_backend import ( from .layerwise_backend import LayerWiseOffloadBackend from .module_residency import BoundedAllocatorCache, PinnedModuleStager from .offload_plan import OffloadPlan, get_offload_plan +from .plan_resolver import ( + BlockStack, + ResolvedComponent, + ResolvedOffloadPlan, + resolve_offload_plan, +) from .sequential_backend import ( ModelLevelOffloadBackend, apply_sequential_offload, @@ -42,6 +48,10 @@ __all__ = [ "OffloadConfig", "OffloadPlan", "OffloadStrategy", + "BlockStack", + "ResolvedComponent", + "ResolvedOffloadPlan", + "resolve_offload_plan", "SupportsModelCpuOffload", "LayerWiseOffloadBackend", "DistributedLayerwiseOffloadBackend", diff --git a/vllm_omni/diffusion/offloader/component_utils.py b/vllm_omni/diffusion/offloader/component_utils.py index 1000b29d6..cbc63159d 100644 --- a/vllm_omni/diffusion/offloader/component_utils.py +++ b/vllm_omni/diffusion/offloader/component_utils.py @@ -4,7 +4,7 @@ from __future__ import annotations -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Sequence from itertools import chain from operator import attrgetter from typing import TYPE_CHECKING, Any @@ -14,14 +14,12 @@ from torch import nn from torch.distributed.tensor import DTensor from vllm.logger import init_logger -from .block_discovery import get_blocks_from_dit -from .config import DIT_COMPONENT, TEXT_ENCODER_COMPONENT +from .config import TEXT_ENCODER_COMPONENT from .offload_plan import OffloadPlan from .tensor_utils import set_tensor_storage if TYPE_CHECKING: - from .base import OffloadConfig - from .module_collector import PipelineModules + from .plan_resolver import BlockStack, ResolvedComponent, ResolvedOffloadPlan logger = init_logger(__name__) @@ -76,30 +74,25 @@ def get_encoder_block_groups( def iter_streamable_dits( - modules: PipelineModules, - config: OffloadConfig, + resolved: ResolvedOffloadPlan, device: torch.device, - plan: OffloadPlan | None, -) -> Iterator[tuple[str, nn.Module, list[str], list[nn.Module]]]: - """Yield selected DiTs whose block metadata resolves successfully.""" - if not config.offloads(DIT_COMPONENT): - return - for name, module in zip(modules.dit_names, modules.dits): - logger.info("Applying hooks on %s (%s)", name, module.__class__.__name__) - planned_attrs = None if plan is None else plan.block_attrs.get(name) - block_attrs, blocks = get_blocks_from_dit(module, planned_attrs) - if blocks: - yield name, module, block_attrs, blocks +) -> Iterator[tuple[ResolvedComponent, BlockStack]]: + """Yield selected DiTs that resolved to a streamable block stack.""" + for component in resolved.dits: + if not component.selected: continue - if config.components is not None: - raise ValueError(f"Selected DiT {name!r} has no streamable layerwise-offload blocks") - logger.warning("Target layers (blocks) not found. Skipping offloading on %s (%s)", name, type(module).__name__) - module.to(device) + logger.info("Applying hooks on %s (%s)", component.path, component.module.__class__.__name__) + if not component.stacks: + # The resolver already warned; a legacy selection keeps the + # unstreamable DiT resident instead of failing the run. + component.module.to(device) + continue + yield component, component.stacks[0] def move_non_block_state_to_device( module: nn.Module, - block_groups: list[nn.ModuleList], + block_groups: Sequence[Sequence[nn.Module]], device: torch.device, ) -> None: """Keep component state outside streamed block lists resident on device.""" @@ -120,7 +113,7 @@ def move_non_block_state_to_device( def set_encoder_layerwise_state( module: nn.Module, hooks: list[Any], - block_groups: list[nn.ModuleList], + block_groups: Sequence[Sequence[nn.Module]], ) -> None: """Publish the backend-neutral state used by encoder stage lifecycles.""" module._omni_layerwise_hooks = hooks @@ -163,59 +156,38 @@ def prepare_component( def prepare_pipeline_components( - modules: PipelineModules, - config: OffloadConfig, - plan: OffloadPlan | None, + resolved: ResolvedOffloadPlan, *, device: torch.device, staged_components: list[nn.Module], - enable_encoder_blocks: Callable[[nn.Module, str, OffloadPlan | None, bool], bool], + enable_encoder_blocks: Callable[[ResolvedComponent], bool], ) -> None: """Apply the shared encoder/VAE/resident placement policy.""" - if config.components is not None and config.offloads(TEXT_ENCODER_COMPONENT): - selected_encoder_names = [name for name in modules.encoder_names if config.offloads_encoder(name, plan)] - if not selected_encoder_names: - raise ValueError("No text encoder modules found for selected text_encoder offload") - - if plan is not None: - for encoder, name in zip(modules.encoders, modules.encoder_names): - if config.should_offload_encoder(name, plan) and name in plan.on_demand_component_paths: - validate_on_demand_component(encoder, name) - - for encoder, name in zip(modules.encoders, modules.encoder_names): - selected = config.should_offload_encoder(name, plan) - stage_on_demand = bool(selected and plan is not None and name in plan.on_demand_component_paths) - blockwise = selected and enable_encoder_blocks(encoder, name, plan, stage_on_demand) - if stage_on_demand and not blockwise and config.uses_allgather(TEXT_ENCODER_COMPONENT): - raise ValueError( - f"Text encoder {name!r} cannot use AllGather without a model-declared streamable block plan" - ) - if selected and config.components is not None and not (blockwise or stage_on_demand): - raise ValueError(f"Selected text encoder {name!r} requires a model-declared streamable or on-demand plan") + for component in resolved.encoders: + blockwise = bool(component.stacks) and enable_encoder_blocks(component) prepare_component( - encoder, - name, + component.module, + component.path, device=device, - stage_on_demand=stage_on_demand, + stage_on_demand=component.on_demand, blockwise=blockwise, staged_components=staged_components, ) - for vae, name in zip(modules.vaes, modules.vae_names): - legacy_staged = config.components is None and plan is not None and name in plan.on_demand_component_paths + for component in resolved.vaes: prepare_component( - vae, - name, + component.module, + component.path, device=device, - stage_on_demand=legacy_staged, + stage_on_demand=component.on_demand, blockwise=False, staged_components=staged_components, ) - for name, module in zip(modules.resident_names, modules.resident_modules): - module.to(device) - logger.debug("Moved resident module %s to %s", name, device) + for component in resolved.residents: + component.module.to(device) + logger.debug("Moved resident module %s to %s", component.path, device) - if not config.offloads(DIT_COMPONENT): - for dit in modules.dits: - dit.to(device) + for component in resolved.dits: + if not component.selected: + component.module.to(device) diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index e626d87a5..299ec3d3a 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -21,6 +21,7 @@ from __future__ import annotations import threading import time import weakref +from collections.abc import Sequence from itertools import chain from typing import Any @@ -39,7 +40,6 @@ from vllm_omni.platforms import current_omni_platform from .base import OffloadBackend, OffloadConfig, run_cleanup_steps from .component_utils import ( clear_encoder_layerwise_state, - get_encoder_block_groups, iter_streamable_dits, move_non_block_state_to_device, prepare_component, @@ -53,11 +53,8 @@ from .host_registration import ( HostRegistrationError, register_host_mappings, ) -from .module_collector import ModuleDiscovery -from .offload_plan import ( - OffloadPlan, - get_offload_plan, -) +from .offload_plan import OffloadPlan +from .plan_resolver import ResolvedComponent, resolve_offload_plan from .tensor_utils import ( clear_block_storage, clear_tensor_storage, @@ -1437,7 +1434,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): def _install_hook_group( self, - blocks: list[nn.Module] | nn.ModuleList, + blocks: Sequence[nn.Module] | nn.ModuleList, component: str, *, use_dit_mmap: bool = False, @@ -1483,48 +1480,27 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): hooks[1]._is_group_first = True return hooks - def _try_layerwise_offload_encoder( - self, - module: nn.Module, - name: str, - plan: OffloadPlan | None, - stage_on_demand: bool = False, - ) -> bool: - """Apply DLO hooks to a plan-declared text encoder. + def _try_layerwise_offload_encoder(self, component: ResolvedComponent) -> bool: + """Apply DLO hooks to a text encoder's resolved block stacks. Rank-local transfer is always safe because it preserves the tensors produced by the model loader. Multi-rank AllGather is enabled only when the model declares that those tensors are replicated across the - selected DLO group; encoder TP groups contain different shards and - are deliberately not used for this transport. + selected DLO group; the resolver rejects every other layout before + this method installs anything. """ - if plan is None: - return False - - group_size = self._component_transport(TEXT_ENCODER_COMPONENT)[1] - if group_size > 1 and name not in plan.encoder_dlo_weight_replication: - raise ValueError( - f"Text encoder {name!r} cannot use DLO AllGather across the DiT offload group: " - "its loader-produced weights are not declared replicated across that group. " - "Set layer_options.text_encoder.weight_transfer='rank-local' in diffusion_offload_config " - "for encoder-TP or rank-specific layouts." - ) - block_groups = get_encoder_block_groups( - module, - name, - plan, - strict=self.config.components is not None or group_size > 1, - ) + block_groups = [stack.blocks for stack in component.stacks] if not block_groups: return False + module = component.module encoder_hooks: list[DistributedLayerwiseOffloadHook] = [] for blocks in block_groups: encoder_hooks.extend(self._install_hook_group(blocks, TEXT_ENCODER_COMPONENT)) # Track the module before placement, which may fail, so outer rollback # also removes its partially-installed block hooks and marker state. self._encoder_modules.append(module) - if not stage_on_demand: + if not component.on_demand: move_non_block_state_to_device(module, block_groups, self.device) set_encoder_layerwise_state( module, @@ -1534,10 +1510,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): logger.info( "Enabled %s DLO transfer for text encoder %s (%d blocks across %d stacks, group_size=%d)", self.config.transfer_for(TEXT_ENCODER_COMPONENT).value, - name, + component.path, sum(len(blocks) for blocks in block_groups), len(block_groups), - group_size, + self._component_transport(TEXT_ENCODER_COMPONENT)[1], ) return True @@ -1700,31 +1676,15 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): if self.dp_group is None and self._has_multirank_allgather(): self._init_dp_group() - modules = ModuleDiscovery.discover(pipeline) - if not modules.dits and self.config.offloads(DIT_COMPONENT): + # One resolver call validates the whole topology; an explicit selection + # has already failed by the time any of the checks below run. + resolved = resolve_offload_plan(pipeline, self.config) + if not resolved.dits and self.config.offloads(DIT_COMPONENT): if self.host_weight_plan is not None: raise RuntimeError( "DLO received a loader-owned host-weight plan, but no DiT modules were discovered to consume it" ) - message = "No DiT/transformer modules found for selected DiT offload" - if self.config.components is not None: - raise ValueError(message) - logger.warning(message) - - # Retrieve optional declarative OffloadPlan from the pipeline. - # When present, replaces heuristic block discovery. - plan = get_offload_plan(pipeline) - - if self.config.dlo_resident_layers: - resident_paths = frozenset() if plan is None else plan.resident_dit_paths - if not resident_paths.intersection(modules.dit_names): - message = ( - f"resident_layers={self.config.dlo_resident_layers} was requested, but this model declares " - "no matching resident_dit_paths" - ) - if self.config.components is not None: - raise ValueError(message) - logger.warning("%s; all blocks will be streamed.", message) + logger.warning("No DiT/transformer modules found for selected DiT offload") # Storage selection belongs to the loader. DLO consumes the exact # prevalidated plan that caused the loader to skip materialization; @@ -1755,7 +1715,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): elif host_weight_plan.backing_kind == "checkpoint_mmap": self._load_weights_via_mmap( pipeline, - modules, + resolved.modules, host_weight_plan, ) else: @@ -1768,10 +1728,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): else: remaining_meta = [ name - for dit_name, dit_module in zip(modules.dit_names, modules.dits) + for component in resolved.dits for name, tensor in chain( - dit_module.named_parameters(), - dit_module.named_buffers(), + component.module.named_parameters(), + component.module.named_buffers(), ) if getattr(tensor, "is_meta", False) ] @@ -1784,66 +1744,46 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): # Apply each selected encoder transfer while keeping explicit VAEs and # unselected components resident. prepare_pipeline_components( - modules, - self.config, - plan, + resolved, device=self.device, staged_components=self._staged_components, enable_encoder_blocks=self._try_layerwise_offload_encoder, ) if self.config.offloads(DIT_COMPONENT): - logger.info("Applying distributed layer-wise offloading on %s", modules.dit_names) + logger.info( + "Applying distributed layer-wise offloading on %s", + [component.path for component in resolved.dits], + ) # Collect all DiT module objects to detect submodules that are # already handled as a separate DiT module (avoids duplicate hooks). - all_dit_modules = set(id(m) for m in modules.dits) + all_dit_modules = set(id(component.module) for component in resolved.dits) # Apply hooks for each DiT module - for dit_name, dit_module, blocks_attr_names, blocks in iter_streamable_dits( - modules, self.config, self.device, plan - ): - resident_count = 0 - if plan is not None and dit_name in plan.resident_dit_paths: - resident_count = min(self.config.dlo_resident_layers, len(blocks)) - if resident_count: - resident_blocks = blocks[:resident_count] - self._resident_blocks.extend(resident_blocks) - blocks = blocks[resident_count:] + for component, stack in iter_streamable_dits(resolved, self.device): + streaming = list(stack.streaming) + if stack.resident_head: + self._resident_blocks.extend(stack.resident) logger.info( "Keeping %d leading blocks resident on %s; streaming %d tail blocks", - resident_count, - dit_name, - len(blocks), + stack.resident_head, + component.path, + len(streaming), ) self._prepare_dit_non_block_modules( - dit_module, - blocks_attr_names, + component.module, + list(stack.attrs), all_dit_modules, - plan, + resolved.declaration, ) - num_blocks = len(blocks) - if num_blocks == 0: - logger.info("All blocks for %s are resident; no streaming hooks required", dit_name) - continue - if num_blocks <= 1: - if self.config.components is not None: - raise ValueError( - f"Selected DiT {dit_name!r} leaves only one streaming block after " - f"resident_layers={resident_count}; choose a resident count that " - "leaves zero or at least two streaming blocks" - ) - logger.warning( - "#Streaming target layers (blocks) <= 1. Keeping the final block resident on %s (%s)", - dit_name, - dit_module.__class__.__name__, - ) - self._resident_blocks.extend(blocks) + if not streaming: + logger.info("All blocks for %s are resident; no streaming hooks required", component.path) continue - self._install_hook_group(blocks, DIT_COMPONENT, use_dit_mmap=True) + self._install_hook_group(streaming, DIT_COMPONENT, use_dit_mmap=True) if self._resident_blocks: self._resident_layer_group = PinnedResidentLayerGroup( self._resident_blocks, diff --git a/vllm_omni/diffusion/offloader/layerwise_backend.py b/vllm_omni/diffusion/offloader/layerwise_backend.py index 9b63f9162..a1e594f8b 100644 --- a/vllm_omni/diffusion/offloader/layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/layerwise_backend.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations +from collections.abc import Sequence from itertools import chain from typing import Any @@ -20,15 +21,13 @@ from .block_discovery import ( ) from .component_utils import ( clear_encoder_layerwise_state, - get_encoder_block_groups, iter_streamable_dits, move_non_block_state_to_device, prepare_pipeline_components, set_encoder_layerwise_state, ) from .config import DIT_COMPONENT -from .module_collector import ModuleDiscovery -from .offload_plan import OffloadPlan, get_offload_plan +from .plan_resolver import ResolvedComponent, resolve_offload_plan from .tensor_utils import ( clear_block_storage, clear_tensor_storage, @@ -279,7 +278,7 @@ def remove_block_hook(module: nn.Module) -> None: def _install_layerwise_hook_group( - blocks: list[nn.Module] | nn.ModuleList, + blocks: Sequence[nn.Module] | nn.ModuleList, device: torch.device, stream: Any, pin_memory: bool, @@ -327,36 +326,29 @@ def _install_layerwise_hook_group( def enable_plan_encoder_layerwise_offload( - module: nn.Module, - name: str, - plan: OffloadPlan | None, + component: ResolvedComponent, *, device: torch.device, stream: current_omni_platform.Stream, pin_memory: bool, - stage_on_demand: bool = False, - strict: bool = False, ) -> bool: - """Apply rank-local layerwise hooks to plan-declared encoder stacks.""" + """Apply rank-local layerwise hooks to the encoder's resolved stacks.""" + module = component.module if getattr(module, "_omni_layerwise_enabled", False): return True - hooks: list[LayerwiseOffloadHook] = [] - hooked_blocks: list[nn.Module] = [] - block_groups = get_encoder_block_groups( - module, - name, - plan, - strict=strict, - ) + block_groups = [stack.blocks for stack in component.stacks] if not block_groups: return False + + hooks: list[LayerwiseOffloadHook] = [] + hooked_blocks: list[nn.Module] = [] try: for blocks in block_groups: group_hooks = _install_layerwise_hook_group(blocks, device, stream, pin_memory) hooks.extend(group_hooks) hooked_blocks.extend(blocks) - if not stage_on_demand: + if not component.on_demand: move_non_block_state_to_device(module, block_groups, device) except BaseException: run_cleanup_steps( @@ -376,7 +368,7 @@ def enable_plan_encoder_layerwise_offload( ) logger.info( "Enabled rank-local layerwise offload for encoder %s (%d blocks across %d stacks)", - name, + component.path, sum(len(blocks) for blocks in block_groups), len(block_groups), ) @@ -440,41 +432,27 @@ class LayerWiseOffloadBackend(OffloadBackend): logger.warning("LayerWiseOffloadBackend already enabled") return - modules = ModuleDiscovery.discover(pipeline) - plan = get_offload_plan(pipeline) - if not modules.dits and self.config.offloads(DIT_COMPONENT): - message = "No DiT/transformer modules found for selected DiT layerwise offload" - if self.config.components is not None: - raise ValueError(message) - logger.warning(message) + resolved = resolve_offload_plan(pipeline, self.config) + if not resolved.dits and self.config.offloads(DIT_COMPONENT): + # An explicit selection already failed while resolving. + logger.warning("No DiT/transformer modules found for selected DiT layerwise offload") return - def enable_encoder_blocks( - module: nn.Module, - name: str, - component_plan: OffloadPlan | None, - stage_on_demand: bool, - ) -> bool: + def enable_encoder_blocks(component: ResolvedComponent) -> bool: enabled = enable_plan_encoder_layerwise_offload( - module, - name, - component_plan, + component, device=self.device, stream=self.copy_stream, pin_memory=self.config.pin_cpu_memory, - stage_on_demand=stage_on_demand, - strict=self.config.components is not None, ) if enabled: # Record each successful installation immediately so a later # component failure can remove these hooks transactionally. - self._encoder_modules.append(module) + self._encoder_modules.append(component.module) return enabled prepare_pipeline_components( - modules, - self.config, - plan, + resolved, device=self.device, staged_components=self._staged_components, enable_encoder_blocks=enable_encoder_blocks, @@ -489,30 +467,17 @@ class LayerWiseOffloadBackend(OffloadBackend): ) return - logger.info("Applying layer-wise offloading on %s", modules.dit_names) + logger.info("Applying layer-wise offloading on %s", [component.path for component in resolved.dits]) # Apply block-wise offloading hook for each of the blocks in DiT model(s) # Note that there might exist multiple DiT models in specific pipelines - for dit_name, dit_module, blocks_attr_names, blocks in iter_streamable_dits( - modules, self.config, self.device, plan - ): - num_blocks = len(blocks) - if num_blocks <= 1: - if self.config.components is not None: - raise ValueError( - f"Selected DiT {dit_name!r} requires at least two streamable layerwise-offload blocks" - ) - logger.warning( - "#Target layers (blocks) <= 1. Skipping offloading on %s (%s)", - dit_name, - dit_module.__class__.__name__, - ) - dit_module.to(self.device) - continue + for component, stack in iter_streamable_dits(resolved, self.device): + dit_module = component.module + blocks = list(stack.blocks) # Move non-block modules to GPU (they stay resident) for name, m in dit_module.named_children(): - if name not in blocks_attr_names: + if name not in stack.attrs: m.to(self.device) logger.debug(f"Moved {name} to device {self.device}") else: @@ -540,7 +505,7 @@ class LayerWiseOffloadBackend(OffloadBackend): # zero once; later denoising iterations prefetch it from the ring. block_hooks[0].prefetch_layer(non_blocking=False) - logger.info(f"Layer-wise offloading enabled on {num_blocks} layers (blocks)") + logger.info(f"Layer-wise offloading enabled on {len(blocks)} layers (blocks)") # Track hooked blocks for cleanup self._blocks.append(blocks) diff --git a/vllm_omni/diffusion/offloader/offload_plan.py b/vllm_omni/diffusion/offloader/offload_plan.py index 4aa37ae8f..5cc0c24c7 100644 --- a/vllm_omni/diffusion/offloader/offload_plan.py +++ b/vllm_omni/diffusion/offloader/offload_plan.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 """Declarative model capabilities for component layerwise offload. Models declare an ``_offload_plan`` class attribute on the pipeline. Both the @@ -22,9 +22,11 @@ class OffloadPlan: pipeline class. When present, both layerwise backends use it instead of model-specific backend branches, making new integrations data-driven. - If not declared, the offloader falls back to: - 1. ``_layerwise_offload_blocks_attrs`` on each DiT module class. - 2. Heuristic search for ``layers`` / ``blocks`` / ``h`` attributes. + If not declared, the offloader falls back to + ``_layerwise_offload_blocks_attrs`` on each DiT module class. + + :func:`~vllm_omni.diffusion.offloader.plan_resolver.resolve_offload_plan` + is the only consumer; backends read the resolved artifact it returns. Attributes: block_attrs: Maps DiT path → tuple of block-list attribute names. diff --git a/vllm_omni/diffusion/offloader/plan_resolver.py b/vllm_omni/diffusion/offloader/plan_resolver.py new file mode 100644 index 000000000..b2e951b4e --- /dev/null +++ b/vllm_omni/diffusion/offloader/plan_resolver.py @@ -0,0 +1,330 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""One topology resolver shared by every diffusion offload backend. + +Model topology is declared through several mechanisms that grew independently: +:class:`~vllm_omni.diffusion.models.interface.SupportsComponentDiscovery` class +variables, the pipeline-level :class:`OffloadPlan`, per-DiT +``_layerwise_offload_blocks_attrs``, and a leaf-name rule for encoders. This +module runs them in one defined order and returns one backend-neutral artifact. + +The resolver is pure: it moves no tensor, installs no hook, writes no module +attribute, and touches no process group. Everything it can reject is rejected +here, so plan-dependent failures happen before a backend mutates the model. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import cache +from typing import TYPE_CHECKING + +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery + +from .block_discovery import get_blocks_from_dit +from .component_utils import get_encoder_block_groups, validate_on_demand_component +from .config import DIT_COMPONENT, TEXT_ENCODER_COMPONENT, OffloadStrategy +from .module_collector import ModuleDiscovery, PipelineModules +from .offload_plan import OffloadPlan, get_offload_plan + +if TYPE_CHECKING: + from .base import OffloadConfig + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class BlockStack: + """One ring of repeated blocks that a backend streams together. + + ``attrs`` names the owner attributes that contributed blocks; backends use + it to skip those children when they place the non-streamed remainder. It is + filled for DiT stacks only — encoder stacks are hooked as whole stacks and + their remainder is placed by tensor identity. + """ + + blocks: tuple[nn.Module, ...] + attrs: tuple[str, ...] = () + resident_head: int = 0 + + @property + def resident(self) -> tuple[nn.Module, ...]: + """Leading blocks that stay on the device instead of streaming.""" + return self.blocks[: self.resident_head] + + @property + def streaming(self) -> tuple[nn.Module, ...]: + """Blocks a hook ring transfers on demand.""" + return self.blocks[self.resident_head :] + + +@dataclass(frozen=True) +class ResolvedComponent: + """One pipeline component with its resolved, validated offload topology. + + ``selected`` means the active selector covers this component. ``on_demand`` + means the pipeline owns its residency through ``load_to_device`` / + ``offload_to_cpu``; a VAE staged by a legacy model plan is ``on_demand`` + without being selectable through the public component grammar. + """ + + path: str + module: nn.Module + selected: bool + stacks: tuple[BlockStack, ...] = () + on_demand: bool = False + + +@dataclass(frozen=True) +class ResolvedOffloadPlan: + """Backend-neutral topology for one pipeline under one offload config.""" + + dits: tuple[ResolvedComponent, ...] + encoders: tuple[ResolvedComponent, ...] + vaes: tuple[ResolvedComponent, ...] + residents: tuple[ResolvedComponent, ...] + # Raw discovery output and the model declaration remain reachable for the + # loader-owned mmap path and for DLO's nested-submodule staging, which + # still interpret topology themselves. + modules: PipelineModules + declaration: OffloadPlan | None = None + + @property + def components(self) -> tuple[ResolvedComponent, ...]: + return self.dits + self.encoders + self.vaes + self.residents + + +@cache +def _warn_legacy_discovery(pipeline_class: str) -> None: + """Warn once per pipeline class that component discovery guessed. + + The key is the fully qualified class name so two same-named pipeline + classes in different modules each get their own warning. + """ + logger.warning( + "%s does not declare SupportsComponentDiscovery; offload components were " + "discovered by scanning well-known attribute names. Declare _dit_modules, " + "_encoder_modules, and _vae_modules; the attribute scan is deprecated.", + pipeline_class, + ) + + +def _resolve_dit_stacks( + module: nn.Module, + path: str, + declaration: OffloadPlan | None, + config: OffloadConfig, + *, + explicit: bool, +) -> tuple[BlockStack, ...]: + """Resolve one DiT's streamed blocks and its device-resident head.""" + planned_attrs = None if declaration is None else declaration.block_attrs.get(path) + attrs, blocks = get_blocks_from_dit(module, planned_attrs) + if not blocks: + if explicit: + raise ValueError(f"Selected DiT {path!r} has no streamable layerwise-offload blocks") + logger.warning( + "Target layers (blocks) not found. Skipping offloading on %s (%s)", + path, + type(module).__name__, + ) + return () + + distributed = config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE + resident_head = 0 + if ( + distributed + and config.dlo_resident_layers + and declaration is not None + and path in declaration.resident_dit_paths + ): + resident_head = min(config.dlo_resident_layers, len(blocks)) + + if len(blocks) - resident_head == 1: + if explicit and distributed: + raise ValueError( + f"Selected DiT {path!r} leaves only one streaming block after " + f"resident_layers={resident_head}; choose a resident count that " + "leaves zero or at least two streaming blocks" + ) + if explicit: + raise ValueError(f"Selected DiT {path!r} requires at least two streamable layerwise-offload blocks") + if distributed: + # One streamed block cannot form a prefetch ring; keeping it + # resident preserves the model instead of skipping placement. + logger.warning( + "#Streaming target layers (blocks) <= 1. Keeping the final block resident on %s (%s)", + path, + type(module).__name__, + ) + resident_head = len(blocks) + else: + logger.warning( + "#Target layers (blocks) <= 1. Skipping offloading on %s (%s)", + path, + type(module).__name__, + ) + return () + + return (BlockStack(blocks=tuple(blocks), attrs=tuple(attrs), resident_head=resident_head),) + + +def _resolve_encoder_stacks( + module: nn.Module, + path: str, + declaration: OffloadPlan | None, + config: OffloadConfig, + *, + explicit: bool, +) -> tuple[BlockStack, ...]: + """Resolve one encoder's streamable stacks and check transfer safety.""" + if declaration is None: + return () + + group_size = 1 + if config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE and config.uses_allgather(TEXT_ENCODER_COMPONENT): + group_size = config.dp_size + if group_size > 1 and path not in declaration.encoder_dlo_weight_replication: + # Every rank must reject an unsafe group, including ranks whose local + # encoder is an unloaded stub, so the check precedes block discovery. + raise ValueError( + f"Text encoder {path!r} cannot use DLO AllGather across the DiT offload group: " + "its loader-produced weights are not declared replicated across that group. " + "Set layer_options.text_encoder.weight_transfer='rank-local' in diffusion_offload_config " + "for encoder-TP or rank-specific layouts." + ) + + groups = get_encoder_block_groups(module, path, declaration, strict=explicit or group_size > 1) + return tuple(BlockStack(blocks=tuple(blocks)) for blocks in groups) + + +def resolve_offload_plan(pipeline: nn.Module, config: OffloadConfig) -> ResolvedOffloadPlan: + """Resolve and validate one pipeline's offload topology. + + Raises ``ValueError`` for every topology that the requested configuration + cannot serve, before the caller places a module or installs a hook. + """ + modules = ModuleDiscovery.discover(pipeline) + declaration = get_offload_plan(pipeline) + if not isinstance(pipeline, SupportsComponentDiscovery) and (modules.dits or modules.encoders or modules.vaes): + pipeline_class = type(pipeline) + _warn_legacy_discovery(f"{pipeline_class.__module__}.{pipeline_class.__qualname__}") + + explicit = config.components is not None + layerwise = config.strategy in (OffloadStrategy.LAYER_WISE, OffloadStrategy.DISTRIBUTED_LAYER_WISE) + dit_selected = config.offloads(DIT_COMPONENT) + + if explicit: + if config.strategy is OffloadStrategy.MODEL_LEVEL: + # Model-level offload swaps the DiT against an encoder stage, so + # both sides must exist regardless of which side was selected. + if not modules.dits: + raise ValueError("Component-selective model offload requires a DiT/transformer module") + if not modules.encoders: + raise ValueError("Component-selective model offload requires an encoder execution stage") + elif dit_selected and not modules.dits: + raise ValueError("No DiT/transformer modules found for selected DiT offload") + if config.offloads(TEXT_ENCODER_COMPONENT) and not any( + config.offloads_encoder(path, declaration) for path in modules.encoder_names + ): + raise ValueError("No text encoder modules found for selected text_encoder offload") + + if config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE and config.dlo_resident_layers: + resident_paths = frozenset() if declaration is None else declaration.resident_dit_paths + if not resident_paths.intersection(modules.dit_names): + message = ( + f"resident_layers={config.dlo_resident_layers} was requested, but this model declares " + "no matching resident_dit_paths" + ) + if explicit: + raise ValueError(message) + logger.warning("%s; all blocks will be streamed.", message) + + dits: list[ResolvedComponent] = [] + encoders: list[ResolvedComponent] = [] + vaes: list[ResolvedComponent] = [] + residents: list[ResolvedComponent] = [] + + for path, module in zip(modules.dit_names, modules.dits): + stacks: tuple[BlockStack, ...] = () + if dit_selected and layerwise: + stacks = _resolve_dit_stacks(module, path, declaration, config, explicit=explicit) + dits.append(ResolvedComponent(path=path, module=module, selected=dit_selected, stacks=stacks)) + + for path, module in zip(modules.encoder_names, modules.encoders): + selected = config.should_offload_encoder(path, declaration) + # Staged residency belongs to the layer backends; model-level offload + # swaps whole modules through its own hooks. + on_demand = bool( + selected and layerwise and declaration is not None and path in declaration.on_demand_component_paths + ) + if on_demand: + validate_on_demand_component(module, path) + stacks = () + if selected and layerwise: + stacks = _resolve_encoder_stacks(module, path, declaration, config, explicit=explicit) + if on_demand and not stacks and config.uses_allgather(TEXT_ENCODER_COMPONENT): + raise ValueError( + f"Text encoder {path!r} cannot use AllGather without a model-declared streamable block plan" + ) + if explicit and not (stacks or on_demand): + raise ValueError( + f"Selected text encoder {path!r} requires a model-declared streamable or on-demand plan" + ) + encoders.append( + ResolvedComponent( + path=path, + module=module, + selected=selected, + stacks=stacks, + on_demand=on_demand, + ) + ) + + for path, module in zip(modules.vae_names, modules.vaes): + # VAEs are not part of the public selector. A model plan may still own + # their residency, which the compatibility topology preserves. + legacy_staged = ( + layerwise and not explicit and declaration is not None and path in declaration.on_demand_component_paths + ) + if legacy_staged: + validate_on_demand_component(module, path) + vaes.append(ResolvedComponent(path=path, module=module, selected=False, on_demand=legacy_staged)) + + for path, module in zip(modules.resident_names, modules.resident_modules): + residents.append(ResolvedComponent(path=path, module=module, selected=False)) + + resolved = ResolvedOffloadPlan( + dits=tuple(dits), + encoders=tuple(encoders), + vaes=tuple(vaes), + residents=tuple(residents), + modules=modules, + declaration=declaration, + ) + _validate_unique_ownership(resolved) + return resolved + + +def _validate_unique_ownership(resolved: ResolvedOffloadPlan) -> None: + """Reject topologies where two components would hook the same blocks.""" + owner_by_block: dict[int, str] = {} + for component in resolved.components: + for block in (block for stack in component.stacks for block in stack.blocks): + owner = owner_by_block.setdefault(id(block), component.path) + if owner != component.path: + raise ValueError( + f"Offload block is claimed by both {owner!r} and {component.path!r}; " + "declare one owner for each block list" + ) + + +__all__ = [ + "BlockStack", + "ResolvedComponent", + "ResolvedOffloadPlan", + "resolve_offload_plan", +] diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index b6b596070..1187c8cee 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/vllm_omni/diffusion/offloader/sequential_backend.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 from collections.abc import Collection, Iterator from contextlib import contextmanager @@ -13,9 +13,8 @@ from vllm_omni.diffusion.hooks import HookRegistry, ModelHook from vllm_omni.platforms import current_omni_platform from .base import OffloadBackend, OffloadConfig, SupportsModelCpuOffload -from .config import DIT_COMPONENT, TEXT_ENCODER_COMPONENT -from .module_collector import ModuleDiscovery -from .offload_plan import get_offload_plan +from .config import DIT_COMPONENT +from .plan_resolver import resolve_offload_plan logger = init_logger(__name__) @@ -329,59 +328,44 @@ class ModelLevelOffloadBackend(OffloadBackend): ) return - modules = ModuleDiscovery.discover(pipeline) - plan = get_offload_plan(pipeline) - selected_encoders = [ - encoder - for encoder, name in zip(modules.encoders, modules.encoder_names) - if self.config.should_offload_encoder(name, plan) - ] - if self.config.components is not None: - if not modules.dits: - raise ValueError("Component-selective model offload requires a DiT/transformer module") - if not modules.encoders: - raise ValueError("Component-selective model offload requires an encoder execution stage") - if self.config.offloads(TEXT_ENCODER_COMPONENT) and not selected_encoders: - raise ValueError("No text encoder modules found for selected text_encoder module offload") + resolved = resolve_offload_plan(pipeline, self.config) + dits = [component.module for component in resolved.dits] + encoders = [component.module for component in resolved.encoders] + vaes = [component.module for component in resolved.vaes] + residents = [component.module for component in resolved.residents] + selected_encoders = [component.module for component in resolved.encoders if component.selected] - all_modules = [ - *modules.dits, - *modules.encoders, - *modules.vaes, - *modules.resident_modules, - ] + all_modules = [*dits, *encoders, *vaes, *residents] initial_devices = _capture_tensor_devices(all_modules) try: - for encoder in modules.encoders: + for encoder in encoders: encoder.to(self.device) - for vae in modules.vaes: + for vae in vaes: vae.to(self.device, non_blocking=True) - for resident in modules.resident_modules: + for resident in residents: resident.to(self.device) - if not modules.dits: + if not dits: logger.warning("No DiT/transformer modules found, skipping model-level offloading") return - if not modules.encoders: - for dit in modules.dits: + if not encoders: + for dit in dits: dit.to(self.device) logger.warning("No encoder modules found, skipping model-level offloading") return apply_sequential_offload( - dit_modules=modules.dits, - encoder_modules=modules.encoders, + dit_modules=dits, + encoder_modules=encoders, device=self.device, pin_memory=self.config.pin_cpu_memory, use_hsdp=self.config.use_hsdp, - offload_dit_modules=( - modules.dits if self.config.components is None or self.config.offloads(DIT_COMPONENT) else () - ), + offload_dit_modules=(dits if self.config.offloads(DIT_COMPONENT) else ()), offload_encoder_modules=selected_encoders, ) except BaseException: try: - remove_sequential_offload([*modules.dits, *modules.encoders]) + remove_sequential_offload([*dits, *encoders]) except BaseException: logger.exception("Failed to remove every model-level hook during rollback") try: @@ -391,15 +375,19 @@ class ModelLevelOffloadBackend(OffloadBackend): raise # Track modules for cleanup - self._offload_modules = [*modules.dits, *modules.encoders] + self._offload_modules = [*dits, *encoders] self.enabled = True logger.info( "Model-level offloading enabled: %s <-> %s (mutual exclusion)%s", - ", ".join(modules.dit_names), - ", ".join(modules.encoder_names), - f"; resident on GPU: {', '.join(modules.resident_names)}" if modules.resident_names else "", + ", ".join(component.path for component in resolved.dits), + ", ".join(component.path for component in resolved.encoders), + ( + f"; resident on GPU: {', '.join(component.path for component in resolved.residents)}" + if resolved.residents + else "" + ), ) def disable(self) -> None: