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

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

tests/diffusion/offloader/conftest.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Shared runtime fixtures for diffusion offloader tests."""
+
+import gc
+
+import pytest
+import torch.distributed as dist
+
+from tests.helpers.runtime import get_distributed_init_method
+from vllm_omni.platforms import current_omni_platform
+
+
+@pytest.fixture(scope="module")
+def dist_group():
+    dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method())
+    try:
+        yield
+    finally:
+        if dist.is_initialized():
+            dist.destroy_process_group()
+        gc.collect()
+        if current_omni_platform.is_available():
+            current_omni_platform.empty_cache()
+            current_omni_platform.synchronize()

tests/diffusion/offloader/helpers.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Shared model fixtures for diffusion offloader tests."""
+
+from contextlib import contextmanager
+
+import torch
+from torch import nn
+
+
+class DummyStream:
+    def wait_stream(self, _stream) -> None:
+        return None
+
+    def wait_event(self, _event) -> None:
+        return None
+
+
+class DummyEvent:
+    def record(self, _stream) -> None:
+        return None
+
+    def synchronize(self) -> None:
+        return None
+
+
+@contextmanager
+def dummy_stream(_stream):
+    yield None
+
+
+def patch_offload_runtime(monkeypatch, platform, *, synchronize: bool = False) -> None:
+    monkeypatch.setattr(platform, "Stream", DummyStream)
+    monkeypatch.setattr(platform, "Event", DummyEvent)
+    monkeypatch.setattr(platform, "current_stream", lambda: DummyStream())
+    monkeypatch.setattr(platform, "stream", dummy_stream)
+    if synchronize:
+        monkeypatch.setattr(platform, "synchronize", lambda: None)
+        monkeypatch.setattr(platform, "empty_cache", lambda: None)
+
+
+class _DummyBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.randn(10, 10))
+
+
+class _SingleBlockModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["blocks"]
+
+    def __init__(self, num_blocks: int = 3):
+        super().__init__()
+        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
+
+
+class _StagedEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.vision = nn.Module()
+        self.vision.blocks = nn.ModuleList([_DummyBlock(), _DummyBlock()])
+        self.text_model = nn.Module()
+        self.text_model.layers = nn.ModuleList([_DummyBlock(), _DummyBlock()])
+        self.load_calls = 0
+        self.offload_calls = 0
+        self.to_calls = 0
+
+    def load_to_device(self):
+        self.load_calls += 1
+
+    def offload_to_cpu(self):
+        self.offload_calls += 1
+        for hook in getattr(self, "_omni_layerwise_hooks", []):
+            hook.offload_layer()
+
+    def to(self, *args, **kwargs):
+        self.to_calls += 1
+        return super().to(*args, **kwargs)
+
+
+class _StagedVAE(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.proj = nn.Linear(2, 2)
+        self.offload_calls = 0
+        self.to_calls = 0
+
+    def load_to_device(self):
+        return None
+
+    def offload_to_cpu(self):
+        self.offload_calls += 1
+        return self.to("cpu")
+
+    def to(self, *args, **kwargs):
+        self.to_calls += 1
+        return super().to(*args, **kwargs)
+
+
+class _PlainEncoder(nn.Module):
+    """Standard encoder with no offload-specific lifecycle methods."""
+
+    def __init__(self):
+        super().__init__()
+        self.encoder = nn.Module()
+        self.encoder.block = nn.ModuleList([_DummyBlock(), _DummyBlock()])
+        self.final_norm = nn.Linear(2, 2)

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -6,8 +6,8 @@
 import gc
 import json
 import weakref
-from contextlib import contextmanager
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
@@ -17,7 +17,15 @@ from torch import nn
 from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
 
 import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module
-from tests.helpers.runtime import get_distributed_init_method
+from tests.diffusion.offloader.helpers import (
+    DummyStream,
+    _DummyBlock,
+    _PlainEncoder,
+    _SingleBlockModel,
+    _StagedEncoder,
+    _StagedVAE,
+    patch_offload_runtime,
+)
 from vllm_omni.diffusion.data import validate_dlo_host_registration_options
 from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
@@ -51,53 +59,9 @@ from vllm_omni.platforms import current_omni_platform
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
 
 
-class DummyStream:
-    def wait_stream(self, _stream) -> None:
-        return None
-
-    def wait_event(self, _event) -> None:
-        return None
-
-
-class DummyEvent:
-    def record(self, _stream) -> None:
-        return None
-
-    def synchronize(self) -> None:
-        return None
-
-
-@contextmanager
-def dummy_stream(_stream):
-    yield None
-
-
-def _cleanup_distributed() -> None:
-    if dist.is_initialized():
-        dist.destroy_process_group()
-
-    gc.collect()
-    if current_omni_platform.is_available():
-        current_omni_platform.empty_cache()
-        current_omni_platform.synchronize()
-
-
-@pytest.fixture(scope="module")
-def dist_group():
-    dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method())
-    try:
-        yield
-    finally:
-        _cleanup_distributed()
-
-
 @pytest.fixture
 def patched_offload_runtime(monkeypatch):
-    monkeypatch.setattr(dist_backend_module.current_omni_platform, "Stream", DummyStream)
-    monkeypatch.setattr(dist_backend_module.current_omni_platform, "Event", DummyEvent)
-    monkeypatch.setattr(dist_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
-    monkeypatch.setattr(dist_backend_module.current_omni_platform, "stream", dummy_stream)
-    monkeypatch.setattr(dist_backend_module.current_omni_platform, "synchronize", lambda: None)
+    patch_offload_runtime(monkeypatch, dist_backend_module.current_omni_platform, synchronize=True)
 
 
 class TinyBlock(nn.Module):
@@ -113,6 +77,93 @@ def _make_values(start: float) -> torch.Tensor:
 
 
 class TestDistributedLayerwiseOffloadHook:
+    def test_buffer_only_block_reports_offloaded_state(self, patched_offload_runtime):
+        current_block = nn.Module()
+        current_block.register_buffer("state", torch.ones(2))
+        next_block = nn.Module()
+        next_block.register_buffer("state", torch.ones(2))
+        hook = DistributedLayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            dp_group=None,
+            dp_size=1,
+            rank=0,
+            pin_memory=False,
+        )
+        hook.initialize_hook(current_block)
+
+        assert hook.is_materialized
+        hook.offload_layer()
+        assert not hook.is_materialized
+
+    def test_initialize_failure_keeps_next_block_materialized(self, monkeypatch):
+        current_block = nn.Linear(2, 2)
+        next_block = nn.Linear(2, 2)
+        expected = {name: tensor.detach().clone() for name, tensor in next_block.state_dict().items()}
+        hook = DistributedLayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            dp_group=None,
+            dp_size=1,
+            rank=0,
+            pin_memory=False,
+        )
+
+        def fail_buffer_allocation():
+            raise RuntimeError("injected buffer allocation failure")
+
+        monkeypatch.setattr(hook, "_allocate_device_buffers", fail_buffer_allocation)
+
+        with pytest.raises(RuntimeError, match="buffer allocation failure"):
+            hook.initialize_hook(current_block)
+
+        for name, tensor in next_block.state_dict().items():
+            torch.testing.assert_close(tensor, expected[name])
+
+    @pytest.mark.parametrize(
+        ("kwargs", "message"),
+        [
+            ({"dp_size": 0, "rank": 0}, "dp_size must be a positive integer"),
+            ({"dp_size": 2, "rank": 2, "dp_group": object()}, "rank must satisfy"),
+            ({"dp_size": 2, "rank": 0}, "dp_group is required"),
+            (
+                {"dp_size": 2, "rank": 0, "dp_group": object(), "rank_local_mmap": True},
+                "rank_local_mmap requires dp_size=1",
+            ),
+        ],
+    )
+    def test_constructor_rejects_invalid_transport_topology(self, kwargs, message):
+        options = {"dp_group": None, "dp_size": 1, "rank": 0, **kwargs}
+
+        with pytest.raises(ValueError, match=message):
+            DistributedLayerwiseOffloadHook(
+                next_block=nn.Linear(2, 2),
+                device=torch.device("cpu"),
+                pin_memory=False,
+                **options,
+            )
+
+    def test_ring_probes_are_captured_before_any_block_is_cleared(self, patched_offload_runtime):
+        blocks = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(3)])
+        expected_middle = blocks[1].weight.detach().clone()
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+            ),
+            torch.device("cpu"),
+        )
+
+        hooks = backend._install_hook_group(blocks, "dit")
+        shared_buffers = backend._allocate_shared_buffers(hooks)
+        for hook in hooks:
+            hook.gpu_buffers = shared_buffers
+
+        assert all(not hook.is_materialized for hook in hooks)
+        hooks[2].pre_forward(blocks[1])
+        torch.testing.assert_close(blocks[1].weight, expected_middle)
+
     def test_shard_and_pin_single_rank(self, dist_group, patched_offload_runtime):
         """With dp_size=1, the shard should equal the full weights."""
         current_block = TinyBlock(_make_values(1.0))
