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

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

vllm_omni/diffusion/offloader/base.py

@@ -2,8 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from abc import ABC, abstractmethod
+from collections.abc import Callable, Iterable
 from dataclasses import dataclass
-from enum import Enum
 from typing import Protocol, runtime_checkable
 
 import torch
@@ -12,9 +12,33 @@ from vllm.logger import init_logger
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig, validate_dlo_host_registration_options
 
+from .component_utils import encoder_component_type
+from .config import (
+    DIT_COMPONENT,
+    TEXT_ENCODER_COMPONENT,
+    DLOTransfer,
+    OffloadStrategy,
+    parse_dlo_transfer,
+    parse_offload_components,
+    resolve_offload,
+)
+from .offload_plan import OffloadPlan
+
 logger = init_logger(__name__)
 
 
+def run_cleanup_steps(steps: Iterable[tuple[str, Callable[[], None]]]) -> BaseException | None:
+    """Run every cleanup step and return the first failure, if any."""
+    first_error: BaseException | None = None
+    for description, step in steps:
+        try:
+            step()
+        except BaseException as exc:
+            logger.exception("Cleanup failed while %s", description)
+            first_error = first_error or exc
+    return first_error
+
+
 @runtime_checkable
 class SupportsModelCpuOffload(Protocol):
     """Pipeline-owned lifecycle for model-level CPU offload.
@@ -30,18 +54,12 @@ class SupportsModelCpuOffload(Protocol):
         device: torch.device,
         pin_memory: bool,
         use_hsdp: bool,
+        offload_components: frozenset[str] | None = None,
     ) -> None: ...
 
     def disable_omni_model_cpu_offload(self) -> None: ...
 
 
-class OffloadStrategy(Enum):
-    NONE = "none"
-    MODEL_LEVEL = "model_level"  # Sequential offloading between DiT and encoders
-    LAYER_WISE = "layer_wise"  # Block-level
-    DISTRIBUTED_LAYER_WISE = "distributed_layer_wise"  # Block-level with DP sharding + H2D/AllGather overlap
-
-
 @dataclass
 class OffloadConfig:
     strategy: OffloadStrategy
@@ -55,14 +73,59 @@ class OffloadConfig:
     # Optional per-worker ceiling for registering an HWR mmap. Zero means no
     # additional ceiling; pin_cpu_memory controls whether registration is tried.
     dlo_host_registration_limit_gib: float = 0.0
+    # ``None`` preserves the model's legacy plan-driven component topology;
+    # a frozenset is an explicit compact-API selection.
+    components: frozenset[str] | None = None
+    dlo_transfers: dict[str, DLOTransfer] | None = None
+
+    def __post_init__(self) -> None:
+        if self.components is not None:
+            self.components = parse_offload_components(self.components)
+        if self.dlo_transfers is None:
+            self.dlo_transfers = {
+                DIT_COMPONENT: DLOTransfer.ALLGATHER if self.dlo_use_allgather else DLOTransfer.RANK_LOCAL,
+                TEXT_ENCODER_COMPONENT: DLOTransfer.RANK_LOCAL,
+            }
+        self.dlo_transfers = parse_dlo_transfer(self.dlo_transfers)
+        # Preserve the old field as the DiT transfer compatibility view.
+        self.dlo_use_allgather = self.uses_allgather(DIT_COMPONENT)
+
+    def offloads(self, component: str) -> bool:
+        # Historically the generic selector covered DiT directly while model
+        # plans opted auxiliary components into their lifecycle.
+        return component == DIT_COMPONENT if self.components is None else component in self.components
+
+    def transfer_for(self, component: str) -> DLOTransfer:
+        if self.dlo_transfers is None:
+            raise RuntimeError("DLO transfers were not initialized")
+        try:
+            return self.dlo_transfers[component]
+        except KeyError as exc:
+            raise ValueError(f"Unknown DLO component {component!r}") from exc
+
+    def uses_allgather(self, component: str) -> bool:
+        return self.transfer_for(component) is DLOTransfer.ALLGATHER
+
+    def offloads_encoder(self, name: str, plan: OffloadPlan | None = None) -> bool:
+        """Return whether the selector covers a discovered encoder path.
+
+        Plans declare non-standard encoder names explicitly. The name-based
+        fallback preserves compatibility with pipelines that predate OffloadPlan.
+        """
+        component = encoder_component_type(name, plan)
+        return component is not None and self.offloads(component)
+
+    def should_offload_encoder(self, name: str, plan: OffloadPlan | None = None) -> bool:
+        """Apply explicit selection while preserving the legacy encoder topology."""
+        return self.components is None or self.offloads_encoder(name, plan)
 
     @classmethod
     def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
         """Extract and validate offload settings from OmniDiffusionConfig.
 
-        Enforces mutual exclusion among the three offload strategies.
-        Distributed layer-wise takes the highest priority, then layer-wise,
-        then model-level.
+        ``diffusion_offload_config`` is the canonical public selector. The
+        historical ``enable_*_offload`` booleans remain compatibility aliases;
+        ambiguous combinations fail instead of using silent precedence.
 
         The ``dp_size`` is automatically derived from ``parallel_config`` —
         it is NOT a user-configurable parameter. The distributed layerwise
