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

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

docs/design/feature/offloader/README.md

@@ -5,23 +5,27 @@ The three strategies have separate user and design pages:
 
 | Strategy | User guide | Design contract |
 | --- | --- | --- |
-| Model-level | [Guide](../../../user_guide/diffusion/offloader/cpu_offload.md) | [Design](cpu_offload.md) |
+| Model-level | [Guide](../../../user_guide/diffusion/offloader/module_offload.md) | [Design](module_offload.md) |
 | Layerwise | [Guide](../../../user_guide/diffusion/offloader/layerwise_offload.md) | [Design](layerwise_offload.md) |
 | Distributed layerwise | [Guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md) | [Design](distributed_layerwise_offload.md) |
 
 ## Strategy selection
 
 `OffloadConfig.from_od_config()` converts diffusion configuration into one
-`OffloadStrategy`. If callers set multiple flags, selection is deterministic:
+`OffloadStrategy`. The public config separates component selection from
+policy:
 
-1. `DISTRIBUTED_LAYERWISE`;
-2. `LAYERWISE`;
-3. `MODEL_LEVEL`; and
-4. `NONE`.
+- `mode="module"` selects model-level offload;
+- `mode="layer"` with rank-local transfers selects ordinary layerwise
+  offload; and
+- AllGather transfer or resident layers selects the distributed layerwise
+  backend that implements those capabilities.
 
-The factory logs when a higher-priority strategy disables another. It derives
-parallel and HSDP state from `DiffusionParallelConfig`; callers do not provide
-a separate offload group size.
+The compatibility boolean flags retain their historical priority (distributed
+layerwise, layerwise, then model-level). A compact config rejects a conflicting
+legacy strategy or non-default legacy DLO tuning. The factory derives parallel
+and HSDP state from `DiffusionParallelConfig`; callers do not provide a
+separate offload group size.
 
 `get_offload_backend()` then validates platform offload support, resolves the
 device, and creates exactly one backend. Returning `None` means offloading is
@@ -39,6 +43,11 @@ Every backend implements `OffloadBackend`:
 `disable()` does not promise to restore the pipeline's original device
 placement. The caller owns any subsequent rematerialization.
 
+An enable failure must remove every partially installed hook. Rank-local
+backends restore ordinary tensors so the pipeline can be retried. A failed
+multi-rank AllGather startup must not enter unmatched recovery collectives;
+that worker is safely discarded after local resource cleanup.
+
 Hooks are registered through `HookRegistry` and `ModelHook`; offload backends
 must use distinct hook names and remove only hooks they own.
 
@@ -55,13 +64,13 @@ Dotted paths are supported. Legacy pipelines may use the fallback scan of
 well-known attribute names, but new integrations should declare components
 explicitly.
 
-Ordinary layerwise offload obtains topology from declared
-`_layerwise_offload_blocks_attrs` (with a deprecated singular-name
-compatibility path). Distributed layerwise offload may instead use a pipeline
-`OffloadPlan`, then falls back to declared block attributes and its documented
-compatibility discovery. Discovery metadata describes structure only; the
-backend remains responsible for transfer, synchronization, and storage
-ownership.
+Both layerwise backends first consume pipeline `OffloadPlan` metadata. A
+plan's `block_attrs` maps each DiT path to its ordered block containers, while
+`encoder_block_attrs` declares streamable encoder stacks. DiTs absent from the
+plan fall back to `_layerwise_offload_blocks_attrs` (including the deprecated
+singular-name compatibility path). Discovery metadata describes structure
+only; the backend remains responsible for transfer, synchronization, and
+storage ownership.
 
 ## Cross-strategy invariants
 

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -1,10 +1,11 @@
 # Distributed Layerwise Offload
 
 This document describes distributed layerwise offload (DLO) for diffusion
