From 85da94339130337b49e610d08276aedcf84f6b96 Mon Sep 17 00:00:00 2001 From: Anjie Hou <149605198+specture724@users.noreply.github.com> Date: Tue, 15 Sep 2026 08:40:04 +0800 Subject: [PATCH] [Refactor][Diffusion] Complete generic backend plan consumption (#7313) Signed-off-by: specture724 Co-authored-by: Claude Opus 5 --- docs/design/feature/offloader/README.md | 33 +++- .../models/flux2/test_flux2_transformer_tp.py | 7 +- .../offloader/test_backend_plan_contract.py | 163 ++++++++++++++++++ .../offloader/test_layerwise_backend.py | 23 ++- .../diffusion/offloader/layerwise_backend.py | 44 +---- .../diffusion/offloader/plan_resolver.py | 10 ++ .../diffusion/offloader/sequential_backend.py | 4 +- 7 files changed, 228 insertions(+), 56 deletions(-) create mode 100644 tests/diffusion/offloader/test_backend_plan_contract.py diff --git a/docs/design/feature/offloader/README.md b/docs/design/feature/offloader/README.md index 96cba8f60..0ec81f3b5 100644 --- a/docs/design/feature/offloader/README.md +++ b/docs/design/feature/offloader/README.md @@ -41,7 +41,8 @@ Every backend implements `OffloadBackend`: - `is_enabled()` reports lifecycle state. `disable()` does not promise to restore the pipeline's original device -placement. The caller owns any subsequent rematerialization. +placement. It must restore usable parameter and buffer storage before releasing +backend-owned host weights. The caller owns subsequent device placement. An enable failure must remove every partially installed hook. Rank-local backends restore ordinary tensors so the pipeline can be retried. A failed @@ -54,9 +55,19 @@ 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. +The generic model-level and ordinary layerwise paths call it once when +`enable()` starts and consume its component selection and block stacks. +Ordinary layerwise offload also respects a resolved `skip_reason`, preserving +the legacy missing-DiT no-op without placing auxiliary components. Non-block +DiT and encoder state is placed by resolved block tensor identity, so the +backend does not rediscover block containers by attribute name. + +Two compatibility paths remain outside this completed generic cutover: +model-level pipelines implementing `SupportsModelCpuOffload` still own their +custom phase lifecycle (RFC #6648 J4), and distributed layerwise offload still +reads declarations for nested-submodule staging and loader-owned mmap paths +(J3). The custom model-level protocol is delegated to before generic plan +resolution; it must not be subjected to the generic encoder/DiT swap checks. Pipeline component discovery prefers `SupportsComponentDiscovery` declarations: @@ -102,3 +113,17 @@ remains responsible for transfer, synchronization, and storage ownership. The diffusion [Offloader module design](../../module/diffusion/offloader.md) describes how these feature contracts fit into the larger diffusion runtime. + +## Backend acceptance tests + +`tests/diffusion/offloader/test_backend_plan_contract.py` covers the two +behaviors the generic cutover changed. A DiT attribute aliasing a streamed +block must keep the ring's host residency, so placement cannot follow attribute +names; the CUDA case fails against the previous name-based placement. Both +backends must then run with their selector helpers and the model declaration +made unreadable, proving execution consumes the resolved plan alone. + +Selection errors, rollback, residency, and enable/disable cycles stay in +`test_plan_resolver.py`, `test_layerwise_backend.py` and +`test_sequential_backend.py`. Model-specific phases and executed tied aliases +remain the J4 contract suite. diff --git a/tests/diffusion/models/flux2/test_flux2_transformer_tp.py b/tests/diffusion/models/flux2/test_flux2_transformer_tp.py index c935871a0..489a948d5 100644 --- a/tests/diffusion/models/flux2/test_flux2_transformer_tp.py +++ b/tests/diffusion/models/flux2/test_flux2_transformer_tp.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import pytest import torch from pytest_mock import MockerFixture @@ -265,7 +268,7 @@ class TestFlux2TransformerLayerwiseOffload: ] def test_get_blocks_from_dit(self): - from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend + from vllm_omni.diffusion.offloader.block_discovery import get_blocks_from_dit # Block discovery only needs the ModuleList structure, not production dims. model = Flux2Transformer2DModel( @@ -275,6 +278,6 @@ class TestFlux2TransformerLayerwiseOffload: attention_head_dim=4, joint_attention_dim=16, ) - attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model) + attr_names, blocks = get_blocks_from_dit(model) assert attr_names == ["transformer_blocks", "single_transformer_blocks"] assert len(blocks) == 4 diff --git a/tests/diffusion/offloader/test_backend_plan_contract.py b/tests/diffusion/offloader/test_backend_plan_contract.py new file mode 100644 index 000000000..3c406f5d5 --- /dev/null +++ b/tests/diffusion/offloader/test_backend_plan_contract.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Executable contracts for the two behaviors the generic cutover changed. + +Placement now follows resolved block identity, and the backends run on the +resolved plan alone. Declaration parsing, selection errors, rollback, +residency and enable/disable cycles are covered by ``test_plan_resolver.py``, +``test_layerwise_backend.py`` and ``test_sequential_backend.py``. +""" + +from typing import ClassVar + +import pytest +import torch +from torch import nn + +from tests.diffusion.offloader.helpers import patch_offload_runtime +from vllm_omni.diffusion.offloader import layerwise_backend, sequential_backend +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 +from vllm_omni.diffusion.offloader.sequential_backend import ModelLevelOffloadBackend +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.diffusion, pytest.mark.core_model] + +STRATEGIES = [OffloadStrategy.MODEL_LEVEL, OffloadStrategy.LAYER_WISE] +CPU = torch.device("cpu") + + +class _Block(nn.Linear): + def __init__(self): + super().__init__(4, 4) + self.register_buffer("scale", torch.tensor(0.5)) + + def forward(self, x): + return super().forward(x).tanh() * self.scale + + +class _Stack(nn.Module): + """Two block containers plus non-block state that must stay resident.""" + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([_Block() for _ in range(3)]) + self.tail = nn.ModuleList([_Block() for _ in range(2)]) + self.bias = nn.Parameter(torch.ones(4)) + self.register_buffer("scale", torch.tensor(0.25)) + self.proj = nn.Linear(4, 4) + + def forward(self, x): + for block in (*self.blocks, *self.tail): + x = block(x) + return self.proj(x) * self.scale + self.bias + + +class _Pipeline(nn.Module): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder", "image_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = ["resident"] + _offload_plan = OffloadPlan( + block_attrs={"transformer": ("blocks", "tail")}, + encoder_block_attrs={"text_encoder": ("blocks", "tail")}, + ) + + def __init__(self): + super().__init__() + self.transformer = _Stack() + self.text_encoder = _Stack() + self.image_encoder = nn.Linear(4, 4) + self.vae = nn.Linear(4, 4) + self.resident = nn.Linear(4, 4) + + def forward(self, x): + return self.resident(self.vae(self.transformer(self.text_encoder(self.image_encoder(x))))) + + +@pytest.fixture(params=[pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=pytest.mark.cuda)]) +def execution_device(request, monkeypatch): + if request.param == "cuda": + if not torch.cuda.is_available(): + pytest.skip("CUDA required for real transfer/residency coverage") + return torch.device("cuda:0") + patch_offload_runtime(monkeypatch, current_omni_platform, synchronize=True) + monkeypatch.setattr(current_omni_platform, "get_free_memory", lambda: 0) + return CPU + + +def _config(strategy, components): + return OffloadConfig(strategy=strategy, components=components, pin_cpu_memory=False) + + +def _backend(config, device): + kind = ModelLevelOffloadBackend if config.strategy is OffloadStrategy.MODEL_LEVEL else LayerWiseOffloadBackend + return kind(config, device) + + +def _assert_no_offload_hooks(pipeline): + for module in pipeline.modules(): + registry = getattr(module, "_hook_registry", None) + if registry is not None: + assert registry.get_hook("sequential_offload") is None + assert registry.get_hook("layerwise_offload") is None + assert not getattr(module, "_omni_layerwise_enabled", False) + + +@torch.inference_mode() +def test_alias_of_a_streamed_block_is_not_placed(execution_device): + """Placement moves DiT state by resolved block identity, not by attribute. + + An attribute aliasing a streamed block must stay with its ring: copying it + to the device makes the ring record the device as that block's home and + strands those weights there after teardown. + """ + pipeline = _Pipeline() + pipeline.transformer.proj.weight = pipeline.transformer.blocks[1].weight + original = pipeline.transformer.proj.weight.clone() + backend = _backend(_config(OffloadStrategy.LAYER_WISE, frozenset({"dit"})), execution_device) + + backend.enable(pipeline) + try: + assert pipeline.transformer.proj.weight is pipeline.transformer.blocks[1].weight + assert pipeline.transformer.proj.weight.device == CPU + # Non-block state around the alias is still placed. + assert pipeline.transformer.proj.bias.device == execution_device + assert pipeline.transformer.bias.device == execution_device + finally: + backend.disable() + + assert pipeline.transformer.proj.weight.device == CPU + torch.testing.assert_close(pipeline.transformer.proj.weight, original) + + +@pytest.mark.parametrize("strategy", STRATEGIES) +@torch.inference_mode() +def test_backends_execute_the_resolved_plan_alone(execution_device, strategy, monkeypatch): + pipeline = _Pipeline().to(execution_device) + x = torch.randn(2, 4, device=execution_device) + expected = pipeline(x) + config = _config(strategy, frozenset({"dit", "text_encoder"})) + resolved = resolve_offload_plan(pipeline, config) + + def forbidden_read(*args, **kwargs): + pytest.fail("Backend reinterpreted topology after plan resolution") + + # Once resolved, neither declarations nor selector helpers are an input to + # backend execution. Transport options stay readable. + monkeypatch.setattr(config, "offloads", forbidden_read) + monkeypatch.setattr(config, "should_offload_encoder", forbidden_read) + monkeypatch.setattr(pipeline, "_offload_plan", None) + backend_module = sequential_backend if strategy is OffloadStrategy.MODEL_LEVEL else layerwise_backend + monkeypatch.setattr(backend_module, "resolve_offload_plan", lambda *_: resolved) + + backend = _backend(config, execution_device) + try: + backend.enable(pipeline) + torch.testing.assert_close(pipeline(x), expected) + finally: + backend.disable() + _assert_no_offload_hooks(pipeline) diff --git a/tests/diffusion/offloader/test_layerwise_backend.py b/tests/diffusion/offloader/test_layerwise_backend.py index 099e93f79..4c965adaa 100644 --- a/tests/diffusion/offloader/test_layerwise_backend.py +++ b/tests/diffusion/offloader/test_layerwise_backend.py @@ -24,6 +24,11 @@ from tests.diffusion.offloader.helpers import ( patch_offload_runtime, ) from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy +from vllm_omni.diffusion.offloader.block_discovery import ( + get_blocks_attr_names, + get_blocks_from_dit, + set_blocks_attr_names, +) from vllm_omni.diffusion.offloader.layerwise_backend import ( LayerWiseOffloadBackend, LayerwiseOffloadHook, @@ -209,21 +214,21 @@ class _NoAttrsModel(nn.Module): class TestGetBlocksFromDit: def test_get_blocks_from_dit_single_block_attr(self): model = _SingleBlockModel(num_blocks=3) - attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model) + attr_names, blocks = get_blocks_from_dit(model) assert attr_names == ["blocks"] assert len(blocks) == 3 assert all(isinstance(b, _DummyBlock) for b in blocks) def test_get_blocks_from_dit_multi_block_attrs(self): model = _MultiBlockModel(num_transformer=2, num_single=3) - attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model) + attr_names, blocks = get_blocks_from_dit(model) assert set(attr_names) == {"transformer_blocks", "single_transformer_blocks"} assert len(blocks) == 5 assert all(isinstance(b, _DummyBlock) for b in blocks) def test_get_blocks_from_dit_empty_blocks(self): model = _EmptyBlocksModel() - attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model) + attr_names, blocks = get_blocks_from_dit(model) assert attr_names == [] assert blocks == [] @@ -233,17 +238,17 @@ class TestGetBlocksFromDit: AttributeError, match="Attribute 'nonexistent_blocks' declared in _layerwise_offload_blocks_attrs does not exist", ): - LayerWiseOffloadBackend.get_blocks_from_dit(model) + get_blocks_from_dit(model) def test_get_blocks_from_dit_no_attrs_defined(self): model = _NoAttrsModel(num_blocks=3) - attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model) + attr_names, blocks = get_blocks_from_dit(model) assert attr_names == [] assert blocks == [] def test_get_blocks_from_dit_deprecated_single_attr(self): model = _DeprecatedSingleAttrModel(num_blocks=2) - attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model) + attr_names, blocks = get_blocks_from_dit(model) assert attr_names == ["blocks"] assert len(blocks) == 2 @@ -251,17 +256,17 @@ class TestGetBlocksFromDit: class TestGetBlocksAttrNames: def test_get_blocks_attr_names_new_format(self): model = _MultiBlockModel() - attrs = LayerWiseOffloadBackend.get_blocks_attr_names(model) + attrs = get_blocks_attr_names(model) assert attrs == ["transformer_blocks", "single_transformer_blocks"] def test_get_blocks_attr_names_no_attrs(self): model = _NoAttrsModel() - attrs = LayerWiseOffloadBackend.get_blocks_attr_names(model) + attrs = get_blocks_attr_names(model) assert attrs == [] def test_set_blocks_attr_names(self): model = _NoAttrsModel() - LayerWiseOffloadBackend.set_blocks_attr_names(model, ["new_blocks"]) + set_blocks_attr_names(model, ["new_blocks"]) assert hasattr(model.__class__, "_layerwise_offload_blocks_attrs") assert model.__class__._layerwise_offload_blocks_attrs == ["new_blocks"] diff --git a/vllm_omni/diffusion/offloader/layerwise_backend.py b/vllm_omni/diffusion/offloader/layerwise_backend.py index a1e594f8b..e9f8cf5be 100644 --- a/vllm_omni/diffusion/offloader/layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/layerwise_backend.py @@ -14,11 +14,6 @@ from vllm_omni.diffusion.hooks import HookRegistry, ModelHook from vllm_omni.platforms import current_omni_platform from .base import OffloadBackend, OffloadConfig, run_cleanup_steps -from .block_discovery import ( - get_blocks_attr_names, - get_blocks_from_dit, - set_blocks_attr_names, -) from .component_utils import ( clear_encoder_layerwise_state, iter_streamable_dits, @@ -26,7 +21,6 @@ from .component_utils import ( prepare_pipeline_components, set_encoder_layerwise_state, ) -from .config import DIT_COMPONENT from .plan_resolver import ResolvedComponent, resolve_offload_plan from .tensor_utils import ( clear_block_storage, @@ -433,9 +427,8 @@ class LayerWiseOffloadBackend(OffloadBackend): return 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") + if resolved.skip_reason is not None: + logger.warning("%s", resolved.skip_reason) return def enable_encoder_blocks(component: ResolvedComponent) -> bool: @@ -458,15 +451,6 @@ class LayerWiseOffloadBackend(OffloadBackend): enable_encoder_blocks=enable_encoder_blocks, ) - if not self.config.offloads(DIT_COMPONENT): - self.enabled = bool(self._encoder_modules or self._staged_components) - if not self.enabled: - raise ValueError( - "None of the selected layerwise offload components have " - "a model-declared streamable or on-demand plan" - ) - return - 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) @@ -475,22 +459,9 @@ class LayerWiseOffloadBackend(OffloadBackend): 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 stack.attrs: - m.to(self.device) - logger.debug(f"Moved {name} to device {self.device}") - else: - logger.debug(f"Skipped blocks module {name}") - - # Move top-level params/buffers to GPU (dit_module's own, not sub-modules) - for param in dit_module._parameters.values(): - if param is not None: - param.data = param.data.to(self.device, non_blocking=True) - - for buffer in dit_module._buffers.values(): - if buffer is not None: - buffer.data = buffer.data.to(self.device, non_blocking=True) + # Place the remainder by resolved block tensor identity, just as + # for encoders. Attribute aliases must not move streamed weights. + move_non_block_state_to_device(dit_module, (stack.blocks,), self.device) block_hooks = _install_layerwise_hook_group( blocks, @@ -550,8 +521,3 @@ class LayerWiseOffloadBackend(OffloadBackend): def disable(self) -> None: self._disable(restore_weights=True) - - # Compatibility aliases for existing model integrations. - get_blocks_attr_names = staticmethod(get_blocks_attr_names) - set_blocks_attr_names = staticmethod(set_blocks_attr_names) - get_blocks_from_dit = staticmethod(get_blocks_from_dit) diff --git a/vllm_omni/diffusion/offloader/plan_resolver.py b/vllm_omni/diffusion/offloader/plan_resolver.py index b2e951b4e..08527318f 100644 --- a/vllm_omni/diffusion/offloader/plan_resolver.py +++ b/vllm_omni/diffusion/offloader/plan_resolver.py @@ -91,6 +91,11 @@ class ResolvedOffloadPlan: # still interpret topology themselves. modules: PipelineModules declaration: OffloadPlan | None = None + # Ordinary layerwise offload turns a legacy selection it cannot serve into + # a complete no-op, placement included. Distributed layerwise offload warns + # and keeps preparing its other components instead, so this stays scoped to + # the strategy that honors it. + skip_reason: str | None = None @property def components(self) -> tuple[ResolvedComponent, ...]: @@ -304,6 +309,11 @@ def resolve_offload_plan(pipeline: nn.Module, config: OffloadConfig) -> Resolved residents=tuple(residents), modules=modules, declaration=declaration, + skip_reason=( + "No DiT/transformer modules found for selected DiT layerwise offload" + if config.strategy is OffloadStrategy.LAYER_WISE and dit_selected and not dits + else None + ), ) _validate_unique_ownership(resolved) return resolved diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index 1187c8cee..b03be3687 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/vllm_omni/diffusion/offloader/sequential_backend.py @@ -13,7 +13,6 @@ 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 from .plan_resolver import resolve_offload_plan logger = init_logger(__name__) @@ -333,6 +332,7 @@ class ModelLevelOffloadBackend(OffloadBackend): 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_dits = [component.module for component in resolved.dits if component.selected] selected_encoders = [component.module for component in resolved.encoders if component.selected] all_modules = [*dits, *encoders, *vaes, *residents] @@ -360,7 +360,7 @@ class ModelLevelOffloadBackend(OffloadBackend): device=self.device, pin_memory=self.config.pin_cpu_memory, use_hsdp=self.config.use_hsdp, - offload_dit_modules=(dits if self.config.offloads(DIT_COMPONENT) else ()), + offload_dit_modules=selected_dits, offload_encoder_modules=selected_encoders, ) except BaseException: