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

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

examples/offline_inference/image_to_video/README.md

@@ -152,9 +152,6 @@ python image_to_video.py \
 | `--vae-use-tiling` | flag | off | Enable VAE tiling for memory optimization |
 | `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
 | `--enable-layerwise-offload` | flag | off | Enable layerwise offloading on DiT modules |
-| `--enable-distributed-layerwise-offload` | flag | off | Enable distributed layerwise offload |
-| `--dlo-use-allgather` / `--dlo-no-use-allgather` | flag | on / off | Select sharded AllGather reconstruction (default) or rank-local streaming |
-| `--dlo-resident-layers` | int | `0` | Leading main-DiT blocks kept device-resident during distributed layerwise offload |
 | `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
 | `--tensor-parallel-size` | int | `1` | Tensor parallel size (effective for models that support TP, e.g. LTX2) |
 | `--ulysses-degree` | int | `1` | Ulysses sequence parallel degree |
@@ -347,6 +344,8 @@ Key arguments:
 - `--hsdp-shard-size`: Number of GPUs to shard model weights across within each replica group. -1 (default) auto-calculates as world_size / replicate_size.
 - `--hsdp-replicate-size`: Number of replica groups for HSDP. Each replica holds a full sharded copy. Default 1 means pure sharding (no replication).
 
+
+
 > ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
 
 ## Wan2.1 VACE Conditional Tasks

examples/offline_inference/minimax_h3/dlo_lifecycle.py

@@ -121,7 +121,6 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
             data_parallel_size=1,
             text_encoder_tp_size=args.tp_size,
             vae_patch_parallel_size=args.tp_size,
-            enable_distributed_layerwise_offload=False,
         )
     return common
 

examples/offline_inference/text_to_video/text_to_video.md

@@ -262,16 +262,14 @@ python text_to_video.py \
 - `--num-inference-steps`: sampling steps. Default depends on model.
 - `--fps`: frames per second for the saved MP4.
 - `--output`: path to save the generated video.
-- `--extra-body`: JSON dict of model-specific knobs (declared in `vllm_omni/model_extras/`), merged into sampling `extra_args`. See the Helios recipes above.
+- `--extra-body`: JSON object of model-specific generation knobs, filtered against the model's declared `extra_body_params` and merged into sampling `extra_args` (see the Helios and Cosmos3 examples above and [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)).
 - `--vae-use-slicing`: enable VAE slicing for memory optimization.
 - `--vae-use-tiling`: enable VAE tiling for memory optimization.
 - `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
 - `--tensor-parallel-size`: tensor parallel size (effective for models that support TP, e.g. LTX2).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 - `--enable-layerwise-offload`: enable layerwise offloading on DiT modules.
-- `--enable-distributed-layerwise-offload`: enable distributed layerwise offload.
-- `--dlo-use-allgather` / `--dlo-no-use-allgather`: use sharded AllGather reconstruction (the default) or stream rank-local weights without AllGather.
-- `--dlo-resident-layers`: keep this many leading main-DiT blocks device-resident during distributed layerwise offload (default: `0`).
+- `--diffusion-offload-config`: component-selective offload JSON; do not combine it with the legacy flags above.
 - `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
 - `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio.
 - `--quantization`: quantization method (such as `fp8` for FP8).
@@ -279,7 +277,6 @@ python text_to_video.py \
 - `--lora-path`: path to PEFT LoRA adapter folder or checkpoint file.
 - `--lora-scale`: scale factor for LoRA weights.
 - `--lora-backend`: backend for loading LoRA adapters. Default: peft. Available options: peft, distill.
-- `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above).
 
 ### Wan2.2-specific
 
@@ -299,4 +296,6 @@ python text_to_video.py \
 | 720p I2V | 7.0 | 6.0 | 50 |
 | CFG-distilled | (same) | 1.0 | 50 |
 
-> If you encounter OOM errors, try `--vae-use-slicing`, `--vae-use-tiling`, `--enable-cpu-offload`, or `--quantization fp8`.
+> If you encounter OOM errors, try `--vae-use-slicing`, `--vae-use-tiling`,
+> module/layer offload through `--diffusion-offload-config`, or
+> `--quantization fp8`.

examples/offline_inference/text_to_video/text_to_video.py

@@ -337,6 +337,12 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
+    parser.add_argument(
+        "--diffusion-offload-config",
+        type=json.loads,
+        default=None,
+        help="Component-selective diffusion offload config as JSON.",
+    )
     parser.add_argument(
         "--enable-distributed-layerwise-offload",
         action="store_true",
@@ -534,6 +540,7 @@ def main():
     omni_kwargs = dict(
         model=args.model,
         enable_layerwise_offload=args.enable_layerwise_offload,
+        diffusion_offload_config=args.diffusion_offload_config,
         vae_use_slicing=args.vae_use_slicing,
         vae_use_tiling=args.vae_use_tiling,
         enable_cpu_offload=args.enable_cpu_offload,

examples/online_serving/diffusers_pipeline_adapter/README.md

@@ -13,6 +13,7 @@ implemented with clear Diffusers or vLLM-Omni behavior.
 - Sequence parallel execution
 - TeaCache / Cache-DiT acceleration
 - Step-wise execution (continuous batching)
+- Component-selective `diffusion_offload_config` (the legacy whole-pipeline CPU-offload options remain supported)
 
 For these features, it is recommended to use natively supported pipelines instead.