-models. DLO keeps only a small number of DiT blocks on the accelerator and
-streams the remaining blocks from host memory. The distributed backend can
-either shard those host-side weights across an existing parallel group or keep
-complete rank-local block sources and avoid an additional collective.
+models. DLO keeps only a small number of selected DiT or text-encoder blocks
+on the accelerator and streams the remaining blocks from host memory. The
+distributed backend can either shard those host-side weights across an
+existing parallel group or keep complete rank-local block sources and avoid
+an additional collective.
 
 For user-facing commands, see the
 [distributed layerwise offloading guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md)
@@ -38,10 +39,11 @@ detailed contracts and validation boundaries.
 
 ## Status
 
-DLO is implemented for multi-device diffusion execution. The default
-AllGather path is the primary path for DP and SP deployments. The
-`--dlo-no-use-allgather` path streams complete blocks independently and adds no
-DLO weight collective.
+DLO is implemented for multi-device diffusion execution. The default public
+transfer is `rank-local`; setting a selected component's
+`layer_options.NAME.weight_transfer` to `allgather` enables weight sharding across a
+compatible DP or SP group. Transfers are configured independently for `dit`
+and `text_encoder`.
 
 Host storage is selected separately from the transfer protocol. The loader can
 produce a direct-checkpoint mmap plan for a proven-compatible runtime layout;
@@ -109,7 +111,7 @@ checkpoint adapter beside the ordinary loader.
 
 ### AllGather path
 
-With the default `dlo_use_allgather=True`, each rank stores approximately
+For a component configured with `weight_transfer="allgather"`, each rank stores approximately
 `1 / group_size` of each streamable block in pinned host memory. The next
 block's shard is copied to a device buffer and reconstructed with
 `all_gather_into_tensor` on a communication stream while the current block is
@@ -132,9 +134,8 @@ source used to prepare each rank's persistent shard. They can be closed after
 shard preparation. Across the AllGather group, those private shards total
 approximately one runtime model copy.
 
-An effective DLO group size of one performs no collective, even when
-`dlo_use_allgather=True`; it follows the rank-local transfer path described
-below.
+An effective DLO group size of one performs no collective; it follows the
+rank-local transfer path described below.
 
 When DP is greater than one, the engine can process one request per DP rank in
 the same denoising wave. Because AllGather is a collective, all participating
@@ -142,9 +143,10 @@ requests must take the same execution path at every denoising step.
 
 ### Rank-local path without DLO AllGather
 
-With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
-one and streams complete blocks using H2D copies only. The host backing may be
-either a loader-approved checkpoint mapping or ordinary runtime tensors.
+With `weight_transfer="rank-local"`, DLO forces the affected component's internal
+offload shard size to one and streams complete blocks using H2D copies only.
+The host backing may be either a loader-approved checkpoint mapping or
+ordinary runtime tensors.
 
 For direct mmap, each process retains immutable safetensors views and uses two
 bounded pinned host staging slots. Processes on the same node that map the same
@@ -204,7 +206,7 @@ adopt it.
 > optimizations without changing AllGather behavior.
 
 Final-layout Host Weight Runtime (HWR) backing is opt-in and currently applies
-only to no-AllGather DLO. Enable it with
+only to a selected rank-local DiT. Enable it with
 `--host-weight-runtime-mode preferred` (or `required`) and
 `--host-weight-runtime-root <node-local-root>`.
 
@@ -422,8 +424,8 @@ failure.
   collective. When the model declares the final-layout contract, HWR keys the
   reusable artifact by TP rank/size and SP layout.
 - **HSDP + SP:** the general parallel configuration permits HSDP over SP, but
-  DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight
-  materialization and synchronization.
+  each selected component must use `weight_transfer="rank-local"`. HSDP remains
+  responsible for weight materialization and synchronization.
 - **HSDP + DP or TP:** rejected independently by the diffusion parallel
   configuration.
 