@@ -511,7 +562,7 @@ class TestDistributedLayerwiseOffloadHook:
         hook = DistributedLayerwiseOffloadHook(
             next_block=next_block,
             device=torch.device("cpu"),
-            dp_group=None,
+            dp_group=object(),
             dp_size=4,
             rank=1,
             pin_memory=False,
@@ -537,7 +588,7 @@ class TestDistributedLayerwiseOffloadHook:
         hook0 = DistributedLayerwiseOffloadHook(
             next_block=next_block,
             device=torch.device("cpu"),
-            dp_group=None,
+            dp_group=object(),
             dp_size=2,
             rank=0,
             pin_memory=False,
@@ -559,7 +610,7 @@ class TestDistributedLayerwiseOffloadHook:
         hook1 = DistributedLayerwiseOffloadHook(
             next_block=next_block1,
             device=torch.device("cpu"),
-            dp_group=None,
+            dp_group=object(),
             dp_size=2,
             rank=1,
             pin_memory=False,
@@ -583,7 +634,7 @@ class TestDistributedLayerwiseOffloadHook:
         hook = DistributedLayerwiseOffloadHook(
             next_block=next_block,
             device=torch.device("cpu"),
-            dp_group=None,
+            dp_group=object(),
             dp_size=2,
             rank=0,
             pin_memory=False,
@@ -602,7 +653,7 @@ class TestDistributedLayerwiseOffloadHook:
         hook = DistributedLayerwiseOffloadHook(
             next_block=next_block,
             device=torch.device("cpu"),
-            dp_group=None,
+            dp_group=object(),
             dp_size=2,
             rank=0,
             pin_memory=False,
@@ -689,6 +740,32 @@ class TestDistributedLayerwiseOffloadHook:
 
 
 class TestPinnedResidentLayerGroup:
+    def test_constructor_failure_leaves_every_block_materialized(self, monkeypatch):
+        blocks = [nn.Linear(2, 2), nn.Linear(2, 2)]
+        expected = [block.weight.detach().clone() for block in blocks]
+        shard_and_pin = DistributedLayerwiseOffloadHook._shard_and_pin
+        calls = 0
+
+        def fail_second_block(*args, **kwargs):
+            nonlocal calls
+            calls += 1
+            if calls == 2:
+                raise RuntimeError("injected resident setup failure")
+            return shard_and_pin(*args, **kwargs)
+
+        monkeypatch.setattr(DistributedLayerwiseOffloadHook, "_shard_and_pin", fail_second_block)
+
+        with pytest.raises(RuntimeError, match="resident setup failure"):
+            PinnedResidentLayerGroup(
+                blocks,
+                device=torch.device("cpu"),
+                copy_stream=DummyStream(),
+                pin_memory=False,
+            )
+
+        for block, weight in zip(blocks, expected, strict=True):
+            torch.testing.assert_close(block.weight, weight)
+
     def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
         blocks = [nn.Linear(2, 2), nn.Linear(2, 2)]
         expected = []
@@ -843,20 +920,6 @@ class TestPinnedModuleStager:
         stager.offload()
 
 
-class _DummyBlock(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.weight = nn.Parameter(torch.randn(10, 10))
-
-
-class _SingleBlockModel(nn.Module):
-    _layerwise_offload_blocks_attrs = ["blocks"]
-
-    def __init__(self, num_blocks: int = 3):
-        super().__init__()
-        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
-
-
 class _HWRPipeline(nn.Module):
     def __init__(self):
         super().__init__()
@@ -1285,6 +1348,31 @@ class TestMmapWeightLoading:
         backend.disable()
         assert lease.closed
 
+        backend.enable(pipeline)
+        assert backend.host_weight_plan is None
+        assert backend.enabled
+        backend.disable()
+
+    def test_rank_local_disable_drains_pending_transfers(self, patched_offload_runtime, monkeypatch):
+        pipeline = _HWRPipeline()
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=1,
+                dlo_use_allgather=False,
+            ),
+            torch.device("cpu"),
+        )
+        backend.enable(pipeline)
+        backend._all_hook_groups[0][0].prefetch_layer(0)
+        synchronize = Mock()
+        monkeypatch.setattr(current_omni_platform, "synchronize", synchronize)
+
+        backend.disable()
+
+        synchronize.assert_called()
+
     def test_hwr_registration_failure_falls_back_to_bounded_staging(
         self,
         monkeypatch,
@@ -1563,6 +1651,7 @@ class TestOffloadPlan:
         assert plan.block_attrs == {}
         assert plan.offload_submodules == {}
         assert plan.resident_dit_paths == frozenset()
+        assert plan.encoder_dlo_weight_replication == frozenset()
 
     def test_offload_plan_is_frozen(self):
         """OffloadPlan should be immutable (frozen=True)."""
@@ -1595,6 +1684,7 @@ class TestOffloadPlan:
                 self.transformer = Transformer()
 
         pipeline = Pipeline()
+        expected_parameters = {name: parameter.detach().clone() for name, parameter in pipeline.named_parameters()}
         backend = DistributedLayerwiseOffloadBackend(
             OffloadConfig(
                 strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
@@ -1610,10 +1700,22 @@ class TestOffloadPlan:
         assert len(backend._resident_blocks) == 2
         assert len(backend._all_hook_groups) == 1
         assert len(backend._all_hook_groups[0]) == 2
+        assert pipeline._dlo_residency_controller is backend
         assert backend.enabled
 
         backend.disable()
 
+        assert pipeline._dlo_residency_controller is None
+
+        for name, parameter in pipeline.named_parameters():
+            torch.testing.assert_close(parameter, expected_parameters[name])
+
+        backend.enable(pipeline)
+        assert len(backend._resident_blocks) == 2
+        backend.disable()
+        for name, parameter in pipeline.named_parameters():
+            torch.testing.assert_close(parameter, expected_parameters[name])
+
 
 class TestMmapValidation:
     """Tests for loader preflight and backend plan realization."""
@@ -2024,8 +2126,8 @@ class TestConfigValidation:
             (2, 1, 1, True, 2),
             (1, 4, 1, True, 4),
             (2, 2, 2, True, 2),
-            (2, 2, 2, False, 1),
-            (1, 4, 2, False, 1),
+            (2, 2, 2, False, 2),
+            (1, 4, 2, False, 4),
         ],
     )
     def test_dlo_group_selection_covers_dp_sp_tp_matrix(
@@ -2199,7 +2301,7 @@ class TestConfigValidation:
             model = "/fake/path"
 
         config = OffloadConfig.from_od_config(FakeODConfig())
-        assert config.dp_size == 1  # forced to 1 when no AllGather
+        assert config.dp_size == 2  # topology is retained; the hook selects group size 1
         assert config.dlo_resident_layers == 20
 
     def test_resident_layers_with_allgather_rejected(self):
@@ -2218,96 +2320,9 @@ class TestConfigValidation:
             parallel_config = FakePC()
             model = "/fake/path"
 
-        with pytest.raises(ValueError, match="requires --dlo-no-use-allgather"):
+        with pytest.raises(ValueError, match="requires the DiT DLO transfer to be rank-local"):
             OffloadConfig.from_od_config(FakeODConfig())
 
-    def test_num_inference_steps_none_rejected(self):
-        """DP multi-concurrency should reject None num_inference_steps."""
-        from types import SimpleNamespace
-
-        # Mock requests with None steps
-        reqs = [
-            SimpleNamespace(
-                req=SimpleNamespace(
-                    request_id=f"req-{i}",
-                    sampling_params=SimpleNamespace(num_inference_steps=None),
-                )
-            )
-            for i in range(2)
-        ]
-
-        # We can't easily instantiate the full executor, but we can test
-        # the validation logic by checking that the code path raises.
-        # The validation is in execute_request, which needs self._ensure_open().
-        # Instead, test the validation logic directly:
-        step_counts = {
-            r.req.sampling_params.num_inference_steps
-            for r in reqs
-            if r.req.sampling_params.num_inference_steps is not None
-        }
-        has_none = any(r.req.sampling_params.num_inference_steps is None for r in reqs)
-        assert has_none, "Test setup: should have None steps"
-        assert len(step_counts) == 0, "Test setup: no explicit steps"
-
-        # The validation condition: (len(step_counts) > 1) or has_none → should reject
-        should_reject = (len(step_counts) > 1) or has_none
-        assert should_reject, "None steps should trigger rejection"
-
-    def test_num_inference_steps_same_explicit_allowed(self):
-        """DP multi-concurrency should allow same explicit num_inference_steps."""
-        from types import SimpleNamespace
-
-        reqs = [
-            SimpleNamespace(
-                req=SimpleNamespace(
-                    request_id=f"req-{i}",
-                    sampling_params=SimpleNamespace(num_inference_steps=35),
-                )
-            )
-            for i in range(4)
-        ]
-
-        step_counts = {
-            r.req.sampling_params.num_inference_steps
-            for r in reqs
-            if r.req.sampling_params.num_inference_steps is not None
-        }
-        has_none = any(r.req.sampling_params.num_inference_steps is None for r in reqs)
-
-        should_reject = (len(step_counts) > 1) or has_none
-        assert not should_reject, "Same explicit steps should be allowed"
-        assert step_counts == {35}
-
-    def test_num_inference_steps_different_explicit_rejected(self):
-        """DP multi-concurrency should reject different explicit steps."""
-        from types import SimpleNamespace
-
-        reqs = [
-            SimpleNamespace(
-                req=SimpleNamespace(
-                    request_id="req-0",
-                    sampling_params=SimpleNamespace(num_inference_steps=35),
-                )
-            ),
-            SimpleNamespace(
-                req=SimpleNamespace(
-                    request_id="req-1",
-                    sampling_params=SimpleNamespace(num_inference_steps=30),
-                )
-            ),
-        ]
-
-        step_counts = {
-            r.req.sampling_params.num_inference_steps
-            for r in reqs
-            if r.req.sampling_params.num_inference_steps is not None
-        }
-        has_none = any(r.req.sampling_params.num_inference_steps is None for r in reqs)
-
-        should_reject = (len(step_counts) > 1) or has_none
-        assert should_reject, "Different steps should trigger rejection"
-        assert len(step_counts) == 2
-
 
 class TestDynamicSlotTracking:
     """Tests for dynamic slot tracking with odd and even block counts.
@@ -2472,3 +2487,344 @@ class TestDynamicSlotTracking:
         # pre_forward should override to slot 1
         hook_b.pre_forward(block_b)
         assert hook_b.current_slot == 1, "pre_forward should read _prev_hook._prefetched_slot=1, not keep initial 0"
+
+
+class _DistributedComponentPipeline(nn.Module):
+    _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"}),
+    )
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _SingleBlockModel(num_blocks=2)
+        self.text_encoder = _StagedEncoder()
+        self.vae = _StagedVAE()
+
+
+class _GenericDistributedEncoderPipeline(nn.Module):
+    _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=2)
+        self.text_encoder = _PlainEncoder()
+
+
+class TestDistributedComponentSelection:
+    def test_on_demand_only_encoder_cannot_claim_allgather(self, patched_offload_runtime):
+        class Pipeline(nn.Module):
+            _offload_plan = OffloadPlan(
+                encoder_component_types={"text_encoder": "text_encoder"},
+                on_demand_component_paths=frozenset({"text_encoder"}),
+                encoder_dlo_weight_replication=frozenset({"text_encoder"}),
+            )
+
+            def __init__(self):
+                super().__init__()
+                self.text_encoder = _StagedEncoder()
+
+        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"),
+        )
+        backend.dp_group = object()
+
+        with pytest.raises(ValueError, match="cannot use AllGather without.*streamable block plan"):
+            backend.enable(Pipeline())
+
+    def test_explicit_single_block_dit_is_rejected(self, patched_offload_runtime):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=1)
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                components=frozenset({"dit"}),
+            ),
+            torch.device("cpu"),
+        )
+
+        with pytest.raises(ValueError, match="leaves only one streaming block"):
+            backend.enable(pipeline)
+
+    def test_enable_defers_collectives_until_first_forward(self, patched_offload_runtime, monkeypatch):
+        pipeline = _GenericDistributedEncoderPipeline()
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=2,
+                components=frozenset({"dit", "text_encoder"}),
+                dlo_transfers={"dit": "allgather", "text_encoder": "rank-local"},
+            ),
+            torch.device("cpu"),
+        )
+        backend.dp_group = object()
+
+        def reject_startup_collective(*_args, **_kwargs):
+            raise AssertionError("enable() entered a rank-asymmetric collective")
+
+        monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", reject_startup_collective)
+
+        backend.enable(pipeline)
+
+        assert all(hook._prefetched_slot is None for group in backend._all_hook_groups for hook in group)
+        encoder_hook = pipeline.text_encoder.encoder.block[0]._hook_registry.get_hook("distributed_layerwise_offload")
+        dit_hook = pipeline.transformer.blocks[0]._hook_registry.get_hook("distributed_layerwise_offload")
+        assert encoder_hook.gpu_shard_buffers == [None, None]
+        assert all(buffer is not None for buffer in dit_hook.gpu_shard_buffers)
+
+        monkeypatch.setattr(
+            torch.distributed,
+            "all_gather_into_tensor",
+            lambda output, local, group: output.copy_(local.repeat(2)),
+        )
+        backend.disable()
+
+    def test_partial_dit_enable_failure_restores_weights_and_hooks(
+        self,
+        patched_offload_runtime,
+        monkeypatch,
+    ):
+        pipeline = _DistributedComponentPipeline()
+        expected_weights = [block.weight.detach().clone() for block in pipeline.transformer.blocks]
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+                components=frozenset({"dit"}),
+            ),
+            torch.device("cpu"),
+        )
+        original_apply = dist_backend_module.apply_distributed_block_hook
+        calls = 0
+
+        def fail_second_hook(*args, **kwargs):
+            nonlocal calls
+            calls += 1
+            if calls == 2:
+                raise RuntimeError("injected DLO hook failure")
+            return original_apply(*args, **kwargs)
+
+        monkeypatch.setattr(dist_backend_module, "apply_distributed_block_hook", fail_second_hook)
+
+        with pytest.raises(RuntimeError, match="injected DLO hook failure"):
+            backend.enable(pipeline)
+
+        assert not backend.enabled
+        for block, expected in zip(pipeline.transformer.blocks, expected_weights, strict=True):
+            registry = getattr(block, "_hook_registry", None)
+            assert registry is None or registry.get_hook("distributed_layerwise_offload") is None
+            torch.testing.assert_close(block.weight, expected)
+
+    def test_multirank_enable_failure_cleanup_skips_restore_collective(self, monkeypatch, mocker):
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=2,
+                dlo_use_allgather=True,
+                components=frozenset({"dit"}),
+            ),
+            torch.device("cpu"),
+        )
+        monkeypatch.setattr(backend, "_enable", mocker.Mock(side_effect=RuntimeError("injected startup failure")))
+        cleanup = mocker.Mock()
+        monkeypatch.setattr(backend, "_disable", cleanup)
+
+        with pytest.raises(RuntimeError, match="injected startup failure"):
+            backend.enable(_DistributedComponentPipeline())
+
+        cleanup.assert_called_once_with(restore_allgather_weights=False)
+
+    def test_multirank_failure_restores_only_rank_local_hooks(self, patched_offload_runtime):
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=2,
+                components=frozenset({"dit", "text_encoder"}),
+                dlo_transfers={"dit": "allgather", "text_encoder": "rank-local"},
+            ),
+            torch.device("cpu"),
+        )
+        allgather_hook = Mock(dp_size=2, next_block=nn.Linear(1, 1))
+        rank_local_hook = Mock(dp_size=1, next_block=nn.Linear(1, 1))
+        backend._all_hook_groups = [[allgather_hook], [rank_local_hook]]
+
+        backend._disable(restore_allgather_weights=False)
+
+        allgather_hook.restore_next_block_to_cpu.assert_not_called()
+        rank_local_hook.restore_next_block_to_cpu.assert_called_once_with()
+
+        with pytest.raises(RuntimeError, match="recreate the backend and reload the pipeline"):
+            backend.enable(nn.Module())
+
+    def test_disable_runs_collectives_before_best_effort_local_cleanup(
+        self,
+        patched_offload_runtime,
+        monkeypatch,
+    ):
+        events: list[str] = []
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=2,
+            ),
+            torch.device("cpu"),
+        )
+        rank_local_hook = Mock(dp_size=1, next_block=nn.Linear(1, 1))
+        allgather_hook = Mock(dp_size=2, next_block=nn.Linear(1, 1))
+        allgather_hook.restore_next_block_to_cpu.side_effect = lambda: events.append("allgather")
+
+        def fail_rank_local_restore():
+            events.append("rank-local")
+            raise RuntimeError("injected rank-local restore failure")
+
+        rank_local_hook.restore_next_block_to_cpu.side_effect = fail_rank_local_restore
+        backend._all_hook_groups = [[rank_local_hook], [allgather_hook]]
+        backend._blocks = [[nn.Module()]]
+        backend.enabled = True
+        monkeypatch.setattr(
+            dist_backend_module,
+            "remove_distributed_block_hook",
+            lambda _block: events.append("remove"),
+        )
+        monkeypatch.setattr(backend, "_release_registered_mmap", lambda: events.append("unregister"))
+        release_handles = Mock()
+        monkeypatch.setattr(backend, "_release_mmap_handles", release_handles)
+
+        with pytest.raises(RuntimeError, match="injected rank-local restore failure"):
+            backend.disable()
+
+        assert events == ["allgather", "rank-local", "remove", "unregister"]
+        release_handles.assert_not_called()
+        assert backend.enabled
+        assert backend._all_hook_groups
+
+    def test_all_streams_encoder_and_keeps_vae_resident(self, patched_offload_runtime, monkeypatch):
+        pipeline = _DistributedComponentPipeline()
+        move_non_block_state = Mock()
+        monkeypatch.setattr(dist_backend_module, "move_non_block_state_to_device", move_non_block_state)
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+                components=frozenset({"dit", "text_encoder"}),
+            ),
+            torch.device("cpu"),
+        )
+
+        backend.enable(pipeline)
+
+        assert pipeline.text_encoder._omni_layerwise_enabled
+        assert len(pipeline.text_encoder._omni_layerwise_hooks) == 4
+        # The pipeline-owned stage lifecycle releases the encoder's non-block
+        # state after DLO has installed blockwise hooks.
+        assert pipeline.text_encoder.offload_calls == 1
+        assert pipeline.vae.offload_calls == 0
+        assert pipeline.vae.to_calls == 1
+        move_non_block_state.assert_not_called()
+
+        backend.disable()
+
+        assert not pipeline.text_encoder._omni_layerwise_enabled
+        encoder_blocks = [
+            *pipeline.text_encoder.vision.blocks,
+            *pipeline.text_encoder.text_model.layers,
+        ]
+        assert all(block._hook_registry.get_hook("distributed_layerwise_offload") is None for block in encoder_blocks)
+
+    def test_dit_and_encoder_choose_transfer_independently(self, patched_offload_runtime, monkeypatch):
+        pipeline = _GenericDistributedEncoderPipeline()
+        expected_parameters = {name: parameter.detach().clone() for name, parameter in pipeline.named_parameters()}
+        full_weights_by_rank_zero_shard = {
+            tuple(parameter.flatten()[: (parameter.numel() + 1) // 2].tolist()): parameter.flatten()
+            for parameter in expected_parameters.values()
+            if parameter.numel() > 4
+        }
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=2,
+                components=frozenset({"dit", "text_encoder"}),
+                dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"},
+            ),
+            torch.device("cpu"),
+        )
+        backend.dp_group = object()
+
+        def fake_allgather(output, local, group):
+            del group
+            output.copy_(full_weights_by_rank_zero_shard[tuple(local.tolist())])
+
+        monkeypatch.setattr(
+            torch.distributed,
+            "all_gather_into_tensor",
+            fake_allgather,
+        )
+
+        backend.enable(pipeline)
+
+        encoder_hook = pipeline.text_encoder.encoder.block[0]._hook_registry.get_hook("distributed_layerwise_offload")
+        dit_hook = pipeline.transformer.blocks[0]._hook_registry.get_hook("distributed_layerwise_offload")
+        assert encoder_hook.dp_size == 2
+        assert dit_hook.dp_size == 1
+        assert all(buffer is not None for buffer in encoder_hook.gpu_shard_buffers)
+        assert dit_hook.gpu_shard_buffers == [None, None]
+
+        backend.disable()
+        actual_parameters = dict(pipeline.named_parameters())
+        for name, expected in expected_parameters.items():
+            torch.testing.assert_close(actual_parameters[name], expected)
+
+    def test_encoder_allgather_rejects_undeclared_replication(self, patched_offload_runtime):
+        pipeline = _DistributedComponentPipeline()
+        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"),
+        )
+        backend.dp_group = object()
+
+        with pytest.raises(ValueError, match="not declared replicated"):
+            backend.enable(pipeline)
+
+    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"),
+        )
+        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)

tests/diffusion/offloader/test_layerwise_backend.py

@@ -3,66 +3,39 @@
 
 """Unit tests for LayerwiseOffloadHook and LayerWiseOffloadBackend utilities."""
 
-import gc
-from contextlib import contextmanager
+from multiprocessing.reduction import ForkingPickler
+from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
-import torch.distributed as dist
 from torch import nn
 from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
 
 import vllm_omni.diffusion.offloader.layerwise_backend as layerwise_backend_module
-from tests.helpers.runtime import get_distributed_init_method
-from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend, LayerwiseOffloadHook
-from vllm_omni.platforms import current_omni_platform
+import vllm_omni.diffusion.offloader.tensor_utils as tensor_utils_module
+from tests.diffusion.offloader.helpers import (
+    DummyStream,
+    _DummyBlock,
+    _PlainEncoder,
+    _SingleBlockModel,
+    _StagedEncoder,
+    _StagedVAE,
+    patch_offload_runtime,
+)
+from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+from vllm_omni.diffusion.offloader.layerwise_backend import (
+    LayerWiseOffloadBackend,
+    LayerwiseOffloadHook,
+)
+from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
 
 
-class DummyStream:
-    def wait_stream(self, _stream) -> None:
-        return None
-
-    def wait_event(self, _event) -> None:
-        return None
-
-
-class DummyEvent:
-    def record(self, _stream) -> None:
-        return None
-
-
-@contextmanager
-def dummy_stream(_stream):
-    yield None
-
-
-def _cleanup_distributed() -> None:
-    if dist.is_initialized():
-        dist.destroy_process_group()
-
-    gc.collect()
-    if current_omni_platform.is_available():
-        current_omni_platform.empty_cache()
-        current_omni_platform.synchronize()
-
-
-@pytest.fixture(scope="module")
-def dist_group():
-    dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method())
-    try:
-        yield
-    finally:
-        _cleanup_distributed()
-
-
 @pytest.fixture
 def patched_offload_runtime(monkeypatch):
-    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Stream", DummyStream)
-    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Event", DummyEvent)
-    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
-    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "stream", dummy_stream)
+    patch_offload_runtime(monkeypatch, layerwise_backend_module.current_omni_platform)
 
 
 class TinyBlock(nn.Module):
@@ -77,7 +50,66 @@ def _make_values(start: float) -> torch.Tensor:
     return torch.arange(start, start + 4, dtype=torch.float32)
 
 
+def test_clear_tensor_storage_rolls_back_partial_commit(monkeypatch):
+    first = nn.Parameter(torch.tensor([1.0, 2.0]))
+    second = nn.Parameter(torch.tensor([3.0, 4.0]))
+    expected_first = first.detach().clone()
+    expected_second = second.detach().clone()
+    second_data_ptr = second.data_ptr()
+    original_set_storage = tensor_utils_module.set_tensor_storage
+    calls = 0
+
+    def fail_second_commit(target, value):
+        nonlocal calls
+        calls += 1
+        if calls == 2:
+            raise RuntimeError("injected storage clear failure")
+        original_set_storage(target, value)
+
+    monkeypatch.setattr(tensor_utils_module, "set_tensor_storage", fail_second_commit)
+
+    with pytest.raises(RuntimeError, match="injected storage clear failure"):
+        tensor_utils_module.clear_tensor_storage([first, second])
+
+    assert calls == 3  # first commit, failed second commit, first rollback
+    torch.testing.assert_close(first, expected_first)
+    torch.testing.assert_close(second, expected_second)
+    assert second.data_ptr() == second_data_ptr
+
+
 class TestLayerwiseOffloadHook:
+    def test_buffer_only_block_reports_offloaded_state(self, patched_offload_runtime):
+        current_block = nn.Module()
+        current_block.register_buffer("state", torch.ones(2))
+        next_block = nn.Module()
+        next_block.register_buffer("state", torch.ones(2))
+        hook = LayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            stream=DummyStream(),
+            pin_memory=False,
+        )
+        hook.initialize_hook(current_block)
+
+        assert hook.is_materialized
+        hook.offload_layer()
+        assert not hook.is_materialized
+
+    def test_ring_probes_are_captured_before_any_block_is_cleared(self, patched_offload_runtime):
+        blocks = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(3)])
+        expected_middle = blocks[1].weight.detach().clone()
+
+        hooks = layerwise_backend_module._install_layerwise_hook_group(
+            blocks,
+            torch.device("cpu"),
+            DummyStream(),
+            pin_memory=False,
+        )
+
+        assert all(not hook.is_materialized for hook in hooks)
+        hooks[2].pre_forward(blocks[1])
+        torch.testing.assert_close(blocks[1].weight, expected_middle)
+
     def test_dtensor_wrapper_is_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime):
         current_block = TinyBlock(_make_values(1.0))
         next_block = TinyBlock(_make_values(10.0))
@@ -135,20 +167,6 @@ class TestLayerwiseOffloadHook:
         assert torch.equal(next_block.weight, expected)
 
 
-class _DummyBlock(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.weight = nn.Parameter(torch.randn(10, 10))
-
-
-class _SingleBlockModel(nn.Module):
-    _layerwise_offload_blocks_attrs = ["blocks"]
-
-    def __init__(self, num_blocks: int = 3):
-        super().__init__()
-        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
-
-
 class _MultiBlockModel(nn.Module):
     _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
 
@@ -246,3 +264,732 @@ class TestGetBlocksAttrNames:
         LayerWiseOffloadBackend.set_blocks_attr_names(model, ["new_blocks"])
         assert hasattr(model.__class__, "_layerwise_offload_blocks_attrs")
         assert model.__class__._layerwise_offload_blocks_attrs == ["new_blocks"]
+
+
+class _ComponentPipeline(nn.Module):
+    _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"}),
+    )
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _SingleBlockModel()
+        self.text_encoder = _StagedEncoder()
+        self.vae = _StagedVAE()
+
+
+class _LegacyComponentPipeline(nn.Module):
+    """Pipeline with DiT block metadata but no auxiliary OffloadPlan."""
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _SingleBlockModel()
+        self.text_encoder = _StagedEncoder()
+        self.vae = _StagedVAE()
+
+
+class _LegacyEncoderOnlyPipeline(nn.Module):
+    _offload_plan = _ComponentPipeline._offload_plan
+
+    def __init__(self):
+        super().__init__()
+        self.text_encoder = _StagedEncoder()
+        self.vae = _StagedVAE()
+
+
+class _PlannedDitPipeline(nn.Module):
+    _offload_plan = OffloadPlan(block_attrs={"transformer": ("blocks",)})
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _NoAttrsModel()
+
+
+class _GenericEncoderPipeline(nn.Module):
+    _offload_plan = OffloadPlan(
+        encoder_component_types={"text_encoder": "text_encoder"},
+        encoder_block_attrs={"text_encoder": ("encoder.block",)},
+    )
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = _SingleBlockModel()
+        self.text_encoder = _PlainEncoder()
+
+
+def _layer_backend(components: set[str] | None = None) -> LayerWiseOffloadBackend:
+    options = {}
+    if components is not None:
+        options = {"components": frozenset(components)}
+    return LayerWiseOffloadBackend(
+        OffloadConfig(
+            strategy=OffloadStrategy.LAYER_WISE,
+            pin_cpu_memory=False,
+            **options,
+        ),
+        torch.device("cpu"),
+    )
+
+
+class TestLayerwiseComponentSelection:
+    def test_plan_block_attrs_drive_dit_discovery(self, patched_offload_runtime):
+        pipeline = _PlannedDitPipeline()
+        backend = _layer_backend({"dit"})
+
+        backend.enable(pipeline)
+
+        assert hasattr(pipeline.transformer.blocks[0], "_hook_registry")
+        backend.disable()
+
+    def test_legacy_missing_dit_preserves_noop(self, patched_offload_runtime):
+        pipeline = _LegacyEncoderOnlyPipeline()
+        backend = _layer_backend()
+
+        backend.enable(pipeline)
+
+        assert not backend.enabled
+        assert pipeline.text_encoder.offload_calls == 0
+        assert pipeline.vae.offload_calls == 0
+
+    def test_encoder_only_streams_planned_blocks(self, patched_offload_runtime, monkeypatch):
+        pipeline = _ComponentPipeline()
+        backend = _layer_backend({"text_encoder"})
+        move_non_block_state = Mock()
+        monkeypatch.setattr(layerwise_backend_module, "move_non_block_state_to_device", move_non_block_state)
+
+        backend.enable(pipeline)
+
+        assert pipeline.text_encoder._omni_layerwise_enabled
+        assert len(pipeline.text_encoder._omni_layerwise_hooks) == 4
+        assert pipeline.text_encoder.offload_calls == 1
+        assert pipeline.vae.offload_calls == 0
+        assert pipeline.vae.to_calls == 1
+        assert not hasattr(pipeline.transformer.blocks[0], "_hook_registry")
+        assert backend.enabled
+        move_non_block_state.assert_not_called()
+
+        backend.disable()
+        assert not pipeline.text_encoder._omni_layerwise_enabled
+
+    def test_single_gpu_dit_only_keeps_encoder_and_vae_resident(self, patched_offload_runtime):
+        pipeline = _ComponentPipeline()
+        expected_weights = [block.weight.detach().clone() for block in pipeline.transformer.blocks]
+        backend = _layer_backend({"dit"})
+
+        backend.enable(pipeline)
+
+        assert not hasattr(pipeline.text_encoder, "_omni_layerwise_enabled")
+        assert pipeline.text_encoder.to_calls == 1
+        assert pipeline.text_encoder.offload_calls == 0
+        assert pipeline.vae.to_calls == 1
+        assert pipeline.vae.offload_calls == 0
+        assert hasattr(pipeline.transformer.blocks[0], "_hook_registry")
+
+        backend.disable()
+
+        for block, expected in zip(pipeline.transformer.blocks, expected_weights, strict=True):
+            torch.testing.assert_close(block.weight, expected)
+
+        backend.enable(pipeline)
+        backend.disable()
+        for block, expected in zip(pipeline.transformer.blocks, expected_weights, strict=True):
+            torch.testing.assert_close(block.weight, expected)
+
+    def test_partial_dit_enable_failure_restores_weights_and_hooks(self, patched_offload_runtime, monkeypatch):
+        pipeline = _ComponentPipeline()
+        expected_weights = [block.weight.detach().clone() for block in pipeline.transformer.blocks]
+        backend = _layer_backend({"dit"})
+        original_apply = layerwise_backend_module.apply_block_hook
+        calls = 0
+
+        def fail_second_hook(*args, **kwargs):
+            nonlocal calls
+            calls += 1
+            if calls == 2:
+                raise RuntimeError("injected hook failure")
+            return original_apply(*args, **kwargs)
+
+        monkeypatch.setattr(layerwise_backend_module, "apply_block_hook", fail_second_hook)
+
+        with pytest.raises(RuntimeError, match="injected hook failure"):
+            backend.enable(pipeline)
+
+        assert not backend.enabled
+        for block, expected in zip(pipeline.transformer.blocks, expected_weights, strict=True):
+            registry = getattr(block, "_hook_registry", None)
+            assert registry is None or registry.get_hook("layerwise_offload") is None
+            torch.testing.assert_close(block.weight, expected)
+
+    def test_later_component_failure_removes_installed_encoder_hooks(
+        self,
+        patched_offload_runtime,
+        monkeypatch,
+    ):
+        pipeline = _ComponentPipeline()
+        blocks = [*pipeline.text_encoder.vision.blocks, *pipeline.text_encoder.text_model.layers]
+        expected = [block.weight.detach().clone() for block in blocks]
+        backend = _layer_backend({"text_encoder"})
+
+        def fail_vae_placement(*_args, **_kwargs):
+            raise RuntimeError("injected VAE placement failure")
+
+        monkeypatch.setattr(pipeline.vae, "to", fail_vae_placement)
+
+        with pytest.raises(RuntimeError, match="injected VAE placement failure"):
+            backend.enable(pipeline)
+
+        assert not backend.enabled
+        assert not pipeline.text_encoder._omni_layerwise_enabled
+        assert not backend._encoder_modules
+        for block, original in zip(blocks, expected, strict=True):
+            assert block._hook_registry.get_hook("layerwise_offload") is None
+            torch.testing.assert_close(block.weight, original)
+
+    def test_encoder_only_requires_streamable_offload_plan(self, patched_offload_runtime):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel()
+        pipeline.text_encoder = _StagedEncoder()
+        backend = _layer_backend({"text_encoder"})
+
+        with pytest.raises(ValueError, match="Selected text encoder 'text_encoder' requires"):
+            backend.enable(pipeline)
+
+    def test_explicit_text_encoder_selection_requires_a_matching_module(self, patched_offload_runtime):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=2)
+        backend = _layer_backend({"dit", "text_encoder"})
+
+        with pytest.raises(ValueError, match="No text encoder modules found"):
+            backend.enable(pipeline)
+
+        assert not hasattr(pipeline.transformer.blocks[0], "_hook_registry")
+
+    def test_every_selected_encoder_requires_its_own_plan(self, patched_offload_runtime):
+        class Pipeline(nn.Module):
+            _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()
+                self.text_encoder = _PlainEncoder()
+                self.text_encoder_2 = _PlainEncoder()
+
+        pipeline = Pipeline()
+        backend = _layer_backend({"text_encoder"})
+
+        with pytest.raises(ValueError, match="Selected text encoder 'text_encoder_2' requires"):
+            backend.enable(pipeline)
+
+        assert not pipeline.text_encoder._omni_layerwise_enabled
+
+    def test_default_selection_preserves_unplanned_auxiliaries(self, patched_offload_runtime):
+        pipeline = _LegacyComponentPipeline()
+        backend = _layer_backend()
+
+        backend.enable(pipeline)
+
+        assert hasattr(pipeline.transformer.blocks[0], "_hook_registry")
+        assert not hasattr(pipeline.text_encoder, "_omni_layerwise_enabled")
+        assert pipeline.text_encoder.to_calls == 1
+        assert pipeline.text_encoder.offload_calls == 0
+        assert pipeline.vae.to_calls == 1
+        assert pipeline.vae.offload_calls == 0
+
+        backend.disable()
+
+    def test_disable_failure_keeps_host_masters_for_retry(self, patched_offload_runtime, monkeypatch):
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=3)
+        expected = {name: tensor.detach().clone() for name, tensor in pipeline.state_dict().items()}
+        backend = _layer_backend()
+        backend.enable(pipeline)
+        failing_hook = backend._dit_hooks[0]
+        restore = failing_hook.restore_next_block
+        attempts = 0
+
+        def fail_once():
+            nonlocal attempts
+            attempts += 1
+            if attempts == 1:
+                raise RuntimeError("injected restore failure")
+            restore()
+
+        monkeypatch.setattr(failing_hook, "restore_next_block", fail_once)
+
+        with pytest.raises(RuntimeError, match="Failed to fully disable"):
+            backend.disable()
+
+        assert backend.enabled
+        assert backend._dit_hooks
+
+        backend.disable()
+
+        assert not backend.enabled
+        for name, tensor in pipeline.state_dict().items():
+            torch.testing.assert_close(tensor, expected[name])
+
+    def test_encoder_disable_failure_keeps_host_masters_for_retry(
+        self,
+        patched_offload_runtime,
+        monkeypatch,
+    ):
+        pipeline = _ComponentPipeline()
+        encoder = pipeline.text_encoder
+        expected = {name: tensor.detach().clone() for name, tensor in encoder.state_dict().items()}
+        backend = _layer_backend({"text_encoder"})
+        backend.enable(pipeline)
+        failing_hook = encoder._omni_layerwise_hooks[0]
+        restore = failing_hook.restore_next_block
+        attempts = 0
+
+        def fail_once():
+            nonlocal attempts
+            attempts += 1
+            if attempts == 1:
+                raise RuntimeError("injected encoder restore failure")
+            restore()
+
+        monkeypatch.setattr(failing_hook, "restore_next_block", fail_once)
+
+        with pytest.raises(RuntimeError, match="Failed to fully disable"):
+            backend.disable()
+
+        assert backend.enabled
+        assert backend._encoder_modules == [encoder]
+        assert encoder._omni_layerwise_enabled
+        assert encoder._omni_layerwise_hooks
+
+        backend.disable()
+
+        assert not backend.enabled
+        assert not encoder._omni_layerwise_enabled
+        for name, tensor in encoder.state_dict().items():
+            torch.testing.assert_close(tensor, expected[name])
+
+    def test_legacy_selector_preserves_planned_encoder_and_vae_lifecycle(self, patched_offload_runtime):
+        pipeline = _ComponentPipeline()
+        backend = _layer_backend()
+
+        backend.enable(pipeline)
+
+        assert pipeline.text_encoder._omni_layerwise_enabled
+        assert pipeline.text_encoder.offload_calls == 1
+        assert pipeline.vae.offload_calls == 1
+
+        backend.disable()
+
+    def test_standard_encoder_needs_only_declared_block_paths(self, patched_offload_runtime):
+        pipeline = _GenericEncoderPipeline()
+        backend = _layer_backend({"text_encoder"})
+
+        backend.enable(pipeline)
+
+        blocks = pipeline.text_encoder.encoder.block
+        assert pipeline.text_encoder._omni_layerwise_enabled
+        assert all(block._hook_registry.get_hook("layerwise_offload") is not None for block in blocks)
+        assert pipeline.text_encoder.final_norm.weight.numel() == 4
+
+        backend.disable()
+
+        assert not pipeline.text_encoder._omni_layerwise_enabled
+        assert all(block._hook_registry.get_hook("layerwise_offload") is None for block in blocks)
+        assert all(block.weight.numel() == 100 for block in blocks)
+
+
+def _offload_od_config(**overrides):
+    values = {
+        "diffusion_offload_config": None,
+        "enable_cpu_offload": False,
+        "enable_layerwise_offload": False,
+        "enable_distributed_layerwise_offload": False,
+        "dlo_use_allgather": True,
+        "dlo_resident_layers": 0,
+        "pin_cpu_memory": True,
+        "parallel_config": None,
+        "model": "/fake/model",
+    }
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+def _resolve_offload_config(public=None, **overrides) -> OffloadConfig:
+    return OffloadConfig.from_od_config(_offload_od_config(diffusion_offload_config=public, **overrides))
+
+
+class TestLayerwiseComponentConfig:
+    def test_layer_mode_selects_components_without_backend_jargon(self):
+        config = _resolve_offload_config(
+            {
+                "mode": "layer",
+                "components": ["dit", "text_encoder"],
+            }
+        )
+
+        assert config.strategy is OffloadStrategy.LAYER_WISE
+        assert config.components == frozenset({"dit", "text_encoder"})
+        assert not config.uses_allgather("dit")
+        assert not config.uses_allgather("text_encoder")
+        plan = OffloadPlan(encoder_component_types={"mllm": "text_encoder"})
+        assert config.offloads_encoder("mllm", plan)
+
+    @pytest.mark.parametrize("component", ["image_encoder", "vae", "scheduler", "text-encoder"])
+    def test_unknown_or_noncanonical_component_is_rejected(self, component):
+        with pytest.raises(ValueError, match="Unknown diffusion offload component"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(diffusion_offload_config={"mode": "layer", "components": [component]})
+            )
+
+    def test_components_is_selection_only_not_an_options_mapping(self):
+        with pytest.raises(TypeError, match="components must be a non-empty list"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={
+                        "mode": "layer",
+                        "components": {"dit": {}},
+                    }
+                )
+            )
+
+    def test_layer_options_requires_component_name_keys(self):
+        with pytest.raises(TypeError, match="layer_options keys must be strings"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={
+                        "mode": "layer",
+                        "components": ["dit"],
+                        "layer_options": {0: {}},
+                    }
+                )
+            )
+
+    @pytest.mark.parametrize(
+        ("config", "message"),
+        [
+            ({0: "layer"}, "diffusion_offload_config keys must be strings"),
+            (
+                {
+                    "mode": "layer",
+                    "components": ["dit"],
+                    "layer_options": {"dit": {0: "allgather"}},
+                },
+                "layer_options\\['dit'\\] keys must be strings",
+            ),
+        ],
+    )
+    def test_all_public_mapping_levels_require_string_keys(self, config, message):
+        with pytest.raises(TypeError, match=message):
+            OffloadConfig.from_od_config(_offload_od_config(diffusion_offload_config=config))
+
+    @pytest.mark.parametrize(
+        "config,match",
+        [
+            ({"components": ["dit"]}, "requires 'mode'"),
+            ({"mode": "layer"}, "requires 'components'"),
+            ({"mode": "layer", "components": []}, "must not be empty"),
+            ({"mode": "layer", "components": [" dit "]}, "Unknown diffusion offload component"),
+            ({"mode": "layerwise", "components": ["dit"]}, "Unknown diffusion offload mode"),
+            (
+                {
+                    "mode": "layer",
+                    "components": ["dit"],
+                    "layer_options": {"dit": {"weight_transfer": "rank_local"}},
+                },
+                "Unknown offload transfer",
+            ),
+            (
+                {
+                    "mode": "layer",
+                    "components": ["dit"],
+                    "layer_options": {"dit": {"prefetch": 2}},
+                },
+                "Unknown diffusion offload setting",
+            ),
+            (
+                {
+                    "mode": "layer",
+                    "components": ["text_encoder"],
+                    "layer_options": {"text_encoder": {"resident_layers": 1}},
+                },
+                "supports only the 'dit'",
+            ),
+            (
+                {
+                    "mode": "layer",
+                    "components": ["dit"],
+                    "layer_options": {"text_encoder": {}},
+                },
+                "requires selecting the same component",
+            ),
+        ],
+    )
+    def test_schema_rejects_ambiguous_or_unsupported_values(self, config, match):
+        with pytest.raises(ValueError, match=match):
+            OffloadConfig.from_od_config(_offload_od_config(diffusion_offload_config=config))
+
+    def test_module_mode_accepts_component_selection(self):
+        config = _resolve_offload_config(
+            {
+                "mode": "module",
+                "components": ["dit", "text_encoder"],
+            }
+        )
+
+        assert config.strategy is OffloadStrategy.MODEL_LEVEL
+        assert config.components == frozenset({"dit", "text_encoder"})
+
+    @pytest.mark.parametrize("setting", [{"weight_transfer": "rank-local"}, {"resident_layers": 1}])
+    def test_module_mode_rejects_layer_settings(self, setting):
+        with pytest.raises(ValueError, match="layer_options requires mode='layer'"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={
+                        "mode": "module",
+                        "components": ["dit"],
+                        "layer_options": {"dit": setting},
+                    }
+                )
+            )
+
+    def test_each_component_can_choose_its_transfer(self):
+        config = _resolve_offload_config(
+            {
+                "mode": "layer",
+                "components": ["dit", "text_encoder"],
+                "layer_options": {
+                    "dit": {"weight_transfer": "rank-local"},
+                    "text_encoder": {"weight_transfer": "allgather"},
+                },
+            }
+        )
+
+        assert config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+        assert not config.uses_allgather("dit")
+        assert config.uses_allgather("text_encoder")
+        assert config.dlo_use_allgather is False
+
+    def test_dit_residency_selects_capable_backend(self):
+        config = _resolve_offload_config(
+            {
+                "mode": "layer",
+                "components": ["dit"],
+                "layer_options": {"dit": {"resident_layers": 20}},
+            }
+        )
+
+        assert config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+        assert config.dlo_resident_layers == 20
+        assert not config.uses_allgather("dit")
+
+    def test_dit_residency_rejects_allgather(self):
+        with pytest.raises(ValueError, match="resident_layers requires dit.weight_transfer='rank-local'"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={
+                        "mode": "layer",
+                        "components": ["dit"],
+                        "layer_options": {"dit": {"weight_transfer": "allgather", "resident_layers": 1}},
+                    }
+                )
+            )
+
+    def test_host_weight_runtime_remains_separate_from_compact_config(self):
+        with pytest.raises(ValueError, match="cannot be combined with Host Weight Runtime"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={"mode": "layer", "components": ["dit"]},
+                    host_weight_runtime_mode="preferred",
+                )
+            )
+
+    def test_pin_memory_override_is_scoped_to_offload_config(self):
+        config = _resolve_offload_config(
+            {
+                "mode": "layer",
+                "components": ["dit"],
+                "pin_memory": False,
+            },
+            pin_cpu_memory=True,
+        )
+
+        assert config.pin_cpu_memory is False
+
+    def test_legacy_omitted_selector_preserves_dit_only_behavior(self):
+        config = _resolve_offload_config(enable_layerwise_offload=True)
+
+        assert config.strategy is OffloadStrategy.LAYER_WISE
+        assert config.components is None
+
+    def test_offload_resolution_is_cached_for_runtime_accessors(self, monkeypatch):
+        import vllm_omni.diffusion.offloader.config as config_module
+
+        od_config = _offload_od_config(
+            diffusion_offload_config={
+                "mode": "layer",
+                "components": ["dit", "text_encoder"],
+            }
+        )
+        parse_calls = 0
+        original_parse = config_module.parse_diffusion_offload_config
+
+        def track_parse(value):
+            nonlocal parse_calls
+            parse_calls += 1
+            return original_parse(value)
+
+        monkeypatch.setattr(config_module, "parse_diffusion_offload_config", track_parse)
+
+        assert config_module.resolve_offload_strategy(od_config) is OffloadStrategy.LAYER_WISE
+        assert config_module.selected_offload_components(od_config) == {"dit", "text_encoder"}
+        assert not config_module.component_uses_allgather(od_config, "dit")
+        assert not config_module.any_selected_component_uses_allgather(od_config)
+        assert parse_calls == 1
+
+    def test_cached_policy_is_deeply_immutable_and_pickle_safe(self):
+        import vllm_omni.diffusion.offloader.config as config_module
+
+        od_config = _offload_od_config(
+            diffusion_offload_config={
+                "mode": "layer",
+                "components": ["dit"],
+                "layer_options": {"dit": {"weight_transfer": "rank-local"}},
+            }
+        )
+        resolved = config_module.resolve_offload(od_config)
+
+        with pytest.raises(TypeError):
+            resolved.transfers["dit"] = config_module.DLOTransfer.ALLGATHER  # type: ignore[index]
+        assert resolved.public is not None
+        with pytest.raises(TypeError):
+            resolved.public.layer_options["dit"] = config_module.LayerOffloadOptions()  # type: ignore[index]
+
+        restored = ForkingPickler.loads(ForkingPickler.dumps(resolved))
+        assert not restored.uses_allgather("dit")
+
+    def test_dp_text_encoder_allgather_rejects_rank_local_prompt_cache(self):
+        config = _offload_od_config(
+            diffusion_offload_config={
+                "mode": "layer",
+                "components": ["text_encoder"],
+                "layer_options": {"text_encoder": {"weight_transfer": "allgather"}},
+            },
+            enable_prompt_embed_cache=True,
+            parallel_config=SimpleNamespace(
+                data_parallel_size=2,
+                sequence_parallel_size=1,
+                use_hsdp=False,
+            ),
+        )
+
+        with pytest.raises(ValueError, match="rank-local cache hits would skip different encoder collectives"):
+            OffloadConfig.from_od_config(config)
+
+    def test_dp_allgather_rejects_rank_local_dit_cache_decisions(self):
+        config = _offload_od_config(
+            diffusion_offload_config={
+                "mode": "layer",
+                "components": ["dit"],
+                "layer_options": {"dit": {"weight_transfer": "allgather"}},
+            },
+            cache_backend="tea_cache",
+            parallel_config=SimpleNamespace(
+                data_parallel_size=2,
+                sequence_parallel_size=1,
+                use_hsdp=False,
+            ),
+        )
+
+        with pytest.raises(ValueError, match="rank-local cache decisions can skip different weight collectives"):
+            OffloadConfig.from_od_config(config)
+
+    def test_internal_transfer_map_requires_every_component(self):
+        with pytest.raises(ValueError, match="missing: text_encoder"):
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                dlo_transfers={"dit": "rank-local"},
+            )
+
+    def test_legacy_distributed_settings_still_work(self):
+        config = _resolve_offload_config(
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=False,
+            dlo_resident_layers=3,
+        )
+
+        assert config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+        assert config.dlo_resident_layers == 3
+        assert not config.uses_allgather("dit")
+        assert not config.uses_allgather("text_encoder")
+
+    def test_legacy_allgather_remains_dit_only(self):
+        config = _resolve_offload_config(
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+
+        assert config.uses_allgather("dit")
+        assert not config.uses_allgather("text_encoder")
+
+    @pytest.mark.parametrize(
+        ("legacy_flags", "expected_strategy"),
+        [
+            (
+                {"enable_cpu_offload": True, "enable_layerwise_offload": True},
+                OffloadStrategy.LAYER_WISE,
+            ),
+            (
+                {
+                    "enable_cpu_offload": True,
+                    "enable_layerwise_offload": True,
+                    "enable_distributed_layerwise_offload": True,
+                },
+                OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+            ),
+        ],
+    )
+    def test_legacy_aliases_preserve_existing_priority(self, legacy_flags, expected_strategy):
+        config = _resolve_offload_config(**legacy_flags)
+
+        assert config.strategy is expected_strategy
+
+    def test_compact_config_cannot_mix_with_legacy_alias(self):
+        with pytest.raises(ValueError, match="cannot be combined"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={"mode": "layer", "components": ["dit"]},
+                    enable_cpu_offload=True,
+                )
+            )
+
+    @pytest.mark.parametrize(
+        "legacy_options",
+        [
+            {"dlo_use_allgather": False},
+            {"dlo_resident_layers": 12},
+        ],
+    )
+    def test_compact_config_cannot_mix_with_legacy_dlo_options(self, legacy_options):
+        with pytest.raises(ValueError, match="cannot be combined with legacy DLO option"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    diffusion_offload_config={"mode": "layer", "components": ["dit"]},
+                    **legacy_options,
+                )
+            )
+
+    def test_legacy_allgather_resident_layers_fail_during_config_resolution(self):
+        with pytest.raises(ValueError, match="requires the DiT DLO transfer to be rank-local"):
+            OffloadConfig.from_od_config(
+                _offload_od_config(
+                    enable_distributed_layerwise_offload=True,
+                    dlo_use_allgather=True,
+                    dlo_resident_layers=12,
+                )
+            )

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()