@@ -74,10 +137,11 @@ class OffloadConfig:
         Returns:
             OffloadConfig with validated settings
         """
-        enable_cpu_offload = getattr(od_config, "enable_cpu_offload", False)
-        enable_layerwise_offload = getattr(od_config, "enable_layerwise_offload", False)
-        enable_distributed_layerwise_offload = getattr(od_config, "enable_distributed_layerwise_offload", False)
-        pin_cpu_memory = getattr(od_config, "pin_cpu_memory", True)
+        resolved = resolve_offload(od_config)
+        strategy = resolved.strategy
+        public_config = resolved.public
+        enable_distributed_layerwise_offload = strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE
+        pin_cpu_memory = resolved.pin_memory
 
         parallel_config = getattr(od_config, "parallel_config", None)
         use_hsdp = getattr(parallel_config, "use_hsdp", False) if parallel_config else False
@@ -102,61 +166,40 @@ class OffloadConfig:
                 if sp_size and sp_size > 1:
                     dp_size = sp_size
 
-        # Determine strategy (mutual exclusion, distributed layer-wise takes priority)
-        if enable_distributed_layerwise_offload:
-            strategy = OffloadStrategy.DISTRIBUTED_LAYER_WISE
-            if enable_layerwise_offload or enable_cpu_offload:
-                logger.info("Distributed layer-wise offloading takes priority, disabling other offloading strategies.")
-        elif enable_layerwise_offload:
-            strategy = OffloadStrategy.LAYER_WISE
-            if enable_cpu_offload:
-                logger.info(
-                    "Both model-level and layer-wise offloading enabled. "
-                    "Layer-wise takes priority, disabling model-level offloading."
-                )
-        elif enable_cpu_offload:
-            strategy = OffloadStrategy.MODEL_LEVEL
+        if public_config is not None:
+            components: frozenset[str] | None = resolved.components
         else:
-            strategy = OffloadStrategy.NONE
+            components = None
 
-        # With dlo_use_allgather=False, do not add another DP shard. Each rank
-        # streams the tensors produced by the standard loader, which may
-        # already be TP-local shards. This avoids AllGather synchronization
-        # requirements (concurrent requests, dummy run skip).
-        dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True)
-        dlo_resident_layers = int(getattr(od_config, "dlo_resident_layers", 0))
+        dlo_transfers = dict(resolved.transfers)
+        dlo_resident_layers = resolved.resident_layers
+        dit_uses_allgather = resolved.uses_allgather(DIT_COMPONENT)
         dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
             limit_gib=getattr(od_config, "dlo_host_registration_limit_gib", 0.0),
             enable_dlo=enable_distributed_layerwise_offload,
-            use_allgather=dlo_use_allgather,
+            use_allgather=dit_uses_allgather,
             hwr_mode=getattr(od_config, "host_weight_runtime_mode", "disabled"),
         )
-        if dlo_resident_layers < 0:
-            raise ValueError(f"dlo_resident_layers must be >= 0, got {dlo_resident_layers}")
-        if dlo_resident_layers and dlo_use_allgather:
-            raise ValueError(
-                "dlo_resident_layers currently requires --dlo-no-use-allgather so "
-                "resident blocks use weights prepared by the standard TP-aware loader"
-            )
 
-        # If dlo_use_allgather=False, force dp_size=1 (each rank independent)
-        if enable_distributed_layerwise_offload and not dlo_use_allgather:
-            dp_size = 1
+        if enable_distributed_layerwise_offload and all(
+            transfer is DLOTransfer.RANK_LOCAL
+            for component, transfer in dlo_transfers.items()
+            if component in resolved.components
+        ):
             logger.info(
-                "Distributed layerwise offload: dlo_use_allgather=False, "
-                "streaming complete rank-local blocks (no DLO shard or AllGather); "
-                "the backend will select mmap or standard-loader host storage"
+                "Distributed layerwise offload: all selected components use "
+                "rank-local transfer (no DLO shard or AllGather)"
             )
 
         # HSDP already shards parameters into DTensors.  Running distributed
         # layerwise offload on top would shard each to_local() again, producing
         # incorrect reconstruction after AllGather.  Reject this combination.
-        if enable_distributed_layerwise_offload and use_hsdp and dlo_use_allgather:
+        if enable_distributed_layerwise_offload and use_hsdp and resolved.any_allgather:
             raise ValueError(
                 "Distributed layerwise offload with AllGather is incompatible with "
                 "HSDP: HSDP parameters are already sharded DTensors, and the offloader "
-                "would double-shard them. Use --dlo-no-use-allgather (standard-loader "
-                "rank-local weights) or disable HSDP."
+                "would double-shard them. Set weight_transfer='rank-local' for the affected "
+                "component in diffusion_offload_config, or disable HSDP."
             )
 
         return cls(
@@ -164,9 +207,11 @@ class OffloadConfig:
             pin_cpu_memory=pin_cpu_memory,
             use_hsdp=use_hsdp,
             dp_size=dp_size,
-            dlo_use_allgather=dlo_use_allgather,
+            dlo_use_allgather=dit_uses_allgather,
             dlo_resident_layers=dlo_resident_layers,
             dlo_host_registration_limit_gib=dlo_host_registration_limit_gib,
+            components=components,
+            dlo_transfers=dlo_transfers,
         )