[diffusion][feature] Add component-selective offload policies (#5929)

Signed-off-by: lishunyang12 <lishunyang12@163.com>

tests/diffusion/offloader/test_sequential_backend.py

@@ -7,7 +7,9 @@ import pytest
 import torch
 from torch import nn
 
+from vllm_omni.diffusion.hooks import HookRegistry
 from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan
 from vllm_omni.diffusion.offloader.sequential_backend import (
     ModelLevelOffloadBackend,
     SequentialOffloadHook,
@@ -48,20 +50,21 @@ def _track_pin_memory_calls():
     return tracker, mock
 
 
+class _CustomPipeline(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.enable_args: dict[str, object] | None = None
+        self.disable_called = False
+
+    def enable_omni_model_cpu_offload(self, **kwargs) -> None:
+        self.enable_args = kwargs
+
+    def disable_omni_model_cpu_offload(self) -> None:
+        self.disable_called = True
+
+
 def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None:
-    class CustomPipeline(nn.Module):
-        def __init__(self) -> None:
-            super().__init__()
-            self.enable_args = None
-            self.disable_called = False
-
-        def enable_omni_model_cpu_offload(self, **kwargs) -> None:
-            self.enable_args = kwargs
-
-        def disable_omni_model_cpu_offload(self) -> None:
-            self.disable_called = True
-
-    pipeline = CustomPipeline()
+    pipeline = _CustomPipeline()
     backend = ModelLevelOffloadBackend(
         OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL, pin_cpu_memory=False),
         torch.device("cpu"),
@@ -82,6 +85,205 @@ def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None:
     assert pipeline.disable_called is True
 
 
+def test_model_level_backend_passes_explicit_component_selection() -> None:
+    pipeline = _CustomPipeline()
+    backend = ModelLevelOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.MODEL_LEVEL,
+            components=frozenset({"dit", "text_encoder"}),
+        ),
+        torch.device("cpu"),
+    )
+
+    backend.enable(pipeline)
+
+    assert pipeline.enable_args is not None
+    assert pipeline.enable_args["offload_components"] == frozenset({"dit", "text_encoder"})
+
+
+def test_model_level_backend_rolls_back_partial_custom_enable() -> None:
+    class CustomPipeline(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.disable_called = False
+
+        def enable_omni_model_cpu_offload(self, **kwargs) -> None:
+            del kwargs
+            raise RuntimeError("injected custom enable failure")
+
+        def disable_omni_model_cpu_offload(self) -> None:
+            self.disable_called = True
+
+    pipeline = CustomPipeline()
+    backend = ModelLevelOffloadBackend(
+        OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL),
+        torch.device("cpu"),
+    )
+
+    with pytest.raises(RuntimeError, match="injected custom enable failure"):
+        backend.enable(pipeline)
+
+    assert pipeline.disable_called
+    assert not backend.enabled
+
+
+def test_sequential_offload_rolls_back_partial_hook_registration(monkeypatch: pytest.MonkeyPatch) -> None:
+    dit = _create_simple_module()
+    encoder = _create_simple_module()
+    original_register = HookRegistry.register_hook
+    calls = 0
+
+    def fail_second_registration(self, name, hook):
+        nonlocal calls
+        calls += 1
+        if calls == 2:
+            raise RuntimeError("injected registration failure")
+        return original_register(self, name, hook)
+
+    monkeypatch.setattr(HookRegistry, "register_hook", fail_second_registration)
+
+    with pytest.raises(RuntimeError, match="injected registration failure"):
+        apply_sequential_offload(
+            dit_modules=[dit],
+            encoder_modules=[encoder],
+            device=torch.device("cpu"),
+        )
+
+    for module in (dit, encoder):
+        registry = getattr(module, "_hook_registry", None)
+        assert registry is None or registry.get_hook(SequentialOffloadHook._HOOK_NAME) is None
+
+
+def test_model_level_disable_cleans_remaining_hooks_and_can_retry(monkeypatch: pytest.MonkeyPatch) -> None:
+    pipeline = nn.Module()
+    pipeline.transformer = _create_simple_module()
+    pipeline.text_encoder = _create_simple_module()
+    backend = ModelLevelOffloadBackend(
+        OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL, pin_cpu_memory=False),
+        torch.device("cpu"),
+    )
+    backend.enable(pipeline)
+
+    original_remove = HookRegistry.remove_hook
+    calls = 0
+
+    def fail_first_removal(self, name):
+        nonlocal calls
+        calls += 1
+        if calls == 1:
+            raise RuntimeError("injected hook removal failure")
+        original_remove(self, name)
+
+    monkeypatch.setattr(HookRegistry, "remove_hook", fail_first_removal)
+
+    with pytest.raises(RuntimeError, match="Failed to remove one or more sequential offload hooks"):
+        backend.disable()
+
+    assert backend.enabled
+    assert backend._offload_modules == [pipeline.transformer, pipeline.text_encoder]
+    assert pipeline.transformer._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) is not None
+    assert pipeline.text_encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) is None
+
+    backend.disable()
+
+    assert not backend.enabled
+    assert not backend._offload_modules
+    for module in (pipeline.transformer, pipeline.text_encoder):
+        assert module._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) is None
+
+
+def test_sequential_offload_filters_cpu_eligible_components() -> None:
+    dit = _create_simple_module()
+    encoder = _create_simple_module()
+    resident_stage = _create_simple_module()
+
+    apply_sequential_offload(
+        dit_modules=[dit],
+        encoder_modules=[encoder, resident_stage],
+        device=torch.device("cpu"),
+        offload_dit_modules=[dit],
+        offload_encoder_modules=[encoder],
+    )
+
+    dit_hook = dit._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    encoder_hook = encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    resident_hook = resident_stage._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    assert dit_hook.offload_targets == [encoder]
+    assert encoder_hook.offload_targets == [dit]
+    assert encoder_hook.offload_after_context is True
+    assert resident_hook.offload_targets == [dit]
+    assert resident_hook.offload_after_context is False
+
+    remove_sequential_offload([dit, encoder, resident_stage])
+
+
+def test_move_params_handles_buffer_only_modules() -> None:
+    module = nn.Module()
+    module.register_buffer("state", torch.ones(2))
+
+    # An indexed CPU device compares differently while remaining usable on
+    # CPU-only test hosts, so the helper must inspect the buffer itself.
+    moved = SequentialOffloadHook._move_params(module, torch.device("cpu:1"))
+
+    assert moved
+    assert module.state.device.type == "cpu"
+
+
+def test_model_level_backend_keeps_declared_image_encoder_resident() -> None:
+    class MixedEncoderPipeline(nn.Module):
+        _offload_plan = OffloadPlan(
+            encoder_component_types={
+                "text_encoder": "text_encoder",
+            }
+        )
+
+        def __init__(self) -> None:
+            super().__init__()
+            self.transformer = _create_simple_module()
+            self.text_encoder = _create_simple_module()
+            self.image_encoder = _create_simple_module()
+
+    pipeline = MixedEncoderPipeline()
+    backend = ModelLevelOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.MODEL_LEVEL,
+            components=frozenset({"dit", "text_encoder"}),
+        ),
+        torch.device("cpu"),
+    )
+
+    backend.enable(pipeline)
+
+    text_hook = pipeline.text_encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    image_hook = pipeline.image_encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+    assert text_hook.offload_after_context is True
+    assert image_hook.offload_after_context is False
+
+    backend.disable()
+
+
+@pytest.mark.parametrize(
+    ("component", "attribute", "message"),
+    [
+        ("text_encoder", "text_encoder", "requires a DiT/transformer"),
+        ("dit", "transformer", "requires an encoder execution stage"),
+    ],
+)
+def test_component_selective_model_offload_requires_swap_counterpart(component, attribute, message) -> None:
+    pipeline = nn.Module()
+    setattr(pipeline, attribute, _create_simple_module())
+    backend = ModelLevelOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.MODEL_LEVEL,
+            components=frozenset({component}),
+        ),
+        torch.device("cpu"),
+    )
+
+    with pytest.raises(ValueError, match=message):
+        backend.enable(pipeline)
+
+
 def test_sequential_offload_can_begin_with_dit_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None:
     dit = _create_simple_module()
     encoder = _create_simple_module()