[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>

docs/design/feature/offloader/README.md

@@ -84,10 +84,22 @@ Block topology comes from the pipeline `OffloadPlan` first: `block_attrs` maps
 each DiT path to its ordered block containers, `encoder_block_attrs` declares
 streamable encoder stacks, `on_demand_component_paths` marks pipeline-managed
 residency, `resident_dit_paths` marks the DiTs that may hold resident layers,
-and `encoder_dlo_weight_replication` marks the encoders whose loader-produced
-weights are safe for AllGather. DiTs absent from the plan fall back to
-`_layerwise_offload_blocks_attrs` (including the deprecated singular-name
-compatibility path).
+`encoder_dlo_weight_replication` marks the encoders whose loader-produced
+weights are safe for AllGather, and `offload_submodules` maps a DiT submodule
+that owns its residency to its own block container. DiTs absent from the plan
+fall back to `_layerwise_offload_blocks_attrs` (including the deprecated
+singular-name compatibility path).
+
+An undeclared DiT submodule joins the plan only when it is large enough that
+keeping it resident would defeat streaming; its block container is then found
+by scanning well-known attribute names, which warns once per submodule class.
+Size is computed from parameter shapes and dtypes, including meta parameters,
+so resolving before mmap loading preserves the same residency decision.
+Block ownership validation includes nested components: a block shared with a
+parent, sibling, or encoder is rejected before any hooks or storage are changed.
+A submodule with no block container must implement the `load_to_device` /
+`offload_to_cpu` lifecycle, and the resolver rejects it before any component of
+the pipeline is placed or hooked.
 
 The resolver is pure: it moves no tensor, installs no hook, writes no module
 attribute, and reads no process group — multi-rank facts come from

tests/diffusion/models/magi2/test_native_dlo_allgather.py

@@ -36,6 +36,7 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
     DistributedLayerwiseOffloadBackend,
     DistributedLayerwiseOffloadHook,
 )
+from vllm_omni.diffusion.offloader.plan_resolver import ResolvedComponent
 
 
 @dataclass(frozen=True)
@@ -46,12 +47,6 @@ class _WeightSourceStub:
     prefix: str
 
 
-@dataclass(frozen=True)
-class _PipelineModulesStub:
-    dits: list[nn.Module]
-    dit_names: list[str]
-
-
 _WORLD_SIZE = 4
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
 