@@ -447,8 +449,9 @@ recorded layouts (they all keep plain 1-byte dtypes over ordinary strided
 views: online INT8 has an int8 weight plus fp32 scale, online MXFP8 has an
 fp8 weight plus e8m0 block scale, each with a contiguous or transposed-view
 finalized tensor, all covered by the same physical-order packing as online
-FP8). Other online methods must use `--dlo-no-use-allgather` or disable
-online quantization until their runtime layouts are validated.
+FP8). Other online methods must set
+`layer_options.dit.weight_transfer="rank-local"` or disable online quantization
+until their runtime layouts are validated.
 
 The Host Weight Runtime representation, publication, and no-AllGather consumer
 contracts are merged; see

docs/design/feature/offloader/layerwise_offload.md

@@ -36,12 +36,14 @@ dependencies, not incidental global synchronization, define correctness.
 ## Topology contract
 
 A DiT declares ordered block-container attributes using
-`_layerwise_offload_blocks_attrs` or an `OffloadPlan`. Each resolved item must
-be an executable `nn.Module`; ordering must match forward execution. Multiple
-containers are concatenated in execution order.
+`_layerwise_offload_blocks_attrs` or an `OffloadPlan`. A text encoder declares
+its block-list paths through `OffloadPlan.encoder_block_attrs`. Each resolved
+item must be an executable `nn.Module`; ordering must match forward execution.
+Multiple containers are streamed as separate hook rings.
 
-Non-block DiT modules, encoders, VAEs, and declared resident modules are not
-streamed by this backend.
+Non-block state of a selected DiT or text encoder remains device resident.
+Unselected components, VAEs, and declared resident modules are not streamed by
+this backend.
 
 ## Invariants and limitations
 
@@ -50,7 +52,8 @@ streamed by this backend.
 - Blocks with heterogeneous shapes and dtypes use their recorded metadata;
   implementations must not assume identical blocks.
 - The ring is initialized before the first forward so block zero is available.
-- This backend is single-device; distributed weight reconstruction belongs to
+- This backend uses rank-local host weights. Multi-device compute may remain
+  active, but cross-rank weight reconstruction belongs to
   [Distributed Layerwise Offload](distributed_layerwise_offload.md).
 
 Shared selection and lifecycle behavior is defined in the

docs/design/feature/offloader/module_offload.md

@@ -2,7 +2,7 @@
 
 Model-level offload enforces mutual exclusion between pipeline component
 groups. For commands and model support, see the
-[model-level user guide](../../../user_guide/diffusion/offloader/cpu_offload.md).
+[model-level user guide](../../../user_guide/diffusion/offloader/module_offload.md).
 
 ## Generic hook path
 
@@ -10,10 +10,15 @@ groups. For commands and model support, see the
 makes encoders and VAEs device resident for initialization, then installs
 `SequentialOffloadHook` instances:
 
-- each DiT hook offloads all encoders and other DiTs before loading itself;
-- each encoder hook offloads every DiT before loading itself; and
+- each DiT hook offloads selected encoders and other selected DiTs before
+  loading itself;
+- each encoder hook offloads selected DiTs before loading itself; and
 - VAEs are not part of the mutual-exclusion hook set.
 
+All execution-stage modules still receive a hook so an unselected resident
+stage can evict a selected component before it runs. Only components listed in
+`diffusion_offload_config.components` are eligible to move to CPU.
+
 The hook moves parameters and buffers without recursively calling
 `module.to()`. This avoids recursion through modules that retain references to
 other transformer blocks. Host tensors are pinned when configured; DTensors
@@ -36,7 +41,9 @@ contract:
 
 - only the executing component is device resident;
 - enable and disable are idempotent;
-- transfers use the supplied device, pinning, and HSDP settings; and
+- transfers use the supplied device, pinning, and HSDP settings;
+- explicit component selection is either honored or rejected before partial
+  setup; and
 - disabling removes all model-local offload hooks or contexts.
 
 The delegation is currently duck typed. If another independent implementation