[Refactor][Diffusion] Resolve offload topology through one plan resolver (#7209)

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

vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

@@ -21,6 +21,7 @@ from __future__ import annotations
 import threading
 import time
 import weakref
+from collections.abc import Sequence
 from itertools import chain
 from typing import Any
 
@@ -39,7 +40,6 @@ from vllm_omni.platforms import current_omni_platform
 from .base import OffloadBackend, OffloadConfig, run_cleanup_steps
 from .component_utils import (
     clear_encoder_layerwise_state,
-    get_encoder_block_groups,
     iter_streamable_dits,
     move_non_block_state_to_device,
     prepare_component,
@@ -53,11 +53,8 @@ from .host_registration import (
     HostRegistrationError,
     register_host_mappings,
 )
-from .module_collector import ModuleDiscovery
-from .offload_plan import (
-    OffloadPlan,
-    get_offload_plan,
-)
+from .offload_plan import OffloadPlan
+from .plan_resolver import ResolvedComponent, resolve_offload_plan
 from .tensor_utils import (
     clear_block_storage,
     clear_tensor_storage,
@@ -1437,7 +1434,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
     def _install_hook_group(
         self,
-        blocks: list[nn.Module] | nn.ModuleList,
+        blocks: Sequence[nn.Module] | nn.ModuleList,
         component: str,
         *,
         use_dit_mmap: bool = False,
@@ -1483,48 +1480,27 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         hooks[1]._is_group_first = True
         return hooks
 
-    def _try_layerwise_offload_encoder(
-        self,
-        module: nn.Module,
-        name: str,
-        plan: OffloadPlan | None,
-        stage_on_demand: bool = False,
-    ) -> bool:
-        """Apply DLO hooks to a plan-declared text encoder.
+    def _try_layerwise_offload_encoder(self, component: ResolvedComponent) -> bool:
+        """Apply DLO hooks to a text encoder's resolved block stacks.
 
         Rank-local transfer is always safe because it preserves the tensors
         produced by the model loader. Multi-rank AllGather is enabled only
         when the model declares that those tensors are replicated across the
-        selected DLO group; encoder TP groups contain different shards and
-        are deliberately not used for this transport.
+        selected DLO group; the resolver rejects every other layout before
+        this method installs anything.
         """
-        if plan is None:
-            return False
-
-        group_size = self._component_transport(TEXT_ENCODER_COMPONENT)[1]
-        if group_size > 1 and name not in plan.encoder_dlo_weight_replication:
-            raise ValueError(
-                f"Text encoder {name!r} cannot use DLO AllGather across the DiT offload group: "
-                "its loader-produced weights are not declared replicated across that group. "
-                "Set layer_options.text_encoder.weight_transfer='rank-local' in diffusion_offload_config "
-                "for encoder-TP or rank-specific layouts."
-            )
-        block_groups = get_encoder_block_groups(
-            module,
-            name,
-            plan,
-            strict=self.config.components is not None or group_size > 1,
-        )
+        block_groups = [stack.blocks for stack in component.stacks]
         if not block_groups:
             return False
 
+        module = component.module
         encoder_hooks: list[DistributedLayerwiseOffloadHook] = []
         for blocks in block_groups:
             encoder_hooks.extend(self._install_hook_group(blocks, TEXT_ENCODER_COMPONENT))
         # Track the module before placement, which may fail, so outer rollback
         # also removes its partially-installed block hooks and marker state.
         self._encoder_modules.append(module)
-        if not stage_on_demand:
+        if not component.on_demand:
             move_non_block_state_to_device(module, block_groups, self.device)
         set_encoder_layerwise_state(
             module,
@@ -1534,10 +1510,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         logger.info(
             "Enabled %s DLO transfer for text encoder %s (%d blocks across %d stacks, group_size=%d)",
             self.config.transfer_for(TEXT_ENCODER_COMPONENT).value,
-            name,
+            component.path,
             sum(len(blocks) for blocks in block_groups),
             len(block_groups),
-            group_size,
+            self._component_transport(TEXT_ENCODER_COMPONENT)[1],
         )
         return True
 
@@ -1700,31 +1676,15 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         if self.dp_group is None and self._has_multirank_allgather():
             self._init_dp_group()
 
-        modules = ModuleDiscovery.discover(pipeline)
-        if not modules.dits and self.config.offloads(DIT_COMPONENT):
+        # One resolver call validates the whole topology; an explicit selection
+        # has already failed by the time any of the checks below run.
+        resolved = resolve_offload_plan(pipeline, self.config)
+        if not resolved.dits and self.config.offloads(DIT_COMPONENT):
             if self.host_weight_plan is not None:
                 raise RuntimeError(
                     "DLO received a loader-owned host-weight plan, but no DiT modules were discovered to consume it"
                 )
-            message = "No DiT/transformer modules found for selected DiT offload"
-            if self.config.components is not None:
-                raise ValueError(message)
-            logger.warning(message)
-
-        # Retrieve optional declarative OffloadPlan from the pipeline.
-        # When present, replaces heuristic block discovery.
-        plan = get_offload_plan(pipeline)
-
-        if self.config.dlo_resident_layers:
-            resident_paths = frozenset() if plan is None else plan.resident_dit_paths
-            if not resident_paths.intersection(modules.dit_names):
-                message = (
-                    f"resident_layers={self.config.dlo_resident_layers} was requested, but this model declares "
-                    "no matching resident_dit_paths"
-                )
-                if self.config.components is not None:
-                    raise ValueError(message)
-                logger.warning("%s; all blocks will be streamed.", message)
+            logger.warning("No DiT/transformer modules found for selected DiT offload")
 
         # Storage selection belongs to the loader.  DLO consumes the exact
         # prevalidated plan that caused the loader to skip materialization;
@@ -1755,7 +1715,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
             elif host_weight_plan.backing_kind == "checkpoint_mmap":
                 self._load_weights_via_mmap(
                     pipeline,
-                    modules,
+                    resolved.modules,
                     host_weight_plan,
                 )
             else:
@@ -1768,10 +1728,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         else:
             remaining_meta = [
                 name
-                for dit_name, dit_module in zip(modules.dit_names, modules.dits)
+                for component in resolved.dits
                 for name, tensor in chain(
-                    dit_module.named_parameters(),
-                    dit_module.named_buffers(),
+                    component.module.named_parameters(),
+                    component.module.named_buffers(),
                 )
                 if getattr(tensor, "is_meta", False)
             ]
@@ -1784,66 +1744,46 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         # Apply each selected encoder transfer while keeping explicit VAEs and
         # unselected components resident.
         prepare_pipeline_components(
-            modules,
-            self.config,
-            plan,
+            resolved,
             device=self.device,
             staged_components=self._staged_components,
             enable_encoder_blocks=self._try_layerwise_offload_encoder,
         )
 
         if self.config.offloads(DIT_COMPONENT):
-            logger.info("Applying distributed layer-wise offloading on %s", modules.dit_names)
+            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(m) for m in modules.dits)
+        all_dit_modules = set(id(component.module) for component in resolved.dits)
 
         # Apply hooks for each DiT module
-        for dit_name, dit_module, blocks_attr_names, blocks in iter_streamable_dits(
-            modules, self.config, self.device, plan
-        ):
-            resident_count = 0
-            if plan is not None and dit_name in plan.resident_dit_paths:
-                resident_count = min(self.config.dlo_resident_layers, len(blocks))
-            if resident_count:
-                resident_blocks = blocks[:resident_count]
-                self._resident_blocks.extend(resident_blocks)
-                blocks = blocks[resident_count:]
+        for component, stack in iter_streamable_dits(resolved, self.device):
+            streaming = list(stack.streaming)
+            if stack.resident_head:
+                self._resident_blocks.extend(stack.resident)
                 logger.info(
                     "Keeping %d leading blocks resident on %s; streaming %d tail blocks",
-                    resident_count,
-                    dit_name,
-                    len(blocks),
+                    stack.resident_head,
+                    component.path,
+                    len(streaming),
                 )
 
             self._prepare_dit_non_block_modules(
-                dit_module,
-                blocks_attr_names,
+                component.module,
+                list(stack.attrs),
                 all_dit_modules,
-                plan,
+                resolved.declaration,
             )
 
-            num_blocks = len(blocks)
-            if num_blocks == 0:
-                logger.info("All blocks for %s are resident; no streaming hooks required", dit_name)
-                continue
-            if num_blocks <= 1:
-                if self.config.components is not None:
-                    raise ValueError(
-                        f"Selected DiT {dit_name!r} leaves only one streaming block after "
-                        f"resident_layers={resident_count}; choose a resident count that "
-                        "leaves zero or at least two streaming blocks"
-                    )
-                logger.warning(
-                    "#Streaming target layers (blocks) <= 1. Keeping the final block resident on %s (%s)",
-                    dit_name,
-                    dit_module.__class__.__name__,
-                )
-                self._resident_blocks.extend(blocks)
+            if not streaming:
+                logger.info("All blocks for %s are resident; no streaming hooks required", component.path)
                 continue
 
-            self._install_hook_group(blocks, DIT_COMPONENT, use_dit_mmap=True)
+            self._install_hook_group(streaming, DIT_COMPONENT, use_dit_mmap=True)
         if self._resident_blocks:
             self._resident_layer_group = PinnedResidentLayerGroup(
                 self._resident_blocks,