@@ -189,7 +184,7 @@ def _load_mmap_transform_and_reconstruct(
         backend._mmap_transforms_by_tensor_id = {}
         backend._load_weights_via_mmap(
             pipeline,
-            _PipelineModulesStub(dits=[target], dit_names=["transformer"]),
+            (ResolvedComponent(path="transformer", module=target, selected=True),),
             plan_result.plan,
         )
 

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"
 
 

vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

@@ -53,8 +53,7 @@ from .host_registration import (
     HostRegistrationError,
     register_host_mappings,
 )
-from .offload_plan import OffloadPlan
-from .plan_resolver import ResolvedComponent, resolve_offload_plan
+from .plan_resolver import BlockStack, ResolvedComponent, resolve_offload_plan
 from .tensor_utils import (
     clear_block_storage,
     clear_tensor_storage,
@@ -111,7 +110,6 @@ def _forget_active_hwr_registration(
 # use layerwise offload (streaming hooks) or be moved to GPU as a resident
 # module.  Submodules larger than this are offloaded to save HBM; smaller
 # ones stay resident for lower latency.
-_ON_DEMAND_THRESHOLD_MB = 1024
 
 
 class DistributedLayerwiseOffloadHook(ModelHook):
@@ -1182,7 +1180,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
     def _load_weights_via_mmap(
         self,
         pipeline: nn.Module,
-        modules,
+        dits: tuple[ResolvedComponent, ...],
         plan: HostWeightPlan,
     ) -> None:
         """Load DiT checkpoint tensors as file-backed safetensors views.
@@ -1226,7 +1224,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         # converts the child's buffers too.  So we must save ALL buffers
         # from ALL DiT modules BEFORE any to_empty call.
         saved_buffers: dict[int, dict[str, torch.Tensor]] = {}
-        for dit_module in modules.dits:
+        for dit_module in (component.module for component in dits):
             bufs: dict[str, torch.Tensor] = {}
             for name, buf in dit_module.named_buffers():
                 # Check if this buffer is non-persistent on its OWNING module
@@ -1246,7 +1244,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         # Skip modules that are already meta (happens when a parent DiT
         # module contains a child DiT module — to_empty on the parent
         # already converted the child).
-        for dit_module in modules.dits:
+        for dit_module in (component.module for component in dits):
             if any(p.is_meta for p in dit_module.parameters()):
                 logger.info(
                     "%s already on meta device (skipping to_empty, %d buffers saved)",
@@ -1307,7 +1305,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         # assigning checkpoint tensors. The mmap path bypasses that loader, so
         # preserve the same lifecycle for transforms such as Cosmos3's fp32
         # timestep embedder.
-        for dit_name, dit_module in zip(modules.dit_names, modules.dits):
+        for dit_name, dit_module in ((component.path, component.module) for component in dits):
             # Restore non-persistent buffers before post-load hooks and strict
             # validation. They are constructor-derived and intentionally have
             # no checkpoint binding.
@@ -1352,7 +1350,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
             self._mmap_transforms_by_tensor_id[id(target)] = binding.transform
 
         remaining_meta: list[str] = []
-        for dit_name, dit_module in zip(modules.dit_names, modules.dits):
+        for dit_name, dit_module in ((component.path, component.module) for component in dits):
             remaining_meta.extend(
                 f"{dit_name}.{name}" for name, tensor in dit_module.named_parameters() if tensor.is_meta
             )
@@ -1517,73 +1515,26 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         )
         return True
 
-    def _try_layerwise_offload_submodule(self, module: nn.Module, name: str, plan: OffloadPlan | None = None) -> bool:
-        """Try to apply layerwise offload to a large submodule's blocks.
-        Resolution order:
-        1. OffloadPlan.offload_submodules (declarative, if plan is provided)
-        2. Heuristic search for common block-list attributes
-
-        Returns True if layerwise offload was applied, False otherwise.
-        """
-        from operator import attrgetter
-
-        blocks = None
-        blocks_attr = None
-
-        # 1. Check OffloadPlan first (declarative — no guessing)
-        if plan is not None and name in plan.offload_submodules:
-            attr_name = plan.offload_submodules[name]
-            try:
-                candidate = attrgetter(attr_name)(module)
-                if isinstance(candidate, nn.ModuleList) and len(candidate) > 1:
-                    blocks = candidate
-                    blocks_attr = attr_name
-            except AttributeError:
-                logger.warning(
-                    "OffloadPlan declared block attr '%s' for submodule '%s' "
-                    "but attribute not found — falling back to heuristic",
-                    attr_name,
-                    name,
-                )
-
-        # 2. Fallback: heuristic search
-        if blocks is None:
-            for attr_name in ("layers", "blocks", "h", "model.layers"):
-                try:
-                    candidate = attrgetter(attr_name)(module)
-                except AttributeError:
-                    continue
-                if isinstance(candidate, nn.ModuleList) and len(candidate) > 1:
-                    blocks = candidate
-                    blocks_attr = attr_name
-                    break
-
-        if blocks is None:
-            return False
-
+    def _offload_nested_submodule(self, component: ResolvedComponent) -> None:
+        """Stream one resolved DiT submodule through its own hook ring."""
+        stack = component.stacks[0]
         logger.info(
             "Distributed layerwise offload for submodule '%s.%s' (%d blocks, %.0f MB total, group_size=%d)",
-            name,
-            blocks_attr,
-            len(blocks),
-            sum(p.nelement() * p.element_size() for p in module.parameters()) / 1048576,
+            component.path,
+            stack.attrs[0],
+            len(stack.blocks),
+            sum(p.nelement() * p.element_size() for p in component.module.parameters()) / 1048576,
             self._component_transport(DIT_COMPONENT)[1],
         )
-
-        # Move non-block parts of the submodule to GPU (small: embeddings, norms)
-        for child_name, child in module.named_children():
-            if child_name != blocks_attr:
-                child.to(self.device)
-
-        self._install_hook_group(blocks, DIT_COMPONENT, use_dit_mmap=True)
-        return True
+        # Keep the submodule's own non-block state resident by block identity,
+        # so an attribute aliasing a streamed block is not placed with it.
+        move_non_block_state_to_device(component.module, (stack.blocks,), self.device)
+        self._install_hook_group(stack.blocks, DIT_COMPONENT, use_dit_mmap=True)
 
     def _prepare_dit_non_block_modules(
         self,
-        dit_module: nn.Module,
-        blocks_attr_names: list[str],
-        all_dit_modules: set[int],
-        plan: OffloadPlan | None,
+        component: ResolvedComponent,
+        stack: BlockStack,
     ) -> None:
         """Place or hook the DiT parts that are outside its repeated blocks.
 
@@ -1591,28 +1542,23 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         an all-resident stage skips placement for modules such as H3's token
         refiner and enters the forward pass with CPU or meta tensors.
         """
+        dit_module = component.module
+        nested = {child.path.rsplit(".", 1)[-1]: child for child in component.children}
         for name, module in dit_module.named_children():
-            if name in blocks_attr_names:
+            if name in stack.attrs:
                 logger.debug("Skipped blocks module %s", name)
                 continue
 
-            module_mb = (
-                sum(
-                    param.nelement() * param.element_size() if not getattr(param, "is_meta", False) else 0
-                    for param in module.parameters()
-                )
-                / 1048576
-            )
-            explicitly_planned = plan is not None and name in plan.offload_submodules
-            if explicitly_planned or module_mb > _ON_DEMAND_THRESHOLD_MB:
-                if id(module) in all_dit_modules:
+            child = nested.get(name)
+            if child is not None:
+                if not child.selected:
                     logger.info("Submodule '%s' is already a DiT module, skipping layerwise offload", name)
-                elif self._try_layerwise_offload_submodule(module, name, plan):
-                    pass
+                elif child.stacks:
+                    self._offload_nested_submodule(child)
                 else:
                     prepare_component(
-                        module,
-                        name,
+                        child.module,
+                        child.path,
                         device=self.device,
                         stage_on_demand=True,
                         blockwise=False,
@@ -1715,7 +1661,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
             elif host_weight_plan.backing_kind == "checkpoint_mmap":
                 self._load_weights_via_mmap(
                     pipeline,
-                    resolved.modules,
+                    resolved.dits,
                     host_weight_plan,
                 )
             else:
@@ -1750,16 +1696,12 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
             enable_encoder_blocks=self._try_layerwise_offload_encoder,
         )
 
-        if self.config.offloads(DIT_COMPONENT):
+        if any(component.selected for component in resolved.dits):
             logger.info(
                 "Applying distributed layer-wise offloading on %s",
                 [component.path for component in resolved.dits],
             )
 
-        # Collect all DiT module objects to detect submodules that are
-        # already handled as a separate DiT module (avoids duplicate hooks).
-        all_dit_modules = set(id(component.module) for component in resolved.dits)
-
         # Apply hooks for each DiT module
         for component, stack in iter_streamable_dits(resolved, self.device):
             streaming = list(stack.streaming)
@@ -1772,12 +1714,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                     len(streaming),
                 )
 
-            self._prepare_dit_non_block_modules(
-                component.module,
-                list(stack.attrs),
-                all_dit_modules,
-                resolved.declaration,
-            )
+            self._prepare_dit_non_block_modules(component, stack)
 
             if not streaming:
                 logger.info("All blocks for %s are resident; no streaming hooks required", component.path)
@@ -1804,11 +1741,6 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
             self.enabled = bool(self._resident_blocks or self._encoder_modules or self._staged_components)
             if self._using_mmap and not self.enabled:
                 self._release_mmap_handles()
-            if not self.enabled and not self.config.offloads(DIT_COMPONENT):
-                raise ValueError(
-                    "None of the selected distributed layerwise offload components have "
-                    "a model-declared streamable or on-demand plan"
-                )
             return
 
         # Unified allocation: 2 shared output buffers + 2 shared shard buffers

vllm_omni/diffusion/offloader/plan_resolver.py

@@ -17,6 +17,7 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 from functools import cache
+from operator import attrgetter
 from typing import TYPE_CHECKING
 
 from torch import nn
@@ -27,7 +28,7 @@ from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from .block_discovery import get_blocks_from_dit
 from .component_utils import get_encoder_block_groups, validate_on_demand_component
 from .config import DIT_COMPONENT, TEXT_ENCODER_COMPONENT, OffloadStrategy
-from .module_collector import ModuleDiscovery, PipelineModules
+from .module_collector import ModuleDiscovery
 from .offload_plan import OffloadPlan, get_offload_plan
 
 if TYPE_CHECKING:
@@ -76,6 +77,9 @@ class ResolvedComponent:
     selected: bool
     stacks: tuple[BlockStack, ...] = ()
     on_demand: bool = False
+    # Submodules of a DiT that are large enough to own their residency instead
+    # of staying resident with the rest of its non-block state.
+    children: tuple[ResolvedComponent, ...] = ()
 
 
 @dataclass(frozen=True)
@@ -86,11 +90,6 @@ class ResolvedOffloadPlan:
     encoders: tuple[ResolvedComponent, ...]
     vaes: tuple[ResolvedComponent, ...]
     residents: tuple[ResolvedComponent, ...]
-    # Raw discovery output and the model declaration remain reachable for the
-    # loader-owned mmap path and for DLO's nested-submodule staging, which
-    # still interpret topology themselves.
-    modules: PipelineModules
-    declaration: OffloadPlan | None = None
     # Ordinary layerwise offload turns a legacy selection it cannot serve into
     # a complete no-op, placement included. Distributed layerwise offload warns
     # and keeps preparing its other components instead, so this stays scoped to
@@ -102,6 +101,22 @@ class ResolvedOffloadPlan:
         return self.dits + self.encoders + self.vaes + self.residents
 
 
+# A DiT submodule this large owns its residency even when no plan declares it.
+_NESTED_OFFLOAD_THRESHOLD_MB = 1024
+_NESTED_BLOCK_ATTRS = ("layers", "blocks", "h", "model.layers")
+
+
+@cache
+def _warn_nested_block_scan(module_class: str, attr: str) -> None:
+    """Warn once per class that a nested block list was found by guessing."""
+    logger.warning(
+        "%s block list was discovered as %r by scanning well-known attribute names. "
+        "Declare it in OffloadPlan.offload_submodules; the attribute scan is deprecated.",
+        module_class,
+        attr,
+    )
+
+
 @cache
 def _warn_legacy_discovery(pipeline_class: str) -> None:
     """Warn once per pipeline class that component discovery guessed.
@@ -117,6 +132,77 @@ def _warn_legacy_discovery(pipeline_class: str) -> None:
     )
 
 
+def _nested_block_stack(module: nn.Module, declared_attr: str | None) -> BlockStack | None:
+    """Resolve one DiT submodule's own block ring, if it has one."""
+    candidates = (declared_attr,) if declared_attr is not None else _NESTED_BLOCK_ATTRS
+    for attr in candidates:
+        try:
+            blocks = attrgetter(attr)(module)
+        except AttributeError:
+            continue
+        if isinstance(blocks, nn.ModuleList) and len(blocks) > 1:
+            if declared_attr is None:
+                _warn_nested_block_scan(type(module).__name__, attr)
+            return BlockStack(attrs=(attr,), blocks=tuple(blocks))
+    if declared_attr is not None:
+        logger.warning(
+            "OffloadPlan declared block attr %r for submodule %s but it is not a streamable block list",
+            declared_attr,
+            type(module).__name__,
+        )
+    return None
+
+
+def _resolve_dit_children(
+    dit_module: nn.Module,
+    dit_path: str,
+    stack: BlockStack,
+    declaration: OffloadPlan | None,
+    dit_module_ids: set[int],
+) -> tuple[ResolvedComponent, ...]:
+    """Resolve the DiT submodules that own their residency.
+
+    A submodule qualifies when the model declares it or when it is large enough
+    that keeping it resident would defeat streaming. Submodules that are
+    themselves discovered DiTs are left to their own component.
+    """
+    children: list[ResolvedComponent] = []
+    for name, module in dit_module.named_children():
+        if name in stack.attrs:
+            continue
+        declared_attr = None if declaration is None else declaration.offload_submodules.get(name)
+        if declared_attr is None and _module_size_mb(module) <= _NESTED_OFFLOAD_THRESHOLD_MB:
+            continue
+        path = f"{dit_path}.{name}"
+        if id(module) in dit_module_ids:
+            # Its own component owns this submodule; the parent must not place it.
+            children.append(ResolvedComponent(path=path, module=module, selected=False))
+            continue
+        nested_stack = _nested_block_stack(module, declared_attr)
+        if nested_stack is None:
+            # Without its own ring the submodule needs the pipeline lifecycle.
+            validate_on_demand_component(module, path)
+        children.append(
+            ResolvedComponent(
+                path=path,
+                module=module,
+                selected=True,
+                stacks=() if nested_stack is None else (nested_stack,),
+                on_demand=nested_stack is None,
+            )
+        )
+    return tuple(children)
+
+
+def _module_size_mb(module: nn.Module) -> float:
+    """Expected parameter bytes in MiB, including not-yet-loaded meta tensors.
+
+    Resolution precedes mmap materialization, so residency must depend on the
+    tensor metadata rather than whether the loader has allocated storage yet.
+    """
+    return sum(parameter.nelement() * parameter.element_size() for parameter in module.parameters()) / 1048576
+
+
 def _resolve_dit_stacks(
     module: nn.Module,
     path: str,
@@ -253,11 +339,25 @@ def resolve_offload_plan(pipeline: nn.Module, config: OffloadConfig) -> Resolved
     vaes: list[ResolvedComponent] = []
     residents: list[ResolvedComponent] = []
 
+    dit_module_ids = {id(module) for module in modules.dits}
     for path, module in zip(modules.dit_names, modules.dits):
         stacks: tuple[BlockStack, ...] = ()
+        children: tuple[ResolvedComponent, ...] = ()
         if dit_selected and layerwise:
             stacks = _resolve_dit_stacks(module, path, declaration, config, explicit=explicit)
-        dits.append(ResolvedComponent(path=path, module=module, selected=dit_selected, stacks=stacks))
+            # Only distributed layerwise offload streams or stages submodules;
+            # the ordinary backend keeps all non-block state resident.
+            if stacks and config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE:
+                children = _resolve_dit_children(module, path, stacks[0], declaration, dit_module_ids)
+        dits.append(
+            ResolvedComponent(
+                path=path,
+                module=module,
+                selected=dit_selected,
+                stacks=stacks,
+                children=children,
+            )
+        )
 
     for path, module in zip(modules.encoder_names, modules.encoders):
         selected = config.should_offload_encoder(path, declaration)
@@ -307,8 +407,6 @@ def resolve_offload_plan(pipeline: nn.Module, config: OffloadConfig) -> Resolved
         encoders=tuple(encoders),
         vaes=tuple(vaes),
         residents=tuple(residents),
-        modules=modules,
-        declaration=declaration,
         skip_reason=(
             "No DiT/transformer modules found for selected DiT layerwise offload"
             if config.strategy is OffloadStrategy.LAYER_WISE and dit_selected and not dits
@@ -322,7 +420,10 @@ def resolve_offload_plan(pipeline: nn.Module, config: OffloadConfig) -> Resolved
 def _validate_unique_ownership(resolved: ResolvedOffloadPlan) -> None:
     """Reject topologies where two components would hook the same blocks."""
     owner_by_block: dict[int, str] = {}
-    for component in resolved.components:
+    pending = list(reversed(resolved.components))
+    while pending:
+        component = pending.pop()
+        pending.extend(reversed(component.children))
         for block in (block for stack in component.stacks for block in stack.blocks):
             owner = owner_by_block.setdefault(id(block), component.path)
             if owner != component.path: