[Refactor][Diffusion] Resolve distributed layerwise topology in the plan (#7326)

Signed-off-by: specture724 <specture724@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -32,6 +32,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import (
     TensorBinding,
     build_checkpoint_mmap_plan,
 )
+from vllm_omni.diffusion.offloader import plan_resolver
 from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
 from vllm_omni.diffusion.offloader.block_discovery import (
     get_blocks_attr_names,
@@ -52,7 +53,7 @@ from vllm_omni.diffusion.offloader.offload_plan import (
     OffloadPlan,
     get_offload_plan,
 )
-from vllm_omni.diffusion.offloader.plan_resolver import resolve_offload_plan
+from vllm_omni.diffusion.offloader.plan_resolver import ResolvedComponent, resolve_offload_plan
 from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state
 from vllm_omni.host_weight_runtime import MappedHostRegion
 from vllm_omni.platforms import current_omni_platform
@@ -1235,6 +1236,78 @@ class _MmapPostLoadPipeline(nn.Module):
 
 
 class TestMmapWeightLoading:
+    @pytest.mark.parametrize("initial_device", ["cpu", "meta"])
+    @pytest.mark.parametrize("use_allgather", [False, True])
+    def test_mmap_keeps_undeclared_nested_blocks_streamed(
+        self, tmp_path, patched_offload_runtime, monkeypatch, initial_device, use_allgather
+    ):
+        class Refiner(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
+
+            def forward(self, x):
+                for block in self.layers:
+                    x = block(x)
+                return x
+
+        class Transformer(nn.Module):
+            _layerwise_offload_blocks_attrs = ["blocks"]
+
+            def __init__(self):
+                super().__init__()
+                self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
+                self.refiner = Refiner()
+
+            def forward(self, x):
+                x = self.refiner(x)
+                for block in self.blocks:
+                    x = block(x)
+                return x
+
+        pipeline = nn.Module()
+        pipeline.transformer = Transformer()
+        weights = {name: param.detach().clone() for name, param in pipeline.named_parameters()}
+        x = torch.randn(1, 2)
+        with torch.no_grad():
+            expected = pipeline.transformer(x)
+        checkpoint_file = tmp_path / "model.safetensors"
+        save_file(weights, str(checkpoint_file))
+        # Exercise the real mmap path with a small model instead of allocating
+        # the 1 GiB needed to trigger the production compatibility threshold.
+        monkeypatch.setattr(plan_resolver, "_NESTED_OFFLOAD_THRESHOLD_MB", 0)
+        if initial_device == "meta":
+            pipeline.transformer.to_empty(device="meta")
+        plan = HostWeightPlan(
+            backing_kind="checkpoint_mmap",
+            bindings={name: TensorBinding(checkpoint_key=name, file_path=str(checkpoint_file)) for name in weights},
+        )
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=use_allgather,
+            ),
+            torch.device("cpu"),
+            host_weight_plan=plan,
+        )
+
+        for _ in range(2):
+            try:
+                backend.enable(pipeline)
+                assert len(backend._all_hook_groups) == 2
+                with torch.no_grad():
+                    torch.testing.assert_close(pipeline.transformer(x), expected)
+            finally:
+                backend.disable()
+            for name, parameter in pipeline.named_parameters():
+                assert not parameter.is_meta
+                torch.testing.assert_close(parameter, weights[name])
+            for module in pipeline.modules():
+                registry = getattr(module, "_hook_registry", None)
+                if registry is not None:
+                    assert registry.get_hook("distributed_layerwise_offload") is None
+
     def test_runs_model_post_load_hook(self, tmp_path, patched_offload_runtime):
         pipeline = _MmapPostLoadPipeline()
         weights = {name: torch.ones(param.shape, dtype=torch.bfloat16) for name, param in pipeline.named_parameters()}
@@ -1250,10 +1323,7 @@ class TestMmapWeightLoading:
             ),
             torch.device("cpu"),
         )
-        modules = SimpleNamespace(
-            dits=[pipeline.transformer],
-            dit_names=["transformer"],
-        )
+        dits = (ResolvedComponent(path="transformer", module=pipeline.transformer, selected=True),)
         plan = HostWeightPlan(
             backing_kind="checkpoint_mmap",
             bindings={
@@ -1265,7 +1335,7 @@ class TestMmapWeightLoading:
             },
         )
 
-        backend._load_weights_via_mmap(pipeline, modules, plan)
+        backend._load_weights_via_mmap(pipeline, dits, plan)
 
         assert pipeline.transformer.post_load_calls == 1
         assert pipeline.transformer.time_embedder.weight.dtype == torch.float32
@@ -1717,6 +1787,128 @@ class TestOffloadPlan:
         for name, parameter in pipeline.named_parameters():
             torch.testing.assert_close(parameter, expected_parameters[name])
 
+    @pytest.mark.parametrize("shared_owner", ["parent", "encoder", "sibling"])
+    def test_nested_duplicate_ownership_is_rejected_before_any_hook(self, patched_offload_runtime, shared_owner):
+        class Pipeline(nn.Module):
+            _offload_plan = OffloadPlan(
+                offload_submodules={"refiner": "blocks", "other_refiner": "blocks"},
+                encoder_block_attrs={"text_encoder": ("encoder.block",)},
+            )
+
+            def __init__(self):
+                super().__init__()
+                self.transformer = _SingleBlockModel(num_blocks=2)
+                self.transformer.refiner = _SingleBlockModel(num_blocks=2)
+                self.transformer.other_refiner = _SingleBlockModel(num_blocks=2)
+                self.text_encoder = _PlainEncoder()
+                owners = {
+                    "parent": self.transformer.blocks,
+                    "encoder": self.text_encoder.encoder.block,
+                    "sibling": self.transformer.other_refiner.blocks,
+                }
+                self.transformer.refiner.blocks = owners[shared_owner]
+
+        pipeline = Pipeline()
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False),
+            torch.device("cpu"),
+        )
+        storage = {name: parameter.data_ptr() for name, parameter in pipeline.named_parameters()}
+        try:
+            with pytest.raises(ValueError, match="claimed by both"):
+                backend.enable(pipeline)
+            assert not backend.enabled
+            assert not backend._all_hook_groups
+            assert {name: parameter.data_ptr() for name, parameter in pipeline.named_parameters()} == storage
+            assert not getattr(pipeline.text_encoder, "_omni_layerwise_enabled", False)
+            for module in pipeline.modules():
+                assert getattr(module, "_hook_registry", None) is None
+        finally:
+            backend.disable()
+
+    def test_unstageable_submodule_is_rejected_before_any_hook(self, patched_offload_runtime):
+        """Plan-dependent rejection happens before earlier components are hooked."""
+
+        class Pipeline(nn.Module):
+            _offload_plan = OffloadPlan(
+                offload_submodules={"refiner": "missing"},
+                encoder_component_types={"text_encoder": "text_encoder"},
+                encoder_block_attrs={"text_encoder": ("encoder.block",)},
+            )
+
+            def __init__(self):
+                super().__init__()
+                self.transformer = _SingleBlockModel(num_blocks=2)
+                # Declared for submodule offload but without a block list or the
+                # load_to_device/offload_to_cpu lifecycle.
+                self.transformer.refiner = nn.Linear(2, 2)
+                self.text_encoder = _PlainEncoder()
+
+        pipeline = Pipeline()
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False),
+            torch.device("cpu"),
+        )
+
+        # Storage identity proves the encoder was never staged and rolled back.
+        encoder_storage = [block.weight.data_ptr() for block in pipeline.text_encoder.encoder.block]
+
+        with pytest.raises(ValueError, match="must implement load_to_device"):
+            backend.enable(pipeline)
+
+        assert not backend.enabled
+        assert [block.weight.data_ptr() for block in pipeline.text_encoder.encoder.block] == encoder_storage
+        assert not getattr(pipeline.text_encoder, "_omni_layerwise_enabled", False)
+        for block in pipeline.transformer.blocks:
+            assert getattr(block, "_hook_registry", None) is None
+
+    def test_undeclared_submodule_block_scan_warns(self, patched_offload_runtime, monkeypatch):
+        """The size-triggered attribute scan still works and is deprecated."""
+
+        class Refiner(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
+
+        class Transformer(nn.Module):
+            _layerwise_offload_blocks_attrs = ["blocks"]
+
+            def __init__(self):
+                super().__init__()
+                self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
+                self.refiner = Refiner()
+
+        class Pipeline(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.transformer = Transformer()
+
+        warnings: list[str] = []
+
+        class _Recorder:
+            def warning(self, message, *args):
+                warnings.append(message % args if args else message)
+
+            def __getattr__(self, _name):
+                return lambda *args, **kwargs: None
+
+        monkeypatch.setattr(plan_resolver, "logger", _Recorder())
+        monkeypatch.setattr(plan_resolver, "_NESTED_OFFLOAD_THRESHOLD_MB", 0)
+        plan_resolver._warn_nested_block_scan.cache_clear()
+        plan_resolver._warn_legacy_discovery.cache_clear()
+        pipeline = Pipeline()
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False),
+            torch.device("cpu"),
+        )
+
+        backend.enable(pipeline)
+
+        assert len(backend._all_hook_groups) == 2
+        assert sum("OffloadPlan.offload_submodules" in message for message in warnings) == 1
+
+        backend.disable()
+
 
 class TestMmapValidation:
     """Tests for loader preflight and backend plan realization."""
@@ -2101,7 +2293,7 @@ class TestMmapValidation:
             ),
             torch.device("cpu"),
         )
-        modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
+        dits = (ResolvedComponent(path="transformer", module=pipeline.transformer, selected=True),)
         result = build_checkpoint_mmap_plan(
             pipeline,
             dit_modules=(("transformer", pipeline.transformer),),
@@ -2113,7 +2305,7 @@ class TestMmapValidation:
         )
 
         assert result.plan is not None
-        backend._load_weights_via_mmap(pipeline, modules, result.plan)
+        backend._load_weights_via_mmap(pipeline, dits, result.plan)
         assert pipeline.transformer.validate_called, "validate_loaded_weights should be called"