From ea3cf99894f379c7732e6d62aa464a6f4285fbd9 Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Sun, 6 Sep 2026 05:40:09 +0800 Subject: [PATCH] [diffusion][feature] Add component-selective offload policies (#5929) Signed-off-by: lishunyang12 --- docs/.nav.yml | 4 +- docs/design/feature/offloader/README.md | 41 +- .../distributed_layerwise_offload.md | 43 +- .../feature/offloader/layerwise_offload.md | 15 +- .../{cpu_offload.md => module_offload.md} | 15 +- docs/design/index.md | 2 +- docs/user_guide/diffusion/cpu_offload.md | 6 +- .../distributed_layerwise_offload.md | 188 +++- .../diffusion/offloader/layerwise_offload.md | 62 +- .../{cpu_offload.md => module_offload.md} | 28 +- .../diffusers_pipeline_adapter.md | 1 + .../image_to_video/README.md | 5 +- .../minimax_h3/dlo_lifecycle.py | 1 - .../text_to_video/text_to_video.md | 11 +- .../text_to_video/text_to_video.py | 7 + .../diffusers_pipeline_adapter/README.md | 1 + mkdocs.yml | 4 +- recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md | 3 +- recipes/MiniMaxAI/MiniMax-H3.md | 59 +- tests/config/test_omni_config.py | 60 ++ .../test_diffusers_backend.py | 7 +- .../model_loader/test_diffusers_loader.py | 153 ++- .../models/cosmos3/test_cosmos3_pipeline.py | 30 + .../models/magi2/test_native_dlo_allgather.py | 4 +- .../models/magi2/test_pipeline_magi2.py | 10 +- .../minimax_h3/test_minimax_h3_contract.py | 141 +-- .../test_minimax_h3_dlo_lifecycle.py | 42 + .../models/minimax_h3/test_minimax_h3_lora.py | 49 + .../minimax_h3/test_minimax_h3_offload.py | 67 +- .../models/sana_video/test_cache_offload.py | 4 +- .../test_sensenova_u1_paged_decode.py | 8 + .../models/wan2_2/test_wan22_i2v_pipeline.py | 5 + .../wan2_2/test_wan22_pipeline_helpers.py | 9 + .../models/wan2_2/test_wan22_s2v_pipeline.py | 65 +- .../models/wan2_2/test_wan22_vace_pipeline.py | 41 + tests/diffusion/offloader/conftest.py | 26 + tests/diffusion/offloader/helpers.py | 107 ++ .../test_distributed_layerwise_backend.py | 670 +++++++++--- .../offloader/test_layerwise_backend.py | 869 ++++++++++++++-- .../offloader/test_sequential_backend.py | 228 ++++- tests/diffusion/test_diffusion_engine.py | 56 + .../test_diffusion_engine_dummy_run.py | 17 + .../diffusion/test_diffusion_model_runner.py | 54 +- .../test_multiproc_engine_concurrency.py | 114 ++- .../test_minimax_h3_dlo_dp2_t2va.py | 13 +- .../test_async_omni_diffusion_config.py | 59 +- vllm_omni/config/omni_config.py | 31 +- vllm_omni/config/stage_config.py | 3 + vllm_omni/diffusion/data.py | 27 +- vllm_omni/diffusion/diffusion_engine.py | 44 +- .../diffusion/executor/multiproc_executor.py | 57 +- .../model_loader/diffusers_loader.py | 105 +- .../models/cosmos3/pipeline_cosmos3.py | 13 + .../pipeline_diffusers_adapter.py | 7 + .../diffusion/models/minimax_h3/encoder.py | 49 - .../models/minimax_h3/pipeline_minimax_h3.py | 91 +- .../models/wan2_2/pipeline_wan2_2.py | 9 + .../models/wan2_2/pipeline_wan2_2_i2v.py | 2 + .../models/wan2_2/pipeline_wan2_2_s2v.py | 8 +- .../models/wan2_2/pipeline_wan2_2_vace.py | 7 +- vllm_omni/diffusion/offloader/__init__.py | 3 +- vllm_omni/diffusion/offloader/base.py | 153 ++- .../diffusion/offloader/block_discovery.py | 14 +- .../diffusion/offloader/component_utils.py | 221 ++++ vllm_omni/diffusion/offloader/config.py | 429 ++++++++ .../distributed_layerwise_backend.py | 954 +++++++++--------- .../diffusion/offloader/layerwise_backend.py | 564 ++++++----- vllm_omni/diffusion/offloader/offload_plan.py | 28 +- .../diffusion/offloader/sequential_backend.py | 274 +++-- vllm_omni/diffusion/offloader/tensor_utils.py | 150 +++ .../worker/diffusion_model_runner.py | 25 +- vllm_omni/engine/arg_utils.py | 3 + vllm_omni/engine/async_omni_engine.py | 1 + vllm_omni/entrypoints/cli/serve.py | 29 +- 74 files changed, 5121 insertions(+), 1584 deletions(-) rename docs/design/feature/offloader/{cpu_offload.md => module_offload.md} (81%) rename docs/user_guide/diffusion/offloader/{cpu_offload.md => module_offload.md} (65%) create mode 100644 tests/diffusion/offloader/conftest.py create mode 100644 tests/diffusion/offloader/helpers.py create mode 100644 vllm_omni/diffusion/offloader/component_utils.py create mode 100644 vllm_omni/diffusion/offloader/config.py diff --git a/docs/.nav.yml b/docs/.nav.yml index b5a83c338..6477323c4 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -72,7 +72,7 @@ nav: - Feature Compatibility: user_guide/feature_compatibility.md - CPU Offloading: - Overview: user_guide/diffusion/cpu_offload.md - - Model-Level: user_guide/diffusion/offloader/cpu_offload.md + - Model-Level: user_guide/diffusion/offloader/module_offload.md - Layerwise: user_guide/diffusion/offloader/layerwise_offload.md - Distributed Layerwise: user_guide/diffusion/offloader/distributed_layerwise_offload.md - Cache Acceleration: @@ -151,7 +151,7 @@ nav: - Skip-Softmax: design/feature/skip_softmax.md - CPU Offloading: - Overview: design/feature/offloader/README.md - - Model-Level: design/feature/offloader/cpu_offload.md + - Model-Level: design/feature/offloader/module_offload.md - Layerwise: design/feature/offloader/layerwise_offload.md - Distributed Layerwise: design/feature/offloader/distributed_layerwise_offload.md - Cache-DiT: design/feature/cache_dit.md diff --git a/docs/design/feature/offloader/README.md b/docs/design/feature/offloader/README.md index f071196b1..7b767f41f 100644 --- a/docs/design/feature/offloader/README.md +++ b/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 diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index 2b5f929ed..730d55472 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/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 `. @@ -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 diff --git a/docs/design/feature/offloader/layerwise_offload.md b/docs/design/feature/offloader/layerwise_offload.md index 3fff501da..344fd3c8c 100644 --- a/docs/design/feature/offloader/layerwise_offload.md +++ b/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 diff --git a/docs/design/feature/offloader/cpu_offload.md b/docs/design/feature/offloader/module_offload.md similarity index 81% rename from docs/design/feature/offloader/cpu_offload.md rename to docs/design/feature/offloader/module_offload.md index 2318611cc..43ef49194 100644 --- a/docs/design/feature/offloader/cpu_offload.md +++ b/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 diff --git a/docs/design/index.md b/docs/design/index.md index 81eac4736..c5b28130c 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -66,7 +66,7 @@ The design contracts separate selection mechanics from backend algorithms: #### CPU offloading - [Overview and Shared Contracts](feature/offloader/README.md) -- [Model-Level Offload](feature/offloader/cpu_offload.md) +- [Model-Level Offload](feature/offloader/module_offload.md) - [Layerwise Offload](feature/offloader/layerwise_offload.md) - [Distributed Layerwise Offload](feature/offloader/distributed_layerwise_offload.md) diff --git a/docs/user_guide/diffusion/cpu_offload.md b/docs/user_guide/diffusion/cpu_offload.md index 93ac5e48e..bd34bf3bf 100644 --- a/docs/user_guide/diffusion/cpu_offload.md +++ b/docs/user_guide/diffusion/cpu_offload.md @@ -11,7 +11,7 @@ For the shared factory, discovery, and lifecycle contract, see the | Strategy | Device residency | Parallel scope | Guide | | --- | --- | --- | --- | -| Model-level (sequential) | One pipeline component group at a time | Single device | [Model-Level Offloading](offloader/cpu_offload.md) | +| Model-level (sequential) | One pipeline component group at a time | Single device | [Model-Level Offloading](offloader/module_offload.md) | | Layerwise (blockwise) | One transformer block, with next-block prefetch | Single device | [Layerwise Offloading](offloader/layerwise_offload.md) | | Distributed layerwise | Fixed two-block device buffer; optional host-weight sharding and AllGather | Multiple GPU/NPU ranks | [Distributed Layerwise Offloading](offloader/distributed_layerwise_offload.md) | @@ -27,7 +27,7 @@ higher-priority strategy is selected. ## Quick selection -- Use [model-level offloading](offloader/cpu_offload.md) when swapping whole +- Use [model-level offloading](offloader/module_offload.md) when swapping whole encoders and DiTs is enough to fit and phase-boundary transfers are acceptable. - Use [layerwise offloading](offloader/layerwise_offload.md) for compute-heavy video @@ -63,7 +63,7 @@ These headings preserve links to sections that moved into dedicated guides. ## Model-level (Sequential) Offloading -See [Model-Level Offloading](offloader/cpu_offload.md). +See [Model-Level Offloading](offloader/module_offload.md). ## Layerwise (Blockwise) Offloading diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md index 45ff8df2d..b6e67c5ca 100644 --- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md +++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md @@ -1,11 +1,13 @@ # Distributed Layerwise Offloading Distributed layerwise offloading (DLO) extends block streaming to multi-device -deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of -the host weights and reconstructs each layer at runtime. Without AllGather, -each rank streams a complete block independently. Compatible TP1 deployments -can share checkpoint-backed host pages among processes on the same node; -otherwise DLO streams the ordinary loader's rank-local tensors. +deployments. Each selected component chooses its transfer independently. With +`allgather`, each rank stores roughly one shard per member of the effective DLO +group (DP, or SP when DP is one) and reconstructs each layer at runtime. With +`rank-local`, each rank streams a complete loader-produced block independently. +Compatible TP1 DiT deployments can share checkpoint-backed host pages among +processes on the same node; otherwise DLO streams the ordinary loader's +rank-local tensors. See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md) for the implementation contract and compatibility matrix. @@ -24,52 +26,120 @@ Slots: slot 0: Layer N slot 1: Layer N+1 AllGather communicates only request-independent weight shards, so data- parallel ranks may process different requests concurrently. +The public API describes intent in one nested configuration: + +| Level | Setting | Question answered | +| --- | --- | --- | +| Granularity | `mode` | move a complete component (`module`) or stream its blocks (`layer`) | +| Selection | `components` list | whether `dit` and/or `text_encoder` may move to CPU | +| Per-component layer options | `layer_options` | choose rank-local or AllGather weight transfer and optionally retain leading DiT blocks | + +`weight_transfer` applies only to model weights. It does not change encoder +activations or the encoder's compute parallelism. `rank-local` copies each +rank's loader-produced block to its device. `allgather` stores one shard of +that block per rank, copies each shard to its device, and reconstructs the +complete block with a device-side AllGather before compute. + +The historical `enable_cpu_offload`, `enable_layerwise_offload`, and +`enable_distributed_layerwise_offload` fields remain compatibility aliases +for existing callers and model-specific stage lifecycles. Existing +combinations retain their priority (`distributed layerwise` > `layerwise` > +`module`). A compact config rejects conflicting compatibility strategies and +non-default legacy DLO tuning. +`distributed-layerwise` remains an internal backend name; users select +`mode="layer"`, and the runtime chooses the capable backend. + ## Usage ```bash -# Four ranks with sharded host weights and AllGather -vllm serve /path/to/model --omni \ - --enable-distributed-layerwise-offload \ - --data-parallel-size 4 +# Four SP ranks; Wan declares replicated text-encoder weights, so both +# selected components can use AllGather safely. +vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni \ + --diffusion-offload-config \ + '{"mode":"layer","components":["dit","text_encoder"],"layer_options":{"dit":{"weight_transfer":"allgather"},"text_encoder":{"weight_transfer":"allgather"}}}' \ + --usp 4 -# Standard-loader rank-local weights, without DLO AllGather +# Mix transfers independently: sharded DiT, loader-produced encoder layout vllm serve /path/to/model --omni \ - --enable-distributed-layerwise-offload \ - --data-parallel-size 4 \ - --dlo-no-use-allgather + --diffusion-offload-config \ + '{"mode":"layer","components":["dit","text_encoder"],"layer_options":{"dit":{"weight_transfer":"allgather"},"text_encoder":{"weight_transfer":"rank-local"}}}' \ + --usp 4 -# Sequence parallel deployment +# Full-topology rank-local DLO for an existing model-specific lifecycle vllm serve /path/to/model --omni \ --enable-distributed-layerwise-offload \ + --dlo-no-use-allgather \ --usp 4 ``` +A compact layer config whose selected components are all `rank-local` and +whose `resident_layers` is zero resolves to the ordinary layerwise backend. It +uses the ordinary loader and does not request DLO's direct-checkpoint mmap. +Keep the compatibility DLO flags for a model-specific full-topology rank-local +lifecycle; do not add resident layers solely to force backend selection. + ```python from vllm_omni import Omni omni = Omni( model="/path/to/model", - enable_distributed_layerwise_offload=True, - dlo_use_allgather=True, + diffusion_offload_config={ + "mode": "layer", + "components": ["dit", "text_encoder"], + "layer_options": { + "dit": { + "weight_transfer": "rank-local", + "resident_layers": 20, + }, + "text_encoder": { + "weight_transfer": "allgather", + }, + }, + "pin_memory": True, + }, ) ``` ## Flags -| Flag | Meaning | Default | +| Setting | Meaning | Default | | --- | --- | --- | -| `--enable-distributed-layerwise-offload` | Enable DLO | `false` | -| `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` | -| `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` | -| `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` | -| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` | +| `diffusion_offload_config.mode` | `module` or `layer` granularity | required | +| `diffusion_offload_config.components` | Non-empty list containing `dit`, `text_encoder`, or both | required | +| `layer_options.NAME.weight_transfer` | `rank-local` or `allgather` | `rank-local` | +| `layer_options.dit.resident_layers` | Leading main-DiT blocks kept on device; requires `rank-local` and model-declared resident paths | `0` | +| `diffusion_offload_config.pin_memory` | Pin streamed host memory for faster H2D copies | `true` | +| `--data-parallel-size N` | DP ranks; DP is the DLO group when greater than one, otherwise SP is used | `1` | | `--host-weight-runtime-mode {disabled,preferred,required}` | HWR policy: no interaction, populate on a miss, or require an exact hit | `disabled` | | `--host-weight-runtime-root PATH` | Writable node-local HWR store shared by workers in one storage domain; required for `preferred` and `required` | unset | | `--dlo-host-registration-limit-gib N` | Optional per-worker ceiling for registering an HWR mmap; zero adds no ceiling | `0` | +HWR and host-registration tuning remain on the compatibility DLO interface and +cannot be combined with `diffusion_offload_config` in this release. + +## Component and weight-transfer matrix + +| Component | `allgather` | `rank-local` | +| --- | --- | --- | +| DiT | Shard each loader-visible block across the DLO DP group, or the SP group when DP is one | Stream each rank's complete loader-produced block | +| Text encoder | Same shard + AllGather path when the model declares identical encoder weights across the DLO group | Stream each rank's complete encoder block, including encoder-TP shards | + +The text encoder's compute group and weight-transfer group are separate +concepts. An encoder TP group owns different parameter shards, so it is not a +valid AllGather offload group. Models opt in through +`OffloadPlan.encoder_dlo_weight_replication`; otherwise multi-rank encoder +AllGather fails at startup with guidance to use `rank-local`. This makes SP +AllGather available to replicated encoders without silently corrupting an +encoder-TP layout. + +`resident_layers` remains a DiT-only setting and currently requires the DiT +transfer to be `rank-local`. + ## Host-weight loading -The diffusion loader chooses host storage before DLO is enabled. It first +When the distributed backend is selected by AllGather, resident DiT layers, or +the compatibility DLO flag, the diffusion loader chooses host storage before +DLO is enabled. It first attempts to build a complete, validated direct-checkpoint mmap plan. If names, coverage, shape, dtype, topology, or loader-callback compatibility cannot be proven, it runs the ordinary model loader instead. DLO consumes that result and @@ -97,20 +167,20 @@ With direct checkpoint mmap, the loader: 6. preserves `post_load_weights()` and `validate_loaded_weights()` lifecycle hooks. -For AllGather with a group larger than one, each process copies only its -persistent shard and then releases the source mapping. For no-AllGather, each -process keeps the mapping open and packs complete blocks through two bounded -pinned staging slots. Processes mapping the same files on one node share the -immutable pages; no-AllGather still performs a complete-block H2D copy in each -process. +For a DiT using `allgather` with a group larger than one, each process copies +only its persistent shard and then releases the source mapping. For a DiT using +`rank-local`, each process keeps the mapping open and packs complete blocks +through two bounded pinned staging slots. Processes mapping the same files on +one node share the immutable pages; rank-local transfer still performs a +complete-block H2D copy in each process. -When the effective DLO group size is one, `dlo_use_allgather=True` does not +When the effective DLO group size is one, `weight_transfer="allgather"` does not perform a collective and uses the same rank-local transfer behavior. ### Final-layout Host Weight Runtime HWR is an opt-in startup optimization for models that declare the final-layout -BF16 restore contract. Use it only with no-AllGather DLO: +BF16 restore contract. Use it only when the selected DiT uses `rank-local`: The validated BF16 model contracts currently cover MiniMax H3 and `black-forest-labs/FLUX.2-klein-4B`. FLUX.2-klein-9B shares the same model @@ -204,7 +274,7 @@ H2D behavior. #### Registered direct H2D Registration is attempted automatically only for an eligible warm HWR hit when -`pin_cpu_memory` is enabled. A successful path is: +the existing `pin_cpu_memory` policy is enabled. A successful path is: ```text shared read-only HWR mmap -> existing rotating HBM block buffer -> GPU kernel @@ -217,10 +287,10 @@ kernels access host memory directly. CUDA requires read-only host-registration capability for the immutable HWR mapping. Unsupported capability, a positive registration limit smaller than the complete page-aligned mapping, or a safely rolled-back registration error -falls back to bounded staging. Programmatic configurations can set -`pin_cpu_memory=False` to disable registration explicitly. A successful -registration locks the complete mapped range in host memory for that worker's -lifetime, so use +falls back to bounded staging. Programmatic compatibility configurations can +set `pin_cpu_memory=False` to disable registration explicitly. +A successful registration locks the complete mapped range in host memory for +that worker's lifetime, so use `--dlo-host-registration-limit-gib` when an operator-enforced ceiling is required. @@ -228,8 +298,8 @@ Shutdown drains H2D work, releases hook/source references, unregisters every range, and only then closes the HWR lease. AllGather and direct checkpoint mmap retain their existing transfer paths. -When HWR is disabled, DLO is disabled, or AllGather is enabled, the loader does -not resolve HWR sources or construct its store. +When HWR is disabled, DLO is disabled, or the DiT uses `allgather`, the loader +does not resolve HWR sources or construct its store. ## Declarative topology @@ -241,38 +311,54 @@ from vllm_omni.diffusion.offloader import OffloadPlan class MyPipeline(nn.Module): _dit_modules = ["transformer"] + _encoder_modules = ["prompt_model"] _offload_plan = OffloadPlan( block_attrs={"transformer": ("blocks",)}, offload_submodules={"context_encoder": "layers"}, + encoder_component_types={"prompt_model": "text_encoder"}, + encoder_block_attrs={"prompt_model": ("encoder.layers",)}, + encoder_dlo_weight_replication=frozenset({"prompt_model"}), + on_demand_component_paths=frozenset({"prompt_model"}), ) ``` -When no plan exists, discovery falls back to -`_layerwise_offload_blocks_attrs` and then heuristic attribute lookup. +`encoder_component_types` maps arbitrary encoder paths to the public +`text_encoder` selector. `encoder_dlo_weight_replication` is the explicit +safety declaration for reusing the DiT DLO group. On-demand encoder non-block +state is loaded and offloaded by its pipeline phase. + +When no plan exists, DiT discovery falls back to +`_layerwise_offload_blocks_attrs` and then heuristic attribute lookup; +undeclared auxiliary components remain resident. ## Data-parallel concurrency With `data_parallel_size > 1` and AllGather enabled, the scheduler can process -up to `dp_size` requests per denoising step. Every concurrent request must set -the same explicit `num_inference_steps`; `None` is rejected because every rank -must enter each collective. +up to `dp_size` requests per denoising step. Concurrent requests must resolve +to the same denoise schedule, so every request must provide the same explicit +`num_inference_steps`. Pipeline defaults are rejected because different request +modes can resolve the same `None` value to different schedules. ## Limitations - Direct checkpoint mmap currently requires TP1. TP greater than one is outside the Phase A shared-mmap support scope and falls back before model mutation to the ordinary TP-aware loader. DLO can stream that runtime layout, - and eligible no-AllGather configurations may use HWR final-layout artifacts, + and eligible DiT rank-local configurations may use HWR final-layout artifacts, but direct checkpoint mmap provides no shared-mmap host-memory guarantee for TP greater than one. -- HSDP plus AllGather is rejected to avoid double sharding. HSDP without - AllGather has limited end-to-end validation. +- HSDP with `allgather` on any selected component is rejected to avoid double + sharding. HSDP with rank-local transfers has limited end-to-end validation. +- With data parallelism, `allgather` cannot be combined with TeaCache or + Cache-DiT on any selected component. Prompt-embedding cache is also + incompatible when the text encoder uses `allgather`. Rank-local cache hits + could otherwise make ranks enter different weight collectives. - Per-tensor online FP8 linears use the ordinary loader and can run with either - DLO transfer path. With AllGather, every rank temporarily materializes the - complete FP8 model in host memory before DLO retains only its shard. Other - online quantization methods require no-AllGather until their runtime layouts - are validated. -- Resident leading layers require `--dlo-no-use-allgather` and a model + DiT transfer path. With DiT `allgather`, every rank temporarily materializes + the complete FP8 model in host memory before DLO retains only its shard. + Other online quantization methods require rank-local transfer for the + affected component until their runtime layouts are validated. +- Resident leading layers require DiT `rank-local` transfer and a model `OffloadPlan` that declares eligible `resident_dit_paths`. - DP concurrency requires an explicit, identical inference-step count. diff --git a/docs/user_guide/diffusion/offloader/layerwise_offload.md b/docs/user_guide/diffusion/offloader/layerwise_offload.md index c88b67fcc..e948eda67 100644 --- a/docs/user_guide/diffusion/offloader/layerwise_offload.md +++ b/docs/user_guide/diffusion/offloader/layerwise_offload.md @@ -18,8 +18,9 @@ stream. | ... | ... | ... | ... | | last block | Prefetch block 0 | Compute last block | Free last block | -Encoders, VAE modules, and non-block DiT modules such as embeddings and norms -remain device resident. +Selected, plan-declared text-encoder blocks can use the same rank-local +streaming mechanism. Image encoders, VAEs, and non-block DiT modules remain +device resident. ## Usage @@ -28,15 +29,44 @@ from vllm_omni import Omni omni = Omni( model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", - enable_layerwise_offload=True, + diffusion_offload_config={ + "mode": "layer", + "components": ["dit"], + }, ) ``` ```bash vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \ - --omni --enable-layerwise-offload + --omni \ + --diffusion-offload-config '{"mode":"layer","components":["dit"]}' ``` +## Component selection + +Add `dit`, `text_encoder`, or both to the `components` list. Omitting +`layer_options` uses safe rank-local defaults: + +- `["dit"]` streams only DiT blocks. +- `["text_encoder"]` streams only declared text-encoder blocks. +- Listing both streams both components. + +```bash +# DiT-only layer offload +vllm serve /path/to/model --omni \ + --diffusion-offload-config \ + '{"mode":"layer","components":["dit"]}' + +# Stream a model-declared text encoder while keeping the DiT resident +vllm serve /path/to/model --omni \ + --diffusion-offload-config \ + '{"mode":"layer","components":["text_encoder"]}' +``` + +Encoder categories are resolved from `OffloadPlan.encoder_component_types` +first. A name-based fallback is retained for pipelines that predate +`OffloadPlan`. + ## Model integration Transformer classes declare containers of executable blocks: @@ -53,13 +83,29 @@ class Flux2Transformer2DModel(nn.Module): ] ``` -See the [layerwise design](../../../design/feature/offloader/layerwise_offload.md) for -the discovery and hook invariants. `OffloadPlan` is a separate declarative -topology path for distributed layerwise offload. +Auxiliary components use declarative pipeline metadata: + +```python +from vllm_omni.diffusion.offloader import OffloadPlan + + +class MyPipeline(nn.Module): + _encoder_modules = ["prompt_model"] + _offload_plan = OffloadPlan( + encoder_component_types={"prompt_model": "text_encoder"}, + encoder_block_attrs={"prompt_model": ("encoder.layers",)}, + ) +``` + +See the [layerwise design](../../../design/feature/offloader/layerwise_offload.md) +for the discovery and hook invariants. Both ordinary and distributed layerwise +offload consume the same `OffloadPlan` metadata. ## Limitations -- Single device only. +- The default weight transfer is `rank-local`. Set a selected component's + `weight_transfer` to `allgather` to shard its host weights across a + compatible multi-device group; backend selection is automatic. - Setup consolidates and pins block parameters, increasing cold-start time. - Performance depends on block compute time and host-to-device bandwidth; lightweight blocks may not hide transfers. diff --git a/docs/user_guide/diffusion/offloader/cpu_offload.md b/docs/user_guide/diffusion/offloader/module_offload.md similarity index 65% rename from docs/user_guide/diffusion/offloader/cpu_offload.md rename to docs/user_guide/diffusion/offloader/module_offload.md index 558985e66..0e1cb3fac 100644 --- a/docs/user_guide/diffusion/offloader/cpu_offload.md +++ b/docs/user_guide/diffusion/offloader/module_offload.md @@ -2,14 +2,14 @@ Model-level, or sequential, offloading keeps only the pipeline component group currently executing on the accelerator. It is the simplest offload strategy -and is selected with `--enable-cpu-offload`. +and is selected with `diffusion_offload_config.mode="module"`. ## How it works Pre-forward hooks enforce mutual exclusion between DiT and encoder modules: -- before an encoder runs, the DiT moves to CPU; -- before a DiT runs, encoders and other DiTs move to CPU; and +- before an encoder runs, selected DiTs move to CPU; +- before a DiT runs, selected encoders and other selected DiTs move to CPU; and - VAE modules remain on the accelerator. Pinned host memory reduces transfer overhead. Transfers occur at phase @@ -22,15 +22,28 @@ from vllm_omni import Omni omni = Omni( model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", - enable_cpu_offload=True, + diffusion_offload_config={ + "mode": "module", + "components": ["dit", "text_encoder"], + }, ) ``` ```bash vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \ - --omni --enable-cpu-offload + --omni \ + --diffusion-offload-config \ + '{"mode":"module","components":["dit","text_encoder"]}' ``` +List a component to select it for offload. Module mode rejects `layer_options` +such as `weight_transfer` and `resident_layers`. The +`enable_cpu_offload=True` compatibility entry point remains supported. New +integrations should prefer the explicit config; existing model-specific stage +lifecycles do not need to migrate until equivalent component coverage exists. +For example, MiniMax-H3's compatibility lifecycle also stages its VAEs, while +the compact selector intentionally covers only `dit` and `text_encoder`. + ## Model integration Pipelines should implement `SupportsComponentDiscovery`: @@ -73,9 +86,10 @@ while reusing sequential `.to()` movers. ## Limitations -- Single device only. +- Transfers are rank-local; module mode does not shard host payloads or add a + weight AllGather across ranks. - Higher cold-start latency. - Transfers between encoder and denoising phases add latency. -See the [model-level design](../../../design/feature/offloader/cpu_offload.md) +See the [model-level design](../../../design/feature/offloader/module_offload.md) for lifecycle and extension invariants. diff --git a/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md b/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md index 7edf7c61d..7ba8648d1 100644 --- a/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md +++ b/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md @@ -16,6 +16,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. diff --git a/examples/offline_inference/image_to_video/README.md b/examples/offline_inference/image_to_video/README.md index 24bc52164..9386264b5 100644 --- a/examples/offline_inference/image_to_video/README.md +++ b/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 diff --git a/examples/offline_inference/minimax_h3/dlo_lifecycle.py b/examples/offline_inference/minimax_h3/dlo_lifecycle.py index 5630777cb..bcea5c14d 100644 --- a/examples/offline_inference/minimax_h3/dlo_lifecycle.py +++ b/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 diff --git a/examples/offline_inference/text_to_video/text_to_video.md b/examples/offline_inference/text_to_video/text_to_video.md index 99a7a9d9f..f051d8f88 100644 --- a/examples/offline_inference/text_to_video/text_to_video.md +++ b/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`. diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index be18141e6..06179ef9e 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/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, diff --git a/examples/online_serving/diffusers_pipeline_adapter/README.md b/examples/online_serving/diffusers_pipeline_adapter/README.md index 37bf48aa0..56150c52d 100644 --- a/examples/online_serving/diffusers_pipeline_adapter/README.md +++ b/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. diff --git a/mkdocs.yml b/mkdocs.yml index 8a0e322e3..655e71e1f 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -83,7 +83,9 @@ plugins: design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md design/feature/cpu_offloading.md: design/feature/offloader/README.md - design/feature/model_level_offload.md: design/feature/offloader/cpu_offload.md + user_guide/diffusion/offloader/cpu_offload.md: user_guide/diffusion/offloader/module_offload.md + design/feature/model_level_offload.md: design/feature/offloader/module_offload.md + design/feature/offloader/cpu_offload.md: design/feature/offloader/module_offload.md design/feature/layerwise_offload.md: design/feature/offloader/layerwise_offload.md design/feature/distributed_layerwise_offload.md: design/feature/offloader/distributed_layerwise_offload.md - glightbox diff --git a/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md b/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md index debc89d89..f4b1e35ae 100644 --- a/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md +++ b/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md @@ -42,7 +42,8 @@ encoder's BF16 configuration and the video/audio VAEs' FP32 precision. Select one offload strategy per deployment: ```bash -# Stage 1 model-level CPU offload. +# Stage 1 full-topology model offload, including MiniMax-H3's VAEs. Keep the +# compatibility alias because the compact selector covers only DiT/text encoder. --stage-overrides '{"1":{"enable_cpu_offload":true}}' # Stage 1 distributed layerwise offload. Tune resident layers for available RAM. diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index 5ef1e9ebc..2f705812b 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -117,11 +117,12 @@ same time on a host sized for this minimum. The consumer-GPU profiles below are HBM budgets only. They still require the host-RAM budget above. -### Single GPU: accuracy and memory first +### Single GPU: blockwise capacity path -The single-GPU configuration uses model-level CPU offload. -This matches the accuracy-qualified reference path and prevents the Qwen3-VL -encoder and DiT from being resident on the GPU at the same time. +Use ordinary layerwise offload when one GPU cannot keep the Qwen3-VL encoder +and DiT resident together. The component list below streams the active DiT, +the Qwen vision blocks, and the first 50 Qwen text layers. Encoder blocks stay +rank-local; the video/audio VAEs remain resident. ```bash export MODEL=MiniMaxAI/MiniMax-H3 @@ -129,20 +130,37 @@ export PORT=8091 CUDA_VISIBLE_DEVICES=0 \ VLLM_WORKER_MULTIPROC_METHOD=spawn \ -VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \ +VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \ vllm serve "${MODEL}" \ --omni \ --host 0.0.0.0 \ --port "${PORT}" \ --trust-remote-code \ + --task-type fl2va \ --num-gpus 1 \ - --enable-cpu-offload \ + --diffusion-offload-config \ + '{"mode":"layer","components":["dit","text_encoder"]}' \ + --enforce-eager \ --diffusion-attention-backend FLASH_ATTN ``` -Use a GPU with enough memory for the active H3 component and enough system RAM -for both offloaded DiTs plus the shared components. Model-level offload keeps -the two DiTs mutually exclusive on GPU, but adds PCIe/NVLink transfer latency. +This is a capacity profile, not a latency profile: every denoising step streams +DiT blocks over the host link, while encoder blocks are streamed only during +the short conditioning phase. It requires host memory for the complete +checkpoint plus pinned transfer buffers. Re-measure peak HBM on the target +shape; this command does not claim a particular GPU model as validated. + +A one-B300 correctness smoke selected only `text_encoder` (keeping the DiT and +VAEs resident) and completed a 384x672, 5-second T2VA request with two denoise +steps. It installed 77 encoder hooks across the vision and text stacks, used a +77,728 MiB worker peak, and measured 2.803 seconds encode, 2.706 seconds +denoise, and 4.503 seconds decode. These reduced-step numbers validate the +execution path; they are not a quality or production-latency benchmark. + +To use whole-component behavior, change the config to +`{"mode":"module","components":["dit","text_encoder"]}`. Module +offload swaps the complete encoder and DiT and therefore has a higher +encode-phase peak than encoder blockwise offload. ### Two 24/32 GB GPUs: TP2 distributed layerwise offload @@ -483,10 +501,9 @@ Add this option to an existing H3 server command: Use the FL2VA-only partition for this capacity test. Loading the combined service would also load the Ref2VA DiT and would test a different memory -budget. A no-offload capacity check should contain none of -`--enable-cpu-offload`, `--enable-layerwise-offload`, or -`--enable-distributed-layerwise-offload`. VAE tiling changes decode placement -but does not offload model weights to the CPU. +budget. A no-offload capacity check should omit +`--diffusion-offload-config` and all legacy `--enable-*-offload` aliases. VAE +tiling changes decode placement but does not offload model weights to the CPU. The run passes the capacity check when the server initializes, the request finishes without CUDA OOM or Xid errors, `peak_used_mib` remains below the @@ -518,10 +535,11 @@ For example, keep the first main block's attention projections in BF16 with: ``` The structured option replaces `--quantization fp8`. Online FP8 can be used -with H3 layerwise offload and with both DLO transfer paths. The default -AllGather path uses the ordinary loader to finalize FP8 weights and scales -before sharding them across ranks. `--dlo-no-use-allgather` instead retains -complete rank-local tensors and avoids the synchronized request-wave contract. +with H3 layerwise offload and with either DiT DLO transfer. DiT `allgather` +uses the ordinary loader to finalize FP8 weights and scales before sharding +them across ranks. DiT `rank-local` instead retains complete loader-produced +tensors and avoids the synchronized request-wave contract. H3's TP-sharded +text encoder uses `rank-local`. ## AMD ROCm (gfx942 / gfx950) @@ -567,6 +585,13 @@ vllm serve "${MODEL}" \ --diffusion-attention-backend FLASH_ATTN ``` +This validated ROCm capacity recipe intentionally retains the compatibility +full-topology alias because MiniMax-H3 also stages its VAEs on that path. The +compact API in this release selects only `dit` and `text_encoder`, so replacing +the flag would change residency rather than perform a mechanical migration. +No removal deadline is assigned until the compact API offers equivalent +component coverage. + ### ROCm four GPUs The best-practice CUDA four-GPU configuration works on ROCm with the changes above. diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index f99de9bf3..068d90d28 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -6,6 +6,7 @@ from __future__ import annotations from dataclasses import fields from inspect import Parameter, signature +from multiprocessing.reduction import ForkingPickler from pathlib import Path import msgspec @@ -1390,3 +1391,62 @@ def test_diffusion_quantization_mapping_reaches_terminal_config(monkeypatch): assert cfg.quantization_config is not None assert cfg.quantization_config.get_name() == "int8" + + +def test_compact_offload_config_reaches_terminal_config(monkeypatch): + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.offloader.config import OffloadStrategy, resolve_offload_strategy + + compact_config = { + "mode": "layer", + "components": ["dit", "text_encoder"], + "layer_options": { + "dit": {"weight_transfer": "rank-local", "resident_layers": 12}, + "text_encoder": {"weight_transfer": "allgather"}, + }, + } + cfg = omni_config_module._DiffusionConfigProjection.from_kwargs( + diffusion_offload_config=compact_config, + ) + + monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500) + monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda _self: None) + cfg.enrich_config() + + assert cfg.diffusion_offload_config == compact_config + assert cfg.extras == {} + assert resolve_offload_strategy(cfg) is OffloadStrategy.DISTRIBUTED_LAYER_WISE + restored = ForkingPickler.loads(ForkingPickler.dumps(cfg)) + assert resolve_offload_strategy(restored) is OffloadStrategy.DISTRIBUTED_LAYER_WISE + + +def test_global_diffusion_offload_config_targets_only_diffusion_stages(): + compact_config = {"mode": "layer", "components": ["dit"]} + + config = _from_pipeline_key( + "minimax_h3_disaggregated", + cli_overrides={"diffusion_offload_config": compact_config}, + ) + + assert not hasattr(config.stage_by_id(0), "diffusion_config") + assert config.stage_by_id(1).diffusion_config.diffusion_offload_config == compact_config + + +def test_explicit_llm_stage_diffusion_offload_override_is_rejected(): + with pytest.raises(ValueError, match="no structured config owner: diffusion_offload_config"): + _from_pipeline_key( + "minimax_h3_disaggregated", + cli_overrides={ + "stage_0_diffusion_offload_config": {"mode": "layer", "components": ["dit"]}, + }, + ) + + +def test_compact_offload_config_is_validated_during_projection(): + with pytest.raises(ValueError, match="Unknown diffusion offload mode"): + omni_config_module._DiffusionConfigProjection.from_kwargs( + diffusion_offload_config={ + "mode": "layerwise", + "components": ["dit"], + } + ) diff --git a/tests/diffusion/diffusion_backend/test_diffusers_backend.py b/tests/diffusion/diffusion_backend/test_diffusers_backend.py index c5b71f502..082855b43 100644 --- a/tests/diffusion/diffusion_backend/test_diffusers_backend.py +++ b/tests/diffusion/diffusion_backend/test_diffusers_backend.py @@ -224,6 +224,7 @@ class TestPipelineArgumentsHandling: @pytest.mark.parametrize( "feature_id", [ + "compact_offload", "cfg_parallel", "tensor_parallel", "ulysses", @@ -235,7 +236,11 @@ class TestPipelineArgumentsHandling: ], ) def test_adapter_guard_unsupported_feature(self, feature_id): - if feature_id == "cfg_parallel": + if feature_id == "compact_offload": + od_config = _make_od_config( + diffusion_offload_config={"mode": "module", "components": ["dit"]}, + ) + elif feature_id == "cfg_parallel": od_config = _make_od_config( parallel_config=DiffusionParallelConfig(cfg_parallel_size=2, sequence_parallel_size=1), cache_backend="none", diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 4304f5d2e..4f60e379b 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -734,6 +734,153 @@ def test_dlo_transfers_loader_plan_and_skips_ordinary_weight_loading(monkeypatch assert loader.take_host_weight_plan() is None +def test_compact_rank_local_layer_offload_skips_dlo_mmap_planning(monkeypatch): + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + from vllm_omni.diffusion.offloader.config import materialize_legacy_offload_flags + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1), + quantization_config=None, + diffusion_offload_config={"mode": "layer", "components": ["dit"]}, + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, + dlo_use_allgather=True, + dlo_resident_layers=0, + pin_cpu_memory=True, + model="unused", + ) + materialize_legacy_offload_flags(od_config) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + calls: list[str] = [] + + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + loader.load_weights = lambda _model: calls.append("load") # type: ignore[method-assign] + loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + monkeypatch.setattr( + loader_mod, + "build_checkpoint_mmap_plan", + lambda *_args, **_kwargs: pytest.fail("ordinary rank-local offload must not plan DLO mmap"), + ) + + assert loader.load_model(load_device="cpu") is model + assert calls == ["load", "process"] + assert loader.take_host_weight_plan() is None + + +class _UnsupportedOnlineQuantMethod: + uses_meta_device = True + + +def _compact_layer_offload_config( + components: list[str], + layer_options: dict[str, dict[str, str]], +) -> SimpleNamespace: + return SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace( + use_hsdp=False, + tensor_parallel_size=1, + data_parallel_size=2, + sequence_parallel_size=1, + ), + quantization_config=None, + diffusion_offload_config={ + "mode": "layer", + "components": components, + "layer_options": layer_options, + }, + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, + dlo_use_allgather=True, + dlo_resident_layers=0, + pin_cpu_memory=False, + host_weight_runtime_mode="disabled", + model="unused", + ) + + +def _stub_ordinary_loader(loader, model, monkeypatch, loader_mod) -> None: + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + loader.load_weights = lambda _model: None # type: ignore[method-assign] + loader._process_weights_after_loading = lambda *_args: None # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + monkeypatch.setattr( + loader_mod, + "build_checkpoint_mmap_plan", + lambda *_args, **_kwargs: HostWeightPlanResult(None, "test fallback"), + ) + + +def test_text_encoder_allgather_rejects_unsupported_online_quantization(monkeypatch): + config = _compact_layer_offload_config( + ["text_encoder"], + {"text_encoder": {"weight_transfer": "allgather"}}, + ) + loader = DiffusersPipelineLoader(LoadConfig(), config) + model = nn.Module() + model.text_encoder = nn.Linear(2, 2, bias=False) + model.text_encoder.quant_method = _UnsupportedOnlineQuantMethod() + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + + with pytest.raises(ValueError, match="unsupported online methods: _UnsupportedOnlineQuantMethod"): + loader.load_model(load_device="cpu") + + +def test_dit_allgather_ignores_unsupported_quantization_on_unselected_encoder(monkeypatch): + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + + config = _compact_layer_offload_config( + ["dit"], + {"dit": {"weight_transfer": "allgather"}}, + ) + loader = DiffusersPipelineLoader(LoadConfig(), config) + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + model.text_encoder = nn.Linear(2, 2, bias=False) + model.text_encoder.quant_method = _UnsupportedOnlineQuantMethod() + _stub_ordinary_loader(loader, model, monkeypatch, loader_mod) + + assert loader.load_model(load_device="cpu") is model + + +def test_rank_local_encoder_quantization_does_not_mark_dit_mmap_plan_online(monkeypatch): + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + + config = _compact_layer_offload_config( + ["dit", "text_encoder"], + { + "dit": {"weight_transfer": "allgather"}, + "text_encoder": {"weight_transfer": "rank-local"}, + }, + ) + loader = DiffusersPipelineLoader(LoadConfig(), config) + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + model.text_encoder = nn.Linear(2, 2, bias=False) + model.text_encoder.quant_method = _UnsupportedOnlineQuantMethod() + planned_online_quantization: list[bool] = [] + + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + loader.load_weights = lambda _model: None # type: ignore[method-assign] + loader._process_weights_after_loading = lambda *_args: None # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + + def capture_plan(*_args, **kwargs): + planned_online_quantization.append(kwargs["online_quantization"]) + return HostWeightPlanResult(None, "test fallback") + + monkeypatch.setattr(loader_mod, "build_checkpoint_mmap_plan", capture_plan) + + assert loader.load_model(load_device="cpu") is model + assert planned_online_quantization == [False] + + def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch): import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod @@ -866,7 +1013,7 @@ def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch): ) assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model - assert allowlist_models == [model] + assert allowlist_models == [model.transformer] assert calls == [("load", True), "process"] assert loader.take_host_weight_plan() is None @@ -908,7 +1055,7 @@ def test_dlo_allgather_online_int8_uses_ordinary_loader(monkeypatch): ) assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model - assert allowlist_models == [model] + assert allowlist_models == [model.transformer] assert calls == [("load", True), "process"] assert loader.take_host_weight_plan() is None @@ -1024,7 +1171,7 @@ def test_dlo_allgather_online_mxfp8_uses_ordinary_loader(monkeypatch): ) assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model - assert allowlist_models == [model] + assert allowlist_models == [model.transformer] assert calls == [("load", True), "process"] assert loader.take_host_weight_plan() is None diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index 105b012bb..836e87219 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -30,6 +30,36 @@ def test_pipeline_declares_layerwise_offload_components() -> None: assert Cosmos3OmniDiffusersPipeline._resident_modules == [] assert hasattr(Cosmos3OmniDiffusersPipeline, "enable_omni_model_cpu_offload") + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import ( + Cosmos3LanguageModel, + Cosmos3VFMTransformer, + ) + + assert Cosmos3LanguageModel._layerwise_offload_blocks_attrs == ["layers"] + assert Cosmos3VFMTransformer._layerwise_offload_blocks_attrs == ["gen_layers"] + + +def test_component_selective_model_offload_fails_before_component_loading(monkeypatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3 as pipeline_module + + monkeypatch.setattr( + pipeline_module, + "get_local_device", + lambda: pytest.fail("component loading must not start for an unsupported selector"), + ) + config = SimpleNamespace( + diffusion_offload_config={"mode": "module", "components": ["dit"]}, + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, + dlo_use_allgather=True, + dlo_resident_layers=0, + pin_cpu_memory=True, + ) + + with pytest.raises(ValueError, match="does not support the dit/text_encoder component selector"): + pipeline_module.Cosmos3OmniDiffusersPipeline(od_config=config) + class StubScheduler: def __init__( diff --git a/tests/diffusion/models/magi2/test_native_dlo_allgather.py b/tests/diffusion/models/magi2/test_native_dlo_allgather.py index 0c097cfc0..9bea5dc30 100644 --- a/tests/diffusion/models/magi2/test_native_dlo_allgather.py +++ b/tests/diffusion/models/magi2/test_native_dlo_allgather.py @@ -220,7 +220,9 @@ def _load_mmap_transform_and_reconstruct( tensor_transforms=backend._mmap_transforms_by_tensor_id, ) assert all(not shard.requires_grad for shard in cpu_shards.values()) - assert all(parameter.numel() == 0 for parameter in target_block.parameters()) + # Shard preparation is non-destructive; initialize_hook clears source + # storage only after every shard and metadata entry has been prepared. + assert all(parameter.numel() > 0 for parameter in target_block.parameters()) for dtype, local_shard in cpu_shards.items(): gathered = torch.empty( diff --git a/tests/diffusion/models/magi2/test_pipeline_magi2.py b/tests/diffusion/models/magi2/test_pipeline_magi2.py index 7acb6b209..2cda7e66e 100644 --- a/tests/diffusion/models/magi2/test_pipeline_magi2.py +++ b/tests/diffusion/models/magi2/test_pipeline_magi2.py @@ -578,6 +578,8 @@ class _FakeLayerwiseHook: def __init__(self) -> None: self._prev_hook = None self.current_slot = 0 + self.dp_size = 1 + self.rank_local_mmap = False self.prefetch_layer = Mock() self.get_weights = Mock() @@ -625,7 +627,7 @@ def test_dlo_no_allgather_enable_preserves_staged_aux_and_streams_blocks( assert len(hooks) == len(pipeline.transformer.block) assert all(stager.offload.call_count == 1 for stager in stagers) assert all(stager.load.call_count == 0 for stager in stagers) - hooks[-1].prefetch_layer.assert_called_once_with( - slot=hooks[0].current_slot, - non_blocking=False, - ) + # First prefetch stays lazy so ranks with different local groups cannot + # enter collectives in a different order during startup. + for hook in hooks: + hook.prefetch_layer.assert_not_called() diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 91aac2df7..791c8d79d 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -487,6 +487,10 @@ def test_dlo_offload_plan_includes_token_refiner(): assert MiniMaxH3Pipeline._offload_plan.offload_submodules == {"token_refiner": "blocks"} assert MiniMaxH3Pipeline._offload_plan.resident_dit_paths == frozenset({"transformer"}) + assert MiniMaxH3Pipeline._offload_plan.encoder_component_types == {"text_encoder": "text_encoder"} + assert MiniMaxH3Pipeline._offload_plan.encoder_block_attrs == { + "text_encoder": ("vision.blocks", "text_model.layers") + } def test_joint_postprocess_is_multiprocessing_picklable(): @@ -1185,6 +1189,12 @@ def test_text_encoder_stub_constructs_without_group_or_weights(): from vllm_omni.diffusion.models.minimax_h3.encoder import ( MiniMaxH3Qwen3VLEncoder, ) + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + MiniMaxH3Pipeline, + ) + from vllm_omni.diffusion.offloader.component_utils import ( + get_encoder_block_groups, + ) encoder = MiniMaxH3Qwen3VLEncoder( "/nonexistent/text_encoder", @@ -1197,6 +1207,25 @@ def test_text_encoder_stub_constructs_without_group_or_weights(): # The stub has no parameters, so it never contributes to the runner's # strict missing-parameter check on non-encoder ranks. assert list(encoder.named_parameters()) == [] + assert ( + get_encoder_block_groups( + encoder, + "text_encoder", + MiniMaxH3Pipeline._offload_plan, + strict=True, + ) + == [] + ) + + # An arbitrary empty module is not silently treated as a distributed + # stub; explicit plans still validate their declared paths. + with pytest.raises(ValueError, match=r"text_encoder\.vision\.blocks was not found"): + get_encoder_block_groups( + nn.Module(), + "text_encoder", + MiniMaxH3Pipeline._offload_plan, + strict=True, + ) def test_global_quant_config_is_shared_by_dit_and_encoder(): @@ -1303,10 +1332,38 @@ def test_model_offload_uses_hooked_text_encoder_call(): @pytest.mark.parametrize( - "offload_flag", - ["enable_layerwise_offload", "enable_distributed_layerwise_offload"], + ("enable_layerwise", "enable_distributed"), + [(True, False), (False, True)], ) -def test_layerwise_offload_releases_text_encoder(offload_flag): +def test_legacy_layer_offload_preserves_minimax_stage_lifecycle(enable_layerwise, enable_distributed): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + diffusion_offload_config=None, + enable_cpu_offload=False, + enable_layerwise_offload=enable_layerwise, + enable_distributed_layerwise_offload=enable_distributed, + ) + pipeline.text_encoder = Mock() + expected = torch.ones(2, 3) + pipeline.text_encoder.encode_ids.return_value = expected + + actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) + + assert actual is expected + pipeline.text_encoder.load_to_device.assert_called_once_with() + pipeline.text_encoder.offload_to_cpu.assert_called_once_with() + + vae = Mock() + with pipeline._component_on_device(vae): + pass + vae.load_to_device.assert_called_once_with() + vae.offload_to_cpu.assert_called_once_with() + + +def test_layerwise_encoder_selection_releases_text_encoder(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline pipeline = object.__new__(MiniMaxH3Pipeline) @@ -1315,8 +1372,11 @@ def test_layerwise_offload_releases_text_encoder(offload_flag): enable_cpu_offload=False, enable_layerwise_offload=False, enable_distributed_layerwise_offload=False, + diffusion_offload_config={ + "mode": "layer", + "components": ["text_encoder"], + }, ) - setattr(pipeline.od_config, offload_flag, True) pipeline.text_encoder = Mock() expected = torch.ones(2, 3) pipeline.text_encoder.encode_ids.return_value = expected @@ -1328,7 +1388,7 @@ def test_layerwise_offload_releases_text_encoder(offload_flag): pipeline.text_encoder.offload_to_cpu.assert_called_once_with() -def test_distributed_layerwise_offload_releases_text_encoder(): +def test_layerwise_dit_only_keeps_text_encoder_resident(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline pipeline = object.__new__(MiniMaxH3Pipeline) @@ -1336,7 +1396,8 @@ def test_distributed_layerwise_offload_releases_text_encoder(): pipeline.od_config = SimpleNamespace( enable_cpu_offload=False, enable_layerwise_offload=False, - enable_distributed_layerwise_offload=True, + enable_distributed_layerwise_offload=False, + diffusion_offload_config={"mode": "layer", "components": ["dit"]}, ) pipeline.text_encoder = Mock() expected = torch.ones(2, 3) @@ -1346,25 +1407,30 @@ def test_distributed_layerwise_offload_releases_text_encoder(): assert actual is expected pipeline.text_encoder.load_to_device.assert_called_once_with() - pipeline.text_encoder.offload_to_cpu.assert_called_once_with() + pipeline.text_encoder.offload_to_cpu.assert_not_called() -def test_distributed_layerwise_offload_stages_vae_component(): +def test_dit_encoder_selection_keeps_vae_resident(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline pipeline = object.__new__(MiniMaxH3Pipeline) torch.nn.Module.__init__(pipeline) pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, enable_layerwise_offload=False, - enable_distributed_layerwise_offload=True, + enable_distributed_layerwise_offload=False, + diffusion_offload_config={ + "mode": "layer", + "components": ["dit", "text_encoder"], + }, ) component = Mock() with pipeline._component_on_device(component): - component.load_to_device.assert_called_once_with() - component.offload_to_cpu.assert_not_called() + pass - component.offload_to_cpu.assert_called_once_with() + component.load_to_device.assert_not_called() + component.offload_to_cpu.assert_not_called() def test_distributed_layerwise_resident_blocks_are_stage_scoped(): @@ -1413,57 +1479,6 @@ def test_distributed_layerwise_resident_blocks_can_be_skipped(): controller.offload_resident_layers.assert_not_called() -def test_encoder_layerwise_offload_keeps_tp_blocks_rank_local(monkeypatch): - from vllm_omni.diffusion.models.minimax_h3.encoder import ( - MiniMaxH3Qwen3VLEncoder, - ) - - class Stack(torch.nn.Module): - def __init__(self, count): - super().__init__() - self.blocks = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)]) - - class TextStack(torch.nn.Module): - def __init__(self, count): - super().__init__() - self.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)]) - - hooks = [] - - def fake_apply(block, next_block, device, stream, pin_memory): - hook = SimpleNamespace( - block=block, - next_block=next_block, - device=device, - stream=stream, - pin_memory=pin_memory, - _prev_hook=None, - offload_layer=Mock(), - ) - hooks.append(hook) - return hook - - monkeypatch.setattr( - "vllm_omni.diffusion.offloader.layerwise_backend.apply_block_hook", - fake_apply, - ) - monkeypatch.setattr( - "vllm_omni.platforms.current_omni_platform.Stream", - Mock(return_value="copy-stream"), - ) - encoder = object.__new__(MiniMaxH3Qwen3VLEncoder) - torch.nn.Module.__init__(encoder) - encoder.device_target = torch.device("cpu") - encoder.vision = Stack(2) - encoder.text_model = TextStack(3) - - encoder.enable_omni_layerwise_offload(pin_memory=False) - - assert len(hooks) == 5 - assert all(hook._prev_hook is not None for hook in hooks) - assert all(hook.device == torch.device("cpu") for hook in hooks) - - def test_video_vae_keeps_reference_fp32_weights(monkeypatch): from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py index b0ebf1038..faa6b9a8d 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +from types import SimpleNamespace + import pytest import torch @@ -11,6 +13,44 @@ class _OffloadAbort(BaseException): pass +def test_no_allgather_example_selects_distributed_backend(): + from examples.offline_inference.minimax_h3.dlo_lifecycle import engine_kwargs + from vllm_omni.diffusion.offloader.base import OffloadConfig + from vllm_omni.diffusion.offloader.config import OffloadStrategy + + args = SimpleNamespace( + model="/fake/MiniMax-H3/FL2VA", + mode="dlo-no-allgather", + dp_size=2, + tp_size=2, + batch_wait_ms=500.0, + init_timeout=1800.0, + ) + kwargs = engine_kwargs(args) + od_config = SimpleNamespace( + diffusion_offload_config=None, + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=kwargs["enable_distributed_layerwise_offload"], + dlo_use_allgather=kwargs["dlo_use_allgather"], + dlo_resident_layers=kwargs["dlo_resident_layers"], + dlo_host_registration_limit_gib=0.0, + host_weight_runtime_mode="disabled", + pin_cpu_memory=True, + parallel_config=SimpleNamespace( + data_parallel_size=kwargs["data_parallel_size"], + sequence_parallel_size=1, + use_hsdp=False, + ), + ) + + config = OffloadConfig.from_od_config(od_config) + + assert config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE + assert config.dlo_resident_layers == 0 + assert not config.uses_allgather("dit") + + def test_encoder_non_block_children_use_one_shared_snapshot_stager(monkeypatch, mocker): from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module @@ -61,6 +101,7 @@ def test_manual_component_failure_forces_retained_cache_release(mocker): pipeline = object.__new__(MiniMaxH3Pipeline) torch.nn.Module.__init__(pipeline) pipeline.od_config = mocker.Mock() + pipeline.od_config.diffusion_offload_config = None pipeline.od_config.enable_layerwise_offload = False pipeline.od_config.enable_distributed_layerwise_offload = True pipeline._model_cpu_offload_modules = [] @@ -83,6 +124,7 @@ def test_manual_component_offload_failure_forces_retained_cache_release(mocker): pipeline = object.__new__(MiniMaxH3Pipeline) torch.nn.Module.__init__(pipeline) pipeline.od_config = mocker.Mock() + pipeline.od_config.diffusion_offload_config = None pipeline.od_config.enable_layerwise_offload = False pipeline.od_config.enable_distributed_layerwise_offload = True pipeline._model_cpu_offload_modules = [] diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py index 5f98a88d5..c493ac58f 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py @@ -247,6 +247,55 @@ def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch): assert captured["unsupported_offload_mode"] is None +@pytest.mark.parametrize( + ("mode", "components", "expected_mode"), + [ + ("module", ["text_encoder"], None), + ("layer", ["text_encoder"], None), + ("module", ["dit"], "model-level CPU offload"), + ("layer", ["dit"], "layerwise offload"), + ], +) +def test_h3_turbo_offload_compatibility_follows_dit_selection( + monkeypatch, + mode, + components, + expected_mode, +): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline._turbo_lora_adapter_ids = set() + pipeline._native_lora_adapter_ids = set() + pipeline._lora_sigma_schedules = {} + pipeline.od_config = SimpleNamespace( + diffusion_offload_config={"mode": mode, "components": components}, + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, + dlo_use_allgather=True, + dlo_resident_layers=0, + pin_cpu_memory=True, + ) + captured = {} + + def load_turbo(**kwargs): + captured.update(kwargs) + return object(), object() + + monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", load_turbo) + + assert pipeline._load_diffusion_lora_adapter( + lora_request=_request("turbo"), + lora_path="turbo", + dtype=torch.bfloat16, + ) + assert captured["unsupported_offload_mode"] == expected_mode + + def test_h3_turbo_accepts_only_the_declared_v1_artifact(tmp_path): unsupported = tmp_path / "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors" _write_tiny_turbo(unsupported) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py index c993de492..dade4437c 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from contextlib import contextmanager from types import SimpleNamespace +from typing import Any from unittest.mock import Mock import pytest @@ -111,6 +112,66 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch): remove_offload.assert_called_once_with([*dits, *stages]) +def test_h3_model_cpu_offload_keeps_unselected_vaes_resident(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.transformer = torch.nn.Linear(2, 2) + pipeline.transformers_ref = torch.nn.Linear(2, 2) + pipeline.text_encoder = torch.nn.Linear(2, 2) + pipeline.video_vae = torch.nn.Linear(2, 2) + pipeline.audio_vae = torch.nn.Linear(2, 2) + apply_offload = Mock() + monkeypatch.setattr(module, "apply_sequential_offload", apply_offload) + + pipeline.enable_omni_model_cpu_offload( + device=torch.device("cpu"), + pin_memory=False, + use_hsdp=False, + offload_components=frozenset({"dit", "text_encoder"}), + ) + + dits = [pipeline.transformer, pipeline.transformers_ref] + stages = [pipeline.text_encoder, pipeline.video_vae, pipeline.audio_vae] + apply_offload.assert_called_once_with( + dit_modules=dits, + encoder_modules=stages, + device=torch.device("cpu"), + pin_memory=False, + use_hsdp=False, + offload_initial_dits=True, + offload_dit_modules=dits, + offload_encoder_modules=[pipeline.text_encoder], + ) + + +def test_h3_model_cpu_offload_rejects_unloaded_selected_encoder(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.transformer = torch.nn.Linear(2, 2) + pipeline.transformers_ref = torch.nn.Linear(2, 2) + pipeline.text_encoder = None + pipeline.video_vae = torch.nn.Linear(2, 2) + pipeline.audio_vae = torch.nn.Linear(2, 2) + apply_offload = Mock() + monkeypatch.setattr(module, "apply_sequential_offload", apply_offload) + + with pytest.raises(ValueError, match="no loaded text encoder"): + pipeline.enable_omni_model_cpu_offload( + device=torch.device("cpu"), + pin_memory=False, + use_hsdp=False, + offload_components=frozenset({"text_encoder"}), + ) + + apply_offload.assert_not_called() + + @pytest.mark.parametrize("decode_fails", [False, True]) def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline @@ -119,7 +180,7 @@ def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails): pipeline = object.__new__(MiniMaxH3Pipeline) torch.nn.Module.__init__(pipeline) component = torch.nn.Linear(2, 2) - events = [] + events: list[tuple[Any, ...]] = [] @contextmanager def record_component(value): @@ -262,7 +323,7 @@ def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeyp torch.nn.Module.__init__(pipeline) pipeline.audio_vae = Mock() pipeline._model_cpu_offload_modules = [pipeline.audio_vae] - events = [] + events: list[tuple[Any, ...]] = [] @contextmanager def record_component(value): diff --git a/tests/diffusion/models/sana_video/test_cache_offload.py b/tests/diffusion/models/sana_video/test_cache_offload.py index 9ec779f11..52ced37c8 100644 --- a/tests/diffusion/models/sana_video/test_cache_offload.py +++ b/tests/diffusion/models/sana_video/test_cache_offload.py @@ -16,7 +16,7 @@ from tests.diffusion.models.sana_video.test_transformer_sana_video import ( # n _TINY_CONFIG, _init_distributed, ) -from tests.diffusion.offloader.test_layerwise_backend import DummyEvent, DummyStream, dummy_stream +from tests.diffusion.offloader.helpers import DummyEvent, DummyStream, dummy_stream from vllm_omni.diffusion.attention import selector as attention_selector from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend @@ -226,7 +226,7 @@ def test_component_loading_uses_loader_device_not_runtime_device(monkeypatch): monkeypatch.setattr( pipeline_module, "_load_json", - lambda _model, filename, _local: ({"vae": [None, "FakeVAE"]} if filename == "model_index.json" else {}), + lambda _model, filename, _local: {"vae": [None, "FakeVAE"]} if filename == "model_index.json" else {}, ) monkeypatch.setattr( pipeline_module, diff --git a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py index cae40b220..368852a1d 100644 --- a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py +++ b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py @@ -173,9 +173,14 @@ def test_length_lives_in_device_tensors(): # --------------------------------------------------------------------------- cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +vllm_flash_attn_only = pytest.mark.skipif( + current_omni_platform.is_rocm(), + reason="vLLM's bundled FlashAttention extension is CUDA-only", +) @cuda_only +@vllm_flash_attn_only @pytest.mark.cuda @pytest.mark.L4 def test_paged_attention_matches_sdpa_over_the_same_buffer(): @@ -217,6 +222,7 @@ def test_paged_attention_matches_sdpa_over_the_same_buffer(): @cuda_only +@vllm_flash_attn_only @pytest.mark.cuda @pytest.mark.L4 def test_a_captured_graph_follows_a_later_length(): @@ -278,6 +284,7 @@ def test_a_captured_graph_follows_a_later_length(): @cuda_only +@vllm_flash_attn_only @pytest.mark.cuda @pytest.mark.L4 def test_a_captured_graph_writes_each_step_to_its_own_slot(): @@ -339,6 +346,7 @@ def test_a_captured_graph_writes_each_step_to_its_own_slot(): @cuda_only +@vllm_flash_attn_only @pytest.mark.cuda @pytest.mark.L4 def test_a_whole_decode_loop_matches_the_unpaged_path(): diff --git a/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py index 12c0ebf45..26e78ddab 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py +++ b/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py @@ -9,6 +9,7 @@ from PIL import Image from torch import nn from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import _WAN_TEXT_ENCODER_OFFLOAD_PLAN from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import ( Wan22I2VPipeline, get_wan22_i2v_post_process_func, @@ -30,6 +31,10 @@ def test_wan22_i2v_postprocess_honors_request_output_type() -> None: assert output is video +def test_i2v_pipeline_declares_text_encoder_offload_blocks() -> None: + assert Wan22I2VPipeline._offload_plan is _WAN_TEXT_ENCODER_OFFLOAD_PLAN + + def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline: pipeline = object.__new__(Wan22I2VPipeline) nn.Module.__init__(pipeline) diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py index 5a37e14f2..6f564f256 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py +++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py @@ -18,6 +18,15 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] +def test_t2v_pipeline_declares_text_encoder_offload_blocks() -> None: + plan = wan22_module.Wan22Pipeline._offload_plan + + assert plan.encoder_component_types == {"text_encoder": "text_encoder"} + assert plan.encoder_block_attrs == {"text_encoder": ("encoder.block",)} + assert plan.encoder_dlo_weight_replication == frozenset({"text_encoder"}) + assert plan.on_demand_component_paths == frozenset() + + class _LatentDist: def sample(self, generator): assert isinstance(generator, torch.Generator) diff --git a/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py index e5bae65df..3b49d7edb 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py +++ b/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py @@ -392,52 +392,31 @@ def test_encode_audio_skips_unshard_reshard_when_not_fsdp(): assert "audio_emb" in result -def test_s2v_pipeline_skips_cpu_offload_when_hsdp_enabled(): - """Test that transformer.to('cpu') is NOT called when HSDP is active.""" - from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import Wan22S2VPipeline - +@pytest.mark.parametrize( + ("components", "use_hsdp", "expected"), + [ + (None, False, True), + (["dit"], False, True), + (["text_encoder"], False, False), + (["dit"], True, False), + ], +) +def test_s2v_dit_release_honors_component_selection(components, use_hsdp, expected): pipeline = object.__new__(Wan22S2VPipeline) nn.Module.__init__(pipeline) + compact = None if components is None else {"mode": "module", "components": components} + pipeline.od_config = SimpleNamespace( + diffusion_offload_config=compact, + enable_cpu_offload=True, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, + dlo_use_allgather=True, + dlo_resident_layers=0, + pin_cpu_memory=True, + parallel_config=SimpleNamespace(use_hsdp=use_hsdp), + ) - od_config = MagicMock() - od_config.enable_cpu_offload = True - parallel_config = MagicMock() - parallel_config.use_hsdp = True - od_config.parallel_config = parallel_config - pipeline.od_config = od_config - - mock_transformer = MagicMock() - pipeline.transformer = mock_transformer - - # Simulate the offload decision from the forward loop - if pipeline.od_config.enable_cpu_offload and not getattr(pipeline.od_config.parallel_config, "use_hsdp", False): - pipeline.transformer.to("cpu") - - mock_transformer.to.assert_not_called() - - -def test_s2v_pipeline_allows_cpu_offload_when_hsdp_disabled(): - """Test that transformer.to('cpu') IS called when HSDP is not active.""" - from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import Wan22S2VPipeline - - pipeline = object.__new__(Wan22S2VPipeline) - nn.Module.__init__(pipeline) - - od_config = MagicMock() - od_config.enable_cpu_offload = True - parallel_config = MagicMock() - parallel_config.use_hsdp = False - od_config.parallel_config = parallel_config - pipeline.od_config = od_config - - mock_transformer = MagicMock() - pipeline.transformer = mock_transformer - - # Simulate the offload decision from the forward loop - if pipeline.od_config.enable_cpu_offload and not getattr(pipeline.od_config.parallel_config, "use_hsdp", False): - pipeline.transformer.to("cpu") - - mock_transformer.to.assert_called_once_with("cpu") + assert pipeline._should_release_dit_before_decode() is expected def test_s2v_pipeline_hsdp_forward_complete_process(): diff --git a/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py index 6534b7cff..5f712bc8a 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py +++ b/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import dataclass, field from types import SimpleNamespace import pytest @@ -20,6 +21,46 @@ from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] +def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeypatch) -> None: + @dataclass + class RuntimeConfig: + flow_shift: float | None = None + diffusion_offload_config: dict[str, object] = field( + default_factory=lambda: { + "mode": "layer", + "components": ["text_encoder"], + } + ) + enable_layerwise_offload: bool = True + post_init_calls: int = 0 + + def __post_init__(self) -> None: + self.post_init_calls += 1 + + config = RuntimeConfig() + resolved_offload = object() + config._resolved_diffusion_offload = resolved_offload + config._diffusion_offload_flags_materialized = True + captured = {} + + def capture_base_init(self, *, od_config, prefix="") -> None: + captured.update(od_config=od_config, prefix=prefix) + + monkeypatch.setattr(Wan22VACEPipeline.__mro__[1], "__init__", capture_base_init) + + Wan22VACEPipeline(od_config=config, prefix="stage") + + forwarded = captured["od_config"] + assert forwarded is not config + assert config.flow_shift is None + assert forwarded.flow_shift == 3.0 + assert forwarded.diffusion_offload_config is config.diffusion_offload_config + assert forwarded._resolved_diffusion_offload is resolved_offload + assert forwarded._diffusion_offload_flags_materialized is True + assert forwarded.post_init_calls == 1 + assert captured["prefix"] == "stage" + + def _make_vace_pipeline() -> Wan22VACEPipeline: pipeline = object.__new__(Wan22VACEPipeline) nn.Module.__init__(pipeline) diff --git a/tests/diffusion/offloader/conftest.py b/tests/diffusion/offloader/conftest.py new file mode 100644 index 000000000..6a55b7284 --- /dev/null +++ b/tests/diffusion/offloader/conftest.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Shared runtime fixtures for diffusion offloader tests.""" + +import gc + +import pytest +import torch.distributed as dist + +from tests.helpers.runtime import get_distributed_init_method +from vllm_omni.platforms import current_omni_platform + + +@pytest.fixture(scope="module") +def dist_group(): + dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method()) + try: + yield + finally: + if dist.is_initialized(): + dist.destroy_process_group() + gc.collect() + if current_omni_platform.is_available(): + current_omni_platform.empty_cache() + current_omni_platform.synchronize() diff --git a/tests/diffusion/offloader/helpers.py b/tests/diffusion/offloader/helpers.py new file mode 100644 index 000000000..5f65b450e --- /dev/null +++ b/tests/diffusion/offloader/helpers.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Shared model fixtures for diffusion offloader tests.""" + +from contextlib import contextmanager + +import torch +from torch import nn + + +class DummyStream: + def wait_stream(self, _stream) -> None: + return None + + def wait_event(self, _event) -> None: + return None + + +class DummyEvent: + def record(self, _stream) -> None: + return None + + def synchronize(self) -> None: + return None + + +@contextmanager +def dummy_stream(_stream): + yield None + + +def patch_offload_runtime(monkeypatch, platform, *, synchronize: bool = False) -> None: + monkeypatch.setattr(platform, "Stream", DummyStream) + monkeypatch.setattr(platform, "Event", DummyEvent) + monkeypatch.setattr(platform, "current_stream", lambda: DummyStream()) + monkeypatch.setattr(platform, "stream", dummy_stream) + if synchronize: + monkeypatch.setattr(platform, "synchronize", lambda: None) + monkeypatch.setattr(platform, "empty_cache", lambda: None) + + +class _DummyBlock(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.randn(10, 10)) + + +class _SingleBlockModel(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self, num_blocks: int = 3): + super().__init__() + self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)]) + + +class _StagedEncoder(nn.Module): + def __init__(self): + super().__init__() + self.vision = nn.Module() + self.vision.blocks = nn.ModuleList([_DummyBlock(), _DummyBlock()]) + self.text_model = nn.Module() + self.text_model.layers = nn.ModuleList([_DummyBlock(), _DummyBlock()]) + self.load_calls = 0 + self.offload_calls = 0 + self.to_calls = 0 + + def load_to_device(self): + self.load_calls += 1 + + def offload_to_cpu(self): + self.offload_calls += 1 + for hook in getattr(self, "_omni_layerwise_hooks", []): + hook.offload_layer() + + def to(self, *args, **kwargs): + self.to_calls += 1 + return super().to(*args, **kwargs) + + +class _StagedVAE(nn.Module): + def __init__(self): + super().__init__() + self.proj = nn.Linear(2, 2) + self.offload_calls = 0 + self.to_calls = 0 + + def load_to_device(self): + return None + + def offload_to_cpu(self): + self.offload_calls += 1 + return self.to("cpu") + + def to(self, *args, **kwargs): + self.to_calls += 1 + return super().to(*args, **kwargs) + + +class _PlainEncoder(nn.Module): + """Standard encoder with no offload-specific lifecycle methods.""" + + def __init__(self): + super().__init__() + self.encoder = nn.Module() + self.encoder.block = nn.ModuleList([_DummyBlock(), _DummyBlock()]) + self.final_norm = nn.Linear(2, 2) diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index 934fe3299..05d7d93b5 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -6,8 +6,8 @@ import gc import json import weakref -from contextlib import contextmanager from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch @@ -17,7 +17,15 @@ from torch import nn from torch.distributed.tensor import DeviceMesh, DTensor, Replicate import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module -from tests.helpers.runtime import get_distributed_init_method +from tests.diffusion.offloader.helpers import ( + DummyStream, + _DummyBlock, + _PlainEncoder, + _SingleBlockModel, + _StagedEncoder, + _StagedVAE, + patch_offload_runtime, +) from vllm_omni.diffusion.data import validate_dlo_host_registration_options from vllm_omni.diffusion.model_loader.host_weight_plan import ( HostWeightPlan, @@ -51,53 +59,9 @@ from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] -class DummyStream: - def wait_stream(self, _stream) -> None: - return None - - def wait_event(self, _event) -> None: - return None - - -class DummyEvent: - def record(self, _stream) -> None: - return None - - def synchronize(self) -> None: - return None - - -@contextmanager -def dummy_stream(_stream): - yield None - - -def _cleanup_distributed() -> None: - if dist.is_initialized(): - dist.destroy_process_group() - - gc.collect() - if current_omni_platform.is_available(): - current_omni_platform.empty_cache() - current_omni_platform.synchronize() - - -@pytest.fixture(scope="module") -def dist_group(): - dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method()) - try: - yield - finally: - _cleanup_distributed() - - @pytest.fixture def patched_offload_runtime(monkeypatch): - monkeypatch.setattr(dist_backend_module.current_omni_platform, "Stream", DummyStream) - monkeypatch.setattr(dist_backend_module.current_omni_platform, "Event", DummyEvent) - monkeypatch.setattr(dist_backend_module.current_omni_platform, "current_stream", lambda: DummyStream()) - monkeypatch.setattr(dist_backend_module.current_omni_platform, "stream", dummy_stream) - monkeypatch.setattr(dist_backend_module.current_omni_platform, "synchronize", lambda: None) + patch_offload_runtime(monkeypatch, dist_backend_module.current_omni_platform, synchronize=True) class TinyBlock(nn.Module): @@ -113,6 +77,93 @@ def _make_values(start: float) -> torch.Tensor: class TestDistributedLayerwiseOffloadHook: + def test_buffer_only_block_reports_offloaded_state(self, patched_offload_runtime): + current_block = nn.Module() + current_block.register_buffer("state", torch.ones(2)) + next_block = nn.Module() + next_block.register_buffer("state", torch.ones(2)) + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + pin_memory=False, + ) + hook.initialize_hook(current_block) + + assert hook.is_materialized + hook.offload_layer() + assert not hook.is_materialized + + def test_initialize_failure_keeps_next_block_materialized(self, monkeypatch): + current_block = nn.Linear(2, 2) + next_block = nn.Linear(2, 2) + expected = {name: tensor.detach().clone() for name, tensor in next_block.state_dict().items()} + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + pin_memory=False, + ) + + def fail_buffer_allocation(): + raise RuntimeError("injected buffer allocation failure") + + monkeypatch.setattr(hook, "_allocate_device_buffers", fail_buffer_allocation) + + with pytest.raises(RuntimeError, match="buffer allocation failure"): + hook.initialize_hook(current_block) + + for name, tensor in next_block.state_dict().items(): + torch.testing.assert_close(tensor, expected[name]) + + @pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"dp_size": 0, "rank": 0}, "dp_size must be a positive integer"), + ({"dp_size": 2, "rank": 2, "dp_group": object()}, "rank must satisfy"), + ({"dp_size": 2, "rank": 0}, "dp_group is required"), + ( + {"dp_size": 2, "rank": 0, "dp_group": object(), "rank_local_mmap": True}, + "rank_local_mmap requires dp_size=1", + ), + ], + ) + def test_constructor_rejects_invalid_transport_topology(self, kwargs, message): + options = {"dp_group": None, "dp_size": 1, "rank": 0, **kwargs} + + with pytest.raises(ValueError, match=message): + DistributedLayerwiseOffloadHook( + next_block=nn.Linear(2, 2), + device=torch.device("cpu"), + pin_memory=False, + **options, + ) + + def test_ring_probes_are_captured_before_any_block_is_cleared(self, patched_offload_runtime): + blocks = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(3)]) + expected_middle = blocks[1].weight.detach().clone() + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dlo_use_allgather=False, + ), + torch.device("cpu"), + ) + + hooks = backend._install_hook_group(blocks, "dit") + shared_buffers = backend._allocate_shared_buffers(hooks) + for hook in hooks: + hook.gpu_buffers = shared_buffers + + assert all(not hook.is_materialized for hook in hooks) + hooks[2].pre_forward(blocks[1]) + torch.testing.assert_close(blocks[1].weight, expected_middle) + def test_shard_and_pin_single_rank(self, dist_group, patched_offload_runtime): """With dp_size=1, the shard should equal the full weights.""" current_block = TinyBlock(_make_values(1.0)) @@ -511,7 +562,7 @@ class TestDistributedLayerwiseOffloadHook: hook = DistributedLayerwiseOffloadHook( next_block=next_block, device=torch.device("cpu"), - dp_group=None, + dp_group=object(), dp_size=4, rank=1, pin_memory=False, @@ -537,7 +588,7 @@ class TestDistributedLayerwiseOffloadHook: hook0 = DistributedLayerwiseOffloadHook( next_block=next_block, device=torch.device("cpu"), - dp_group=None, + dp_group=object(), dp_size=2, rank=0, pin_memory=False, @@ -559,7 +610,7 @@ class TestDistributedLayerwiseOffloadHook: hook1 = DistributedLayerwiseOffloadHook( next_block=next_block1, device=torch.device("cpu"), - dp_group=None, + dp_group=object(), dp_size=2, rank=1, pin_memory=False, @@ -583,7 +634,7 @@ class TestDistributedLayerwiseOffloadHook: hook = DistributedLayerwiseOffloadHook( next_block=next_block, device=torch.device("cpu"), - dp_group=None, + dp_group=object(), dp_size=2, rank=0, pin_memory=False, @@ -602,7 +653,7 @@ class TestDistributedLayerwiseOffloadHook: hook = DistributedLayerwiseOffloadHook( next_block=next_block, device=torch.device("cpu"), - dp_group=None, + dp_group=object(), dp_size=2, rank=0, pin_memory=False, @@ -689,6 +740,32 @@ class TestDistributedLayerwiseOffloadHook: class TestPinnedResidentLayerGroup: + def test_constructor_failure_leaves_every_block_materialized(self, monkeypatch): + blocks = [nn.Linear(2, 2), nn.Linear(2, 2)] + expected = [block.weight.detach().clone() for block in blocks] + shard_and_pin = DistributedLayerwiseOffloadHook._shard_and_pin + calls = 0 + + def fail_second_block(*args, **kwargs): + nonlocal calls + calls += 1 + if calls == 2: + raise RuntimeError("injected resident setup failure") + return shard_and_pin(*args, **kwargs) + + monkeypatch.setattr(DistributedLayerwiseOffloadHook, "_shard_and_pin", fail_second_block) + + with pytest.raises(RuntimeError, match="resident setup failure"): + PinnedResidentLayerGroup( + blocks, + device=torch.device("cpu"), + copy_stream=DummyStream(), + pin_memory=False, + ) + + for block, weight in zip(blocks, expected, strict=True): + torch.testing.assert_close(block.weight, weight) + def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime): blocks = [nn.Linear(2, 2), nn.Linear(2, 2)] expected = [] @@ -843,20 +920,6 @@ class TestPinnedModuleStager: stager.offload() -class _DummyBlock(nn.Module): - def __init__(self): - super().__init__() - self.weight = nn.Parameter(torch.randn(10, 10)) - - -class _SingleBlockModel(nn.Module): - _layerwise_offload_blocks_attrs = ["blocks"] - - def __init__(self, num_blocks: int = 3): - super().__init__() - self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)]) - - class _HWRPipeline(nn.Module): def __init__(self): super().__init__() @@ -1285,6 +1348,31 @@ class TestMmapWeightLoading: backend.disable() assert lease.closed + backend.enable(pipeline) + assert backend.host_weight_plan is None + assert backend.enabled + backend.disable() + + def test_rank_local_disable_drains_pending_transfers(self, patched_offload_runtime, monkeypatch): + pipeline = _HWRPipeline() + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dp_size=1, + dlo_use_allgather=False, + ), + torch.device("cpu"), + ) + backend.enable(pipeline) + backend._all_hook_groups[0][0].prefetch_layer(0) + synchronize = Mock() + monkeypatch.setattr(current_omni_platform, "synchronize", synchronize) + + backend.disable() + + synchronize.assert_called() + def test_hwr_registration_failure_falls_back_to_bounded_staging( self, monkeypatch, @@ -1563,6 +1651,7 @@ class TestOffloadPlan: assert plan.block_attrs == {} assert plan.offload_submodules == {} assert plan.resident_dit_paths == frozenset() + assert plan.encoder_dlo_weight_replication == frozenset() def test_offload_plan_is_frozen(self): """OffloadPlan should be immutable (frozen=True).""" @@ -1595,6 +1684,7 @@ class TestOffloadPlan: self.transformer = Transformer() pipeline = Pipeline() + expected_parameters = {name: parameter.detach().clone() for name, parameter in pipeline.named_parameters()} backend = DistributedLayerwiseOffloadBackend( OffloadConfig( strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, @@ -1610,10 +1700,22 @@ class TestOffloadPlan: assert len(backend._resident_blocks) == 2 assert len(backend._all_hook_groups) == 1 assert len(backend._all_hook_groups[0]) == 2 + assert pipeline._dlo_residency_controller is backend assert backend.enabled backend.disable() + assert pipeline._dlo_residency_controller is None + + for name, parameter in pipeline.named_parameters(): + torch.testing.assert_close(parameter, expected_parameters[name]) + + backend.enable(pipeline) + assert len(backend._resident_blocks) == 2 + backend.disable() + for name, parameter in pipeline.named_parameters(): + torch.testing.assert_close(parameter, expected_parameters[name]) + class TestMmapValidation: """Tests for loader preflight and backend plan realization.""" @@ -2024,8 +2126,8 @@ class TestConfigValidation: (2, 1, 1, True, 2), (1, 4, 1, True, 4), (2, 2, 2, True, 2), - (2, 2, 2, False, 1), - (1, 4, 2, False, 1), + (2, 2, 2, False, 2), + (1, 4, 2, False, 4), ], ) def test_dlo_group_selection_covers_dp_sp_tp_matrix( @@ -2199,7 +2301,7 @@ class TestConfigValidation: model = "/fake/path" config = OffloadConfig.from_od_config(FakeODConfig()) - assert config.dp_size == 1 # forced to 1 when no AllGather + assert config.dp_size == 2 # topology is retained; the hook selects group size 1 assert config.dlo_resident_layers == 20 def test_resident_layers_with_allgather_rejected(self): @@ -2218,96 +2320,9 @@ class TestConfigValidation: parallel_config = FakePC() model = "/fake/path" - with pytest.raises(ValueError, match="requires --dlo-no-use-allgather"): + with pytest.raises(ValueError, match="requires the DiT DLO transfer to be rank-local"): OffloadConfig.from_od_config(FakeODConfig()) - def test_num_inference_steps_none_rejected(self): - """DP multi-concurrency should reject None num_inference_steps.""" - from types import SimpleNamespace - - # Mock requests with None steps - reqs = [ - SimpleNamespace( - req=SimpleNamespace( - request_id=f"req-{i}", - sampling_params=SimpleNamespace(num_inference_steps=None), - ) - ) - for i in range(2) - ] - - # We can't easily instantiate the full executor, but we can test - # the validation logic by checking that the code path raises. - # The validation is in execute_request, which needs self._ensure_open(). - # Instead, test the validation logic directly: - step_counts = { - r.req.sampling_params.num_inference_steps - for r in reqs - if r.req.sampling_params.num_inference_steps is not None - } - has_none = any(r.req.sampling_params.num_inference_steps is None for r in reqs) - assert has_none, "Test setup: should have None steps" - assert len(step_counts) == 0, "Test setup: no explicit steps" - - # The validation condition: (len(step_counts) > 1) or has_none → should reject - should_reject = (len(step_counts) > 1) or has_none - assert should_reject, "None steps should trigger rejection" - - def test_num_inference_steps_same_explicit_allowed(self): - """DP multi-concurrency should allow same explicit num_inference_steps.""" - from types import SimpleNamespace - - reqs = [ - SimpleNamespace( - req=SimpleNamespace( - request_id=f"req-{i}", - sampling_params=SimpleNamespace(num_inference_steps=35), - ) - ) - for i in range(4) - ] - - step_counts = { - r.req.sampling_params.num_inference_steps - for r in reqs - if r.req.sampling_params.num_inference_steps is not None - } - has_none = any(r.req.sampling_params.num_inference_steps is None for r in reqs) - - should_reject = (len(step_counts) > 1) or has_none - assert not should_reject, "Same explicit steps should be allowed" - assert step_counts == {35} - - def test_num_inference_steps_different_explicit_rejected(self): - """DP multi-concurrency should reject different explicit steps.""" - from types import SimpleNamespace - - reqs = [ - SimpleNamespace( - req=SimpleNamespace( - request_id="req-0", - sampling_params=SimpleNamespace(num_inference_steps=35), - ) - ), - SimpleNamespace( - req=SimpleNamespace( - request_id="req-1", - sampling_params=SimpleNamespace(num_inference_steps=30), - ) - ), - ] - - step_counts = { - r.req.sampling_params.num_inference_steps - for r in reqs - if r.req.sampling_params.num_inference_steps is not None - } - has_none = any(r.req.sampling_params.num_inference_steps is None for r in reqs) - - should_reject = (len(step_counts) > 1) or has_none - assert should_reject, "Different steps should trigger rejection" - assert len(step_counts) == 2 - class TestDynamicSlotTracking: """Tests for dynamic slot tracking with odd and even block counts. @@ -2472,3 +2487,344 @@ class TestDynamicSlotTracking: # pre_forward should override to slot 1 hook_b.pre_forward(block_b) assert hook_b.current_slot == 1, "pre_forward should read _prev_hook._prefetched_slot=1, not keep initial 0" + + +class _DistributedComponentPipeline(nn.Module): + _offload_plan = OffloadPlan( + encoder_component_types={"text_encoder": "text_encoder"}, + encoder_block_attrs={"text_encoder": ("vision.blocks", "text_model.layers")}, + on_demand_component_paths=frozenset({"text_encoder", "vae"}), + ) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + self.text_encoder = _StagedEncoder() + self.vae = _StagedVAE() + + +class _GenericDistributedEncoderPipeline(nn.Module): + _offload_plan = OffloadPlan( + encoder_component_types={"text_encoder": "text_encoder"}, + encoder_block_attrs={"text_encoder": ("encoder.block",)}, + encoder_dlo_weight_replication=frozenset({"text_encoder"}), + ) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + self.text_encoder = _PlainEncoder() + + +class TestDistributedComponentSelection: + def test_on_demand_only_encoder_cannot_claim_allgather(self, patched_offload_runtime): + class Pipeline(nn.Module): + _offload_plan = OffloadPlan( + encoder_component_types={"text_encoder": "text_encoder"}, + on_demand_component_paths=frozenset({"text_encoder"}), + encoder_dlo_weight_replication=frozenset({"text_encoder"}), + ) + + def __init__(self): + super().__init__() + self.text_encoder = _StagedEncoder() + + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dp_size=2, + components=frozenset({"text_encoder"}), + dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"}, + ), + torch.device("cpu"), + ) + backend.dp_group = object() + + with pytest.raises(ValueError, match="cannot use AllGather without.*streamable block plan"): + backend.enable(Pipeline()) + + def test_explicit_single_block_dit_is_rejected(self, patched_offload_runtime): + pipeline = nn.Module() + pipeline.transformer = _SingleBlockModel(num_blocks=1) + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + components=frozenset({"dit"}), + ), + torch.device("cpu"), + ) + + with pytest.raises(ValueError, match="leaves only one streaming block"): + backend.enable(pipeline) + + def test_enable_defers_collectives_until_first_forward(self, patched_offload_runtime, monkeypatch): + pipeline = _GenericDistributedEncoderPipeline() + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dp_size=2, + components=frozenset({"dit", "text_encoder"}), + dlo_transfers={"dit": "allgather", "text_encoder": "rank-local"}, + ), + torch.device("cpu"), + ) + backend.dp_group = object() + + def reject_startup_collective(*_args, **_kwargs): + raise AssertionError("enable() entered a rank-asymmetric collective") + + monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", reject_startup_collective) + + backend.enable(pipeline) + + assert all(hook._prefetched_slot is None for group in backend._all_hook_groups for hook in group) + encoder_hook = pipeline.text_encoder.encoder.block[0]._hook_registry.get_hook("distributed_layerwise_offload") + dit_hook = pipeline.transformer.blocks[0]._hook_registry.get_hook("distributed_layerwise_offload") + assert encoder_hook.gpu_shard_buffers == [None, None] + assert all(buffer is not None for buffer in dit_hook.gpu_shard_buffers) + + monkeypatch.setattr( + torch.distributed, + "all_gather_into_tensor", + lambda output, local, group: output.copy_(local.repeat(2)), + ) + backend.disable() + + def test_partial_dit_enable_failure_restores_weights_and_hooks( + self, + patched_offload_runtime, + monkeypatch, + ): + pipeline = _DistributedComponentPipeline() + expected_weights = [block.weight.detach().clone() for block in pipeline.transformer.blocks] + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dlo_use_allgather=False, + components=frozenset({"dit"}), + ), + torch.device("cpu"), + ) + original_apply = dist_backend_module.apply_distributed_block_hook + calls = 0 + + def fail_second_hook(*args, **kwargs): + nonlocal calls + calls += 1 + if calls == 2: + raise RuntimeError("injected DLO hook failure") + return original_apply(*args, **kwargs) + + monkeypatch.setattr(dist_backend_module, "apply_distributed_block_hook", fail_second_hook) + + with pytest.raises(RuntimeError, match="injected DLO hook failure"): + backend.enable(pipeline) + + assert not backend.enabled + for block, expected in zip(pipeline.transformer.blocks, expected_weights, strict=True): + registry = getattr(block, "_hook_registry", None) + assert registry is None or registry.get_hook("distributed_layerwise_offload") is None + torch.testing.assert_close(block.weight, expected) + + def test_multirank_enable_failure_cleanup_skips_restore_collective(self, monkeypatch, mocker): + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dp_size=2, + dlo_use_allgather=True, + components=frozenset({"dit"}), + ), + torch.device("cpu"), + ) + monkeypatch.setattr(backend, "_enable", mocker.Mock(side_effect=RuntimeError("injected startup failure"))) + cleanup = mocker.Mock() + monkeypatch.setattr(backend, "_disable", cleanup) + + with pytest.raises(RuntimeError, match="injected startup failure"): + backend.enable(_DistributedComponentPipeline()) + + cleanup.assert_called_once_with(restore_allgather_weights=False) + + def test_multirank_failure_restores_only_rank_local_hooks(self, patched_offload_runtime): + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dp_size=2, + components=frozenset({"dit", "text_encoder"}), + dlo_transfers={"dit": "allgather", "text_encoder": "rank-local"}, + ), + torch.device("cpu"), + ) + allgather_hook = Mock(dp_size=2, next_block=nn.Linear(1, 1)) + rank_local_hook = Mock(dp_size=1, next_block=nn.Linear(1, 1)) + backend._all_hook_groups = [[allgather_hook], [rank_local_hook]] + + backend._disable(restore_allgather_weights=False) + + allgather_hook.restore_next_block_to_cpu.assert_not_called() + rank_local_hook.restore_next_block_to_cpu.assert_called_once_with() + + with pytest.raises(RuntimeError, match="recreate the backend and reload the pipeline"): + backend.enable(nn.Module()) + + def test_disable_runs_collectives_before_best_effort_local_cleanup( + self, + patched_offload_runtime, + monkeypatch, + ): + events: list[str] = [] + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dp_size=2, + ), + torch.device("cpu"), + ) + rank_local_hook = Mock(dp_size=1, next_block=nn.Linear(1, 1)) + allgather_hook = Mock(dp_size=2, next_block=nn.Linear(1, 1)) + allgather_hook.restore_next_block_to_cpu.side_effect = lambda: events.append("allgather") + + def fail_rank_local_restore(): + events.append("rank-local") + raise RuntimeError("injected rank-local restore failure") + + rank_local_hook.restore_next_block_to_cpu.side_effect = fail_rank_local_restore + backend._all_hook_groups = [[rank_local_hook], [allgather_hook]] + backend._blocks = [[nn.Module()]] + backend.enabled = True + monkeypatch.setattr( + dist_backend_module, + "remove_distributed_block_hook", + lambda _block: events.append("remove"), + ) + monkeypatch.setattr(backend, "_release_registered_mmap", lambda: events.append("unregister")) + release_handles = Mock() + monkeypatch.setattr(backend, "_release_mmap_handles", release_handles) + + with pytest.raises(RuntimeError, match="injected rank-local restore failure"): + backend.disable() + + assert events == ["allgather", "rank-local", "remove", "unregister"] + release_handles.assert_not_called() + assert backend.enabled + assert backend._all_hook_groups + + def test_all_streams_encoder_and_keeps_vae_resident(self, patched_offload_runtime, monkeypatch): + pipeline = _DistributedComponentPipeline() + move_non_block_state = Mock() + monkeypatch.setattr(dist_backend_module, "move_non_block_state_to_device", move_non_block_state) + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dlo_use_allgather=False, + components=frozenset({"dit", "text_encoder"}), + ), + torch.device("cpu"), + ) + + backend.enable(pipeline) + + assert pipeline.text_encoder._omni_layerwise_enabled + assert len(pipeline.text_encoder._omni_layerwise_hooks) == 4 + # The pipeline-owned stage lifecycle releases the encoder's non-block + # state after DLO has installed blockwise hooks. + assert pipeline.text_encoder.offload_calls == 1 + assert pipeline.vae.offload_calls == 0 + assert pipeline.vae.to_calls == 1 + move_non_block_state.assert_not_called() + + backend.disable() + + assert not pipeline.text_encoder._omni_layerwise_enabled + encoder_blocks = [ + *pipeline.text_encoder.vision.blocks, + *pipeline.text_encoder.text_model.layers, + ] + assert all(block._hook_registry.get_hook("distributed_layerwise_offload") is None for block in encoder_blocks) + + def test_dit_and_encoder_choose_transfer_independently(self, patched_offload_runtime, monkeypatch): + pipeline = _GenericDistributedEncoderPipeline() + expected_parameters = {name: parameter.detach().clone() for name, parameter in pipeline.named_parameters()} + full_weights_by_rank_zero_shard = { + tuple(parameter.flatten()[: (parameter.numel() + 1) // 2].tolist()): parameter.flatten() + for parameter in expected_parameters.values() + if parameter.numel() > 4 + } + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dp_size=2, + components=frozenset({"dit", "text_encoder"}), + dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"}, + ), + torch.device("cpu"), + ) + backend.dp_group = object() + + def fake_allgather(output, local, group): + del group + output.copy_(full_weights_by_rank_zero_shard[tuple(local.tolist())]) + + monkeypatch.setattr( + torch.distributed, + "all_gather_into_tensor", + fake_allgather, + ) + + backend.enable(pipeline) + + encoder_hook = pipeline.text_encoder.encoder.block[0]._hook_registry.get_hook("distributed_layerwise_offload") + dit_hook = pipeline.transformer.blocks[0]._hook_registry.get_hook("distributed_layerwise_offload") + assert encoder_hook.dp_size == 2 + assert dit_hook.dp_size == 1 + assert all(buffer is not None for buffer in encoder_hook.gpu_shard_buffers) + assert dit_hook.gpu_shard_buffers == [None, None] + + backend.disable() + actual_parameters = dict(pipeline.named_parameters()) + for name, expected in expected_parameters.items(): + torch.testing.assert_close(actual_parameters[name], expected) + + def test_encoder_allgather_rejects_undeclared_replication(self, patched_offload_runtime): + pipeline = _DistributedComponentPipeline() + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dp_size=2, + components=frozenset({"text_encoder"}), + dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"}, + ), + torch.device("cpu"), + ) + backend.dp_group = object() + + with pytest.raises(ValueError, match="not declared replicated"): + backend.enable(pipeline) + + def test_encoder_allgather_rejects_stub_rank_before_block_discovery(self): + """Every rank must reject an unsafe encoder group, including stub ranks.""" + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dp_size=2, + components=frozenset({"text_encoder"}), + dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"}, + ), + torch.device("cpu"), + ) + backend.dp_group = object() + stub_plan = OffloadPlan(encoder_block_attrs={"text_encoder": ("missing.blocks",)}) + + with pytest.raises(ValueError, match="not declared replicated"): + backend._try_layerwise_offload_encoder(nn.Module(), "text_encoder", stub_plan) diff --git a/tests/diffusion/offloader/test_layerwise_backend.py b/tests/diffusion/offloader/test_layerwise_backend.py index 9253f6d03..2dd8804cf 100644 --- a/tests/diffusion/offloader/test_layerwise_backend.py +++ b/tests/diffusion/offloader/test_layerwise_backend.py @@ -3,66 +3,39 @@ """Unit tests for LayerwiseOffloadHook and LayerWiseOffloadBackend utilities.""" -import gc -from contextlib import contextmanager +from multiprocessing.reduction import ForkingPickler +from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch -import torch.distributed as dist from torch import nn from torch.distributed.tensor import DeviceMesh, DTensor, Replicate import vllm_omni.diffusion.offloader.layerwise_backend as layerwise_backend_module -from tests.helpers.runtime import get_distributed_init_method -from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend, LayerwiseOffloadHook -from vllm_omni.platforms import current_omni_platform +import vllm_omni.diffusion.offloader.tensor_utils as tensor_utils_module +from tests.diffusion.offloader.helpers import ( + DummyStream, + _DummyBlock, + _PlainEncoder, + _SingleBlockModel, + _StagedEncoder, + _StagedVAE, + patch_offload_runtime, +) +from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy +from vllm_omni.diffusion.offloader.layerwise_backend import ( + LayerWiseOffloadBackend, + LayerwiseOffloadHook, +) +from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] -class DummyStream: - def wait_stream(self, _stream) -> None: - return None - - def wait_event(self, _event) -> None: - return None - - -class DummyEvent: - def record(self, _stream) -> None: - return None - - -@contextmanager -def dummy_stream(_stream): - yield None - - -def _cleanup_distributed() -> None: - if dist.is_initialized(): - dist.destroy_process_group() - - gc.collect() - if current_omni_platform.is_available(): - current_omni_platform.empty_cache() - current_omni_platform.synchronize() - - -@pytest.fixture(scope="module") -def dist_group(): - dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method()) - try: - yield - finally: - _cleanup_distributed() - - @pytest.fixture def patched_offload_runtime(monkeypatch): - monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Stream", DummyStream) - monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Event", DummyEvent) - monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream()) - monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "stream", dummy_stream) + patch_offload_runtime(monkeypatch, layerwise_backend_module.current_omni_platform) class TinyBlock(nn.Module): @@ -77,7 +50,66 @@ def _make_values(start: float) -> torch.Tensor: return torch.arange(start, start + 4, dtype=torch.float32) +def test_clear_tensor_storage_rolls_back_partial_commit(monkeypatch): + first = nn.Parameter(torch.tensor([1.0, 2.0])) + second = nn.Parameter(torch.tensor([3.0, 4.0])) + expected_first = first.detach().clone() + expected_second = second.detach().clone() + second_data_ptr = second.data_ptr() + original_set_storage = tensor_utils_module.set_tensor_storage + calls = 0 + + def fail_second_commit(target, value): + nonlocal calls + calls += 1 + if calls == 2: + raise RuntimeError("injected storage clear failure") + original_set_storage(target, value) + + monkeypatch.setattr(tensor_utils_module, "set_tensor_storage", fail_second_commit) + + with pytest.raises(RuntimeError, match="injected storage clear failure"): + tensor_utils_module.clear_tensor_storage([first, second]) + + assert calls == 3 # first commit, failed second commit, first rollback + torch.testing.assert_close(first, expected_first) + torch.testing.assert_close(second, expected_second) + assert second.data_ptr() == second_data_ptr + + class TestLayerwiseOffloadHook: + def test_buffer_only_block_reports_offloaded_state(self, patched_offload_runtime): + current_block = nn.Module() + current_block.register_buffer("state", torch.ones(2)) + next_block = nn.Module() + next_block.register_buffer("state", torch.ones(2)) + hook = LayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + stream=DummyStream(), + pin_memory=False, + ) + hook.initialize_hook(current_block) + + assert hook.is_materialized + hook.offload_layer() + assert not hook.is_materialized + + def test_ring_probes_are_captured_before_any_block_is_cleared(self, patched_offload_runtime): + blocks = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(3)]) + expected_middle = blocks[1].weight.detach().clone() + + hooks = layerwise_backend_module._install_layerwise_hook_group( + blocks, + torch.device("cpu"), + DummyStream(), + pin_memory=False, + ) + + assert all(not hook.is_materialized for hook in hooks) + hooks[2].pre_forward(blocks[1]) + torch.testing.assert_close(blocks[1].weight, expected_middle) + def test_dtensor_wrapper_is_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime): current_block = TinyBlock(_make_values(1.0)) next_block = TinyBlock(_make_values(10.0)) @@ -135,20 +167,6 @@ class TestLayerwiseOffloadHook: assert torch.equal(next_block.weight, expected) -class _DummyBlock(nn.Module): - def __init__(self): - super().__init__() - self.weight = nn.Parameter(torch.randn(10, 10)) - - -class _SingleBlockModel(nn.Module): - _layerwise_offload_blocks_attrs = ["blocks"] - - def __init__(self, num_blocks: int = 3): - super().__init__() - self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)]) - - class _MultiBlockModel(nn.Module): _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"] @@ -246,3 +264,732 @@ class TestGetBlocksAttrNames: LayerWiseOffloadBackend.set_blocks_attr_names(model, ["new_blocks"]) assert hasattr(model.__class__, "_layerwise_offload_blocks_attrs") assert model.__class__._layerwise_offload_blocks_attrs == ["new_blocks"] + + +class _ComponentPipeline(nn.Module): + _offload_plan = OffloadPlan( + encoder_component_types={"text_encoder": "text_encoder"}, + encoder_block_attrs={"text_encoder": ("vision.blocks", "text_model.layers")}, + on_demand_component_paths=frozenset({"text_encoder", "vae"}), + ) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel() + self.text_encoder = _StagedEncoder() + self.vae = _StagedVAE() + + +class _LegacyComponentPipeline(nn.Module): + """Pipeline with DiT block metadata but no auxiliary OffloadPlan.""" + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel() + self.text_encoder = _StagedEncoder() + self.vae = _StagedVAE() + + +class _LegacyEncoderOnlyPipeline(nn.Module): + _offload_plan = _ComponentPipeline._offload_plan + + def __init__(self): + super().__init__() + self.text_encoder = _StagedEncoder() + self.vae = _StagedVAE() + + +class _PlannedDitPipeline(nn.Module): + _offload_plan = OffloadPlan(block_attrs={"transformer": ("blocks",)}) + + def __init__(self): + super().__init__() + self.transformer = _NoAttrsModel() + + +class _GenericEncoderPipeline(nn.Module): + _offload_plan = OffloadPlan( + encoder_component_types={"text_encoder": "text_encoder"}, + encoder_block_attrs={"text_encoder": ("encoder.block",)}, + ) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel() + self.text_encoder = _PlainEncoder() + + +def _layer_backend(components: set[str] | None = None) -> LayerWiseOffloadBackend: + options = {} + if components is not None: + options = {"components": frozenset(components)} + return LayerWiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.LAYER_WISE, + pin_cpu_memory=False, + **options, + ), + torch.device("cpu"), + ) + + +class TestLayerwiseComponentSelection: + def test_plan_block_attrs_drive_dit_discovery(self, patched_offload_runtime): + pipeline = _PlannedDitPipeline() + backend = _layer_backend({"dit"}) + + backend.enable(pipeline) + + assert hasattr(pipeline.transformer.blocks[0], "_hook_registry") + backend.disable() + + def test_legacy_missing_dit_preserves_noop(self, patched_offload_runtime): + pipeline = _LegacyEncoderOnlyPipeline() + backend = _layer_backend() + + backend.enable(pipeline) + + assert not backend.enabled + assert pipeline.text_encoder.offload_calls == 0 + assert pipeline.vae.offload_calls == 0 + + def test_encoder_only_streams_planned_blocks(self, patched_offload_runtime, monkeypatch): + pipeline = _ComponentPipeline() + backend = _layer_backend({"text_encoder"}) + move_non_block_state = Mock() + monkeypatch.setattr(layerwise_backend_module, "move_non_block_state_to_device", move_non_block_state) + + backend.enable(pipeline) + + assert pipeline.text_encoder._omni_layerwise_enabled + assert len(pipeline.text_encoder._omni_layerwise_hooks) == 4 + assert pipeline.text_encoder.offload_calls == 1 + assert pipeline.vae.offload_calls == 0 + assert pipeline.vae.to_calls == 1 + assert not hasattr(pipeline.transformer.blocks[0], "_hook_registry") + assert backend.enabled + move_non_block_state.assert_not_called() + + backend.disable() + assert not pipeline.text_encoder._omni_layerwise_enabled + + def test_single_gpu_dit_only_keeps_encoder_and_vae_resident(self, patched_offload_runtime): + pipeline = _ComponentPipeline() + expected_weights = [block.weight.detach().clone() for block in pipeline.transformer.blocks] + backend = _layer_backend({"dit"}) + + backend.enable(pipeline) + + assert not hasattr(pipeline.text_encoder, "_omni_layerwise_enabled") + assert pipeline.text_encoder.to_calls == 1 + assert pipeline.text_encoder.offload_calls == 0 + assert pipeline.vae.to_calls == 1 + assert pipeline.vae.offload_calls == 0 + assert hasattr(pipeline.transformer.blocks[0], "_hook_registry") + + backend.disable() + + for block, expected in zip(pipeline.transformer.blocks, expected_weights, strict=True): + torch.testing.assert_close(block.weight, expected) + + backend.enable(pipeline) + backend.disable() + for block, expected in zip(pipeline.transformer.blocks, expected_weights, strict=True): + torch.testing.assert_close(block.weight, expected) + + def test_partial_dit_enable_failure_restores_weights_and_hooks(self, patched_offload_runtime, monkeypatch): + pipeline = _ComponentPipeline() + expected_weights = [block.weight.detach().clone() for block in pipeline.transformer.blocks] + backend = _layer_backend({"dit"}) + original_apply = layerwise_backend_module.apply_block_hook + calls = 0 + + def fail_second_hook(*args, **kwargs): + nonlocal calls + calls += 1 + if calls == 2: + raise RuntimeError("injected hook failure") + return original_apply(*args, **kwargs) + + monkeypatch.setattr(layerwise_backend_module, "apply_block_hook", fail_second_hook) + + with pytest.raises(RuntimeError, match="injected hook failure"): + backend.enable(pipeline) + + assert not backend.enabled + for block, expected in zip(pipeline.transformer.blocks, expected_weights, strict=True): + registry = getattr(block, "_hook_registry", None) + assert registry is None or registry.get_hook("layerwise_offload") is None + torch.testing.assert_close(block.weight, expected) + + def test_later_component_failure_removes_installed_encoder_hooks( + self, + patched_offload_runtime, + monkeypatch, + ): + pipeline = _ComponentPipeline() + blocks = [*pipeline.text_encoder.vision.blocks, *pipeline.text_encoder.text_model.layers] + expected = [block.weight.detach().clone() for block in blocks] + backend = _layer_backend({"text_encoder"}) + + def fail_vae_placement(*_args, **_kwargs): + raise RuntimeError("injected VAE placement failure") + + monkeypatch.setattr(pipeline.vae, "to", fail_vae_placement) + + with pytest.raises(RuntimeError, match="injected VAE placement failure"): + backend.enable(pipeline) + + assert not backend.enabled + assert not pipeline.text_encoder._omni_layerwise_enabled + assert not backend._encoder_modules + for block, original in zip(blocks, expected, strict=True): + assert block._hook_registry.get_hook("layerwise_offload") is None + torch.testing.assert_close(block.weight, original) + + def test_encoder_only_requires_streamable_offload_plan(self, patched_offload_runtime): + pipeline = nn.Module() + pipeline.transformer = _SingleBlockModel() + pipeline.text_encoder = _StagedEncoder() + backend = _layer_backend({"text_encoder"}) + + with pytest.raises(ValueError, match="Selected text encoder 'text_encoder' requires"): + backend.enable(pipeline) + + def test_explicit_text_encoder_selection_requires_a_matching_module(self, patched_offload_runtime): + pipeline = nn.Module() + pipeline.transformer = _SingleBlockModel(num_blocks=2) + backend = _layer_backend({"dit", "text_encoder"}) + + with pytest.raises(ValueError, match="No text encoder modules found"): + backend.enable(pipeline) + + assert not hasattr(pipeline.transformer.blocks[0], "_hook_registry") + + def test_every_selected_encoder_requires_its_own_plan(self, patched_offload_runtime): + class Pipeline(nn.Module): + _offload_plan = OffloadPlan( + encoder_component_types={ + "text_encoder": "text_encoder", + "text_encoder_2": "text_encoder", + }, + encoder_block_attrs={"text_encoder": ("encoder.block",)}, + ) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel() + self.text_encoder = _PlainEncoder() + self.text_encoder_2 = _PlainEncoder() + + pipeline = Pipeline() + backend = _layer_backend({"text_encoder"}) + + with pytest.raises(ValueError, match="Selected text encoder 'text_encoder_2' requires"): + backend.enable(pipeline) + + assert not pipeline.text_encoder._omni_layerwise_enabled + + def test_default_selection_preserves_unplanned_auxiliaries(self, patched_offload_runtime): + pipeline = _LegacyComponentPipeline() + backend = _layer_backend() + + backend.enable(pipeline) + + assert hasattr(pipeline.transformer.blocks[0], "_hook_registry") + assert not hasattr(pipeline.text_encoder, "_omni_layerwise_enabled") + assert pipeline.text_encoder.to_calls == 1 + assert pipeline.text_encoder.offload_calls == 0 + assert pipeline.vae.to_calls == 1 + assert pipeline.vae.offload_calls == 0 + + backend.disable() + + def test_disable_failure_keeps_host_masters_for_retry(self, patched_offload_runtime, monkeypatch): + pipeline = nn.Module() + pipeline.transformer = _SingleBlockModel(num_blocks=3) + expected = {name: tensor.detach().clone() for name, tensor in pipeline.state_dict().items()} + backend = _layer_backend() + backend.enable(pipeline) + failing_hook = backend._dit_hooks[0] + restore = failing_hook.restore_next_block + attempts = 0 + + def fail_once(): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise RuntimeError("injected restore failure") + restore() + + monkeypatch.setattr(failing_hook, "restore_next_block", fail_once) + + with pytest.raises(RuntimeError, match="Failed to fully disable"): + backend.disable() + + assert backend.enabled + assert backend._dit_hooks + + backend.disable() + + assert not backend.enabled + for name, tensor in pipeline.state_dict().items(): + torch.testing.assert_close(tensor, expected[name]) + + def test_encoder_disable_failure_keeps_host_masters_for_retry( + self, + patched_offload_runtime, + monkeypatch, + ): + pipeline = _ComponentPipeline() + encoder = pipeline.text_encoder + expected = {name: tensor.detach().clone() for name, tensor in encoder.state_dict().items()} + backend = _layer_backend({"text_encoder"}) + backend.enable(pipeline) + failing_hook = encoder._omni_layerwise_hooks[0] + restore = failing_hook.restore_next_block + attempts = 0 + + def fail_once(): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise RuntimeError("injected encoder restore failure") + restore() + + monkeypatch.setattr(failing_hook, "restore_next_block", fail_once) + + with pytest.raises(RuntimeError, match="Failed to fully disable"): + backend.disable() + + assert backend.enabled + assert backend._encoder_modules == [encoder] + assert encoder._omni_layerwise_enabled + assert encoder._omni_layerwise_hooks + + backend.disable() + + assert not backend.enabled + assert not encoder._omni_layerwise_enabled + for name, tensor in encoder.state_dict().items(): + torch.testing.assert_close(tensor, expected[name]) + + def test_legacy_selector_preserves_planned_encoder_and_vae_lifecycle(self, patched_offload_runtime): + pipeline = _ComponentPipeline() + backend = _layer_backend() + + backend.enable(pipeline) + + assert pipeline.text_encoder._omni_layerwise_enabled + assert pipeline.text_encoder.offload_calls == 1 + assert pipeline.vae.offload_calls == 1 + + backend.disable() + + def test_standard_encoder_needs_only_declared_block_paths(self, patched_offload_runtime): + pipeline = _GenericEncoderPipeline() + backend = _layer_backend({"text_encoder"}) + + backend.enable(pipeline) + + blocks = pipeline.text_encoder.encoder.block + assert pipeline.text_encoder._omni_layerwise_enabled + assert all(block._hook_registry.get_hook("layerwise_offload") is not None for block in blocks) + assert pipeline.text_encoder.final_norm.weight.numel() == 4 + + backend.disable() + + assert not pipeline.text_encoder._omni_layerwise_enabled + assert all(block._hook_registry.get_hook("layerwise_offload") is None for block in blocks) + assert all(block.weight.numel() == 100 for block in blocks) + + +def _offload_od_config(**overrides): + values = { + "diffusion_offload_config": None, + "enable_cpu_offload": False, + "enable_layerwise_offload": False, + "enable_distributed_layerwise_offload": False, + "dlo_use_allgather": True, + "dlo_resident_layers": 0, + "pin_cpu_memory": True, + "parallel_config": None, + "model": "/fake/model", + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _resolve_offload_config(public=None, **overrides) -> OffloadConfig: + return OffloadConfig.from_od_config(_offload_od_config(diffusion_offload_config=public, **overrides)) + + +class TestLayerwiseComponentConfig: + def test_layer_mode_selects_components_without_backend_jargon(self): + config = _resolve_offload_config( + { + "mode": "layer", + "components": ["dit", "text_encoder"], + } + ) + + assert config.strategy is OffloadStrategy.LAYER_WISE + assert config.components == frozenset({"dit", "text_encoder"}) + assert not config.uses_allgather("dit") + assert not config.uses_allgather("text_encoder") + plan = OffloadPlan(encoder_component_types={"mllm": "text_encoder"}) + assert config.offloads_encoder("mllm", plan) + + @pytest.mark.parametrize("component", ["image_encoder", "vae", "scheduler", "text-encoder"]) + def test_unknown_or_noncanonical_component_is_rejected(self, component): + with pytest.raises(ValueError, match="Unknown diffusion offload component"): + OffloadConfig.from_od_config( + _offload_od_config(diffusion_offload_config={"mode": "layer", "components": [component]}) + ) + + def test_components_is_selection_only_not_an_options_mapping(self): + with pytest.raises(TypeError, match="components must be a non-empty list"): + OffloadConfig.from_od_config( + _offload_od_config( + diffusion_offload_config={ + "mode": "layer", + "components": {"dit": {}}, + } + ) + ) + + def test_layer_options_requires_component_name_keys(self): + with pytest.raises(TypeError, match="layer_options keys must be strings"): + OffloadConfig.from_od_config( + _offload_od_config( + diffusion_offload_config={ + "mode": "layer", + "components": ["dit"], + "layer_options": {0: {}}, + } + ) + ) + + @pytest.mark.parametrize( + ("config", "message"), + [ + ({0: "layer"}, "diffusion_offload_config keys must be strings"), + ( + { + "mode": "layer", + "components": ["dit"], + "layer_options": {"dit": {0: "allgather"}}, + }, + "layer_options\\['dit'\\] keys must be strings", + ), + ], + ) + def test_all_public_mapping_levels_require_string_keys(self, config, message): + with pytest.raises(TypeError, match=message): + OffloadConfig.from_od_config(_offload_od_config(diffusion_offload_config=config)) + + @pytest.mark.parametrize( + "config,match", + [ + ({"components": ["dit"]}, "requires 'mode'"), + ({"mode": "layer"}, "requires 'components'"), + ({"mode": "layer", "components": []}, "must not be empty"), + ({"mode": "layer", "components": [" dit "]}, "Unknown diffusion offload component"), + ({"mode": "layerwise", "components": ["dit"]}, "Unknown diffusion offload mode"), + ( + { + "mode": "layer", + "components": ["dit"], + "layer_options": {"dit": {"weight_transfer": "rank_local"}}, + }, + "Unknown offload transfer", + ), + ( + { + "mode": "layer", + "components": ["dit"], + "layer_options": {"dit": {"prefetch": 2}}, + }, + "Unknown diffusion offload setting", + ), + ( + { + "mode": "layer", + "components": ["text_encoder"], + "layer_options": {"text_encoder": {"resident_layers": 1}}, + }, + "supports only the 'dit'", + ), + ( + { + "mode": "layer", + "components": ["dit"], + "layer_options": {"text_encoder": {}}, + }, + "requires selecting the same component", + ), + ], + ) + def test_schema_rejects_ambiguous_or_unsupported_values(self, config, match): + with pytest.raises(ValueError, match=match): + OffloadConfig.from_od_config(_offload_od_config(diffusion_offload_config=config)) + + def test_module_mode_accepts_component_selection(self): + config = _resolve_offload_config( + { + "mode": "module", + "components": ["dit", "text_encoder"], + } + ) + + assert config.strategy is OffloadStrategy.MODEL_LEVEL + assert config.components == frozenset({"dit", "text_encoder"}) + + @pytest.mark.parametrize("setting", [{"weight_transfer": "rank-local"}, {"resident_layers": 1}]) + def test_module_mode_rejects_layer_settings(self, setting): + with pytest.raises(ValueError, match="layer_options requires mode='layer'"): + OffloadConfig.from_od_config( + _offload_od_config( + diffusion_offload_config={ + "mode": "module", + "components": ["dit"], + "layer_options": {"dit": setting}, + } + ) + ) + + def test_each_component_can_choose_its_transfer(self): + config = _resolve_offload_config( + { + "mode": "layer", + "components": ["dit", "text_encoder"], + "layer_options": { + "dit": {"weight_transfer": "rank-local"}, + "text_encoder": {"weight_transfer": "allgather"}, + }, + } + ) + + assert config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE + assert not config.uses_allgather("dit") + assert config.uses_allgather("text_encoder") + assert config.dlo_use_allgather is False + + def test_dit_residency_selects_capable_backend(self): + config = _resolve_offload_config( + { + "mode": "layer", + "components": ["dit"], + "layer_options": {"dit": {"resident_layers": 20}}, + } + ) + + assert config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE + assert config.dlo_resident_layers == 20 + assert not config.uses_allgather("dit") + + def test_dit_residency_rejects_allgather(self): + with pytest.raises(ValueError, match="resident_layers requires dit.weight_transfer='rank-local'"): + OffloadConfig.from_od_config( + _offload_od_config( + diffusion_offload_config={ + "mode": "layer", + "components": ["dit"], + "layer_options": {"dit": {"weight_transfer": "allgather", "resident_layers": 1}}, + } + ) + ) + + def test_host_weight_runtime_remains_separate_from_compact_config(self): + with pytest.raises(ValueError, match="cannot be combined with Host Weight Runtime"): + OffloadConfig.from_od_config( + _offload_od_config( + diffusion_offload_config={"mode": "layer", "components": ["dit"]}, + host_weight_runtime_mode="preferred", + ) + ) + + def test_pin_memory_override_is_scoped_to_offload_config(self): + config = _resolve_offload_config( + { + "mode": "layer", + "components": ["dit"], + "pin_memory": False, + }, + pin_cpu_memory=True, + ) + + assert config.pin_cpu_memory is False + + def test_legacy_omitted_selector_preserves_dit_only_behavior(self): + config = _resolve_offload_config(enable_layerwise_offload=True) + + assert config.strategy is OffloadStrategy.LAYER_WISE + assert config.components is None + + def test_offload_resolution_is_cached_for_runtime_accessors(self, monkeypatch): + import vllm_omni.diffusion.offloader.config as config_module + + od_config = _offload_od_config( + diffusion_offload_config={ + "mode": "layer", + "components": ["dit", "text_encoder"], + } + ) + parse_calls = 0 + original_parse = config_module.parse_diffusion_offload_config + + def track_parse(value): + nonlocal parse_calls + parse_calls += 1 + return original_parse(value) + + monkeypatch.setattr(config_module, "parse_diffusion_offload_config", track_parse) + + assert config_module.resolve_offload_strategy(od_config) is OffloadStrategy.LAYER_WISE + assert config_module.selected_offload_components(od_config) == {"dit", "text_encoder"} + assert not config_module.component_uses_allgather(od_config, "dit") + assert not config_module.any_selected_component_uses_allgather(od_config) + assert parse_calls == 1 + + def test_cached_policy_is_deeply_immutable_and_pickle_safe(self): + import vllm_omni.diffusion.offloader.config as config_module + + od_config = _offload_od_config( + diffusion_offload_config={ + "mode": "layer", + "components": ["dit"], + "layer_options": {"dit": {"weight_transfer": "rank-local"}}, + } + ) + resolved = config_module.resolve_offload(od_config) + + with pytest.raises(TypeError): + resolved.transfers["dit"] = config_module.DLOTransfer.ALLGATHER # type: ignore[index] + assert resolved.public is not None + with pytest.raises(TypeError): + resolved.public.layer_options["dit"] = config_module.LayerOffloadOptions() # type: ignore[index] + + restored = ForkingPickler.loads(ForkingPickler.dumps(resolved)) + assert not restored.uses_allgather("dit") + + def test_dp_text_encoder_allgather_rejects_rank_local_prompt_cache(self): + config = _offload_od_config( + diffusion_offload_config={ + "mode": "layer", + "components": ["text_encoder"], + "layer_options": {"text_encoder": {"weight_transfer": "allgather"}}, + }, + enable_prompt_embed_cache=True, + parallel_config=SimpleNamespace( + data_parallel_size=2, + sequence_parallel_size=1, + use_hsdp=False, + ), + ) + + with pytest.raises(ValueError, match="rank-local cache hits would skip different encoder collectives"): + OffloadConfig.from_od_config(config) + + def test_dp_allgather_rejects_rank_local_dit_cache_decisions(self): + config = _offload_od_config( + diffusion_offload_config={ + "mode": "layer", + "components": ["dit"], + "layer_options": {"dit": {"weight_transfer": "allgather"}}, + }, + cache_backend="tea_cache", + parallel_config=SimpleNamespace( + data_parallel_size=2, + sequence_parallel_size=1, + use_hsdp=False, + ), + ) + + with pytest.raises(ValueError, match="rank-local cache decisions can skip different weight collectives"): + OffloadConfig.from_od_config(config) + + def test_internal_transfer_map_requires_every_component(self): + with pytest.raises(ValueError, match="missing: text_encoder"): + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + dlo_transfers={"dit": "rank-local"}, + ) + + def test_legacy_distributed_settings_still_work(self): + config = _resolve_offload_config( + enable_distributed_layerwise_offload=True, + dlo_use_allgather=False, + dlo_resident_layers=3, + ) + + assert config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE + assert config.dlo_resident_layers == 3 + assert not config.uses_allgather("dit") + assert not config.uses_allgather("text_encoder") + + def test_legacy_allgather_remains_dit_only(self): + config = _resolve_offload_config( + enable_distributed_layerwise_offload=True, + dlo_use_allgather=True, + ) + + assert config.uses_allgather("dit") + assert not config.uses_allgather("text_encoder") + + @pytest.mark.parametrize( + ("legacy_flags", "expected_strategy"), + [ + ( + {"enable_cpu_offload": True, "enable_layerwise_offload": True}, + OffloadStrategy.LAYER_WISE, + ), + ( + { + "enable_cpu_offload": True, + "enable_layerwise_offload": True, + "enable_distributed_layerwise_offload": True, + }, + OffloadStrategy.DISTRIBUTED_LAYER_WISE, + ), + ], + ) + def test_legacy_aliases_preserve_existing_priority(self, legacy_flags, expected_strategy): + config = _resolve_offload_config(**legacy_flags) + + assert config.strategy is expected_strategy + + def test_compact_config_cannot_mix_with_legacy_alias(self): + with pytest.raises(ValueError, match="cannot be combined"): + OffloadConfig.from_od_config( + _offload_od_config( + diffusion_offload_config={"mode": "layer", "components": ["dit"]}, + enable_cpu_offload=True, + ) + ) + + @pytest.mark.parametrize( + "legacy_options", + [ + {"dlo_use_allgather": False}, + {"dlo_resident_layers": 12}, + ], + ) + def test_compact_config_cannot_mix_with_legacy_dlo_options(self, legacy_options): + with pytest.raises(ValueError, match="cannot be combined with legacy DLO option"): + OffloadConfig.from_od_config( + _offload_od_config( + diffusion_offload_config={"mode": "layer", "components": ["dit"]}, + **legacy_options, + ) + ) + + def test_legacy_allgather_resident_layers_fail_during_config_resolution(self): + with pytest.raises(ValueError, match="requires the DiT DLO transfer to be rank-local"): + OffloadConfig.from_od_config( + _offload_od_config( + enable_distributed_layerwise_offload=True, + dlo_use_allgather=True, + dlo_resident_layers=12, + ) + ) diff --git a/tests/diffusion/offloader/test_sequential_backend.py b/tests/diffusion/offloader/test_sequential_backend.py index ac3d3181d..2d3e0db3b 100644 --- a/tests/diffusion/offloader/test_sequential_backend.py +++ b/tests/diffusion/offloader/test_sequential_backend.py @@ -7,7 +7,9 @@ import pytest import torch from torch import nn +from vllm_omni.diffusion.hooks import HookRegistry from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy +from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan from vllm_omni.diffusion.offloader.sequential_backend import ( ModelLevelOffloadBackend, SequentialOffloadHook, @@ -48,20 +50,21 @@ def _track_pin_memory_calls(): return tracker, mock +class _CustomPipeline(nn.Module): + def __init__(self) -> None: + super().__init__() + self.enable_args: dict[str, object] | None = None + self.disable_called = False + + def enable_omni_model_cpu_offload(self, **kwargs) -> None: + self.enable_args = kwargs + + def disable_omni_model_cpu_offload(self) -> None: + self.disable_called = True + + def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None: - class CustomPipeline(nn.Module): - def __init__(self) -> None: - super().__init__() - self.enable_args = None - self.disable_called = False - - def enable_omni_model_cpu_offload(self, **kwargs) -> None: - self.enable_args = kwargs - - def disable_omni_model_cpu_offload(self) -> None: - self.disable_called = True - - pipeline = CustomPipeline() + pipeline = _CustomPipeline() backend = ModelLevelOffloadBackend( OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL, pin_cpu_memory=False), torch.device("cpu"), @@ -82,6 +85,205 @@ def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None: assert pipeline.disable_called is True +def test_model_level_backend_passes_explicit_component_selection() -> None: + pipeline = _CustomPipeline() + backend = ModelLevelOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.MODEL_LEVEL, + components=frozenset({"dit", "text_encoder"}), + ), + torch.device("cpu"), + ) + + backend.enable(pipeline) + + assert pipeline.enable_args is not None + assert pipeline.enable_args["offload_components"] == frozenset({"dit", "text_encoder"}) + + +def test_model_level_backend_rolls_back_partial_custom_enable() -> None: + class CustomPipeline(nn.Module): + def __init__(self) -> None: + super().__init__() + self.disable_called = False + + def enable_omni_model_cpu_offload(self, **kwargs) -> None: + del kwargs + raise RuntimeError("injected custom enable failure") + + def disable_omni_model_cpu_offload(self) -> None: + self.disable_called = True + + pipeline = CustomPipeline() + backend = ModelLevelOffloadBackend( + OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL), + torch.device("cpu"), + ) + + with pytest.raises(RuntimeError, match="injected custom enable failure"): + backend.enable(pipeline) + + assert pipeline.disable_called + assert not backend.enabled + + +def test_sequential_offload_rolls_back_partial_hook_registration(monkeypatch: pytest.MonkeyPatch) -> None: + dit = _create_simple_module() + encoder = _create_simple_module() + original_register = HookRegistry.register_hook + calls = 0 + + def fail_second_registration(self, name, hook): + nonlocal calls + calls += 1 + if calls == 2: + raise RuntimeError("injected registration failure") + return original_register(self, name, hook) + + monkeypatch.setattr(HookRegistry, "register_hook", fail_second_registration) + + with pytest.raises(RuntimeError, match="injected registration failure"): + apply_sequential_offload( + dit_modules=[dit], + encoder_modules=[encoder], + device=torch.device("cpu"), + ) + + for module in (dit, encoder): + registry = getattr(module, "_hook_registry", None) + assert registry is None or registry.get_hook(SequentialOffloadHook._HOOK_NAME) is None + + +def test_model_level_disable_cleans_remaining_hooks_and_can_retry(monkeypatch: pytest.MonkeyPatch) -> None: + pipeline = nn.Module() + pipeline.transformer = _create_simple_module() + pipeline.text_encoder = _create_simple_module() + backend = ModelLevelOffloadBackend( + OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL, pin_cpu_memory=False), + torch.device("cpu"), + ) + backend.enable(pipeline) + + original_remove = HookRegistry.remove_hook + calls = 0 + + def fail_first_removal(self, name): + nonlocal calls + calls += 1 + if calls == 1: + raise RuntimeError("injected hook removal failure") + original_remove(self, name) + + monkeypatch.setattr(HookRegistry, "remove_hook", fail_first_removal) + + with pytest.raises(RuntimeError, match="Failed to remove one or more sequential offload hooks"): + backend.disable() + + assert backend.enabled + assert backend._offload_modules == [pipeline.transformer, pipeline.text_encoder] + assert pipeline.transformer._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) is not None + assert pipeline.text_encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) is None + + backend.disable() + + assert not backend.enabled + assert not backend._offload_modules + for module in (pipeline.transformer, pipeline.text_encoder): + assert module._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) is None + + +def test_sequential_offload_filters_cpu_eligible_components() -> None: + dit = _create_simple_module() + encoder = _create_simple_module() + resident_stage = _create_simple_module() + + apply_sequential_offload( + dit_modules=[dit], + encoder_modules=[encoder, resident_stage], + device=torch.device("cpu"), + offload_dit_modules=[dit], + offload_encoder_modules=[encoder], + ) + + dit_hook = dit._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) + encoder_hook = encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) + resident_hook = resident_stage._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) + assert dit_hook.offload_targets == [encoder] + assert encoder_hook.offload_targets == [dit] + assert encoder_hook.offload_after_context is True + assert resident_hook.offload_targets == [dit] + assert resident_hook.offload_after_context is False + + remove_sequential_offload([dit, encoder, resident_stage]) + + +def test_move_params_handles_buffer_only_modules() -> None: + module = nn.Module() + module.register_buffer("state", torch.ones(2)) + + # An indexed CPU device compares differently while remaining usable on + # CPU-only test hosts, so the helper must inspect the buffer itself. + moved = SequentialOffloadHook._move_params(module, torch.device("cpu:1")) + + assert moved + assert module.state.device.type == "cpu" + + +def test_model_level_backend_keeps_declared_image_encoder_resident() -> None: + class MixedEncoderPipeline(nn.Module): + _offload_plan = OffloadPlan( + encoder_component_types={ + "text_encoder": "text_encoder", + } + ) + + def __init__(self) -> None: + super().__init__() + self.transformer = _create_simple_module() + self.text_encoder = _create_simple_module() + self.image_encoder = _create_simple_module() + + pipeline = MixedEncoderPipeline() + backend = ModelLevelOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.MODEL_LEVEL, + components=frozenset({"dit", "text_encoder"}), + ), + torch.device("cpu"), + ) + + backend.enable(pipeline) + + text_hook = pipeline.text_encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) + image_hook = pipeline.image_encoder._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) + assert text_hook.offload_after_context is True + assert image_hook.offload_after_context is False + + backend.disable() + + +@pytest.mark.parametrize( + ("component", "attribute", "message"), + [ + ("text_encoder", "text_encoder", "requires a DiT/transformer"), + ("dit", "transformer", "requires an encoder execution stage"), + ], +) +def test_component_selective_model_offload_requires_swap_counterpart(component, attribute, message) -> None: + pipeline = nn.Module() + setattr(pipeline, attribute, _create_simple_module()) + backend = ModelLevelOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.MODEL_LEVEL, + components=frozenset({component}), + ), + torch.device("cpu"), + ) + + with pytest.raises(ValueError, match=message): + backend.enable(pipeline) + + def test_sequential_offload_can_begin_with_dit_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None: dit = _create_simple_module() encoder = _create_simple_module() diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index 04d528af4..1ee7424e9 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/tests/diffusion/test_diffusion_engine.py @@ -600,6 +600,34 @@ class TestDiffusionCompileConfig: ), ({"enable_cpu_offload": True}, "CPU offload"), ({"enable_layerwise_offload": True}, "layerwise offload"), + ( + { + "diffusion_offload_config": { + "mode": "module", + "components": ["dit"], + } + }, + "CPU offload", + ), + ( + { + "diffusion_offload_config": { + "mode": "layer", + "components": ["dit"], + } + }, + "layerwise offload", + ), + ( + { + "diffusion_offload_config": { + "mode": "layer", + "components": ["dit"], + "layer_options": {"dit": {"weight_transfer": "allgather"}}, + } + }, + "distributed layerwise offload", + ), ], ) def test_full_compile_rejects_incompatible_features(self, kwargs, feature) -> None: @@ -610,6 +638,34 @@ class TestDiffusionCompileConfig: **kwargs, ) + def test_raw_parallel_mapping_is_normalized_before_allgather_cache_validation(self) -> None: + with pytest.raises(ValueError, match="rank-local cache decisions"): + OmniDiffusionConfig( + model="test", + num_gpus=2, + parallel_config={"data_parallel_size": 2}, + cache_backend="tea_cache", + diffusion_offload_config={ + "mode": "layer", + "components": ["dit"], + "layer_options": {"dit": {"weight_transfer": "allgather"}}, + }, + ) + + def test_auto_dp_is_resolved_before_encoder_prompt_cache_validation(self) -> None: + with pytest.raises(ValueError, match="rank-local cache hits"): + OmniDiffusionConfig( + model="test", + num_gpus=2, + parallel_config={"data_parallel_size": None}, + enable_prompt_embed_cache=True, + diffusion_offload_config={ + "mode": "layer", + "components": ["text_encoder"], + "layer_options": {"text_encoder": {"weight_transfer": "allgather"}}, + }, + ) + class TestRequestBatchAdmission: pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] diff --git a/tests/diffusion/test_diffusion_engine_dummy_run.py b/tests/diffusion/test_diffusion_engine_dummy_run.py index ea37b2fcc..8d524c572 100644 --- a/tests/diffusion/test_diffusion_engine_dummy_run.py +++ b/tests/diffusion/test_diffusion_engine_dummy_run.py @@ -52,6 +52,23 @@ def test_dummy_run_uses_enough_steps_for_execution_mode( assert captured_requests[0].sampling_params.num_inference_steps == 2 +def test_allgather_startup_runs_broadcast_dummy_request() -> None: + engine = object.__new__(DiffusionEngine) + engine.od_config = SimpleNamespace( + diffusion_offload_config={ + "mode": "layer", + "components": ["dit"], + "layer_options": {"dit": {"weight_transfer": "allgather"}}, + }, + parallel_config=SimpleNamespace(data_parallel_size=2, sequence_parallel_size=1), + ) + engine._dummy_run = Mock() + + engine.run_startup_warmup() + + engine._dummy_run.assert_called_once_with() + + def test_dummy_run_num_frames_uses_explicit_model_setting(monkeypatch: pytest.MonkeyPatch) -> None: class JointAudioVideoModel: dummy_run_num_frames = 2 diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index fee348a3a..9f8d4dec6 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -484,20 +484,23 @@ def test_compile_transformer_uses_full_granularity(monkeypatch): @pytest.mark.core_model @pytest.mark.cpu -def test_compile_transformer_falls_back_after_synchronous_setup_failure(monkeypatch, caplog): +def test_compile_transformer_falls_back_after_synchronous_setup_failure(monkeypatch): model = _CompileTrackingModel() runner = _make_compile_runner(model) + warning = Mock() def _regionally_compile(*args, **kwargs): raise RuntimeError("compile setup failed") monkeypatch.setattr(model_runner_module, "regionally_compile", _regionally_compile) + monkeypatch.setattr(model_runner_module.logger, "warning", warning) DiffusionModelRunner._compile_transformer(runner, "transformer") assert runner.pipeline.transformer is model - assert "failed before activation" in caplog.text - assert "lazy compilation errors" in caplog.text + warning.assert_called_once() + assert "failed before activation" in warning.call_args.args[0] + assert "lazy compilation errors" in warning.call_args.args[0] @pytest.mark.core_model @@ -1061,6 +1064,51 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): assert dummy_cache_backend.enabled is False +@pytest.mark.core_model +@pytest.mark.cpu +def test_load_model_rejects_env_prompt_cache_before_text_encoder_allgather(monkeypatch): + runner = object.__new__(DiffusionModelRunner) + runner.vllm_config = object() + runner.device = torch.device("cpu") + runner.pipeline = None + runner.offload_backend = None + runner.od_config = SimpleNamespace( + diffusion_offload_config={ + "mode": "layer", + "components": ["text_encoder"], + "layer_options": {"text_encoder": {"weight_transfer": "allgather"}}, + }, + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=True, + dlo_use_allgather=True, + dlo_resident_layers=0, + pin_cpu_memory=True, + parallel_config=SimpleNamespace(data_parallel_size=2), + enable_prompt_embed_cache=False, + prompt_embed_cache_size=32, + cache_backend="none", + model_class_name="Wan22Pipeline", + enforce_eager=True, + streaming_output=False, + step_execution=False, + ) + monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE", "1") + runtime_init = Mock(side_effect=pytest.fail) + loader = Mock(side_effect=pytest.fail) + monkeypatch.setattr(model_runner_module.current_omni_platform, "init_diffusion_model_runner_runtime", runtime_init) + monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", loader) + enable_offload = Mock(side_effect=pytest.fail) + monkeypatch.setattr(model_runner_module, "enable_offload_backend", enable_offload) + + with pytest.raises(ValueError, match="Prompt embedding cache cannot be combined"): + DiffusionModelRunner.load_model(runner) + + runtime_init.assert_not_called() + loader.assert_not_called() + enable_offload.assert_not_called() + + def test_set_forward_context_enters_vllm_config_contexts(monkeypatch): """Ensure `with set_forward_context(...):` enters vllm's context managers internally and calls desired vllm functions.""" import vllm.config.vllm as vllm_config_module diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index 9bd7a3e8f..efe6bf9a3 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -254,8 +254,11 @@ class TestRequestModeDispatch: executor.od_config = SimpleNamespace( step_execution=False, parallel_config=SimpleNamespace(data_parallel_size=2), - enable_distributed_layerwise_offload=True, - dlo_use_allgather=True, + diffusion_offload_config={ + "mode": "layer", + "components": ["dit"], + "layer_options": {"dit": {"weight_transfer": "allgather"}}, + }, ) executor.execute_request = Mock(return_value="dlo-dp") executor.collective_rpc = Mock() @@ -358,7 +361,7 @@ class TestRequestModeDispatch: for rank, result_mq in enumerate(executor._result_mqs): assert result_mq.dequeue.call_count == (1 if rank in expected_primary_ranks else 0) - @pytest.mark.parametrize("empty_prompt", ["", {"prompt": ""}]) + @pytest.mark.parametrize("empty_prompt", ["", [], {}, {"prompt": ""}]) def test_dlo_dp_rejects_empty_prompt_before_worker_dispatch(self, empty_prompt): executor, _, _ = _make_executor(num_gpus=2) executor.od_config = SimpleNamespace( @@ -376,6 +379,71 @@ class TestRequestModeDispatch: executor.collective_rpc.assert_not_called() + def test_dlo_dp_accepts_precomputed_embeddings_without_prompt_text(self): + executor, _, _ = _make_executor(num_gpus=2) + executor.od_config = SimpleNamespace( + step_execution=False, + parallel_config=SimpleNamespace(data_parallel_size=2), + diffusion_offload_config={ + "mode": "layer", + "components": ["text_encoder"], + "layer_options": {"text_encoder": {"weight_transfer": "allgather"}}, + }, + ) + executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")]) + scheduler_output = _make_sched_output("A", "B") + for new_req in scheduler_output.scheduled_new_reqs: + new_req.req.prompt = {"prompt": "", "prompt_embeds": object()} + + result = executor.execute_request(scheduler_output) + + assert [output.result.error for output in result.runner_outputs] == ["A", "B"] + + @pytest.mark.parametrize( + "prompt", + [ + [1, 2], + {"type": "token", "prompt_token_ids": [1, 2]}, + {"prompt_ids": [1, 2], "prompt_mask": object()}, + ], + ) + def test_dlo_dp_accepts_tokenized_prompts(self, prompt): + executor, _, _ = _make_executor(num_gpus=2) + executor.od_config = SimpleNamespace( + step_execution=False, + parallel_config=SimpleNamespace(data_parallel_size=2), + enable_distributed_layerwise_offload=True, + dlo_use_allgather=True, + ) + executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")]) + scheduler_output = _make_sched_output("A", "B") + for new_req in scheduler_output.scheduled_new_reqs: + new_req.req.prompt = prompt + + result = executor.execute_request(scheduler_output) + + assert [output.result.error for output in result.runner_outputs] == ["A", "B"] + + def test_text_encoder_allgather_rejects_mismatched_embedding_paths(self): + executor, _, _ = _make_executor(num_gpus=2) + executor.od_config = SimpleNamespace( + step_execution=False, + parallel_config=SimpleNamespace(data_parallel_size=2), + diffusion_offload_config={ + "mode": "layer", + "components": ["text_encoder"], + "layer_options": {"text_encoder": {"weight_transfer": "allgather"}}, + }, + ) + executor.collective_rpc = Mock() + scheduler_output = _make_sched_output("A", "B") + scheduler_output.scheduled_new_reqs[1].req.prompt = {"prompt_embeds": object()} + + with pytest.raises(ValueError, match="same positive/negative prompt embedding fields"): + executor.execute_request(scheduler_output) + + executor.collective_rpc.assert_not_called() + def test_dlo_dp_valid_wave_still_runs_after_rejected_wave(self): executor, _, _ = _make_executor(num_gpus=2) executor.od_config = SimpleNamespace( @@ -396,24 +464,6 @@ class TestRequestModeDispatch: assert [output.result.error for output in result.runner_outputs] == ["A", "B"] executor.collective_rpc.assert_called_once() - def test_dlo_dp_allows_shared_default_denoise_steps(self): - executor, _, _ = _make_executor(num_gpus=2) - executor.od_config = SimpleNamespace( - step_execution=False, - parallel_config=SimpleNamespace(data_parallel_size=2), - enable_distributed_layerwise_offload=True, - dlo_use_allgather=True, - ) - executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")]) - scheduler_output = _make_sched_output("A", "B") - for new_req in scheduler_output.scheduled_new_reqs: - new_req.req.sampling_params.num_inference_steps = None - - result = executor.execute_request(scheduler_output) - - assert [output.result.error for output in result.runner_outputs] == ["A", "B"] - executor.collective_rpc.assert_called_once() - def test_dlo_dp_forwards_request_metadata_envelopes_as_one_wave(self): executor, _, _ = _make_executor(num_gpus=2) executor.od_config = SimpleNamespace( @@ -496,6 +546,28 @@ class TestRequestModeDispatch: executor._fail_closed_on_dp_wave_timeout.assert_called_once() assert isinstance(executor._fail_closed_on_dp_wave_timeout.call_args.args[0], TimeoutError) + def test_single_allgather_request_times_out_and_fails_closed(self): + from vllm_omni.diffusion.executor import multiproc_executor as executor_module + + executor, _, _ = _make_executor(num_gpus=2) + executor.od_config = SimpleNamespace( + step_execution=False, + parallel_config=SimpleNamespace(data_parallel_size=2), + diffusion_offload_config={ + "mode": "layer", + "components": ["dit"], + "layer_options": {"dit": {"weight_transfer": "allgather"}}, + }, + ) + executor.collective_rpc = Mock(side_effect=TimeoutError("timed out")) + executor._fail_closed_on_dp_wave_timeout = Mock() + + result = executor.execute_request(_make_sched_output("A")) + + assert result.runner_outputs[0].result.error == "timed out" + assert executor.collective_rpc.call_args.kwargs["timeout"] == executor_module._DLO_DP_WAVE_TIMEOUT_S + executor._fail_closed_on_dp_wave_timeout.assert_called_once() + # ───────────────── concurrent collective RPC ───────────────── diff --git a/tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py b/tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py index 148520ff0..0bf546683 100644 --- a/tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py +++ b/tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py @@ -51,7 +51,18 @@ SERVER_ARGS = [ "--vae-parallel-mode", "tile", "--vae-use-tiling", - "--enable-distributed-layerwise-offload", + "--diffusion-offload-config", + json.dumps( + { + "mode": "layer", + "components": ["dit", "text_encoder"], + "layer_options": { + "dit": {"weight_transfer": "allgather"}, + "text_encoder": {"weight_transfer": "rank-local"}, + }, + }, + separators=(",", ":"), + ), ] diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 71a6e00a5..616d5486e 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -6,6 +6,7 @@ from types import SimpleNamespace import pytest from vllm_omni.diffusion.data import AttentionConfig +from vllm_omni.engine import stage_init_utils from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.entrypoints.cli.serve import OmniServeCommand from vllm_omni.utils.tracking_parser import TrackingArgumentParser @@ -426,8 +427,8 @@ def test_serve_cli_forwards_distilled_lora_to_diffusion_stage(): ] -def test_serve_cli_forwards_distributed_offload_residency(): - """Ensure the two-GPU DLO placement controls reach the diffusion stage.""" +def test_serve_cli_forwards_compact_diffusion_offload_config(): + """Ensure component-specific layer settings reach the diffusion stage.""" parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="command") OmniServeCommand().subparser_init(subparsers) @@ -437,10 +438,11 @@ def test_serve_cli_forwards_distributed_offload_residency(): "serve", "MiniMaxAI/MiniMax-H3", "--omni", - "--enable-distributed-layerwise-offload", - "--dlo-no-use-allgather", - "--dlo-resident-layers", - "20", + "--diffusion-offload-config", + '{"mode":"layer","components":["dit","text_encoder"],' + '"layer_options":{"dit":{"weight_transfer":"rank-local","resident_layers":20},' + '"text_encoder":{"weight_transfer":"allgather"}},' + '"pin_memory":true}', ] ) @@ -448,12 +450,45 @@ def test_serve_cli_forwards_distributed_offload_residency(): stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] engine_args = stage_cfg["engine_args"] - assert args.enable_distributed_layerwise_offload is True - assert args.dlo_use_allgather is False - assert args.dlo_resident_layers == 20 - assert engine_args["enable_distributed_layerwise_offload"] is True - assert engine_args["dlo_use_allgather"] is False - assert engine_args["dlo_resident_layers"] == 20 + expected = { + "mode": "layer", + "components": ["dit", "text_encoder"], + "layer_options": { + "dit": {"weight_transfer": "rank-local", "resident_layers": 20}, + "text_encoder": {"weight_transfer": "allgather"}, + }, + "pin_memory": True, + } + assert args.diffusion_offload_config == expected + assert engine_args["diffusion_offload_config"] == expected + + +def test_invalid_diffusion_offload_config_fails_before_model_loading(monkeypatch, mocker): + load_model = mocker.patch("vllm_omni.diffusion.model_loader.diffusers_loader.DiffusersPipelineLoader.load_model") + create_client = mocker.patch("vllm_omni.diffusion.stage_diffusion_client.create_diffusion_client") + monkeypatch.setattr( + stage_init_utils, + "build_engine_args_dict", + lambda *_args, **_kwargs: { + "model": "test", + "diffusion_offload_config": { + "mode": "layerwise", + "components": ["dit"], + }, + }, + ) + + with pytest.raises(ValueError, match="Unknown diffusion offload mode"): + stage_init_utils.initialize_diffusion_stage( + stage_id=0, + model="test", + stage_cfg=object(), + metadata=mocker.Mock(), + stage_init_timeout=30, + ) + + create_client.assert_not_called() + load_model.assert_not_called() def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo(): diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index dce775cd3..c7c145fea 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -328,9 +328,11 @@ def _stage_cli_overrides( ) -> dict[str, Any]: runtime_overrides = build_stage_runtime_overrides(stage_id, dict(cli_overrides)) global_stage_fields = _global_stage_cli_fields() + owned_fields = None if execution_type is None else _STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type] result: dict[str, Any] = {} for key, value in runtime_overrides.items(): - if key in global_stage_fields or f"stage_{stage_id}_{key}" in cli_overrides: + stage_specific = f"stage_{stage_id}_{key}" in cli_overrides + if stage_specific or (key in global_stage_fields and (owned_fields is None or key in owned_fields)): result[key] = _copy_value(value) # step_execution is a diffusion execution protocol, not an LLM engine @@ -340,6 +342,25 @@ def _stage_cli_overrides( return result +def _validate_global_stage_cli_ownership( + pipeline: PipelineConfig, + cli_overrides: Mapping[str, Any], +) -> None: + """Reject global stage arguments that no stage in the pipeline owns.""" + explicit_global_fields = { + key for key, value in cli_overrides.items() if value is not None and key in _global_stage_cli_fields() + } + owned_fields = { + field for stage in pipeline.stages for field in _STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE[stage.execution_type] + } + unowned_fields = explicit_global_fields - owned_fields + if unowned_fields: + names = ", ".join(sorted(unowned_fields)) + raise ValueError( + f"Pipeline {pipeline.model_type!r} has explicit engine argument(s) with no structured config owner: {names}" + ) + + def _resolve_deploy_path(deploy_config_path: str) -> Path: deploy_path = Path(deploy_config_path) if not deploy_path.exists() and deploy_path.parent == Path("."): @@ -707,6 +728,9 @@ class _DiffusionConfigProjection: lora_backend: str = "peft" max_cpu_loras: int | None = None output_type: str = "pil" + diffusion_offload_config: dict[str, Any] | None = None + # Compatibility aliases for existing callers and model-specific stage + # lifecycles that are broader than the compact dit/text_encoder selector. enable_cpu_offload: bool = False enable_layerwise_offload: bool = False enable_distributed_layerwise_offload: bool = False @@ -759,8 +783,12 @@ class _DiffusionConfigProjection: @classmethod def from_kwargs(cls, **kwargs: Any) -> _DiffusionConfigProjection: from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs + from vllm_omni.diffusion.offloader.config import parse_diffusion_offload_config normalized_kwargs = normalize_omni_diffusion_kwargs(kwargs) + # Validate before stage construction while retaining the raw mapping + # needed by dataclass/config serialization across process boundaries. + parse_diffusion_offload_config(normalized_kwargs.get("diffusion_offload_config")) valid_fields = {f.name for f in fields(cast(Any, cls))} return cls(**{k: v for k, v in normalized_kwargs.items() if k in valid_fields}) @@ -1865,6 +1893,7 @@ class VllmOmniConfig: if cli_overrides is None: cli_overrides = {} cli_overrides = normalize_pipeline_cli_overrides(pipeline_cfg, cli_overrides) + _validate_global_stage_cli_ownership(pipeline_cfg, cli_overrides) deploy, loaded_deploy_config_path = _get_deploy_config( pipeline_cfg, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index e185b6843..06840e55a 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -472,6 +472,9 @@ class StageDeployConfig: # Runtime optimizations used by diffusion loading/execution. enable_multithread_weight_load: bool | None = None num_weight_load_threads: int | None = None + diffusion_offload_config: dict[str, Any] | None = None + # Compatibility aliases for existing callers and model-specific stage + # lifecycles that are broader than the compact dit/text_encoder selector. enable_cpu_offload: bool | None = None enable_layerwise_offload: bool | None = None diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 10a2f9a19..3794652f6 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -804,7 +804,12 @@ class OmniDiffusionConfig: output_type: str = "pil" - # CPU offload parameters + # CPU offload parameters. Keep the public mapping raw so stage configs can + # serialize it across processes; __post_init__ validates it once and caches + # the internal typed resolution used at runtime. + diffusion_offload_config: dict[str, Any] | None = None + # Compatibility aliases. Some model-specific legacy stage lifecycles are + # intentionally broader than the compact dit/text_encoder selector. # When enabled, DiT and encoders swap GPU access (mutual exclusion): # - Text encoders run on GPU while DiT is on CPU # - DiT runs on GPU while encoders are on CPU @@ -1056,6 +1061,11 @@ class OmniDiffusionConfig: ) def __post_init__(self): + from vllm_omni.diffusion.offloader.config import ( + OffloadStrategy, + materialize_legacy_offload_flags, + ) + if self.diffusion_compile_granularity not in {"regional", "full"}: raise ValueError( "diffusion_compile_granularity must be 'regional' or 'full', " @@ -1128,6 +1138,9 @@ class OmniDiffusionConfig: self.num_gpus = 1 self.parallel_config.resolve_data_parallel_size(self.num_gpus) + # Resolve offload only after DP/SP normalization so cached policy + # validation observes the actual execution topology. + offload_strategy = materialize_legacy_offload_flags(self) if self.diffusion_compile_granularity == "full": incompatible_features = [] @@ -1136,10 +1149,14 @@ class OmniDiffusionConfig: incompatible_features.append("HSDP") if self.parallel_config.sequence_parallel_size > 1: incompatible_features.append("sequence parallelism") - if self.enable_cpu_offload: - incompatible_features.append("CPU offload") - if self.enable_layerwise_offload: - incompatible_features.append("layerwise offload") + if offload_strategy is not OffloadStrategy.NONE: + incompatible_features.append( + { + OffloadStrategy.MODEL_LEVEL: "CPU offload", + OffloadStrategy.LAYER_WISE: "layerwise offload", + OffloadStrategy.DISTRIBUTED_LAYER_WISE: "distributed layerwise offload", + }[offload_strategy] + ) if incompatible_features: features = ", ".join(incompatible_features) raise ValueError( diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index f0dc4b0e4..52a22c0b9 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -41,6 +41,7 @@ from vllm_omni.diffusion.io_support import ( supports_audio_output, supports_multimodal_input, ) +from vllm_omni.diffusion.offloader.config import any_selected_component_uses_allgather from vllm_omni.diffusion.output_formatter import ( format_diffusion_outputs, format_empty_diffusion_outputs, @@ -181,11 +182,7 @@ def _max_num_seqs(od_config: OmniDiffusionConfig) -> int: def _uses_dlo_dp_concurrency(od_config: OmniDiffusionConfig) -> bool: parallel_config = getattr(od_config, "parallel_config", None) dp_size = getattr(parallel_config, "data_parallel_size", 1) - return ( - dp_size > 1 - and getattr(od_config, "enable_distributed_layerwise_offload", False) - and getattr(od_config, "dlo_use_allgather", True) - ) + return dp_size > 1 and any_selected_component_uses_allgather(od_config) def _move_tensor_tree_to_cpu(value: object) -> object: @@ -1012,33 +1009,12 @@ class DiffusionEngine: raise RuntimeError(f"Could not {action} profiler: {e}") from e def run_startup_warmup(self) -> None: - dlo_use_allgather = getattr(self.od_config, "dlo_use_allgather", True) - # Skip dummy run when AllGather is used with more than 1 rank, - # because the dummy run sends only 1 request but AllGather requires - # all ranks to participate simultaneously. This covers both DP > 1 - # and SP > 1 (where dp_size is derived from sp_size in OffloadConfig). - pc = getattr(self.od_config, "parallel_config", None) - dp_size = getattr(pc, "data_parallel_size", 1) if pc else 1 - sp_size = getattr(pc, "sequence_parallel_size", 1) if pc else 1 - effective_shard_size = max(dp_size, sp_size) - skip_dummy = ( - getattr(self.od_config, "enable_distributed_layerwise_offload", False) - and dlo_use_allgather - and effective_shard_size > 1 - ) - if skip_dummy: - logger.info( - "Skipping dummy run (dist_offload with AllGather, dp_size=%d, sp_size=%d)", - dp_size, - sp_size, - ) - return try: self._dummy_run() except Exception as e: logger.error(f"Dummy run failed: {e}") self.close() - raise e + raise def _make_dummy_request( self, @@ -1050,10 +1026,13 @@ class DiffusionEngine: num_inference_steps: int = 1, ) -> OmniDiffusionRequest | None: """Build a minimal model request for startup profiling or warmup.""" - prompt: OmniTextPrompt = {"prompt": "dummy run"} + prompt = OmniTextPrompt(prompt="dummy run") + model_class_name = self.od_config.model_class_name + if model_class_name is None: + raise RuntimeError("Dummy request requires a resolved model_class_name") supports_image_input, supports_audio_input = supports_multimodal_input(self.od_config) if supports_image_input: - color_format = image_color_format(self.od_config.model_class_name) + color_format = image_color_format(model_class_name) images = [PIL.Image.new(color_format, (width, height)) for _ in range(num_image_inputs)] prompt.setdefault("multi_modal_data", {})["image"] = images[0] if len(images) == 1 else images @@ -1061,7 +1040,7 @@ class DiffusionEngine: audio_sr = 16000 prompt.setdefault("multi_modal_data", {})["audio"] = np.random.randn(audio_sr * 2).astype(np.float32) - num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input) + num_frames = get_dummy_run_num_frames(model_class_name, supports_audio_input) if num_frames <= 0: return None return OmniDiffusionRequest( @@ -1101,11 +1080,14 @@ class DiffusionEngine: is not DiffusionKVCacheMode.PAGED_SCHEDULER ): return None + model_class_name = self.od_config.model_class_name + if model_class_name is None: + raise RuntimeError("Diffusion KV profiling requires a resolved model_class_name") request = self._make_dummy_request( height=1024, width=1024, guidance_scale=5.0, - num_image_inputs=get_dummy_run_num_image_inputs(self.od_config.model_class_name), + num_image_inputs=get_dummy_run_num_image_inputs(model_class_name), ) if request is None: raise RuntimeError("paged_scheduler requires a runnable Diffusion KV memory profile request") diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index ff9253962..39e85abcc 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -23,6 +23,11 @@ from vllm.v1.executor.multiproc_executor import set_multiprocessing_worker_envs from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, AsyncDiffusionOutput, AsyncOutputKind, DiffusionOutput from vllm_omni.diffusion.executor.abstract import DiffusionExecutor from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, unpack_diffusion_output_shm +from vllm_omni.diffusion.offloader.config import ( + TEXT_ENCODER_COMPONENT, + any_selected_component_uses_allgather, + resolve_offload, +) from vllm_omni.diffusion.sched.request_scheduler import build_request_batch_sampling_params_key from vllm_omni.diffusion.utils.future_utils import try_set_exception, try_set_result from vllm_omni.diffusion.worker import WorkerProc @@ -44,13 +49,30 @@ def _is_empty_dp_prompt(prompt: object) -> bool: """Return whether a DP request has no usable text prompt.""" if prompt is None: return True - if isinstance(prompt, str): + if isinstance(prompt, (str, list, tuple)): return not prompt - if isinstance(prompt, dict) and "prompt" in prompt: - return not prompt["prompt"] + if isinstance(prompt, dict): + return ( + not prompt.get("prompt") + and not prompt.get("prompt_token_ids") + and not prompt.get("prompt_ids") + and prompt.get("prompt_embeds") is None + ) return False +def _text_encoder_input_signature(prompt: object) -> tuple[bool, bool]: + """Describe precomputed embeddings that change encoder forward counts.""" + if not isinstance(prompt, dict): + return False, False + return prompt.get("prompt_embeds") is not None, prompt.get("negative_prompt_embeds") is not None + + +def _uses_text_encoder_allgather(config: object) -> bool: + resolved = resolve_offload(config) + return resolved.offloads(TEXT_ENCODER_COMPONENT) and resolved.uses_allgather(TEXT_ENCODER_COMPONENT) + + @dataclass class _ExecutorShutdownCleaner: """Finalizer that shuts down executor worker processes.""" @@ -471,16 +493,11 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): # request and Diffusion KV metadata inseparable. # All ranks reply (unique_reply_rank=None) so we collect dp_size # responses and match by dp_rank. - if ( - len(new_reqs) > 1 - and getattr(self.od_config, "enable_distributed_layerwise_offload", False) - and getattr(self.od_config, "dlo_use_allgather", True) - ): + if len(new_reqs) > 1 and any_selected_component_uses_allgather(self.od_config): # Reuse the request scheduler's complete compatibility key. DLO # AllGather requires every DP rank to execute the same collective # schedule, including shape, CFG, denoise steps, output count, - # and LoRA settings. Two default (None) step counts are valid and - # resolve identically inside the same pipeline. + # and LoRA settings. compatibility_keys = [build_request_batch_sampling_params_key(nr.req) for nr in new_reqs] if any(key != compatibility_keys[0] for key in compatibility_keys[1:]): raise ValueError( @@ -498,6 +515,13 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): "share identical extra_args. Different extra_args can change " "the forward schedule and cause AllGather deadlock." ) + if _uses_text_encoder_allgather(self.od_config): + encoder_signatures = {_text_encoder_input_signature(nr.req.prompt) for nr in new_reqs} + if len(encoder_signatures) > 1: + raise ValueError( + "DLO text_encoder AllGather requires every concurrent request " + "to provide the same positive/negative prompt embedding fields." + ) empty_prompt_ids = [nr.request_id for nr in new_reqs if _is_empty_dp_prompt(nr.req.prompt)] if empty_prompt_ids: raise ValueError( @@ -547,11 +571,16 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): args: tuple = (req, self.od_config, scheduler_output.kv_prefetch_job) if new_req.diffusion_kv_metadata is not None: args += (new_req.diffusion_kv_metadata,) + allgather_active = any_selected_component_uses_allgather(self.od_config) + timeout_options: dict[str, Any] = {} + if allgather_active: + timeout_options["timeout"] = _DLO_DP_WAVE_TIMEOUT_S result = self.collective_rpc( "execute_model", args=args, unique_reply_rank=0, exec_all_ranks=True, + **timeout_options, ) if isinstance(result, AsyncDiffusionOutput) and result.kind == AsyncOutputKind.COMPUTE_DONE: runner_outputs.append( @@ -575,6 +604,8 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): else: raise RuntimeError(f"Unexpected response type: {type(result)!r}") except Exception as exc: + if isinstance(exc, TimeoutError): + self._fail_closed_on_dp_wave_timeout(exc) runner_outputs.append( RunnerOutput( request_id=new_req.request_id, @@ -605,11 +636,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): parallel_config = getattr(self.od_config, "parallel_config", None) dp_size = getattr(parallel_config, "data_parallel_size", 1) - if ( - dp_size > 1 - and getattr(self.od_config, "enable_distributed_layerwise_offload", False) - and getattr(self.od_config, "dlo_use_allgather", True) - ): + if dp_size > 1 and any_selected_component_uses_allgather(self.od_config): # DLO DP uses one independent request per DP replica. It is not a # fused pipeline request batch, so models such as MiniMax-H3 do not # need to advertise supports_request_batch=True. diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 30d0f0610..ec0c1b14e 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -42,7 +42,15 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import ( has_online_quantization, ) from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline +from vllm_omni.diffusion.offloader.component_utils import encoder_component_type +from vllm_omni.diffusion.offloader.config import ( + DIT_COMPONENT, + TEXT_ENCODER_COMPONENT, + OffloadStrategy, + resolve_offload, +) from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery +from vllm_omni.diffusion.offloader.offload_plan import get_offload_plan from vllm_omni.diffusion.registry import initialize_model from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific @@ -546,14 +554,51 @@ class DiffusersPipelineLoader(HWRLoaderMixin): else: model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False) - _dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False) - _use_ag = getattr(self.od_config, "dlo_use_allgather", True) - _has_online_quant = self._has_online_quant(model) - _tp_size = int(getattr(self.parallel_config, "tensor_parallel_size", 1)) - _use_hsdp = bool(getattr(self.parallel_config, "use_hsdp", False)) - _dp_size = int(getattr(self.parallel_config, "data_parallel_size", 1)) - _sp_size = int(getattr(self.parallel_config, "sequence_parallel_size", 1)) - _dlo_group_size = _dp_size if _dp_size > 1 else _sp_size + resolved_offload = resolve_offload(self.od_config) + distributed_offload = resolved_offload.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE + dit_distributed_offload = distributed_offload and resolved_offload.offloads(DIT_COMPONENT) + dit_uses_allgather = dit_distributed_offload and resolved_offload.uses_allgather(DIT_COMPONENT) + tensor_parallel_size = int(getattr(self.parallel_config, "tensor_parallel_size", 1)) + use_hsdp = bool(getattr(self.parallel_config, "use_hsdp", False)) + data_parallel_size = int(getattr(self.parallel_config, "data_parallel_size", 1)) + sequence_parallel_size = int(getattr(self.parallel_config, "sequence_parallel_size", 1)) + dlo_group_size = data_parallel_size if data_parallel_size > 1 else sequence_parallel_size + modules = ModuleDiscovery.discover(model) + plan = get_offload_plan(model) + selected_encoders = [ + encoder + for name, encoder in zip(modules.encoder_names, modules.encoders) + if resolved_offload.offloads(TEXT_ENCODER_COMPONENT) + and encoder_component_type(name, plan) == TEXT_ENCODER_COMPONENT + ] + allgather_modules: list[nn.Module] = [] + if dlo_group_size > 1: + if dit_uses_allgather: + allgather_modules.extend(modules.dits) + if ( + distributed_offload + and resolved_offload.offloads(TEXT_ENCODER_COMPONENT) + and resolved_offload.uses_allgather(TEXT_ENCODER_COMPONENT) + ): + allgather_modules.extend(selected_encoders) + allgather_online_quant = any(self._has_online_quant(module) for module in allgather_modules) + if allgather_online_quant: + unsupported_methods = { + method + for module in allgather_modules + for method in self._unsupported_dlo_allgather_online_quant_methods(module) + } + if unsupported_methods: + raise ValueError( + "DLO+AllGather supports online quantization only for " + "per-tensor FP8, INT8, and MXFP8 linears; unsupported " + f"online methods: {', '.join(sorted(unsupported_methods))}. " + "Use rank-local transfer for the affected component or " + "disable online quantization." + ) + logger.info( + "Validated online methods (per-tensor FP8, INT8, MXFP8) for every component using DLO+AllGather" + ) plan_result = None weight_sources = self._get_weight_sources(model) @@ -562,9 +607,9 @@ class DiffusersPipelineLoader(HWRLoaderMixin): try: hwr_state = self._resolve_hwr( model, - ModuleDiscovery.discover(model), - dist_offload=_dist_offload, - use_allgather=_use_ag, + modules, + dist_offload=dit_distributed_offload, + use_allgather=dit_uses_allgather, load_format=load_format, sources=weight_sources, ) @@ -584,16 +629,15 @@ class DiffusersPipelineLoader(HWRLoaderMixin): hwr_active = hwr_state is not None if hwr_active and hwr_state is not None: self.host_weight_plan = cast(HostWeightPlan | None, hwr_state.get("plan")) - if _dist_offload and not hwr_active and not self._force_canonical_load: - modules = ModuleDiscovery.discover(model) + if dit_distributed_offload and not hwr_active and not self._force_canonical_load: plan_result = build_checkpoint_mmap_plan( model, dit_modules=tuple(zip(modules.dit_names, modules.dits)), sources=weight_sources, model_path=str(getattr(self.od_config, "model", "")) or None, - tensor_parallel_size=_tp_size, - use_hsdp=_use_hsdp, - online_quantization=_has_online_quant, + tensor_parallel_size=tensor_parallel_size, + use_hsdp=use_hsdp, + online_quantization=any(self._has_online_quant(dit) for dit in modules.dits), ) self.host_weight_plan = plan_result.plan @@ -602,7 +646,7 @@ class DiffusersPipelineLoader(HWRLoaderMixin): if host_weight_plan is not None: logger.info( "DLO host-weight plan active (%s, %s): skipping ordinary materialization for %s", - "AllGather" if _use_ag and _dlo_group_size > 1 else "rank-local", + "AllGather" if dit_uses_allgather and dlo_group_size > 1 else "rank-local", host_weight_plan.backing_kind, sorted(host_weight_plan.planned_source_prefixes) or "legacy DiT sources", ) @@ -623,24 +667,7 @@ class DiffusersPipelineLoader(HWRLoaderMixin): ) self._maybe_fuse_distilled_lora(model) else: - if _dist_offload and _use_ag and _dlo_group_size > 1 and _has_online_quant: - # An effective DLO group size of one performs no weight - # collective, so the AllGather layout allowlist does not - # apply there even with dlo_use_allgather=True. - unsupported_methods = self._unsupported_dlo_allgather_online_quant_methods(model) - if unsupported_methods: - raise ValueError( - "DLO+AllGather supports online quantization only for " - "per-tensor FP8, INT8, and MXFP8 linears; unsupported online " - f"methods: {', '.join(unsupported_methods)}. Please use " - "--dlo-no-use-allgather or disable online quantization." - ) - logger.info( - "Validated online methods (per-tensor FP8, INT8, MXFP8) with " - "DLO+AllGather: using the ordinary loader before sharding " - "finalized weights and scales" - ) - if _dist_offload and plan_result is not None: + if dit_distributed_offload and plan_result is not None: logger.info( "DLO direct checkpoint mmap unavailable; using ordinary loader: %s", plan_result.fallback_reason, @@ -846,14 +873,6 @@ class DiffusersPipelineLoader(HWRLoaderMixin): continue if has_online_quant: - # finalize_layerwise_processing() and the synchronous online - # loader already processed these layers. Avoid moving their - # quantized CPU weights back to the accelerator merely to call - # an idempotent no-op; doing so rebuilds a large CUDA allocator - # cache and defeats streaming CPU offload's startup-memory bound. - if getattr(module, "_already_called_process_weights_after_loading", False): - continue - # Online quant may leave straggler params on the ``meta`` device. # Move only real (non-meta) params onto the target device for # processing and restore them afterward, mirroring upstream vLLM's diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 2220c88f9..a171e2d3b 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -68,6 +68,7 @@ from vllm_omni.diffusion.models.schedulers.scheduling_flow_match_euler_discrete from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import ( FlowUniPCMultistepScheduler, ) +from vllm_omni.diffusion.offloader.config import OffloadStrategy, resolve_offload from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch @@ -934,6 +935,12 @@ class Cosmos3OmniDiffusersPipeline( ) -> None: super().__init__() self.od_config = od_config + resolved_offload = resolve_offload(od_config) + if resolved_offload.public is not None and resolved_offload.strategy is OffloadStrategy.MODEL_LEVEL: + raise ValueError( + "Cosmos3 model offload uses reasoner/generator topology and does " + "not support the dit/text_encoder component selector" + ) self.device = get_local_device() self.dtype = od_config.dtype @@ -1086,6 +1093,7 @@ class Cosmos3OmniDiffusersPipeline( device: torch.device, pin_memory: bool = True, use_hsdp: bool = False, + offload_components: frozenset[str] | None = None, ) -> None: """Enable Cosmos3 component-level model offload. @@ -1094,6 +1102,11 @@ class Cosmos3OmniDiffusersPipeline( mutual-exclusion swaps. The VAE stays resident on GPU like the generic model-level offloader. """ + if offload_components is not None: + raise ValueError( + "Cosmos3 model offload uses reasoner/generator topology and does not support " + "the dit/text_encoder offload_components selector" + ) self.vae.to(device, non_blocking=True) if isinstance(self._sound_tokenizer, nn.Module): self._sound_tokenizer.to(device) diff --git a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py index 59140073d..85e679715 100644 --- a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py +++ b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py @@ -237,6 +237,13 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin): def _raise_unsupported_features(self) -> None: """Raise an error for incompatible feature switches.""" + if self.od_config.diffusion_offload_config is not None: + raise NotImplementedError( + "diffusion_offload_config is not supported with the diffusers backend. " + "Its component selectors cannot be represented by Diffusers' pipeline-wide " + "CPU-offload hooks. Use the legacy enable_cpu_offload or " + "enable_layerwise_offload option, or use a native pipeline." + ) pc = self.od_config.parallel_config if pc.tensor_parallel_size > 1: raise NotImplementedError( diff --git a/vllm_omni/diffusion/models/minimax_h3/encoder.py b/vllm_omni/diffusion/models/minimax_h3/encoder.py index a91cb180b..43760dc7c 100644 --- a/vllm_omni/diffusion/models/minimax_h3/encoder.py +++ b/vllm_omni/diffusion/models/minimax_h3/encoder.py @@ -1275,55 +1275,6 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module): else: cache.release_if_needed() - def enable_omni_layerwise_offload(self, *, pin_memory: bool = True) -> None: - """Stream the TP-local Qwen vision/text blocks for low-HBM serving. - - The encoder has its own TP process group, so these blocks must remain - rank-local. Reusing the DiT AllGather group would concatenate - different TP shards and corrupt the encoder weights. - """ - if not self.is_loaded or getattr(self, "_omni_layerwise_enabled", False): - return - - from vllm_omni.diffusion.offloader.layerwise_backend import apply_block_hook - from vllm_omni.platforms import current_omni_platform - - self._omni_layerwise_hooks = [] - self._omni_layerwise_block_groups = [] - copy_stream = current_omni_platform.Stream() - for blocks in (self.vision.blocks, self.text_model.layers): - if len(blocks) <= 1: - continue - last_hook = apply_block_hook( - blocks[-1], - blocks[0], - self.device_target, - copy_stream, - pin_memory, - ) - hooks = [last_hook] - for index, block in enumerate(blocks[:-1]): - hooks.append( - apply_block_hook( - block, - blocks[index + 1], - self.device_target, - copy_stream, - pin_memory, - ) - ) - for index, hook in enumerate(hooks): - hook._prev_hook = hooks[index - 1] - self._omni_layerwise_hooks.extend(hooks) - self._omni_layerwise_block_groups.append(blocks) - - self._omni_layerwise_enabled = bool(self._omni_layerwise_hooks) - logger.info( - "MiniMax H3 encoder layerwise offload enabled on %d blocks across %d stacks", - sum(len(blocks) for blocks in self._omni_layerwise_block_groups), - len(self._omni_layerwise_block_groups), - ) - def _encode( self, input_ids: torch.Tensor, diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 702eaf594..22f638e9a 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -52,6 +52,13 @@ from vllm_omni.diffusion.offloader import ( remove_sequential_offload, sequential_offload_component, ) +from vllm_omni.diffusion.offloader.config import ( + DIT_COMPONENT, + TEXT_ENCODER_COMPONENT, + OffloadStrategy, + resolve_offload, + should_offload_component, +) from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( DiffusionPipelineProfilerMixin, @@ -668,6 +675,7 @@ class MiniMaxH3Pipeline( _offload_plan: ClassVar[OffloadPlan] = OffloadPlan( offload_submodules={"token_refiner": "blocks"}, resident_dit_paths=frozenset({"transformer"}), + encoder_component_types={"text_encoder": TEXT_ENCODER_COMPONENT}, encoder_block_attrs={"text_encoder": ("vision.blocks", "text_model.layers")}, on_demand_component_paths=frozenset({"text_encoder", "video_vae", "audio_vae"}), ) @@ -703,10 +711,13 @@ class MiniMaxH3Pipeline( self._lora_sigma_schedules.pop(lora_request.lora_int_id, None) od_config = getattr(self, "od_config", None) offload_modes = [] - if getattr(od_config, "enable_cpu_offload", False): - offload_modes.append("model-level CPU offload (--enable-cpu-offload)") - if getattr(od_config, "enable_layerwise_offload", False): - offload_modes.append("layerwise offload (--enable-layerwise-offload)") + if od_config is not None: + resolved_offload = resolve_offload(od_config) + if resolved_offload.offloads(DIT_COMPONENT): + if resolved_offload.strategy is OffloadStrategy.MODEL_LEVEL: + offload_modes.append("model-level CPU offload") + elif resolved_offload.strategy is OffloadStrategy.LAYER_WISE: + offload_modes.append("layerwise offload") loaded = load_minimax_h3_turbo_lora( partition=self.partition, lora_request=lora_request, @@ -1030,10 +1041,13 @@ class MiniMaxH3Pipeline( self.text_encoder_group = None self.text_encoder = None self._encoder_modules = [] - stage_components = bool( + legacy_manual_components = getattr(od_config, "diffusion_offload_config", None) is None and bool( od_config.enable_layerwise_offload or getattr(od_config, "enable_distributed_layerwise_offload", False) ) - component_load_device = torch.device("cpu") if stage_components else self.device + # Preserve the legacy MiniMax-H3 low-residency path. The compact API + # deliberately limits explicit component selection to dit/text_encoder, + # so VAEs stay resident for new configurations. + component_load_device = torch.device("cpu") if legacy_manual_components else self.device self.video_vae = MiniMaxH3VideoVAE( os.path.join(model_path, "video_vae"), device=self.device, @@ -1048,14 +1062,19 @@ class MiniMaxH3Pipeline( self.vae = self.video_vae self._dlo_component_cache = None - if getattr(od_config, "enable_distributed_layerwise_offload", False): + offloads_text_encoder = should_offload_component(od_config, TEXT_ENCODER_COMPONENT) + needs_component_cache = legacy_manual_components or offloads_text_encoder + if getattr(od_config, "enable_distributed_layerwise_offload", False) and needs_component_cache: self._dlo_component_cache = BoundedAllocatorCache(self.device) - _register_dlo_component_cache( - self._dlo_component_cache, - self.text_encoder, - self.video_vae, - self.audio_vae, - ) + if legacy_manual_components: + _register_dlo_component_cache( + self._dlo_component_cache, + self.text_encoder, + self.video_vae, + self.audio_vae, + ) + elif offloads_text_encoder: + _register_dlo_component_cache(self._dlo_component_cache, self.text_encoder) self._quality_policy = MiniMaxH3QualityPolicy(od_config) self._cache_dit_runtime = RequestScopedCacheDiTRuntime(self) @@ -1519,25 +1538,25 @@ class MiniMaxH3Pipeline( # swaps the resident DiT and encoder. return self.text_encoder(input_ids, **vision_kwargs) - if self.od_config.enable_layerwise_offload or getattr( - self.od_config, "enable_distributed_layerwise_offload", False - ): - # Layerwise DiT offload already provides the low-residency encoder - # phase used by the checkpoint reference. + if self._uses_manual_component_offload(self.text_encoder): with self._component_on_device(self.text_encoder): return self.text_encoder.encode_ids(input_ids, **vision_kwargs) - # Keep both Qwen and DiT resident across requests. Moving either model - # here makes encoder latency include a tens-of-gigabytes PCIe transfer, - # which defeats the no-offload contract. + # Keep Qwen resident when it is not selected for layerwise offload. self.text_encoder.load_to_device() return self.text_encoder.encode_ids(input_ids, **vision_kwargs) - def _uses_manual_component_offload(self) -> bool: + def _uses_manual_component_offload(self, component: nn.Module) -> bool: od_config = getattr(self, "od_config", None) - return bool( - getattr(od_config, "enable_layerwise_offload", False) - or getattr(od_config, "enable_distributed_layerwise_offload", False) + if od_config is None: + return False + if getattr(od_config, "diffusion_offload_config", None) is None: + return bool( + getattr(od_config, "enable_layerwise_offload", False) + or getattr(od_config, "enable_distributed_layerwise_offload", False) + ) + return component is getattr(self, "text_encoder", None) and should_offload_component( + od_config, TEXT_ENCODER_COMPONENT ) def enable_omni_model_cpu_offload( @@ -1546,6 +1565,7 @@ class MiniMaxH3Pipeline( device: torch.device, pin_memory: bool, use_hsdp: bool, + offload_components: frozenset[str] | None = None, ) -> None: if getattr(self, "_model_cpu_offload_modules", None): return @@ -1554,19 +1574,32 @@ class MiniMaxH3Pipeline( dits = components.dits stages = [*components.encoders, *components.vaes] modules = [*dits, *stages] + selection_options: dict[str, Any] = {} + if offload_components is not None: + if DIT_COMPONENT in offload_components and not dits: + raise ValueError("MiniMax-H3 has no loaded DiT for selected module offload") + if TEXT_ENCODER_COMPONENT in offload_components and not components.encoders: + raise ValueError("MiniMax-H3 has no loaded text encoder for selected module offload") + selection_options = { + "offload_dit_modules": dits if DIT_COMPONENT in offload_components else (), + "offload_encoder_modules": ( + components.encoders if TEXT_ENCODER_COMPONENT in offload_components else () + ), + } apply_sequential_offload( dit_modules=dits, encoder_modules=stages, device=device, pin_memory=pin_memory, use_hsdp=use_hsdp, - offload_initial_dits=True, + offload_initial_dits=offload_components is None or DIT_COMPONENT in offload_components, + **selection_options, ) self._model_cpu_offload_modules = modules logger.info( - "MiniMax-H3 model-level CPU offload enabled for %d DiT(s), text encoder, video VAE, and audio VAE", - len(dits), + "MiniMax-H3 model-level CPU offload enabled for selected components: %s", + sorted(offload_components) if offload_components is not None else "legacy full topology", ) def disable_omni_model_cpu_offload(self) -> None: @@ -1582,7 +1615,7 @@ class MiniMaxH3Pipeline( with sequential_offload_component(component): yield return - staged = self._uses_manual_component_offload() + staged = self._uses_manual_component_offload(component) try: if staged: component.load_to_device() diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index f92053e52..2383fc6fd 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -34,6 +34,7 @@ from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_z from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler from vllm_omni.diffusion.models.wan2_2.scheduling_wan_euler import WanEulerScheduler from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention, WanTransformer3DModel +from vllm_omni.diffusion.offloader import OffloadPlan from vllm_omni.diffusion.postprocess import interpolate_video_tensor from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest @@ -308,6 +309,13 @@ def get_wan22_pre_process_func( return pre_process_func +_WAN_TEXT_ENCODER_OFFLOAD_PLAN = OffloadPlan( + encoder_component_types={"text_encoder": "text_encoder"}, + encoder_block_attrs={"text_encoder": ("encoder.block",)}, + encoder_dlo_weight_replication=frozenset({"text_encoder"}), +) + + class Wan22Pipeline( nn.Module, PipelineParallelMixin, @@ -322,6 +330,7 @@ class Wan22Pipeline( _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"] _encoder_modules: ClassVar[list[str]] = ["text_encoder"] _vae_modules: ClassVar[list[str]] = ["vae"] + _offload_plan = _WAN_TEXT_ENCODER_OFFLOAD_PLAN def __init__( self, diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index b0cd35458..908a35020 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -32,6 +32,7 @@ from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComp from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero from vllm_omni.diffusion.models.utils import _load_json from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( + _WAN_TEXT_ENCODER_OFFLOAD_PLAN, build_wan_scheduler, create_transformer_from_config, load_transformer_config, @@ -192,6 +193,7 @@ class Wan22I2VPipeline( _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"] _encoder_modules: ClassVar[list[str]] = ["text_encoder", "image_encoder"] _vae_modules: ClassVar[list[str]] = ["vae"] + _offload_plan = _WAN_TEXT_ENCODER_OFFLOAD_PLAN def __init__( self, diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py index bc688fbc4..ea99e065f 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py @@ -41,6 +41,7 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import ( create_s2v_transformer_from_config, ) +from vllm_omni.diffusion.offloader.config import DIT_COMPONENT, selected_offload_components from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request @@ -1090,6 +1091,11 @@ class Wan22S2VPipeline( # Main forward # ------------------------------------------------------------------ + def _should_release_dit_before_decode(self) -> bool: + if getattr(self.od_config.parallel_config, "use_hsdp", False): + return False + return self.od_config.enable_cpu_offload and DIT_COMPONENT in selected_offload_components(self.od_config) + def forward( self, req: DiffusionRequestBatch, @@ -1401,7 +1407,7 @@ class Wan22S2VPipeline( ) # ---- Decode this clip ---- - if self.od_config.enable_cpu_offload and not getattr(self.od_config.parallel_config, "use_hsdp", False): + if self._should_release_dit_before_decode(): self.transformer.to("cpu") current_omni_platform.empty_cache() diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py index 58cd98c3e..0d48dffb5 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py @@ -16,7 +16,7 @@ determined by which inputs are provided (no explicit mode flag): from __future__ import annotations from collections.abc import Iterable -from dataclasses import replace +from copy import copy import PIL.Image import torch @@ -198,7 +198,10 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): ): # VACE defaults to flow_shift=3.0 for 480p (base WAN T2V uses 5.0 for 720p) if od_config.flow_shift is None: - od_config = replace(od_config, flow_shift=3.0) + # Preserve runtime-only config provenance (including compact + # offload aliases) instead of rerunning dataclass __post_init__. + od_config = copy(od_config) + od_config.flow_shift = 3.0 super().__init__(od_config=od_config, prefix=prefix) diff --git a/vllm_omni/diffusion/offloader/__init__.py b/vllm_omni/diffusion/offloader/__init__.py index 8145ac624..0df4039eb 100644 --- a/vllm_omni/diffusion/offloader/__init__.py +++ b/vllm_omni/diffusion/offloader/__init__.py @@ -8,8 +8,9 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.model_loader.host_weight_plan import HostWeightPlan from vllm_omni.platforms import current_omni_platform -from .base import OffloadBackend, OffloadConfig, OffloadStrategy, SupportsModelCpuOffload +from .base import OffloadBackend, OffloadConfig, SupportsModelCpuOffload from .block_discovery import get_blocks_attr_names, get_blocks_from_dit, set_blocks_attr_names +from .config import OffloadStrategy from .distributed_layerwise_backend import ( DistributedLayerwiseOffloadBackend, DistributedLayerwiseOffloadHook, diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py index c4ad93482..bb047a7eb 100644 --- a/vllm_omni/diffusion/offloader/base.py +++ b/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, ) diff --git a/vllm_omni/diffusion/offloader/block_discovery.py b/vllm_omni/diffusion/offloader/block_discovery.py index 643485ed6..de535e34e 100644 --- a/vllm_omni/diffusion/offloader/block_discovery.py +++ b/vllm_omni/diffusion/offloader/block_discovery.py @@ -35,9 +35,13 @@ def set_blocks_attr_names(model: nn.Module, names: list[str]) -> None: setattr(model.__class__, "_layerwise_offload_blocks_attrs", names) -def get_blocks_from_dit(model: nn.Module) -> tuple[list[str], list[nn.Module]]: - """Retrieve blocks and attribute names from provided DiT model.""" - blocks_attr_names = get_blocks_attr_names(model) +def get_blocks_from_dit( + model: nn.Module, + block_attrs: tuple[str, ...] | None = None, +) -> tuple[list[str], list[nn.Module]]: + """Retrieve blocks from an explicit plan or the DiT class metadata.""" + explicit_attrs = block_attrs is not None + blocks_attr_names = list(block_attrs) if explicit_attrs else get_blocks_attr_names(model) if not blocks_attr_names: logger.warning( f"No _layerwise_offload_blocks_attrs defined for {model.__class__.__name__}, " @@ -49,9 +53,9 @@ def get_blocks_from_dit(model: nn.Module) -> tuple[list[str], list[nn.Module]]: for name in blocks_attr_names: attr = getattr(model, name, None) if attr is None: + source = "OffloadPlan.block_attrs" if explicit_attrs else "_layerwise_offload_blocks_attrs" raise AttributeError( - f"Attribute '{name}' declared in _layerwise_offload_blocks_attrs " - f"does not exist on model {model.__class__.__name__}" + f"Attribute '{name}' declared in {source} does not exist on model {model.__class__.__name__}" ) try: attr_iter = iter(attr) diff --git a/vllm_omni/diffusion/offloader/component_utils.py b/vllm_omni/diffusion/offloader/component_utils.py new file mode 100644 index 000000000..1000b29d6 --- /dev/null +++ b/vllm_omni/diffusion/offloader/component_utils.py @@ -0,0 +1,221 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Shared component-plan helpers for diffusion offload backends.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterator +from itertools import chain +from operator import attrgetter +from typing import TYPE_CHECKING, Any + +import torch +from torch import nn +from torch.distributed.tensor import DTensor +from vllm.logger import init_logger + +from .block_discovery import get_blocks_from_dit +from .config import DIT_COMPONENT, TEXT_ENCODER_COMPONENT +from .offload_plan import OffloadPlan +from .tensor_utils import set_tensor_storage + +if TYPE_CHECKING: + from .base import OffloadConfig + from .module_collector import PipelineModules + +logger = init_logger(__name__) + + +def encoder_component_type(name: str, plan: OffloadPlan | None) -> str | None: + """Map a discovered encoder path to its public offload component.""" + declared = None if plan is None else plan.encoder_component_types.get(name) + if declared is not None: + if declared != TEXT_ENCODER_COMPONENT: + raise ValueError(f"OffloadPlan maps encoder {name!r} to unknown component {declared!r}") + return declared + + leaf_name = name.rsplit(".", 1)[-1] + if leaf_name.startswith(TEXT_ENCODER_COMPONENT) or leaf_name.endswith(TEXT_ENCODER_COMPONENT): + return TEXT_ENCODER_COMPONENT + return None + + +def get_encoder_block_groups( + module: nn.Module, + name: str, + plan: OffloadPlan | None, + *, + strict: bool = False, +) -> list[nn.ModuleList]: + """Resolve the streamable block lists declared for one encoder.""" + if plan is None: + return [] + + # Some distributed models expose an unloaded stub on ranks where the + # component never executes. That is a valid local no-op even when the + # component was selected explicitly. + if getattr(module, "is_loaded", True) is False: + return [] + + groups: list[nn.ModuleList] = [] + for block_path in plan.encoder_block_attrs.get(name, ()): + try: + blocks = attrgetter(block_path)(module) + except AttributeError: + if strict: + raise ValueError(f"Encoder offload path {name}.{block_path} was not found") from None + logger.warning("Encoder offload path %s.%s was not found", name, block_path) + continue + if not isinstance(blocks, nn.ModuleList) or len(blocks) <= 1: + if strict: + raise ValueError(f"Encoder offload path {name}.{block_path} is not a streamable block list") + logger.warning("Encoder offload path %s.%s is not a streamable block list", name, block_path) + continue + groups.append(blocks) + return groups + + +def iter_streamable_dits( + modules: PipelineModules, + config: OffloadConfig, + device: torch.device, + plan: OffloadPlan | None, +) -> Iterator[tuple[str, nn.Module, list[str], list[nn.Module]]]: + """Yield selected DiTs whose block metadata resolves successfully.""" + if not config.offloads(DIT_COMPONENT): + return + for name, module in zip(modules.dit_names, modules.dits): + logger.info("Applying hooks on %s (%s)", name, module.__class__.__name__) + planned_attrs = None if plan is None else plan.block_attrs.get(name) + block_attrs, blocks = get_blocks_from_dit(module, planned_attrs) + if blocks: + yield name, module, block_attrs, blocks + continue + if config.components is not None: + raise ValueError(f"Selected DiT {name!r} has no streamable layerwise-offload blocks") + logger.warning("Target layers (blocks) not found. Skipping offloading on %s (%s)", name, type(module).__name__) + module.to(device) + + +def move_non_block_state_to_device( + module: nn.Module, + block_groups: list[nn.ModuleList], + device: torch.device, +) -> None: + """Keep component state outside streamed block lists resident on device.""" + block_tensors = { + id(tensor) + for blocks in block_groups + for block in blocks + for tensor in chain(block.parameters(), block.buffers()) + } + for tensor in chain(module.parameters(), module.buffers()): + if id(tensor) in block_tensors: + continue + local = tensor.to_local() if isinstance(tensor, DTensor) else tensor + if local.device != device: + set_tensor_storage(tensor, local.to(device, non_blocking=True)) + + +def set_encoder_layerwise_state( + module: nn.Module, + hooks: list[Any], + block_groups: list[nn.ModuleList], +) -> None: + """Publish the backend-neutral state used by encoder stage lifecycles.""" + module._omni_layerwise_hooks = hooks + module._omni_layerwise_block_groups = block_groups + module._omni_layerwise_enabled = True + + +def clear_encoder_layerwise_state(module: nn.Module) -> None: + """Clear encoder layerwise state after its backend hooks are removed.""" + module._omni_layerwise_hooks = [] + module._omni_layerwise_block_groups = [] + module._omni_layerwise_enabled = False + + +def validate_on_demand_component(module: nn.Module, name: str) -> None: + """Require the explicit lifecycle used by pipeline-managed components.""" + if not callable(getattr(module, "load_to_device", None)) or not callable(getattr(module, "offload_to_cpu", None)): + raise ValueError( + f"Component {name!r} declares on-demand offload but must implement load_to_device() and offload_to_cpu()" + ) + + +def prepare_component( + module: nn.Module, + name: str, + *, + device: torch.device, + stage_on_demand: bool, + blockwise: bool, + staged_components: list[nn.Module], +) -> None: + """Stage a selected component or keep its non-streamed form resident.""" + if stage_on_demand: + validate_on_demand_component(module, name) + getattr(module, "offload_to_cpu")() + staged_components.append(module) + logger.info("Prepared %s for pipeline-managed staged offload", name) + elif not blockwise: + module.to(device) + + +def prepare_pipeline_components( + modules: PipelineModules, + config: OffloadConfig, + plan: OffloadPlan | None, + *, + device: torch.device, + staged_components: list[nn.Module], + enable_encoder_blocks: Callable[[nn.Module, str, OffloadPlan | None, bool], bool], +) -> None: + """Apply the shared encoder/VAE/resident placement policy.""" + if config.components is not None and config.offloads(TEXT_ENCODER_COMPONENT): + selected_encoder_names = [name for name in modules.encoder_names if config.offloads_encoder(name, plan)] + if not selected_encoder_names: + raise ValueError("No text encoder modules found for selected text_encoder offload") + + if plan is not None: + for encoder, name in zip(modules.encoders, modules.encoder_names): + if config.should_offload_encoder(name, plan) and name in plan.on_demand_component_paths: + validate_on_demand_component(encoder, name) + + for encoder, name in zip(modules.encoders, modules.encoder_names): + selected = config.should_offload_encoder(name, plan) + stage_on_demand = bool(selected and plan is not None and name in plan.on_demand_component_paths) + blockwise = selected and enable_encoder_blocks(encoder, name, plan, stage_on_demand) + if stage_on_demand and not blockwise and config.uses_allgather(TEXT_ENCODER_COMPONENT): + raise ValueError( + f"Text encoder {name!r} cannot use AllGather without a model-declared streamable block plan" + ) + if selected and config.components is not None and not (blockwise or stage_on_demand): + raise ValueError(f"Selected text encoder {name!r} requires a model-declared streamable or on-demand plan") + prepare_component( + encoder, + name, + device=device, + stage_on_demand=stage_on_demand, + blockwise=blockwise, + staged_components=staged_components, + ) + + for vae, name in zip(modules.vaes, modules.vae_names): + legacy_staged = config.components is None and plan is not None and name in plan.on_demand_component_paths + prepare_component( + vae, + name, + device=device, + stage_on_demand=legacy_staged, + blockwise=False, + staged_components=staged_components, + ) + + for name, module in zip(modules.resident_names, modules.resident_modules): + module.to(device) + logger.debug("Moved resident module %s to %s", name, device) + + if not config.offloads(DIT_COMPONENT): + for dit in modules.dits: + dit.to(device) diff --git a/vllm_omni/diffusion/offloader/config.py b/vllm_omni/diffusion/offloader/config.py new file mode 100644 index 000000000..46f5e95f3 --- /dev/null +++ b/vllm_omni/diffusion/offloader/config.py @@ -0,0 +1,429 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Public diffusion CPU-offload configuration helpers.""" + +from __future__ import annotations + +from collections.abc import Collection, Iterator, Mapping +from dataclasses import dataclass +from enum import Enum +from typing import Any, TypeVar + +DIT_COMPONENT = "dit" +TEXT_ENCODER_COMPONENT = "text_encoder" +OFFLOAD_COMPONENTS = frozenset({DIT_COMPONENT, TEXT_ENCODER_COMPONENT}) +DEFAULT_OFFLOAD_COMPONENTS = frozenset({DIT_COMPONENT}) + +_KeyT = TypeVar("_KeyT") +_ValueT = TypeVar("_ValueT") + + +class _FrozenMapping(Mapping[_KeyT, _ValueT]): + """Small immutable mapping that remains safe to pickle between workers.""" + + def __init__(self, value: Mapping[_KeyT, _ValueT]) -> None: + self._items = tuple(value.items()) + + def __getitem__(self, key: _KeyT) -> _ValueT: + for candidate, value in self._items: + if candidate == key: + return value + raise KeyError(key) + + def __iter__(self) -> Iterator[_KeyT]: + return (key for key, _ in self._items) + + def __len__(self) -> int: + return len(self._items) + + def __reduce__(self): + return type(self), (dict(self._items),) + + +class OffloadMode(str, Enum): + """User-facing offload granularity.""" + + MODULE = "module" + LAYER = "layer" + + +class OffloadStrategy(str, Enum): + """Resolved internal backend strategy.""" + + NONE = "none" + MODEL_LEVEL = "model_level" + LAYER_WISE = "layer_wise" + DISTRIBUTED_LAYER_WISE = "distributed_layer_wise" + + +class DLOTransfer(str, Enum): + """How one component's next block reaches the device.""" + + ALLGATHER = "allgather" + RANK_LOCAL = "rank-local" + + +@dataclass(frozen=True) +class LayerOffloadOptions: + """Layer-mode settings for one selected diffusion component.""" + + weight_transfer: DLOTransfer | None = None + resident_layers: int = 0 + + +@dataclass(frozen=True) +class ParsedDiffusionOffloadConfig: + """Validated internal form of the raw public offload mapping.""" + + mode: OffloadMode + components: frozenset[str] + layer_options: Mapping[str, LayerOffloadOptions] + pin_memory: bool | None = None + + def __post_init__(self) -> None: + object.__setattr__(self, "layer_options", _FrozenMapping(self.layer_options)) + + +@dataclass(frozen=True) +class ResolvedOffload: + """Canonical offload policy derived once from public or legacy inputs.""" + + strategy: OffloadStrategy + public: ParsedDiffusionOffloadConfig | None + components: frozenset[str] + transfers: Mapping[str, DLOTransfer] + pin_memory: bool + resident_layers: int + + def __post_init__(self) -> None: + object.__setattr__(self, "transfers", _FrozenMapping(self.transfers)) + + def offloads(self, component: str) -> bool: + return component in self.components + + def uses_allgather(self, component: str = DIT_COMPONENT) -> bool: + try: + return self.transfers[component] is DLOTransfer.ALLGATHER + except KeyError as exc: + raise ValueError(f"Unknown offload component {component!r}") from exc + + @property + def any_allgather(self) -> bool: + return self.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE and any( + self.uses_allgather(component) for component in self.components + ) + + +_LEGACY_STRATEGY_FIELDS = { + "enable_cpu_offload": OffloadStrategy.MODEL_LEVEL, + "enable_layerwise_offload": OffloadStrategy.LAYER_WISE, + "enable_distributed_layerwise_offload": OffloadStrategy.DISTRIBUTED_LAYER_WISE, +} +_LEGACY_STRATEGY_PRIORITY = ( + OffloadStrategy.DISTRIBUTED_LAYER_WISE, + OffloadStrategy.LAYER_WISE, + OffloadStrategy.MODEL_LEVEL, +) + + +_EnumT = TypeVar("_EnumT", bound=Enum) + + +def _parse_enum(value: Any, enum_type: type[_EnumT], label: str) -> _EnumT: + if isinstance(value, enum_type): + return value + if not isinstance(value, str): + raise TypeError(f"{label} must be a string, got {type(value).__name__}") + try: + return enum_type(value) + except ValueError as exc: + choices = ", ".join(item.value for item in enum_type) + raise ValueError(f"Unknown {label} {value!r}; choose from: {choices}") from exc + + +def _validate_string_keys(value: Mapping[Any, Any], label: str) -> None: + if any(not isinstance(key, str) for key in value): + raise TypeError(f"{label} keys must be strings") + + +def _parse_layer_options(component: str, value: Any) -> LayerOffloadOptions: + if not isinstance(value, Mapping): + raise TypeError( + f"diffusion_offload_config.layer_options[{component!r}] must be a mapping, got {type(value).__name__}" + ) + _validate_string_keys(value, f"diffusion_offload_config.layer_options[{component!r}]") + unknown = sorted(set(value) - {"weight_transfer", "resident_layers"}) + if unknown: + raise ValueError( + f"Unknown diffusion offload setting(s) for {component}: {', '.join(unknown)}; " + "choose from: weight_transfer, resident_layers" + ) + weight_transfer = ( + _parse_enum(value["weight_transfer"], DLOTransfer, "offload transfer") if "weight_transfer" in value else None + ) + resident_layers = value.get("resident_layers", 0) + + if type(resident_layers) is not int or resident_layers < 0: + raise ValueError(f"resident_layers for {component} must be a non-negative integer") + parsed = LayerOffloadOptions( + weight_transfer=weight_transfer, + resident_layers=resident_layers, + ) + if component != DIT_COMPONENT and parsed.resident_layers: + raise ValueError("resident_layers currently supports only the 'dit' component") + return parsed + + +def parse_diffusion_offload_config(value: Any) -> ParsedDiffusionOffloadConfig | None: + """Validate the compact ``diffusion_offload_config`` public API.""" + if value is None: + return None + if not isinstance(value, Mapping): + raise TypeError(f"diffusion_offload_config must be a mapping, got {type(value).__name__}") + + _validate_string_keys(value, "diffusion_offload_config") + unknown = sorted(set(value) - {"mode", "components", "layer_options", "pin_memory"}) + if unknown: + raise ValueError( + f"Unknown diffusion_offload_config field(s): {', '.join(unknown)}; " + "choose from: mode, components, layer_options, pin_memory" + ) + if "mode" not in value: + raise ValueError("diffusion_offload_config requires 'mode'") + if "components" not in value: + raise ValueError("diffusion_offload_config requires 'components'") + + mode = _parse_enum(value["mode"], OffloadMode, "diffusion offload mode") + components = parse_offload_components(value["components"]) + + raw_layer_options = value.get("layer_options", {}) + if not isinstance(raw_layer_options, Mapping): + raise TypeError("diffusion_offload_config.layer_options must be a mapping") + _validate_string_keys(raw_layer_options, "diffusion_offload_config.layer_options") + unselected_options = sorted(set(raw_layer_options) - components) + if unselected_options: + raise ValueError( + "diffusion_offload_config.layer_options requires selecting the same component(s): " + + ", ".join(unselected_options) + ) + layer_options = { + component: _parse_layer_options(component, raw_layer_options.get(component, {})) for component in components + } + + pin_memory = value.get("pin_memory") + if pin_memory is not None and type(pin_memory) is not bool: + raise TypeError("diffusion_offload_config.pin_memory must be a bool") + + if mode is OffloadMode.MODULE and raw_layer_options: + raise ValueError( + "diffusion_offload_config.layer_options requires mode='layer'; " + f"configured for: {', '.join(sorted(raw_layer_options))}" + ) + dit = layer_options.get(DIT_COMPONENT) + if dit is not None and dit.resident_layers and dit.weight_transfer is DLOTransfer.ALLGATHER: + raise ValueError("resident_layers requires dit.weight_transfer='rank-local'") + + return ParsedDiffusionOffloadConfig( + mode=mode, + components=components, + layer_options=layer_options, + pin_memory=pin_memory, + ) + + +def _legacy_strategy(config: Any) -> OffloadStrategy: + enabled = {strategy for field, strategy in _LEGACY_STRATEGY_FIELDS.items() if bool(getattr(config, field, False))} + return next((strategy for strategy in _LEGACY_STRATEGY_PRIORITY if strategy in enabled), OffloadStrategy.NONE) + + +def _public_strategy(public: ParsedDiffusionOffloadConfig) -> OffloadStrategy: + if public.mode is OffloadMode.MODULE: + return OffloadStrategy.MODEL_LEVEL + + needs_distributed_backend = any( + settings.weight_transfer is DLOTransfer.ALLGATHER or settings.resident_layers + for settings in public.layer_options.values() + ) + return OffloadStrategy.DISTRIBUTED_LAYER_WISE if needs_distributed_backend else OffloadStrategy.LAYER_WISE + + +def _validate_legacy_layer_options(config: Any, public: ParsedDiffusionOffloadConfig | None) -> None: + """Reject ambiguous or invalid compatibility DLO tuning before loading.""" + resident_layers = getattr(config, "dlo_resident_layers", 0) + if type(resident_layers) is not int or resident_layers < 0: + raise ValueError(f"dlo_resident_layers must be a non-negative integer, got {resident_layers!r}") + + if public is not None: + if getattr(config, "host_weight_runtime_mode", "disabled") != "disabled": + raise ValueError("diffusion_offload_config cannot be combined with Host Weight Runtime") + conflicting_fields = [] + if getattr(config, "dlo_use_allgather", True) is not True: + conflicting_fields.append("dlo_use_allgather") + if resident_layers: + conflicting_fields.append("dlo_resident_layers") + if float(getattr(config, "dlo_host_registration_limit_gib", 0.0)) > 0: + conflicting_fields.append("dlo_host_registration_limit_gib") + if conflicting_fields: + raise ValueError( + "diffusion_offload_config cannot be combined with legacy DLO option(s): " + + ", ".join(conflicting_fields) + ) + return + + if resident_layers and bool(getattr(config, "dlo_use_allgather", True)): + raise ValueError( + "dlo_resident_layers requires the DiT DLO transfer to be rank-local; set dlo_use_allgather=False" + ) + + +def resolve_offload(config: Any) -> ResolvedOffload: + """Resolve and cache the complete offload policy on its owning config.""" + cached = getattr(config, "_resolved_diffusion_offload", None) + if isinstance(cached, ResolvedOffload): + return cached + + public = parse_diffusion_offload_config(getattr(config, "diffusion_offload_config", None)) + _validate_legacy_layer_options(config, public) + legacy = _legacy_strategy(config) + if public is None: + transfers = { + DIT_COMPONENT: ( + DLOTransfer.ALLGATHER if bool(getattr(config, "dlo_use_allgather", True)) else DLOTransfer.RANK_LOCAL + ), + TEXT_ENCODER_COMPONENT: DLOTransfer.RANK_LOCAL, + } + resolved = ResolvedOffload( + strategy=legacy, + public=None, + components=DEFAULT_OFFLOAD_COMPONENTS, + transfers=transfers, + pin_memory=bool(getattr(config, "pin_cpu_memory", True)), + resident_layers=int(getattr(config, "dlo_resident_layers", 0)), + ) + else: + strategy = _public_strategy(public) + if legacy is not OffloadStrategy.NONE and legacy is not strategy: + raise ValueError("diffusion_offload_config cannot be combined with legacy enable_*_offload flags") + transfers = {component: DLOTransfer.RANK_LOCAL for component in OFFLOAD_COMPONENTS} + for component, options in public.layer_options.items(): + transfers[component] = options.weight_transfer or DLOTransfer.RANK_LOCAL + dit_options = public.layer_options.get(DIT_COMPONENT) + resolved = ResolvedOffload( + strategy=strategy, + public=public, + components=public.components, + transfers=transfers, + pin_memory=( + public.pin_memory if public.pin_memory is not None else bool(getattr(config, "pin_cpu_memory", True)) + ), + resident_layers=0 if dit_options is None else dit_options.resident_layers, + ) + + parallel_config = getattr(config, "parallel_config", None) + data_parallel_size = int(getattr(parallel_config, "data_parallel_size", 1)) + cache_backend = getattr(config, "cache_backend", None) + if data_parallel_size > 1 and resolved.any_allgather and cache_backend not in (None, "none"): + raise ValueError( + "Diffusion cache acceleration cannot be combined with DLO AllGather " + "across data-parallel ranks because rank-local cache decisions can skip " + "different weight collectives. Disable cache_backend or use rank-local transfer." + ) + if ( + data_parallel_size > 1 + and bool(getattr(config, "enable_prompt_embed_cache", False)) + and resolved.offloads(TEXT_ENCODER_COMPONENT) + and resolved.uses_allgather(TEXT_ENCODER_COMPONENT) + ): + raise ValueError( + "Prompt embedding cache cannot be combined with text_encoder AllGather " + "across data-parallel ranks: rank-local cache hits would skip different " + "encoder collectives. Use text_encoder weight_transfer='rank-local' or " + "disable enable_prompt_embed_cache." + ) + + # Actual runtime configs are immutable after construction. Keeping the + # derived object off the public dataclass fields also keeps projections and + # serialization transport-safe while making hot-path accessors constant-time. + try: + setattr(config, "_resolved_diffusion_offload", resolved) + except (AttributeError, TypeError): + pass + return resolved + + +def resolve_offload_strategy(config: Any) -> OffloadStrategy: + """Resolve the compact config or compatibility boolean entry points.""" + return resolve_offload(config).strategy + + +def materialize_legacy_offload_flags(config: Any) -> OffloadStrategy: + """Keep existing strategy readers working after resolving the compact API.""" + resolved = resolve_offload(config) + for field, field_strategy in _LEGACY_STRATEGY_FIELDS.items(): + setattr(config, field, resolved.strategy is field_strategy) + setattr(config, "pin_cpu_memory", resolved.pin_memory) + return resolved.strategy + + +def parse_offload_components(value: Collection[str]) -> frozenset[str]: + """Validate an internal component collection.""" + if isinstance(value, (str, Mapping)) or not isinstance(value, Collection): + raise TypeError("diffusion_offload_config.components must be a non-empty list of component names") + if any(not isinstance(item, str) for item in value): + raise TypeError("offload component entries must be strings") + components = list(value) + if not components: + raise ValueError("offload components must not be empty") + unknown = sorted(set(components) - OFFLOAD_COMPONENTS) + if unknown: + choices = ", ".join(sorted(OFFLOAD_COMPONENTS)) + raise ValueError(f"Unknown diffusion offload component(s): {', '.join(unknown)}; choose from: {choices}") + return frozenset(components) + + +def parse_dlo_transfer(value: Mapping[str, str | DLOTransfer]) -> dict[str, DLOTransfer]: + """Resolve an internal per-component transfer mapping.""" + if not isinstance(value, Mapping): + raise TypeError(f"offload transfers must be a mapping, got {type(value).__name__}") + _validate_string_keys(value, "offload transfers") + missing = sorted(OFFLOAD_COMPONENTS - set(value)) + if missing: + raise ValueError("offload transfers require every component; missing: " + ", ".join(missing)) + resolved: dict[str, DLOTransfer] = {} + for component, raw_transfer in value.items(): + if component not in OFFLOAD_COMPONENTS: + choices = ", ".join(sorted(OFFLOAD_COMPONENTS)) + raise ValueError(f"Unknown offload transfer component {component!r}; choose from: {choices}") + resolved[component] = _parse_enum(raw_transfer, DLOTransfer, "offload transfer") + return resolved + + +def component_uses_allgather(config: Any, component: str = DIT_COMPONENT) -> bool: + """Return whether one selected component uses AllGather transport.""" + resolved = resolve_offload(config) + if not resolved.offloads(component): + raise ValueError(f"Offload component {component!r} is not selected") + return resolved.uses_allgather(component) + + +def selected_offload_components(config: Any) -> frozenset[str]: + """Resolve selected components while preserving legacy topology defaults.""" + return resolve_offload(config).components + + +def should_offload_component(config: Any, component: str) -> bool: + """Return whether an active layer policy selects ``component``.""" + if component not in OFFLOAD_COMPONENTS: + raise ValueError(f"Unknown offload component: {component}") + resolved = resolve_offload(config) + if resolved.strategy not in { + OffloadStrategy.LAYER_WISE, + OffloadStrategy.DISTRIBUTED_LAYER_WISE, + }: + return False + return resolved.offloads(component) + + +def any_selected_component_uses_allgather(config: Any) -> bool: + """Return whether an enabled layer backend requires weight collectives.""" + return resolve_offload(config).any_allgather diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index 7063c2759..e626d87a5 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -20,6 +20,7 @@ from __future__ import annotations import threading import time +import weakref from itertools import chain from typing import Any @@ -35,8 +36,17 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import ( from vllm_omni.host_weight_runtime import HostWeightLease from vllm_omni.platforms import current_omni_platform -from .base import OffloadBackend, OffloadConfig -from .block_discovery import get_blocks_from_dit +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, + prepare_pipeline_components, + set_encoder_layerwise_state, +) +from .config import DIT_COMPONENT, TEXT_ENCODER_COMPONENT from .host_registration import ( HostRegistration, HostRegistrationCleanupError, @@ -49,12 +59,21 @@ from .offload_plan import ( get_offload_plan, ) from .tensor_utils import ( - dtype_size as _dtype_size, -) -from .tensor_utils import ( + clear_block_storage, + clear_tensor_storage, + describe_tensor_storage, + flatten_physical_storage, + group_named_tensors_by_dtype, is_materialized_tensor, make_offload_placeholder, + materialization_probe, + module_materialization_probe, + restore_tensor_storage, set_tensor_storage, + tensor_storage_metadata, +) +from .tensor_utils import ( + dtype_size as _dtype_size, ) logger = init_logger(__name__) @@ -124,8 +143,17 @@ class DistributedLayerwiseOffloadHook(ModelHook): shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] | None = None, rank_local_mmap: bool = False, tensor_transforms: dict[int, Any] | None = None, + materialization_probe_tensor: torch.Tensor | None = None, ): assert isinstance(next_block, nn.Module), "transformer block must be type `torch.nn.Module`" + if type(dp_size) is not int or dp_size < 1: + raise ValueError(f"dp_size must be a positive integer, got {dp_size!r}") + if type(rank) is not int or not 0 <= rank < dp_size: + raise ValueError(f"rank must satisfy 0 <= rank < dp_size, got rank={rank!r}, dp_size={dp_size}") + if dp_size > 1 and dp_group is None: + raise ValueError("dp_group is required when dp_size is greater than one") + if rank_local_mmap and dp_size != 1: + raise ValueError("rank_local_mmap requires dp_size=1") self.next_block = next_block self.device = device @@ -136,6 +164,7 @@ class DistributedLayerwiseOffloadHook(ModelHook): self.rank_local_mmap = rank_local_mmap self.registered_mmap = False self.tensor_transforms = tensor_transforms or {} + self._materialization_probe = materialization_probe_tensor self.copy_stream = copy_stream or current_omni_platform.Stream() self.comm_stream = comm_stream or current_omni_platform.Stream() @@ -200,9 +229,7 @@ class DistributedLayerwiseOffloadHook(ModelHook): # Shared shard (AllGather input) buffers — assigned by backend. self.gpu_shard_buffers: list[dict[torch.dtype, torch.Tensor] | None] = [None, None] - # Pending async AllGather work (prevent GC before completion). - self._pending_work: Any | None = None - self._cached_repoint: list | None = None + self._cached_repoint: tuple[tuple[Any, ...], ...] = () # ------------------------------------------------------------------ # # DTensor helpers (shared with LayerwiseOffloadHook) # @@ -213,6 +240,8 @@ class DistributedLayerwiseOffloadHook(ModelHook): self.block_parameters = dict(module.named_parameters()) self.block_buffers = dict(module.named_buffers()) + if self._materialization_probe is None: + self._materialization_probe = materialization_probe(self.block_parameters, self.block_buffers) self.next_block_parameters = dict(self.next_block.named_parameters()) self.next_block_buffers = dict(self.next_block.named_buffers()) @@ -239,37 +268,31 @@ class DistributedLayerwiseOffloadHook(ModelHook): self._allocate_device_buffers() # Cache parameter re-pointing metadata to avoid per-layer dict lookups. - self._cached_repoint = [] - for slot in range(2): - repoint = [] - for dtype, metas in self.metadata.items(): - for m in metas: - target = ( - self.next_block_parameters[m["name"]] - if m["name"] in self.next_block_parameters - else self.next_block_buffers[m["name"]] - ) - repoint.append( - ( - target, - dtype, - m["offset"], - m["numel"], - m["shape"], - m["stride"], - ) - ) - self._cached_repoint.append(repoint) + self._cached_repoint = tuple( + ( + self.next_block_parameters[m["name"]] + if m["name"] in self.next_block_parameters + else self.next_block_buffers[m["name"]], + dtype, + m["offset"], + m["numel"], + m["shape"], + m["stride"], + ) + for dtype, metas in self.metadata.items() + for m in metas + ) # Pre-compute AG output sizes (avoid sum() per layer). self._ag_output_sizes: dict[torch.dtype, int] = {} - for dtype, metas in self.metadata.items(): - total_numel = sum(m["numel"] for m in metas) - if self.rank_local_mmap: - self._ag_output_sizes[dtype] = total_numel - else: + if self.dp_size > 1: + for dtype in self.metadata: shard_numel = self.cpu_shards[dtype].numel() - self._ag_output_sizes[dtype] = shard_numel * self.dp_size if self.dp_size > 1 else total_numel + self._ag_output_sizes[dtype] = shard_numel * self.dp_size + + # Commit the storage mutation only after every allocation and metadata + # calculation above has succeeded. Until here, hook setup is retryable. + clear_tensor_storage(chain(self.next_block_parameters.values(), self.next_block_buffers.values())) return module @@ -279,7 +302,7 @@ class DistributedLayerwiseOffloadHook(ModelHook): bufs: dict[str, torch.Tensor], tensor_transforms: dict[int, Any] | None = None, ) -> tuple[dict[torch.dtype, list[dict[str, Any]]], dict[torch.dtype, list[dict[str, Any]]]]: - """Retain file-backed tensors and replace module storage with placeholders. + """Retain file-backed tensors and describe their runtime layout. The returned sources preserve safetensors mmap storage. Runtime-layout adapters are applied only while packing a bounded staging slot, so no @@ -288,7 +311,6 @@ class DistributedLayerwiseOffloadHook(ModelHook): cpu_sources: dict[torch.dtype, list[dict[str, Any]]] = {} metadata: dict[torch.dtype, list[dict[str, Any]]] = {} offsets: dict[torch.dtype, int] = {} - for name, target in chain(params.items(), bufs.items()): source = target.to_local() if hasattr(target, "to_local") else target if source.device.type != "cpu": @@ -331,10 +353,6 @@ class DistributedLayerwiseOffloadHook(ModelHook): ) offsets[dtype] = offset + storage_numel - # The detached source above keeps the mmap storage alive while the - # module parameter/buffer is rebound to the rotating device slot. - set_tensor_storage(target, make_offload_placeholder(target)) - return cpu_sources, metadata @staticmethod @@ -351,40 +369,14 @@ class DistributedLayerwiseOffloadHook(ModelHook): Each rank stores only ``1/dp_size`` of the total weights. The full tensor is reconstructed at runtime via AllGather. """ - dtype_grouped: dict[torch.dtype, dict[str, torch.Tensor]] = {} dtype_metadata: dict[torch.dtype, list[dict[str, Any]]] = {} - - for name, param_or_buf in chain(params.items(), bufs.items()): - dtype = param_or_buf.dtype - if dtype not in dtype_grouped: - dtype_grouped[dtype] = {} - dtype_grouped[dtype][name] = param_or_buf - cpu_shards: dict[torch.dtype, torch.Tensor] = {} - for dtype, name2weights in dtype_grouped.items(): - # Resolve local tensors (handle DTensor via to_local) - weights_with_local = [] - for name, t in name2weights.items(): - local_t = t.to_local() if hasattr(t, "to_local") else t - mmap_transform = (tensor_transforms or {}).get(id(t)) - if callable(mmap_transform): - # Some checkpoints use a layout that is converted by the - # regular weight loader (for example MiniMax-H3 grouped - # QKV). Apply that conversion one block at a time while - # copying the rank-local CPU shard. Keeping the raw - # parameter as an mmap view avoids a private full-model - # copy in every worker. - local_t = mmap_transform(local_t) - stride = local_t.stride() - storage_numel = ( - 0 - if local_t.numel() == 0 - else 1 + sum((size - 1) * axis_stride for size, axis_stride in zip(local_t.shape, stride)) - ) - weights_with_local.append((name, t, local_t, storage_numel, stride)) - - total_numel = sum(storage_numel for _, _, _, storage_numel, _ in weights_with_local) + for dtype, named_weights in group_named_tensors_by_dtype(params, bufs).items(): + # Apply loader-declared layout conversions block by block while + # preserving the physical tensor layout used by the kernels. + specs = describe_tensor_storage(named_weights, tensor_transforms) + total_numel = sum(spec.storage_numel for spec in specs) # Equal-sized shards (ceil division) for all_gather_into_tensor shard_size = (total_numel + dp_size - 1) // dp_size # ceil @@ -393,69 +385,31 @@ class DistributedLayerwiseOffloadHook(ModelHook): # Allocate ONLY the shard (1/dp_size), zero-padded to ceil. # Avoids materialising the full block on CPU. - shard = torch.zeros(shard_size, dtype=dtype, device="cpu") + shard = torch.zeros( + shard_size, + dtype=dtype, + device="cpu", + pin_memory=pin_memory, + ) current_offset = 0 - for ( - name, - original_tensor, - local_tensor, - storage_numel, - stride, - ) in weights_with_local: - if dtype not in dtype_metadata: - dtype_metadata[dtype] = [] + for spec in specs: # Offsets remain relative to the FULL flattened buffer # (needed for correct AllGather reconstruction). - dtype_metadata[dtype].append( - { - "name": name, - "offset": current_offset, - "numel": storage_numel, - "shape": local_tensor.shape, - "stride": stride, - } - ) + dtype_metadata.setdefault(dtype, []).append(tensor_storage_metadata(spec, current_offset)) # Copy ONLY the portion within [shard_start, shard_end) overlap_start = max(current_offset, shard_start) - overlap_end = min(current_offset + storage_numel, shard_end) + overlap_end = min(current_offset + spec.storage_numel, shard_end) if overlap_start < overlap_end: - if local_tensor.is_contiguous(): - flat_storage = local_tensor.flatten() - else: - # Online FP8 stores Cutlass weights as transposed views - # (e.g. stride=(1, K)). Flattening such a tensor in - # logical order and later rebuilding it with .view() - # changes its layout and makes scaled_mm reject it. - # Pack the physical storage order and preserve the - # original stride for zero-copy reconstruction. - flat_storage = torch.zeros( - storage_numel, - dtype=dtype, - device=local_tensor.device, - ) - physical_view = torch.as_strided( - flat_storage, - size=local_tensor.shape, - stride=stride, - ) - physical_view.copy_(local_tensor) + flat_storage = flatten_physical_storage(spec.value, spec.storage_numel) src_start = overlap_start - current_offset src_end = overlap_end - current_offset dst_start = overlap_start - shard_start dst_end = overlap_end - shard_start shard[dst_start:dst_end].copy_(flat_storage[src_start:src_end]) - # Replace original tensor with placeholder (frees CPU storage) - set_tensor_storage( - original_tensor, - make_offload_placeholder(original_tensor), - ) - current_offset += storage_numel - - if pin_memory: - shard = shard.pin_memory() + current_offset += spec.storage_numel cpu_shards[dtype] = shard @@ -478,9 +432,7 @@ class DistributedLayerwiseOffloadHook(ModelHook): @property def is_materialized(self) -> bool: """Check whether this block's parameters hold real data on device.""" - for param in self.block_parameters.values(): - return is_materialized_tensor(param) - return True + return self._materialization_probe is None or is_materialized_tensor(self._materialization_probe) # ------------------------------------------------------------------ # # Prefetch: H2D + AllGather (overlapped on dedicated streams) # @@ -651,7 +603,7 @@ class DistributedLayerwiseOffloadHook(ModelHook): self._shared_slot_group[slot] = self._group_id # Re-point using cached metadata (avoids per-layer dict lookups). - for target, dtype, offset, numel, shape, stride in self._cached_repoint[slot]: + for target, dtype, offset, numel, shape, stride in self._cached_repoint: set_tensor_storage( target, torch.as_strided( @@ -661,8 +613,8 @@ class DistributedLayerwiseOffloadHook(ModelHook): ), ) - def get_weights(self, slot: int) -> dict[torch.dtype, torch.Tensor] | None: - """Wait for AllGather completion and return full weights for the slot. + def _wait_for_weights(self, slot: int) -> None: + """Wait until the slot holding this block is ready for compute. The ready event for this slot was set by the *previous* hook's prefetch_layer (which prefetched THIS block's weights into the @@ -676,7 +628,6 @@ class DistributedLayerwiseOffloadHook(ModelHook): evt = self._prev_hook.ready_events[slot] if evt is not None: current_omni_platform.current_stream().wait_event(evt) - return self.gpu_buffers[slot] # ------------------------------------------------------------------ # # Offload: free device memory for current block # @@ -685,15 +636,54 @@ class DistributedLayerwiseOffloadHook(ModelHook): @torch.compiler.disable def offload_layer(self) -> None: """Free GPU memory for current block by replacing tensors with placeholders.""" - evt = self._prefetch_done - if evt is not None: - current_omni_platform.current_stream().wait_event(evt) + clear_block_storage(self.block_parameters, self.block_buffers, self._prefetch_done) self._prefetch_done = None - for _, param in self.block_parameters.items(): - set_tensor_storage(param, make_offload_placeholder(param)) - for _, buf in self.block_buffers.items(): - set_tensor_storage(buf, make_offload_placeholder(buf)) + @torch.compiler.disable + def restore_next_block_to_cpu(self) -> None: + """Restore hook-owned master weights before removing the hook. + + Every circular hook owns the host backing for its ``next_block``. + Dropping that hook while the module points at rotating device buffers + or placeholders would make a later enable shard invalid tensors. + """ + + if self.rank_local_mmap: + for dtype, metas in self.metadata.items(): + for source_info, meta in zip(self.cpu_sources[dtype], metas, strict=True): + source = self._resolve_mmap_source(source_info, meta, dtype) + target = ( + self.next_block_parameters[meta["name"]] + if meta["name"] in self.next_block_parameters + else self.next_block_buffers[meta["name"]] + ) + restore_tensor_storage(target, source, device="cpu") + return + + if self.dp_size <= 1: + for dtype, metas in self.metadata.items(): + flat = self.cpu_shards[dtype] + for meta in metas: + value = torch.as_strided( + flat[meta["offset"] : meta["offset"] + meta["numel"]], + size=meta["shape"], + stride=meta["stride"], + ) + target = ( + self.next_block_parameters[meta["name"]] + if meta["name"] in self.next_block_parameters + else self.next_block_buffers[meta["name"]] + ) + restore_tensor_storage(target, value, device="cpu") + return + + if self.dp_group is None: + raise RuntimeError("Cannot restore distributed offload weights without the DLO process group") + self.prefetch_layer(0, non_blocking=False) + current_omni_platform.synchronize() + for target in chain(self.next_block_parameters.values(), self.next_block_buffers.values()): + local = target.to_local() if hasattr(target, "to_local") else target + restore_tensor_storage(target, local, device="cpu") # ------------------------------------------------------------------ # # ModelHook interface # @@ -720,11 +710,11 @@ class DistributedLayerwiseOffloadHook(ModelHook): # Another group (or no group) wrote to our slot — re-fetch self._prev_hook.prefetch_layer(self.current_slot, non_blocking=False) # Always wait for data to be ready (handles both sync and async paths) - self._prev_hook.get_weights(self.current_slot) + self._prev_hook._wait_for_weights(self.current_slot) elif not self.is_materialized and self._prev_hook is not None: # Previous hook was skipped (e.g. by cache-dit), sync-prefetch self._prev_hook.prefetch_layer(self.current_slot, non_blocking=False) - self._prev_hook.get_weights(self.current_slot) + self._prev_hook._wait_for_weights(self.current_slot) # Prefetch next layer into the other slot (overlapped with compute). # No explicit get_weights() here — offload_layer() in the previous @@ -760,6 +750,7 @@ def apply_distributed_block_hook( shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] | None = None, rank_local_mmap: bool = False, tensor_transforms: dict[int, Any] | None = None, + materialization_probe_tensor: torch.Tensor | None = None, ) -> DistributedLayerwiseOffloadHook: """Register a DistributedLayerwiseOffloadHook on *module*.""" registry = HookRegistry.get_or_create(module) @@ -775,6 +766,7 @@ def apply_distributed_block_hook( shared_buffers=shared_buffers, rank_local_mmap=rank_local_mmap, tensor_transforms=tensor_transforms, + materialization_probe_tensor=materialization_probe_tensor, ) registry.register_hook(DistributedLayerwiseOffloadHook._HOOK_NAME, hook) return hook @@ -867,12 +859,20 @@ class PinnedResidentLayerGroup: for _ in range(2): buffers = {} for dtype, total in max_sizes.items(): - buffer = torch.empty(total, dtype=dtype, device="cpu") - if pin_memory: - buffer = buffer.pin_memory() + buffer = torch.empty( + total, + dtype=dtype, + device="cpu", + pin_memory=pin_memory, + ) buffers[dtype] = buffer self._cpu_staging_buffers.append(buffers) + # All host masters and optional staging buffers are now ready. Clear + # model storage as one final constructor commit so an earlier failure + # leaves every resident block intact. + clear_tensor_storage(target for state in self._states for target in state["targets"].values()) + def load(self) -> None: if self.loaded: return @@ -960,6 +960,34 @@ class PinnedResidentLayerGroup: self._gpu_buffers.clear() self.loaded = False + def restore_to_cpu(self) -> None: + """Materialize the persistent host masters back into the module. + + Stage-scoped ``offload()`` deliberately leaves placeholders in the + module while this group owns the CPU backing. ``disable()`` discards + the group, so it must first restore ordinary CPU tensors to make a + later enable cycle safe. + """ + self.offload() + + for state in self._states: + targets = state["targets"] + for dtype, metas in state["metadata"].items(): + if self.rank_local_mmap: + for source_info, meta in zip(state["cpu_sources"][dtype], metas, strict=True): + source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype) + restore_tensor_storage(targets[meta["name"]], source, device="cpu") + continue + + flat = state["cpu_shards"][dtype] + for meta in metas: + source = torch.as_strided( + flat[meta["offset"] : meta["offset"] + meta["numel"]], + size=meta["shape"], + stride=meta["stride"], + ) + restore_tensor_storage(targets[meta["name"]], source, device="cpu") + # ---------------------------------------------------------------------- # # Backend # @@ -995,6 +1023,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._all_hook_groups: list[list[DistributedLayerwiseOffloadHook]] = [] self._resident_blocks: list[nn.Module] = [] self._resident_layer_group: PinnedResidentLayerGroup | None = None + self._residency_pipeline_ref: weakref.ReferenceType[nn.Module] | None = None + self._encoder_modules: list[nn.Module] = [] + self._staged_components: list[nn.Module] = [] self._using_mmap = False self._using_rank_local_mmap = False self._using_registered_mmap = False @@ -1002,6 +1033,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._host_weight_lease: HostWeightLease | None = None self._host_registration: HostRegistration | None = None self._mmap_transforms_by_tensor_id: dict[int, Any] = {} + self._poisoned_reason: str | None = None def load_resident_layers(self) -> None: """Load the model-declared leading blocks for the denoise stage.""" @@ -1013,6 +1045,13 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): if self._resident_layer_group is not None: self._resident_layer_group.offload() + def _clear_residency_controller(self) -> None: + pipeline_ref = self._residency_pipeline_ref + self._residency_pipeline_ref = None + pipeline = None if pipeline_ref is None else pipeline_ref() + if pipeline is not None and getattr(pipeline, "_dlo_residency_controller", None) is self: + pipeline._dlo_residency_controller = None + def _rank_local_source_tensors( self, hooks: list[DistributedLayerwiseOffloadHook], @@ -1384,69 +1423,121 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): coord.ranks, ) - def _register_on_demand_hook(self, module: nn.Module, label: str, *, stage_on_demand: bool = False) -> None: - """Prepare a pipeline-managed stage component or keep it resident. + def _component_transport( + self, + component: str, + ) -> tuple[torch.distributed.ProcessGroup | None, int, int]: + if self.config.uses_allgather(component): + return self.dp_group, self.dp_size, self.rank + return None, 1, 0 - Components that expose an explicit stage lifecycle are initially - offloaded and loaded by their pipeline only around encode/decode. - Other models retain the conservative resident behavior because a - generic post-forward hook can disrupt the DiT prefetch streams. - """ - offload_to_cpu = getattr(module, "offload_to_cpu", None) - if stage_on_demand and callable(offload_to_cpu): - offload_to_cpu() - logger.info("Prepared %s (%s) for pipeline-managed staged offload", label, module.__class__.__name__) - return - module.to(self.device) - logger.info("Moved %s (%s) to GPU (resident)", label, module.__class__.__name__) + def _has_multirank_allgather(self) -> bool: + components = self.config.components or frozenset({DIT_COMPONENT}) + return self.dp_size > 1 and any(self.config.uses_allgather(component) for component in components) - def _try_layerwise_offload_encoder(self, module: nn.Module, name: str, plan: OffloadPlan | None) -> bool: - """Stream plan-declared encoder blocks on each rank without AllGather.""" - if plan is None or name not in plan.encoder_block_attrs: - return False - if getattr(module, "_omni_layerwise_enabled", False): - return True + def _install_hook_group( + self, + blocks: list[nn.Module] | nn.ModuleList, + component: str, + *, + use_dit_mmap: bool = False, + ) -> list[DistributedLayerwiseOffloadHook]: + """Install one circular block ring with the shared slot protocol.""" + block_list = list(blocks) + if len(block_list) <= 1: + raise ValueError("A distributed layerwise hook group requires at least two blocks") - from operator import attrgetter - - from vllm_omni.diffusion.offloader.layerwise_backend import apply_block_hook - - hooks = [] - block_groups = [] - copy_stream = current_omni_platform.Stream() - for block_path in plan.encoder_block_attrs[name]: - try: - blocks = attrgetter(block_path)(module) - except AttributeError: - logger.warning("Encoder offload path %s.%s was not found", name, block_path) - continue - if not isinstance(blocks, nn.ModuleList) or len(blocks) <= 1: - logger.warning("Encoder offload path %s.%s is not a streamable block list", name, block_path) - continue - group_hooks = [ - apply_block_hook(blocks[-1], blocks[0], self.device, copy_stream, self.config.pin_cpu_memory) - ] - group_hooks.extend( - apply_block_hook(block, blocks[index + 1], self.device, copy_stream, self.config.pin_cpu_memory) - for index, block in enumerate(blocks[:-1]) + group, group_size, group_rank = self._component_transport(component) + hooks: list[DistributedLayerwiseOffloadHook] = [] + self._all_hook_groups.append(hooks) + self._blocks.append(block_list) + probes = {id(block): module_materialization_probe(block) for block in block_list} + for block, next_block in zip( + chain((block_list[-1],), block_list[:-1]), + block_list, + strict=True, + ): + hooks.append( + apply_distributed_block_hook( + block, + next_block, + self.device, + group, + group_size, + group_rank, + self.copy_stream, + self.comm_stream, + self.config.pin_cpu_memory, + shared_buffers=[None, None], + rank_local_mmap=self._using_rank_local_mmap if use_dit_mmap else False, + tensor_transforms=self._mmap_transforms_by_tensor_id if use_dit_mmap else None, + materialization_probe_tensor=probes[id(block)], + ) ) - for index, hook in enumerate(group_hooks): - hook._prev_hook = group_hooks[index - 1] - hooks.extend(group_hooks) - block_groups.append(blocks) - if not hooks: + # hooks = [last -> first, block0 -> block1, ...]. Alternating slots + # keep a prefetch from overwriting the current block for any ring size. + for index, hook in enumerate(hooks): + hook._prev_hook = hooks[index - 1] + hook.current_slot = index % 2 + 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. + + 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. + """ + if plan is None: return False - # The component lifecycle uses these generic attributes to keep only - # non-block encoder state resident during the encode phase. - module._omni_layerwise_hooks = hooks - module._omni_layerwise_block_groups = block_groups - module._omni_layerwise_enabled = True + + 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, + ) + if not block_groups: + return False + + 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: + move_non_block_state_to_device(module, block_groups, self.device) + set_encoder_layerwise_state( + module, + encoder_hooks, + block_groups, + ) logger.info( - "Enabled rank-local layerwise offload for encoder %s (%d blocks across %d stacks)", + "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, sum(len(blocks) for blocks in block_groups), len(block_groups), + group_size, ) return True @@ -1494,14 +1585,13 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): if blocks is None: return False - num_blocks = len(blocks) logger.info( - "Distributed layerwise offload for submodule '%s.%s' (%d blocks, %.0f MB total, dp_size=%d)", + "Distributed layerwise offload for submodule '%s.%s' (%d blocks, %.0f MB total, group_size=%d)", name, blocks_attr, - num_blocks, + len(blocks), sum(p.nelement() * p.element_size() for p in module.parameters()) / 1048576, - self.dp_size, + self._component_transport(DIT_COMPONENT)[1], ) # Move non-block parts of the submodule to GPU (small: embeddings, norms) @@ -1509,66 +1599,12 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): if child_name != blocks_attr: child.to(self.device) - # Apply distributed hooks (1/4 sharding + AllGather, same as DiT) - # Pass shared_buffers=[None,None] to defer per-hook allocation - # (prevents OOM on large models where N hooks × 2 buffers >> HBM) - last_block, first_block = blocks[-1], blocks[0] - last_hook = apply_distributed_block_hook( - last_block, - first_block, - self.device, - self.dp_group, - self.dp_size, - self.rank, - self.copy_stream, - self.comm_stream, - self.config.pin_cpu_memory, - shared_buffers=[None, None], - rank_local_mmap=self._using_rank_local_mmap, - tensor_transforms=self._mmap_transforms_by_tensor_id, - ) - sub_hooks = [last_hook] - for i, block in enumerate(blocks[:-1]): - next_block = blocks[(i + 1) % num_blocks] - hook = apply_distributed_block_hook( - block, - next_block, - self.device, - self.dp_group, - self.dp_size, - self.rank, - self.copy_stream, - self.comm_stream, - self.config.pin_cpu_memory, - shared_buffers=[None, None], - rank_local_mmap=self._using_rank_local_mmap, - tensor_transforms=self._mmap_transforms_by_tensor_id, - ) - sub_hooks.append(hook) - - # Wire backward references + slot alternation - for i in range(len(sub_hooks)): - sub_hooks[i]._prev_hook = sub_hooks[i - 1] - # Assign slots in list order: sub_hooks = [last_hook, block0, ..., blockN-2] - # This ensures last_hook.current_slot != block0_hook.current_slot, - # so the circular prefetch (last_hook -> block0) writes to a - # different slot than block0 reads from. Correct for ALL N. - for i, hook in enumerate(sub_hooks): - hook.current_slot = i % 2 - # Mark block0_hook (index 1) as group-first — it must sync-prefetch - # on entry because another group may have overwritten the shared slot. - if len(sub_hooks) > 1: - sub_hooks[1]._is_group_first = True - - # Defer buffer allocation and prefetch to enable() unified allocation - self._all_hook_groups.append(sub_hooks) - self._blocks.append(blocks) + self._install_hook_group(blocks, DIT_COMPONENT, use_dit_mmap=True) return True def _prepare_dit_non_block_modules( self, dit_module: nn.Module, - dit_name: str, blocks_attr_names: list[str], all_dit_modules: set[int], plan: OffloadPlan | None, @@ -1579,7 +1615,6 @@ 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. """ - _ON_DEMAND_THRESHOLD = _ON_DEMAND_THRESHOLD_MB for name, module in dit_module.named_children(): if name in blocks_attr_names: logger.debug("Skipped blocks module %s", name) @@ -1593,13 +1628,20 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): / 1048576 ) explicitly_planned = plan is not None and name in plan.offload_submodules - if explicitly_planned or module_mb > _ON_DEMAND_THRESHOLD: + if explicitly_planned or module_mb > _ON_DEMAND_THRESHOLD_MB: if id(module) in all_dit_modules: logger.info("Submodule '%s' is already a DiT module, skipping layerwise offload", name) elif self._try_layerwise_offload_submodule(module, name, plan): pass else: - self._register_on_demand_hook(module, name) + prepare_component( + module, + name, + device=self.device, + stage_on_demand=True, + blockwise=False, + staged_components=self._staged_components, + ) continue try: @@ -1608,23 +1650,21 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): # Non-persistent buffers such as RoPE frequencies do not # exist in the checkpoint and must be reconstructed. has_meta_buffer = any(getattr(buffer, "is_meta", False) for buffer in module.buffers(recurse=True)) - if has_meta_buffer: - saved_params = { - param_name: param.data.clone() - for param_name, param in module.named_parameters() - if not getattr(param, "is_meta", False) - } - module.to_empty(device=self.device) - for submodule in module.modules(): - if hasattr(submodule, "reset_parameters"): - submodule.reset_parameters() - for param_name, param in module.named_parameters(): - if param_name in saved_params: - param.data.copy_(saved_params[param_name]) - try: - module.to(self.device) - except Exception: - logger.warning("Module %s still has meta params after mmap load", name) + if not has_meta_buffer: + raise + saved_params = { + param_name: param.data.clone() + for param_name, param in module.named_parameters() + if not getattr(param, "is_meta", False) + } + module.to_empty(device=self.device) + for submodule in module.modules(): + if hasattr(submodule, "reset_parameters"): + submodule.reset_parameters() + for param_name, param in module.named_parameters(): + if param_name in saved_params: + param.data.copy_(saved_params[param_name]) + module.to(self.device) for param in dit_module._parameters.values(): if param is not None and not getattr(param, "is_meta", False): @@ -1635,13 +1675,20 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): def enable(self, pipeline: nn.Module) -> None: """Enable DLO and make partial startup failures transactional.""" + if self._poisoned_reason is not None: + raise RuntimeError(self._poisoned_reason) + + # A rank-local rollback can reconstruct ordinary tensors from its host + # masters. A multi-rank AllGather rollback is unsafe: another rank may + # never enter the matching restoration collective after startup fails. + restore_allgather_weights = not self._has_multirank_allgather() try: self._enable(pipeline) except BaseException: try: - self.disable() + self._disable(restore_allgather_weights=restore_allgather_weights) except BaseException: - logger.exception("DLO cleanup failed while handling an enable failure") + logger.exception("DistributedLayerwiseOffloadBackend cleanup failed while handling an enable failure") raise def _enable(self, pipeline: nn.Module) -> None: @@ -1649,42 +1696,48 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): logger.warning("DistributedLayerwiseOffloadBackend already enabled") return - self._on_demand_shard_infos: list[dict] = [] - self._on_demand_handles: list[Any] = [] - # Initialize DP group (if not already done by early init) - if self.dp_group is None and self.dp_size > 1: + if self.dp_group is None and self._has_multirank_allgather(): self._init_dp_group() modules = ModuleDiscovery.discover(pipeline) - if not modules.dits: + if not modules.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" ) - logger.warning("No DiT/transformer modules found, skipping distributed layer-wise offloading") - return + 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 and (plan is None or not plan.resident_dit_paths): - logger.warning( - "dlo_resident_layers=%d was requested, but this model declares no " - "resident_dit_paths; all blocks will be streamed.", - self.config.dlo_resident_layers, - ) + 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) # Storage selection belongs to the loader. DLO consumes the exact # prevalidated plan that caused the loader to skip materialization; # without a plan, all weights must already come from the ordinary # loader. The transfer protocol is selected independently below. - self._using_mmap = self.host_weight_plan is not None - self._using_rank_local_mmap = self._using_mmap and self.dp_size <= 1 - if self._using_mmap: - if self.host_weight_plan.backing_kind == "host_weight_runtime": - carrier = self.host_weight_plan.lease_carrier + host_weight_plan = self.host_weight_plan + self._using_mmap = host_weight_plan is not None + # A one-rank AllGather transport is rank-local in practice. Preserve + # the mmap source as the host master instead of eagerly closing it. + self._using_rank_local_mmap = self._using_mmap and self._component_transport(DIT_COMPONENT)[1] <= 1 + if host_weight_plan is not None: + if host_weight_plan.backing_kind == "host_weight_runtime": + carrier = host_weight_plan.lease_carrier if carrier is None: raise RuntimeError("DLO received a Host Weight Runtime plan without a lease carrier") self._host_weight_lease = carrier.take() @@ -1699,14 +1752,14 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): "DLO consuming final-layout Host Weight Runtime lease %s", self._host_weight_lease.provenance.resolution_id, ) - elif self.host_weight_plan.backing_kind == "checkpoint_mmap": + elif host_weight_plan.backing_kind == "checkpoint_mmap": self._load_weights_via_mmap( pipeline, modules, - self.host_weight_plan, + host_weight_plan, ) else: - raise ValueError(f"Unsupported DLO host-weight backing: {self.host_weight_plan.backing_kind}") + raise ValueError(f"Unsupported DLO host-weight backing: {host_weight_plan.backing_kind}") if self._using_rank_local_mmap: logger.info( "DLO rank-local host storage enabled: source pages are " @@ -1728,51 +1781,28 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): ) logger.info("DLO is using host tensors materialized by the ordinary loader") - # Keep VAE/encoders on CPU; move to GPU on-demand via hooks. - # This saves ~4.3 GB HBM per card (VAE 1.3 + encoder 1.1 + sound 1.9) - # during the DiT forward pass. They are only needed briefly for - # text-encoding (before DiT) and VAE-decode (after DiT). - for enc, enc_name in zip(modules.encoders, modules.encoder_names): - self._try_layerwise_offload_encoder(enc, enc_name, plan) - self._register_on_demand_hook( - enc, "encoder", stage_on_demand=plan is not None and enc_name in plan.on_demand_component_paths - ) - for vae, vae_name in zip(modules.vaes, modules.vae_names): - self._register_on_demand_hook( - vae, - "vae", - stage_on_demand=(plan is not None and vae_name in plan.on_demand_component_paths), - ) + # Apply each selected encoder transfer while keeping explicit VAEs and + # unselected components resident. + prepare_pipeline_components( + modules, + self.config, + plan, + device=self.device, + staged_components=self._staged_components, + enable_encoder_blocks=self._try_layerwise_offload_encoder, + ) - # Move resident modules to GPU (small modules needed every forward) - for name, module in zip(modules.resident_names, modules.resident_modules): - try: - module.to(self.device) - except Exception as exc: - logger.debug("Failed to move resident module %s to GPU: %s", name, exc) - - logger.info("Applying distributed layer-wise offloading on %s", modules.dit_names) + if self.config.offloads(DIT_COMPONENT): + logger.info("Applying distributed layer-wise offloading on %s", modules.dit_names) # 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) # Apply hooks for each DiT module - for i, dit_module in enumerate(modules.dits): - dit_name = modules.dit_names[i] - logger.info(f"Applying hooks on {dit_name} ({dit_module.__class__.__name__})") - - blocks_attr_names, blocks = get_blocks_from_dit(dit_module) - - if not blocks: - logger.warning( - "Target layers (blocks) not found. Skipping offloading on %s (%s)", - dit_name, - dit_module.__class__.__name__, - ) - dit_module.to(self.device) - continue - + 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)) @@ -1789,7 +1819,6 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._prepare_dit_non_block_modules( dit_module, - dit_name, blocks_attr_names, all_dit_modules, plan, @@ -1800,6 +1829,12 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): 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, @@ -1808,66 +1843,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._resident_blocks.extend(blocks) continue - # Register hooks in a circular sliding window: - # last block prefetches first block, block i prefetches block (i+1) - # All hooks share 2 global device buffers (RFC: "exactly two layers on device") - last_block, first_block = blocks[-1], blocks[0] - - # Pass 1: create hooks with deferred buffer allocation - # (shared_buffers=[None,None] prevents per-hook OOM on large models) - last_hook = apply_distributed_block_hook( - last_block, - first_block, - self.device, - self.dp_group, - self.dp_size, - self.rank, - self.copy_stream, - self.comm_stream, - self.config.pin_cpu_memory, - shared_buffers=[None, None], - rank_local_mmap=self._using_rank_local_mmap, - tensor_transforms=self._mmap_transforms_by_tensor_id, - ) - - block_hooks: list[DistributedLayerwiseOffloadHook] = [last_hook] - for i, block in enumerate(blocks[:-1]): - next_block = blocks[(i + 1) % num_blocks] - hook = apply_distributed_block_hook( - block, - next_block, - self.device, - self.dp_group, - self.dp_size, - self.rank, - self.copy_stream, - self.comm_stream, - self.config.pin_cpu_memory, - shared_buffers=[None, None], - rank_local_mmap=self._using_rank_local_mmap, - tensor_transforms=self._mmap_transforms_by_tensor_id, - ) - block_hooks.append(hook) - - # Wire backward references for cache-dit fallback - for i in range(len(block_hooks)): - block_hooks[i]._prev_hook = block_hooks[i - 1] - - # Assign slots in list order: block_hooks = [last_hook, block0, ..., blockN-2] - # This ensures last_hook.current_slot != block0_hook.current_slot, - # so the circular prefetch (last_hook -> block0) writes to a - # different slot than block0 reads from. Correct for ALL N. - for i, hook in enumerate(block_hooks): - hook.current_slot = i % 2 - - # Mark block0_hook (index 1) as group-first - if len(block_hooks) > 1: - block_hooks[1]._is_group_first = True - - # Defer buffer allocation — collected for unified allocation below - self._all_hook_groups.append(block_hooks) - self._blocks.append(blocks) - + self._install_hook_group(blocks, DIT_COMPONENT, use_dit_mmap=True) if self._resident_blocks: self._resident_layer_group = PinnedResidentLayerGroup( self._resident_blocks, @@ -1879,28 +1855,34 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): tensor_transforms=self._mmap_transforms_by_tensor_id, ) pipeline._dlo_residency_controller = self + self._residency_pipeline_ref = weakref.ref(pipeline) all_hooks = [hook for group in self._all_hook_groups for hook in group] self._configure_hwr_transfer(all_hooks) if not self._all_hook_groups: - self.enabled = bool(self._resident_blocks) + 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 # sized to the max block across ALL module groups (gen_layers + # language_model). Groups execute sequentially, so 2 buffers suffice. unified_buffers = self._allocate_shared_buffers(all_hooks) - unified_shard_buffers = None - if self.dp_size > 1: - unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks) + allgather_hooks = [hook for hook in all_hooks if hook.dp_size > 1] + mmap_hooks = [hook for hook in all_hooks if hook.rank_local_mmap] + unified_shard_buffers = self._allocate_shared_shard_buffers(allgather_hooks) if allgather_hooks else None unified_cpu_staging = None cpu_staging_events = None if self._using_rank_local_mmap and not self._using_registered_mmap: unified_cpu_staging = self._allocate_shared_cpu_staging_buffers( - all_hooks, + mmap_hooks, self._resident_layer_group, ) cpu_staging_events = [None, None] @@ -1921,32 +1903,31 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): for hook in group: hook.gpu_buffers = unified_buffers hook._owns_buffers = False - if unified_shard_buffers is not None: + if unified_shard_buffers is not None and hook.dp_size > 1: hook.gpu_shard_buffers = unified_shard_buffers - if unified_cpu_staging is not None and cpu_staging_events is not None: + if hook.rank_local_mmap and unified_cpu_staging is not None and cpu_staging_events is not None: hook.cpu_staging_buffers = unified_cpu_staging hook.cpu_staging_events = cpu_staging_events hook._group_id = group_idx hook._shared_slot_group = shared_slot_group - # Prefetch first block of the FIRST module group only. - # Subsequent groups share the same 2 device buffers; prefetching - # them now would overwrite the first group's data in the shared - # buffer (both groups default to slot 0). Instead, subsequent - # groups' first blocks remain as meta placeholders, and their - # pre_forward will sync-prefetch on-demand via the is_materialized - # check — by which point the first group's forward has completed - # and its buffer slots are free. - if self._all_hook_groups: - group = self._all_hook_groups[0] - first_slot = group[0].current_slot - group[-1].prefetch_layer(slot=first_slot, non_blocking=False) - group[-1].get_weights(first_slot) + # Defer every group's first prefetch until its first forward. DP ranks + # may intentionally own different rank-local groups (for example only + # rank 0 owns the text encoder) while sharing DiT AllGather groups. + # Entering a collective for each rank's locally first group here would + # give the ranks different collective orders and deadlock startup. The + # group-first hook treats the initially unowned slot as contaminated + # and performs the required synchronous prefetch on first use. total_blocks = sum(len(b) for b in self._blocks) + transfer_summary = ", ".join( + f"{component}: {self.config.transfer_for(component).value}" + for component in sorted(self.config.components or {DIT_COMPONENT}) + ) logger.info( f"Distributed layer-wise offloading enabled on {total_blocks} blocks " - f"across {len(self._all_hook_groups)} group(s), dp_size={self.dp_size}, " + f"across {len(self._all_hook_groups)} group(s), " + f"transfers={{{transfer_summary}}}, " f"unified shared_buffers=2" ) @@ -1958,23 +1939,6 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): # the node-shared host master until disable(). self._release_mmap_handles() - # Assign GPU buffers to sharded on-demand modules (VAE/encoders). - # Each module gets a dedicated input (shard-sized) and output - # (full-sized) buffer. VAE/encoders run before/after DiT, never - # concurrently, so peak HBM = max(DiT, VAE) not sum. - for si in self._on_demand_shard_infos: - out_bufs: dict[torch.dtype, torch.Tensor] = {} - in_bufs: dict[torch.dtype, torch.Tensor] = {} - for dtype, shard in si["cpu_shards"].items(): - # AllGather output = dp_size * shard_size - full_size = shard.numel() * self.dp_size if self.dp_size > 1 else shard.numel() - out_bufs[dtype] = torch.empty(full_size, dtype=dtype, device=self.device) - in_bufs[dtype] = torch.empty(shard.shape, dtype=dtype, device=self.device) - si["gpu_output"] = out_bufs - si["gpu_input"] = in_bufs - _mb = sum(t.nelement() * t.element_size() for t in out_bufs.values()) / 1048576 - logger.info("Allocated %.0f MB GPU buffer for sharded on-demand module", _mb) - self._cleanup_after_loading() def _release_mmap_handles(self) -> None: @@ -2007,7 +1971,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): except Exception: pass - def disable(self) -> None: + def _disable(self, *, restore_allgather_weights: bool) -> None: has_open_lease = self._host_weight_lease is not None and not self._host_weight_lease.closed has_registration = self._host_registration is not None has_carrier = ( @@ -2018,8 +1982,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): has_partial_hooks = bool( self._blocks or self._all_hook_groups - or getattr(self, "_on_demand_handles", ()) + or self._encoder_modules + or self._staged_components or self._resident_layer_group is not None + or self._residency_pipeline_ref is not None ) if ( not self.enabled @@ -2031,30 +1997,117 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): ): return - if self._using_rank_local_mmap or has_registration: - current_omni_platform.synchronize() + self._clear_residency_controller() - for blocks in self._blocks: - for block in blocks: - remove_distributed_block_hook(block) + # A hook can leave the circular tail prefetch queued after the final + # forward. Drain every transport before releasing hook-owned host or + # device buffers, including the ordinary rank-local path. + sync_error = run_cleanup_steps([("synchronizing pending DLO transfers", current_omni_platform.synchronize)]) - for h in getattr(self, "_on_demand_handles", []): - h.remove() - self._on_demand_handles = [] - self._on_demand_shard_infos = [] + unique_hooks: list[DistributedLayerwiseOffloadHook] = [] + seen_blocks: set[int] = set() + for hook in chain.from_iterable(self._all_hook_groups): + block_id = id(hook.next_block) + if block_id not in seen_blocks: + seen_blocks.add(block_id) + unique_hooks.append(hook) + + allgather_hooks = [hook for hook in unique_hooks if hook.dp_size > 1] + rank_local_hooks = [hook for hook in unique_hooks if hook.dp_size <= 1] + skipped_allgather = bool(allgather_hooks) and not restore_allgather_weights + if skipped_allgather: + # Startup rollback cannot safely enter a collective that a failed + # peer may never reach. Those blocks cannot be reconstructed, so + # make accidental reuse explicit instead of accepting zero weights. + self._poisoned_reason = ( + "Distributed layerwise offload startup skipped AllGather weight restoration; " + "recreate the backend and reload the pipeline before retrying" + ) + + collective_error = None + if restore_allgather_weights: + # Run collective-bearing restores before any rank-local operation: + # a local failure must never keep this rank out of a later AllGather. + collective_error = run_cleanup_steps( + ("restoring an AllGather block", hook.restore_next_block_to_cpu) for hook in allgather_hooks + ) + rank_local_error = run_cleanup_steps( + ("restoring a rank-local block", hook.restore_next_block_to_cpu) for hook in rank_local_hooks + ) + removal_error = run_cleanup_steps( + ( + "removing a distributed block hook", + lambda block=block: remove_distributed_block_hook(block), + ) + for blocks in self._blocks + for block in blocks + ) + encoder_error = run_cleanup_steps( + ( + "clearing distributed encoder state", + lambda module=module: clear_encoder_layerwise_state(module), + ) + for module in self._encoder_modules + ) + + resident_error = None + if self._resident_layer_group is not None: + # Resident layers are always rank-local (the public parser rejects + # resident DiT layers combined with AllGather). + resident_error = run_cleanup_steps( + [("restoring resident DLO blocks", self._resident_layer_group.restore_to_cpu)] + ) + + lifecycle_error = next( + ( + error + for error in ( + sync_error, + collective_error, + rank_local_error, + removal_error, + encoder_error, + resident_error, + ) + if error is not None + ), + None, + ) + + # Unregistration is independent once streams have drained, but the + # lease/file mappings must stay alive if restoration or hook removal + # needs a retry. + registration_error = None + if sync_error is None: + registration_error = run_cleanup_steps( + [("releasing registered HWR mappings", self._release_registered_mmap)] + ) + cleanup_error = lifecycle_error or registration_error + if cleanup_error is not None: + raise cleanup_error + + release_error = run_cleanup_steps([("releasing DLO mmap handles", self._release_mmap_handles)]) - self.offload_resident_layers() self._blocks.clear() self._all_hook_groups.clear() self._resident_blocks.clear() self._resident_layer_group = None - self._release_registered_mmap() - self._release_mmap_handles() + self._encoder_modules.clear() + self._staged_components.clear() + # Loader plans are single-use (HWR carriers in particular). A clean + # disable can rebuild from restored CPU tensors; a poisoned backend is + # guarded above and must be recreated. + self.host_weight_plan = None self._using_mmap = False self._using_rank_local_mmap = False self._using_registered_mmap = False self.enabled = False logger.info("Distributed layer-wise offloading disabled") + if release_error is not None: + raise release_error + + def disable(self) -> None: + self._disable(restore_allgather_weights=True) @staticmethod def _allocate_shared_buffers( @@ -2114,9 +2167,12 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): for slot in range(2): buffers: dict[torch.dtype, torch.Tensor] = {} for dtype, total_numel in max_sizes.items(): - buffer = torch.empty(total_numel, dtype=dtype, device="cpu") - if pin_memory: - buffer = buffer.pin_memory() + buffer = torch.empty( + total_numel, + dtype=dtype, + device="cpu", + pin_memory=pin_memory, + ) buffers[dtype] = buffer shared_staging[slot] = buffers diff --git a/vllm_omni/diffusion/offloader/layerwise_backend.py b/vllm_omni/diffusion/offloader/layerwise_backend.py index 662e1304f..9b63f9162 100644 --- a/vllm_omni/diffusion/offloader/layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/layerwise_backend.py @@ -7,14 +7,41 @@ from typing import Any import torch from torch import nn -from torch.distributed.tensor import DTensor from vllm.logger import init_logger from vllm_omni.diffusion.hooks import HookRegistry, ModelHook from vllm_omni.platforms import current_omni_platform -from .base import OffloadBackend, OffloadConfig +from .base import OffloadBackend, OffloadConfig, run_cleanup_steps +from .block_discovery import ( + get_blocks_attr_names, + get_blocks_from_dit, + set_blocks_attr_names, +) +from .component_utils import ( + clear_encoder_layerwise_state, + get_encoder_block_groups, + iter_streamable_dits, + move_non_block_state_to_device, + prepare_pipeline_components, + set_encoder_layerwise_state, +) +from .config import DIT_COMPONENT from .module_collector import ModuleDiscovery +from .offload_plan import OffloadPlan, get_offload_plan +from .tensor_utils import ( + clear_block_storage, + clear_tensor_storage, + describe_tensor_storage, + flatten_physical_storage, + group_named_tensors_by_dtype, + is_materialized_tensor, + materialization_probe, + module_materialization_probe, + restore_tensor_storage, + set_tensor_storage, + tensor_storage_metadata, +) logger = init_logger(__name__) @@ -41,6 +68,7 @@ class LayerwiseOffloadHook(ModelHook): device: torch.device, stream: current_omni_platform.Stream | None = None, pin_memory: bool = True, + materialization_probe_tensor: torch.Tensor | None = None, ): assert isinstance(next_block, nn.Module), "transformer block must be type `torch.nn.Module`" @@ -59,31 +87,7 @@ class LayerwiseOffloadHook(ModelHook): self.next_block_buffers: dict[str, torch.Tensor] = {} self.dtype_cpu_flattened_weights: dict[torch.dtype, torch.Tensor] = {} self.dtype_metadata: dict[torch.dtype, list[dict[str, Any]]] = {} - - @staticmethod - def _is_dtensor(t: torch.Tensor) -> bool: - return isinstance(t, DTensor) - - @staticmethod - def _set_tensor_storage(target: torch.Tensor, value: torch.Tensor) -> None: - if LayerwiseOffloadHook._is_dtensor(target): - target._local_tensor = value - else: - target.data = value - - @staticmethod - def _make_offload_placeholder(tensor: torch.Tensor) -> torch.Tensor: - if LayerwiseOffloadHook._is_dtensor(tensor): - local_shape = tuple(tensor.to_local().shape) - return torch.empty(local_shape, device="meta", dtype=tensor.dtype) - return torch.empty((0,), device=tensor.device, dtype=tensor.dtype) - - @staticmethod - def _is_materialized_tensor(t: torch.Tensor) -> bool: - if LayerwiseOffloadHook._is_dtensor(t): - local_t = t.to_local() - return not local_t.is_meta - return not t.is_meta and t.data.numel() > 0 + self._materialization_probe = materialization_probe_tensor def initialize_hook(self, module: nn.Module) -> nn.Module: # This all happen during the hook instance being registered to hook registry; @@ -92,6 +96,8 @@ class LayerwiseOffloadHook(ModelHook): self.block_parameters: dict[str, nn.Parameter] = dict(module.named_parameters()) self.block_buffers: dict[str, torch.Tensor] = dict(module.named_buffers()) + if self._materialization_probe is None: + self._materialization_probe = materialization_probe(self.block_parameters, self.block_buffers) self.next_block_parameters: dict[str, nn.Parameter] = dict(self.next_block.named_parameters()) self.next_block_buffers: dict[str, torch.Tensor] = dict(self.next_block.named_buffers()) @@ -100,7 +106,6 @@ class LayerwiseOffloadHook(ModelHook): self.dtype_cpu_flattened_weights, self.dtype_metadata = LayerwiseOffloadHook._to_cpu( self.next_block_parameters, self.next_block_buffers, - self.device, self.pin_memory, ) @@ -110,7 +115,6 @@ class LayerwiseOffloadHook(ModelHook): def _to_cpu( params: dict[str, nn.Parameter], bufs: dict[str, torch.Tensor], - device: torch.device, pin_memory: bool = True, ) -> tuple[dict[torch.dtype, torch.Tensor], dict[torch.dtype, list[dict[str, Any]]]]: """Helper method to move block parameters and buffers to CPU, flattening by dtype. @@ -123,85 +127,39 @@ class LayerwiseOffloadHook(ModelHook): flattened CPU tensors by dtype, metadata for reconstruction by dtype """ - dtype_grouped_weights: dict[torch.dtype, dict[str, torch.Tensor]] = {} dtype_cpu_flattened_weights: dict[torch.dtype, torch.Tensor] = {} # NOTE: order does matter dtype_metadata: dict[torch.dtype, list[dict[str, Any]]] = {} + targets_to_offload: list[torch.Tensor] = [] - for name, param_or_buf in chain(params.items(), bufs.items()): - dtype = param_or_buf.dtype - if dtype not in dtype_grouped_weights: - dtype_grouped_weights[dtype] = {} - dtype_grouped_weights[dtype][name] = param_or_buf - - for dtype, name2weights in dtype_grouped_weights.items(): + for dtype, named_weights in group_named_tensors_by_dtype(params, bufs).items(): # total # of parameters + buffers - weights_with_local = [] - for name, t in name2weights.items(): - local_t = t.to_local() if hasattr(t, "to_local") else t - stride = local_t.stride() - storage_numel = ( - 0 - if local_t.numel() == 0 - else 1 + sum((size - 1) * axis_stride for size, axis_stride in zip(local_t.shape, stride)) - ) - weights_with_local.append((name, t, local_t, storage_numel, stride)) - total_numel = sum(storage_numel for _, _, _, storage_numel, _ in weights_with_local) + specs = describe_tensor_storage(named_weights) + total_numel = sum(spec.storage_numel for spec in specs) cpu_tensor = torch.empty(total_numel, dtype=dtype, device="cpu", pin_memory=pin_memory) current_offset = 0 - for ( - name, - original_tensor, - local_tensor, - storage_numel, - stride, - ) in weights_with_local: - if local_tensor.is_contiguous(): - flat_storage = local_tensor.flatten() - else: - # Cutlass FP8 weights use a transposed physical layout. - # Preserve it across the flattened CPU staging buffer. - flat_storage = torch.zeros( - storage_numel, - dtype=dtype, - device=local_tensor.device, - ) - physical_view = torch.as_strided( - flat_storage, - size=local_tensor.shape, - stride=stride, - ) - physical_view.copy_(local_tensor) - cpu_tensor[current_offset : current_offset + storage_numel].copy_(flat_storage) - if dtype not in dtype_metadata: - dtype_metadata[dtype] = [] - dtype_metadata[dtype].append( - { - "name": name, - "offset": current_offset, - "numel": storage_numel, - "shape": local_tensor.shape, - "stride": stride, - } + for spec in specs: + flat_storage = flatten_physical_storage(spec.value, spec.storage_numel) + cpu_tensor[current_offset : current_offset + spec.storage_numel].copy_(flat_storage) + dtype_metadata.setdefault(dtype, []).append( + tensor_storage_metadata(spec, current_offset, include_device=True) ) - - LayerwiseOffloadHook._set_tensor_storage( - original_tensor, LayerwiseOffloadHook._make_offload_placeholder(original_tensor) - ) - current_offset += storage_numel + targets_to_offload.append(spec.target) + current_offset += spec.storage_numel dtype_cpu_flattened_weights[dtype] = cpu_tensor + # Do not mutate the module until every host master has been built. + # Allocation/copy failures therefore leave the input module intact. + clear_tensor_storage(targets_to_offload) + return dtype_cpu_flattened_weights, dtype_metadata @property def is_materialized(self) -> bool: """Check whether this block's parameters hold real data on device.""" - for param in self.block_parameters.values(): - return LayerwiseOffloadHook._is_materialized_tensor(param) - - return True + return self._materialization_probe is None or is_materialized_tensor(self._materialization_probe) @torch.compiler.disable def prefetch_layer(self, non_blocking: bool = True) -> None: @@ -236,7 +194,7 @@ class LayerwiseOffloadHook(ModelHook): layer_params[target_name] if target_name in layer_params else layer_bufs[target_name] ) - LayerwiseOffloadHook._set_tensor_storage( + set_tensor_storage( target_param_or_buf, torch.as_strided( gpu_weight[metadata["offset"] : metadata["offset"] + metadata["numel"]], @@ -252,17 +210,27 @@ class LayerwiseOffloadHook(ModelHook): """Free GPU memory for layer by replacing tensors with empty placeholders. This function does not actually offload weights from GPU back to CPU. """ - evt = self._prefetch_done - if evt is not None: - current_omni_platform.current_stream().wait_event(evt) - + clear_block_storage(self.block_parameters, self.block_buffers, self._prefetch_done) self._prefetch_done = None - # free GPU residency - for _, param in self.block_parameters.items(): - LayerwiseOffloadHook._set_tensor_storage(param, LayerwiseOffloadHook._make_offload_placeholder(param)) - for _, buf in self.block_buffers.items(): - LayerwiseOffloadHook._set_tensor_storage(buf, LayerwiseOffloadHook._make_offload_placeholder(buf)) + @torch.compiler.disable + def restore_next_block(self) -> None: + """Detach the next block from this hook's host backing store.""" + for dtype, ordered_metadata in self.dtype_metadata.items(): + flat = self.dtype_cpu_flattened_weights[dtype] + for metadata in ordered_metadata: + value = torch.as_strided( + flat[metadata["offset"] : metadata["offset"] + metadata["numel"]], + size=metadata["shape"], + stride=metadata["stride"], + ) + target_name = metadata["name"] + target = ( + self.next_block_parameters[target_name] + if target_name in self.next_block_parameters + else self.next_block_buffers[target_name] + ) + restore_tensor_storage(target, value, device=metadata["device"]) def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]: # if the previous hook was skipped and the weights are not on device, @@ -287,9 +255,17 @@ def apply_block_hook( device: torch.device, stream: current_omni_platform.Stream | None = None, pin_memory: bool = True, + *, + materialization_probe_tensor: torch.Tensor | None = None, ) -> LayerwiseOffloadHook: registry = HookRegistry.get_or_create(module) - hook = LayerwiseOffloadHook(next_block, device, stream, pin_memory) + hook = LayerwiseOffloadHook( + next_block, + device, + stream, + pin_memory, + materialization_probe_tensor, + ) registry.register_hook(LayerwiseOffloadHook._HOOK_NAME, hook) return hook @@ -302,6 +278,135 @@ def remove_block_hook(module: nn.Module) -> None: logger.debug("Removed offload hook from %s", module.__class__.__name__) +def _install_layerwise_hook_group( + blocks: list[nn.Module] | nn.ModuleList, + device: torch.device, + stream: Any, + pin_memory: bool, +) -> list[LayerwiseOffloadHook]: + """Install one circular hook ring and roll it back transactionally.""" + block_list = list(blocks) + if len(block_list) <= 1: + raise ValueError("A layerwise hook group requires at least two blocks") + + hooks: list[LayerwiseOffloadHook] = [] + hooked_blocks: list[nn.Module] = [] + probes = {id(block): module_materialization_probe(block) for block in block_list} + try: + for block, next_block in zip( + chain((block_list[-1],), block_list[:-1]), + block_list, + strict=True, + ): + hooks.append( + apply_block_hook( + block, + next_block, + device, + stream, + pin_memory, + materialization_probe_tensor=probes[id(block)], + ) + ) + hooked_blocks.append(block) + except BaseException: + run_cleanup_steps( + [ + *(("restoring a partially installed layerwise block", hook.restore_next_block) for hook in hooks), + *( + ("removing a partially installed layerwise hook", lambda block=block: remove_block_hook(block)) + for block in hooked_blocks + ), + ] + ) + raise + + for index, hook in enumerate(hooks): + hook._prev_hook = hooks[index - 1] + return hooks + + +def enable_plan_encoder_layerwise_offload( + module: nn.Module, + name: str, + plan: OffloadPlan | None, + *, + device: torch.device, + stream: current_omni_platform.Stream, + pin_memory: bool, + stage_on_demand: bool = False, + strict: bool = False, +) -> bool: + """Apply rank-local layerwise hooks to plan-declared encoder stacks.""" + if getattr(module, "_omni_layerwise_enabled", False): + return True + + hooks: list[LayerwiseOffloadHook] = [] + hooked_blocks: list[nn.Module] = [] + block_groups = get_encoder_block_groups( + module, + name, + plan, + strict=strict, + ) + if not block_groups: + return False + try: + for blocks in block_groups: + group_hooks = _install_layerwise_hook_group(blocks, device, stream, pin_memory) + hooks.extend(group_hooks) + hooked_blocks.extend(blocks) + if not stage_on_demand: + move_non_block_state_to_device(module, block_groups, device) + except BaseException: + run_cleanup_steps( + [ + *(("restoring a partially installed encoder block", hook.restore_next_block) for hook in hooks), + *( + ("removing a partially installed encoder hook", lambda block=block: remove_block_hook(block)) + for block in hooked_blocks + ), + ] + ) + raise + set_encoder_layerwise_state( + module, + hooks, + block_groups, + ) + logger.info( + "Enabled rank-local layerwise offload for encoder %s (%d blocks across %d stacks)", + name, + sum(len(blocks) for blocks in block_groups), + len(block_groups), + ) + return True + + +def disable_plan_encoder_layerwise_offload( + module: nn.Module, + *, + restore_weights: bool = True, +) -> None: + """Remove hooks installed by :func:`enable_plan_encoder_layerwise_offload`.""" + if not getattr(module, "_omni_layerwise_enabled", False): + return + hooks = getattr(module, "_omni_layerwise_hooks", []) + block_groups = getattr(module, "_omni_layerwise_block_groups", []) + steps = [] + if restore_weights: + steps.extend(("restoring an encoder block", hook.restore_next_block) for hook in hooks) + steps.extend( + ("removing an encoder hook", lambda block=block: remove_block_hook(block)) + for blocks in block_groups + for block in blocks + ) + cleanup_error = run_cleanup_steps(steps) + if cleanup_error is not None: + raise RuntimeError("Failed to fully disable encoder layerwise offload") from cleanup_error + clear_encoder_layerwise_state(module) + + class LayerWiseOffloadBackend(OffloadBackend): """Layer-wise (block-level) offloading backend. @@ -315,56 +420,88 @@ class LayerWiseOffloadBackend(OffloadBackend): self.copy_stream = current_omni_platform.Stream() self._blocks: list[list[nn.Module]] = [] + self._dit_hooks: list[LayerwiseOffloadHook] = [] + self._hooked_dit_blocks: list[nn.Module] = [] + self._encoder_modules: list[nn.Module] = [] + self._staged_components: list[nn.Module] = [] def enable(self, pipeline: nn.Module) -> None: + try: + self._enable(pipeline) + except BaseException: + try: + self.disable() + except BaseException: + logger.exception("LayerWiseOffloadBackend cleanup failed while handling an enable failure") + raise + + def _enable(self, pipeline: nn.Module) -> None: if self.enabled: logger.warning("LayerWiseOffloadBackend already enabled") return modules = ModuleDiscovery.discover(pipeline) - if not modules.dits: - logger.warning("No DiT/transformer modules found, skipping layer-wise offloading") + plan = get_offload_plan(pipeline) + if not modules.dits and self.config.offloads(DIT_COMPONENT): + message = "No DiT/transformer modules found for selected DiT layerwise offload" + if self.config.components is not None: + raise ValueError(message) + logger.warning(message) return - # Move encoders to GPU (they stay resident) - for enc in modules.encoders: - enc.to(self.device) + def enable_encoder_blocks( + module: nn.Module, + name: str, + component_plan: OffloadPlan | None, + stage_on_demand: bool, + ) -> bool: + enabled = enable_plan_encoder_layerwise_offload( + module, + name, + component_plan, + device=self.device, + stream=self.copy_stream, + pin_memory=self.config.pin_cpu_memory, + stage_on_demand=stage_on_demand, + strict=self.config.components is not None, + ) + if enabled: + # Record each successful installation immediately so a later + # component failure can remove these hooks transactionally. + self._encoder_modules.append(module) + return enabled - # Move VAE(s) to GPU if available - for vae in modules.vaes: - try: - vae.to(self.device, non_blocking=True) - except Exception as exc: - logger.debug("Failed to move VAE to GPU: %s", exc) + prepare_pipeline_components( + modules, + self.config, + plan, + device=self.device, + staged_components=self._staged_components, + enable_encoder_blocks=enable_encoder_blocks, + ) - # Move resident modules to GPU (small modules needed every forward) - for name, module in zip(modules.resident_names, modules.resident_modules): - try: - module.to(self.device) - except Exception as exc: - logger.debug("Failed to move resident module %s to GPU: %s", name, exc) + if not self.config.offloads(DIT_COMPONENT): + self.enabled = bool(self._encoder_modules or self._staged_components) + if not self.enabled: + raise ValueError( + "None of the selected layerwise offload components have " + "a model-declared streamable or on-demand plan" + ) + return logger.info("Applying layer-wise offloading on %s", modules.dit_names) # Apply block-wise offloading hook for each of the blocks in DiT model(s) # Note that there might exist multiple DiT models in specific pipelines - for i, dit_module in enumerate(modules.dits): - dit_name = modules.dit_names[i] - logger.info(f"Applying hooks on {dit_name} ({dit_module.__class__.__name__})") - - blocks_attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(dit_module) - - if not blocks: - logger.warning( - "Target layers (blocks) not found. Skipping offloading on %s (%s)", - dit_name, - dit_module.__class__.__name__, - ) - dit_module.to(self.device) - continue - + for dit_name, dit_module, blocks_attr_names, blocks in iter_streamable_dits( + modules, self.config, self.device, plan + ): num_blocks = len(blocks) if num_blocks <= 1: + if self.config.components is not None: + raise ValueError( + f"Selected DiT {dit_name!r} requires at least two streamable layerwise-offload blocks" + ) logger.warning( "#Target layers (blocks) <= 1. Skipping offloading on %s (%s)", dit_name, @@ -390,137 +527,66 @@ class LayerWiseOffloadBackend(OffloadBackend): if buffer is not None: buffer.data = buffer.data.to(self.device, non_blocking=True) - # Pre-fetch the first layer by manually calling the hook function on the last layer; - # For subsequent requests, the first layer/block will be pre-fetched - # during the last layer compute of the previous request. - last_block, first_block = blocks[-1], blocks[0] - last_hook = apply_block_hook( - last_block, - first_block, + block_hooks = _install_layerwise_hook_group( + blocks, self.device, self.copy_stream, self.config.pin_cpu_memory, ) - last_hook.prefetch_layer(non_blocking=False) + self._dit_hooks.extend(block_hooks) + self._hooked_dit_blocks.extend(blocks) - block_hooks: list[LayerwiseOffloadHook] = [last_hook] - # Register hook for each of blocks - for i, block in enumerate(blocks[:-1]): - next_block = blocks[(i + 1) % num_blocks] - hook = apply_block_hook( - block, - next_block, - self.device, - self.copy_stream, - self.config.pin_cpu_memory, - ) - block_hooks.append(hook) - - # NOTE(yuanheng-zhao): We make each hook gets a backward reference to the hook - # that is responsible for prefetching its block's weights. This is specifically a - # workaround for that arbitrary blocks are skipped by caching systems (e.g., cache-dit) - for i in range(len(block_hooks)): - block_hooks[i]._prev_hook = block_hooks[i - 1] + # The last block owns block zero's host backing. Materialize block + # zero once; later denoising iterations prefetch it from the ring. + block_hooks[0].prefetch_layer(non_blocking=False) logger.info(f"Layer-wise offloading enabled on {num_blocks} layers (blocks)") # Track hooked blocks for cleanup self._blocks.append(blocks) - if len(self._blocks) > 0 and len(self._blocks[0]) > 0: - self.enabled = True + self.enabled = bool(self._blocks or self._encoder_modules or self._staged_components) - def disable(self) -> None: - if not self.enabled: + def _disable(self, *, restore_weights: bool) -> None: + if not self.enabled and not ( + self._dit_hooks or self._hooked_dit_blocks or self._encoder_modules or self._staged_components + ): return - for blocks in self._blocks: - for block in blocks: - remove_block_hook(block) - + steps = [] + if restore_weights: + steps.extend(("restoring a DiT block", hook.restore_next_block) for hook in self._dit_hooks) + steps.extend( + ("removing a DiT block hook", lambda block=block: remove_block_hook(block)) + for block in self._hooked_dit_blocks + ) + steps.extend( + ( + "disabling encoder layerwise offload", + lambda module=module: disable_plan_encoder_layerwise_offload( + module, + restore_weights=restore_weights, + ), + ) + for module in self._encoder_modules + ) + cleanup_error = run_cleanup_steps(steps) + if cleanup_error is not None: + # Keep every host-master reference so a transient teardown error + # can be retried without losing weights behind placeholders. + raise RuntimeError("Failed to fully disable layerwise offload") from cleanup_error self._blocks.clear() + self._dit_hooks.clear() + self._hooked_dit_blocks.clear() + self._encoder_modules.clear() + self._staged_components.clear() self.enabled = False logger.info("Layer-wise offloading disabled") - @staticmethod - def get_blocks_attr_names(model: nn.Module) -> list[str]: - """Get block attribute names from model class.""" - attrs: list[str] = getattr(model.__class__, "_layerwise_offload_blocks_attrs", []) + def disable(self) -> None: + self._disable(restore_weights=True) - if not attrs: - old_attr = getattr(model.__class__, "_layerwise_offload_blocks_attr", None) - if old_attr is not None: - logger.warning( - "'_layerwise_offload_blocks_attr' is deprecated, " - "please use '_layerwise_offload_blocks_attrs' instead. " - "Example: _layerwise_offload_blocks_attrs = ['blocks']" - ) - attrs = [old_attr] if isinstance(old_attr, str) else list(old_attr) - - return attrs - - @staticmethod - def set_blocks_attr_names(model: nn.Module, names: list[str]) -> None: - if not hasattr(model.__class__, "_layerwise_offload_blocks_attrs"): - setattr(model.__class__, "_layerwise_offload_blocks_attrs", names) - - @staticmethod - def get_blocks_from_dit(model: nn.Module) -> tuple[list[str], list[nn.Module]]: - """ - Retrieve blocks and attribute names from provided DiT model. Blocks attribute names - are found by `_layerwise_offload_blocks_attrs` set to DiT models. For example, - - ``` - class WanTransformer3DModel(nn.Module): - _layerwise_offload_blocks_attrs = ["blocks"] - ``` - - Returns: - Tuple of (blocks_attr_names, blocks) - """ - blocks_attr_names = LayerWiseOffloadBackend.get_blocks_attr_names(model) - if not blocks_attr_names: - logger.warning( - f"No _layerwise_offload_blocks_attrs defined for {model.__class__.__name__}, " - "skipping layerwise offloading" - ) - return [], [] - - blocks = [] - for name in blocks_attr_names: - attr = getattr(model, name, None) - if attr is None: - raise AttributeError( - f"Attribute '{name}' declared in _layerwise_offload_blocks_attrs " - f"does not exist on model {model.__class__.__name__}" - ) - try: - attr_iter = iter(attr) - except TypeError: - if isinstance(attr, nn.Module): - logger.warning( - "Attribute '%s' on %s is not iterable; treating it as one block.", - name, - model.__class__.__name__, - ) - blocks.append(attr) - continue - - logger.warning( - "Attribute '%s' on %s is not iterable (got %s); skipping it.", - name, - model.__class__.__name__, - type(attr).__name__, - ) - else: - blocks.extend(attr_iter) - - if not blocks: - logger.warning( - "No blocks found in %s for %s, skipping layerwise offloading", - blocks_attr_names, - model.__class__.__name__, - ) - return [], [] - - return blocks_attr_names, blocks + # Compatibility aliases for existing model integrations. + get_blocks_attr_names = staticmethod(get_blocks_attr_names) + set_blocks_attr_names = staticmethod(set_blocks_attr_names) + get_blocks_from_dit = staticmethod(get_blocks_from_dit) diff --git a/vllm_omni/diffusion/offloader/offload_plan.py b/vllm_omni/diffusion/offloader/offload_plan.py index 2e1443df7..4aa37ae8f 100644 --- a/vllm_omni/diffusion/offloader/offload_plan.py +++ b/vllm_omni/diffusion/offloader/offload_plan.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Declarative OffloadPlan for distributed layerwise offload. +"""Declarative model capabilities for component layerwise offload. -Models declare this as a class attribute ``_offload_plan`` on the -pipeline class. When present, the offloader uses it instead of -heuristic block discovery. +Models declare an ``_offload_plan`` class attribute on the pipeline. Both the +ordinary and distributed layerwise backends consume the same plan instead of +embedding model-specific block paths in backend code. """ from __future__ import annotations @@ -16,12 +16,11 @@ from torch import nn @dataclass(frozen=True) class OffloadPlan: - """Optional declarative metadata for distributed layerwise offload. + """Optional declarative metadata for component layerwise offload. Models declare this as a class attribute ``_offload_plan`` on the - pipeline class. When present, the offloader uses it instead of - heuristic block discovery, making it easier to support new models - without offload-specific logic in constructor or forward code. + pipeline class. When present, both layerwise backends use it instead of + model-specific backend branches, making new integrations data-driven. If not declared, the offloader falls back to: 1. ``_layerwise_offload_blocks_attrs`` on each DiT module class. @@ -39,9 +38,14 @@ class OffloadPlan: device when ``dlo_resident_layers`` is nonzero. Keeping this model-declared avoids applying a consumer-GPU tuning knob to auxiliary or dual DiTs unintentionally. - encoder_block_attrs: Maps encoder paths to rank-local block-list paths. - These blocks are streamed with ordinary layerwise hooks, never - with the DiT AllGather group. + encoder_component_types: Maps encoder paths to public selector types + (currently text_encoder). This declaration is used before the + compatibility name heuristic. + encoder_block_attrs: Maps encoder paths to streamable block-list paths. + encoder_dlo_weight_replication: Encoder paths whose loader-produced + block tensors are identical across the DiT DLO group. Only these + encoders may use multi-rank AllGather transfer; this must not be + declared for encoder-TP shards. """ on_demand_component_paths: frozenset[str] = field(default_factory=frozenset) @@ -49,7 +53,9 @@ class OffloadPlan: block_attrs: dict[str, tuple[str, ...]] = field(default_factory=dict) offload_submodules: dict[str, str] = field(default_factory=dict) resident_dit_paths: frozenset[str] = field(default_factory=frozenset) + encoder_component_types: dict[str, str] = field(default_factory=dict) encoder_block_attrs: dict[str, tuple[str, ...]] = field(default_factory=dict) + encoder_dlo_weight_replication: frozenset[str] = field(default_factory=frozenset) def get_offload_plan(pipeline: nn.Module) -> OffloadPlan | None: diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index 36e0f7b22..b6b596070 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/vllm_omni/diffusion/offloader/sequential_backend.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Iterator +from collections.abc import Collection, Iterator from contextlib import contextmanager import torch @@ -13,11 +13,38 @@ from vllm_omni.diffusion.hooks import HookRegistry, ModelHook from vllm_omni.platforms import current_omni_platform from .base import OffloadBackend, OffloadConfig, SupportsModelCpuOffload +from .config import DIT_COMPONENT, TEXT_ENCODER_COMPONENT from .module_collector import ModuleDiscovery +from .offload_plan import get_offload_plan logger = init_logger(__name__) +def _capture_tensor_devices(modules: Collection[nn.Module]) -> list[tuple[torch.Tensor, torch.device]]: + seen: set[int] = set() + devices: list[tuple[torch.Tensor, torch.device]] = [] + for module in modules: + for tensor in (*module.parameters(), *module.buffers()): + if id(tensor) in seen: + continue + seen.add(id(tensor)) + devices.append((tensor, tensor.device)) + return devices + + +def _restore_tensor_devices(devices: list[tuple[torch.Tensor, torch.device]]) -> None: + first_error: BaseException | None = None + for tensor, device in devices: + try: + if tensor.device != device: + tensor.data = tensor.data.to(device) + except BaseException as exc: + logger.exception("Failed to restore tensor placement to %s", device) + first_error = first_error or exc + if first_error is not None: + raise RuntimeError("Failed to restore one or more tensor placements") from first_error + + class SequentialOffloadHook(ModelHook): """Hook for sequential offloading with mutual exclusion on encoder and DiT modules. @@ -34,12 +61,14 @@ class SequentialOffloadHook(ModelHook): device: torch.device, pin_memory: bool = True, use_hsdp: bool = False, + offload_after_context: bool = True, ): # Modules to offload to CPU before this module runs self.offload_targets = offload_targets self.device = device self.pin_memory = pin_memory self.use_hsdp = use_hsdp + self.offload_after_context = offload_after_context @staticmethod def _move_params( @@ -48,7 +77,7 @@ class SequentialOffloadHook(ModelHook): *, non_blocking: bool = False, pin_memory: bool = False, - ) -> None: + ) -> bool: """Move module parameters and buffers to device. This cls method specifically prevents recursion device movement, @@ -57,47 +86,38 @@ class SequentialOffloadHook(ModelHook): refer to https://github.com/vipshop/cache-dit/blob/v1.2.3/src/cache_dit/caching/cache_blocks/__init__.py#L83 """ + moved = False for p in module.parameters(): if p.data.device != target_device: data = p.data.to(target_device, non_blocking=non_blocking) if pin_memory and target_device.type == "cpu" and not isinstance(data, DTensor): data = data.pin_memory() p.data = data + moved = True for b in module.buffers(): if b.device != target_device: data = b.data.to(target_device, non_blocking=non_blocking) if pin_memory and target_device.type == "cpu" and not isinstance(data, DTensor): data = data.pin_memory() b.data = data + moved = True + return moved def _to_cpu(self, module: nn.Module) -> None: - try: - param = next(module.parameters()) - except StopIteration: - return - - if param.device.type == "cpu": - return - # XPU's allocator doesn't respect stream dependencies in empty_cache, # so non-blocking copies can race with cache eviction. Use blocking # copies on XPU to avoid NULL pointer errors during DMA. non_blocking = not self.use_hsdp and not current_omni_platform.is_xpu() - self._move_params( + moved = self._move_params( module, torch.device("cpu"), non_blocking=non_blocking, pin_memory=self.pin_memory, ) - current_omni_platform.empty_cache() + if moved: + current_omni_platform.empty_cache() def _to_gpu(self, module: nn.Module) -> None: - try: - if next(module.parameters()).device == self.device: - return - except StopIteration: - return - self._move_params(module, self.device, non_blocking=False) def pre_forward(self, module: nn.Module, *args, **kwargs) -> tuple[tuple, dict]: @@ -127,6 +147,8 @@ def apply_sequential_offload( pin_memory: bool = True, use_hsdp: bool = False, offload_initial_dits: bool = False, + offload_dit_modules: Collection[nn.Module] | None = None, + offload_encoder_modules: Collection[nn.Module] | None = None, ) -> None: """Apply sequential offloading hooks to DiT and encoder modules. @@ -141,6 +163,10 @@ def apply_sequential_offload( pin_memory: Whether to pin CPU memory for faster transfers use_hsdp: Whether HSDP is enabled (affects non_blocking behavior) offload_initial_dits: Whether to begin with all DiT modules on CPU. + offload_dit_modules: DiT modules allowed to move to CPU. None selects + every DiT for backward compatibility. + offload_encoder_modules: Encoder/stage modules allowed to move to CPU. + None selects every supplied module for backward compatibility. Example: >>> apply_sequential_offload( @@ -150,38 +176,59 @@ def apply_sequential_offload( ... ) >>> # Modules of pipeline now automatically swap between CPU and GPU """ - # Register hooks on DiT modules (offload encoders AND other DiTs when a DiT runs) - for i, dit_mod in enumerate(dit_modules): - other_dits = [d for j, d in enumerate(dit_modules) if j != i] - registry = HookRegistry.get_or_create(dit_mod) - hook = SequentialOffloadHook( - offload_targets=encoder_modules + other_dits, - device=device, - pin_memory=pin_memory, - use_hsdp=use_hsdp, - ) - registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook) - logger.debug("Registered offload hook for %s", dit_mod.__class__.__name__) + selected_dit_ids = {id(module) for module in (dit_modules if offload_dit_modules is None else offload_dit_modules)} + selected_encoder_ids = { + id(module) for module in (encoder_modules if offload_encoder_modules is None else offload_encoder_modules) + } - # Register hooks on encoders (offload DiTs when encoder runs) - for enc in encoder_modules: - registry = HookRegistry.get_or_create(enc) - hook = SequentialOffloadHook( - offload_targets=dit_modules, - device=device, - pin_memory=pin_memory, - use_hsdp=use_hsdp, - ) - registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook) - logger.debug("Registered offload hook for %s", enc.__class__.__name__) + all_modules = [*dit_modules, *encoder_modules] + initial_devices = _capture_tensor_devices(dit_modules) if offload_initial_dits else [] + try: + selected_encoders = [encoder for encoder in encoder_modules if id(encoder) in selected_encoder_ids] + # Register hooks on DiT modules (offload selected encoders and other selected DiTs). + for i, dit_mod in enumerate(dit_modules): + other_dits = [d for j, d in enumerate(dit_modules) if j != i and id(d) in selected_dit_ids] + registry = HookRegistry.get_or_create(dit_mod) + hook = SequentialOffloadHook( + offload_targets=selected_encoders + other_dits, + device=device, + pin_memory=pin_memory, + use_hsdp=use_hsdp, + offload_after_context=id(dit_mod) in selected_dit_ids, + ) + registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook) + logger.debug("Registered offload hook for %s", dit_mod.__class__.__name__) - if offload_initial_dits: - try: + # Register hooks on all execution stages so unselected resident modules can + # still evict selected DiTs before they run. + selected_dits = [dit for dit in dit_modules if id(dit) in selected_dit_ids] + for enc in encoder_modules: + registry = HookRegistry.get_or_create(enc) + hook = SequentialOffloadHook( + offload_targets=selected_dits, + device=device, + pin_memory=pin_memory, + use_hsdp=use_hsdp, + offload_after_context=id(enc) in selected_encoder_ids, + ) + registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook) + logger.debug("Registered offload hook for %s", enc.__class__.__name__) + + if offload_initial_dits: for dit_mod in dit_modules: + if id(dit_mod) not in selected_dit_ids: + continue _get_sequential_offload_hook(dit_mod)._to_cpu(dit_mod) - except Exception: - remove_sequential_offload([*dit_modules, *encoder_modules]) - raise + except BaseException: + try: + remove_sequential_offload(all_modules) + except BaseException: + logger.exception("Failed to remove every sequential hook during rollback") + try: + _restore_tensor_devices(initial_devices) + except BaseException: + logger.exception("Failed to restore initial DiT placement during rollback") + raise def remove_sequential_offload(modules: list[nn.Module]) -> None: @@ -194,11 +241,18 @@ def remove_sequential_offload(modules: list[nn.Module]) -> None: >>> all_modules = [*dit_modules, *encoder_modules] >>> remove_sequential_offload(all_modules) """ + first_error: BaseException | None = None for module in modules: - registry: HookRegistry | None = getattr(module, "_hook_registry", None) - if registry is not None: - registry.remove_hook(SequentialOffloadHook._HOOK_NAME) - logger.debug("Removed offload hook from %s", module.__class__.__name__) + try: + registry: HookRegistry | None = getattr(module, "_hook_registry", None) + if registry is not None: + registry.remove_hook(SequentialOffloadHook._HOOK_NAME) + logger.debug("Removed offload hook from %s", module.__class__.__name__) + except BaseException as exc: + logger.exception("Failed to remove offload hook from %s", module.__class__.__name__) + first_error = first_error or exc + if first_error is not None: + raise RuntimeError("Failed to remove one or more sequential offload hooks") from first_error def _get_sequential_offload_hook(module: nn.Module) -> SequentialOffloadHook: @@ -218,12 +272,14 @@ def sequential_offload_component(module: nn.Module) -> Iterator[None]: yield except BaseException: try: - hook._to_cpu(module) + if hook.offload_after_context: + hook._to_cpu(module) except Exception: logger.exception("Failed to release %s after component failure", module.__class__.__name__) raise else: - hook._to_cpu(module) + if hook.offload_after_context: + hook._to_cpu(module) class ModelLevelOffloadBackend(OffloadBackend): @@ -245,11 +301,26 @@ class ModelLevelOffloadBackend(OffloadBackend): # Pipelines with non-forward component entry points own their complete # mutual-exclusion lifecycle. Delegate through the explicit protocol. if isinstance(pipeline, SupportsModelCpuOffload): - pipeline.enable_omni_model_cpu_offload( - device=self.device, - pin_memory=self.config.pin_cpu_memory, - use_hsdp=self.config.use_hsdp, - ) + try: + if self.config.components is not None: + pipeline.enable_omni_model_cpu_offload( + device=self.device, + pin_memory=self.config.pin_cpu_memory, + use_hsdp=self.config.use_hsdp, + offload_components=self.config.components, + ) + else: + pipeline.enable_omni_model_cpu_offload( + device=self.device, + pin_memory=self.config.pin_cpu_memory, + use_hsdp=self.config.use_hsdp, + ) + except BaseException: + try: + pipeline.disable_omni_model_cpu_offload() + except BaseException: + logger.exception("Model-level cleanup failed while handling an enable failure") + raise self._custom_pipeline = pipeline self.enabled = True logger.info( @@ -259,44 +330,65 @@ class ModelLevelOffloadBackend(OffloadBackend): return modules = ModuleDiscovery.discover(pipeline) + plan = get_offload_plan(pipeline) + selected_encoders = [ + encoder + for encoder, name in zip(modules.encoders, modules.encoder_names) + if self.config.should_offload_encoder(name, plan) + ] + if self.config.components is not None: + if not modules.dits: + raise ValueError("Component-selective model offload requires a DiT/transformer module") + if not modules.encoders: + raise ValueError("Component-selective model offload requires an encoder execution stage") + if self.config.offloads(TEXT_ENCODER_COMPONENT) and not selected_encoders: + raise ValueError("No text encoder modules found for selected text_encoder module offload") - # Move encoders to GPU - for enc in modules.encoders: - enc.to(self.device) - - # Move VAE(s) to GPU if available - for vae in modules.vaes: - try: + all_modules = [ + *modules.dits, + *modules.encoders, + *modules.vaes, + *modules.resident_modules, + ] + initial_devices = _capture_tensor_devices(all_modules) + try: + for encoder in modules.encoders: + encoder.to(self.device) + for vae in modules.vaes: vae.to(self.device, non_blocking=True) - except Exception as exc: - logger.debug("Failed to move VAE to GPU: %s", exc) + for resident in modules.resident_modules: + resident.to(self.device) - # Pin resident modules on GPU (small hot submodules called inside the DiT loop). - for res, name in zip(modules.resident_modules, modules.resident_names): + if not modules.dits: + logger.warning("No DiT/transformer modules found, skipping model-level offloading") + return + if not modules.encoders: + for dit in modules.dits: + dit.to(self.device) + logger.warning("No encoder modules found, skipping model-level offloading") + return + + apply_sequential_offload( + dit_modules=modules.dits, + encoder_modules=modules.encoders, + device=self.device, + pin_memory=self.config.pin_cpu_memory, + use_hsdp=self.config.use_hsdp, + offload_dit_modules=( + modules.dits if self.config.components is None or self.config.offloads(DIT_COMPONENT) else () + ), + offload_encoder_modules=selected_encoders, + ) + except BaseException: try: - res.to(self.device) - except Exception as exc: - logger.warning("Failed to move resident module '%s' to GPU: %s", name, exc) - - if not modules.dits: - logger.warning("No DiT/transformer modules found, skipping model-level offloading") - return - - if not modules.encoders: - # Nothing to swap against — move DiTs to GPU and skip hooks. - for dit in modules.dits: - dit.to(self.device) - logger.warning("No encoder modules found, skipping model-level offloading") - return - - # Apply sequential offloading hooks - apply_sequential_offload( - dit_modules=modules.dits, - encoder_modules=modules.encoders, - device=self.device, - pin_memory=self.config.pin_cpu_memory, - use_hsdp=self.config.use_hsdp, - ) + remove_sequential_offload([*modules.dits, *modules.encoders]) + except BaseException: + logger.exception("Failed to remove every model-level hook during rollback") + try: + _restore_tensor_devices(initial_devices) + except BaseException: + logger.exception("Failed to restore model placement during rollback") + raise # Track modules for cleanup self._offload_modules = [*modules.dits, *modules.encoders] diff --git a/vllm_omni/diffusion/offloader/tensor_utils.py b/vllm_omni/diffusion/offloader/tensor_utils.py index ef2d96011..ffc9ef3a0 100644 --- a/vllm_omni/diffusion/offloader/tensor_utils.py +++ b/vllm_omni/diffusion/offloader/tensor_utils.py @@ -9,15 +9,100 @@ offload backends. from __future__ import annotations +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass +from itertools import chain +from typing import Any + import torch +from torch import nn from torch.distributed.tensor import DTensor +@dataclass(frozen=True) +class TensorStorageSpec: + """One tensor's target and physical-storage layout.""" + + name: str + target: torch.Tensor + value: torch.Tensor + storage_numel: int + stride: tuple[int, ...] + + def dtype_size(dtype: torch.dtype) -> int: """Return element size in bytes for a torch.dtype.""" return torch.empty(1, dtype=dtype).element_size() +def group_named_tensors_by_dtype( + params: Mapping[str, torch.Tensor], + buffers: Mapping[str, torch.Tensor], +) -> dict[torch.dtype, list[tuple[str, torch.Tensor]]]: + """Group parameters and buffers by dtype while preserving their order.""" + grouped: dict[torch.dtype, list[tuple[str, torch.Tensor]]] = {} + for name, tensor in chain(params.items(), buffers.items()): + grouped.setdefault(tensor.dtype, []).append((name, tensor)) + return grouped + + +def physical_storage_numel(tensor: torch.Tensor) -> int: + """Return the storage span needed to preserve ``tensor``'s stride.""" + if tensor.numel() == 0: + return 0 + return 1 + sum((size - 1) * stride for size, stride in zip(tensor.shape, tensor.stride())) + + +def flatten_physical_storage( + tensor: torch.Tensor, + storage_numel: int | None = None, +) -> torch.Tensor: + """Pack a possibly strided tensor in physical storage order.""" + if tensor.is_contiguous(): + return tensor.flatten() + storage_numel = physical_storage_numel(tensor) if storage_numel is None else storage_numel + storage = torch.zeros(storage_numel, dtype=tensor.dtype, device=tensor.device) + torch.as_strided(storage, size=tensor.shape, stride=tensor.stride()).copy_(tensor) + return storage + + +def describe_tensor_storage( + named_tensors: Iterable[tuple[str, torch.Tensor]], + transforms: Mapping[int, Callable[[torch.Tensor], torch.Tensor]] | None = None, +) -> list[TensorStorageSpec]: + """Resolve local tensors and their transport-preserving layouts.""" + specs: list[TensorStorageSpec] = [] + for name, target in named_tensors: + value = target.to_local() if hasattr(target, "to_local") else target + transform = (transforms or {}).get(id(target)) + if callable(transform): + value = transform(value) + specs.append( + TensorStorageSpec( + name=name, + target=target, + value=value, + storage_numel=physical_storage_numel(value), + stride=value.stride(), + ) + ) + return specs + + +def tensor_storage_metadata(spec: TensorStorageSpec, offset: int, *, include_device: bool = False) -> dict[str, Any]: + """Build reconstruction metadata shared by offload transports.""" + metadata: dict[str, Any] = { + "name": spec.name, + "offset": offset, + "numel": spec.storage_numel, + "shape": spec.value.shape, + "stride": spec.stride, + } + if include_device: + metadata["device"] = spec.value.device + return metadata + + def is_dtensor(t: torch.Tensor) -> bool: """Check if tensor is a DTensor.""" return isinstance(t, DTensor) @@ -39,9 +124,74 @@ def make_offload_placeholder(tensor: torch.Tensor) -> torch.Tensor: return torch.empty((0,), device=tensor.device, dtype=tensor.dtype) +def clear_tensor_storage(tensors: Iterable[torch.Tensor]) -> None: + """Atomically replace tensor storage with offload placeholders.""" + targets = list(tensors) + originals = [tensor.to_local() if is_dtensor(tensor) else tensor.data for tensor in targets] + placeholders = [make_offload_placeholder(tensor) for tensor in targets] + committed = 0 + try: + for target, placeholder in zip(targets, placeholders, strict=True): + set_tensor_storage(target, placeholder) + committed += 1 + except BaseException: + rollback_error: BaseException | None = None + for target, original in zip(targets[:committed], originals[:committed], strict=True): + try: + set_tensor_storage(target, original) + except BaseException as exc: + rollback_error = rollback_error or exc + if rollback_error is not None: + raise RuntimeError("Tensor storage clear failed and rollback was incomplete") from rollback_error + raise + + +def clear_block_storage( + params: Mapping[str, torch.Tensor], + buffers: Mapping[str, torch.Tensor], + ready_event: Any | None, +) -> None: + """Wait for an outstanding transfer, then release one block's residency.""" + if ready_event is not None: + from vllm_omni.platforms import current_omni_platform + + current_omni_platform.current_stream().wait_event(ready_event) + clear_tensor_storage(chain(params.values(), buffers.values())) + + def is_materialized_tensor(t: torch.Tensor) -> bool: """Check if tensor holds real data (not meta or empty placeholder).""" if is_dtensor(t): local_t = t.to_local() return not local_t.is_meta return not t.is_meta and t.data.numel() > 0 + + +def materialization_probe( + params: Mapping[str, torch.Tensor], + buffers: Mapping[str, torch.Tensor], +) -> torch.Tensor | None: + """Pick one non-empty tensor that tracks a block's atomic storage state.""" + return next((tensor for tensor in chain(params.values(), buffers.values()) if tensor.numel()), None) + + +def module_materialization_probe(module: nn.Module) -> torch.Tensor | None: + """Capture a probe before another hook releases the module's storage.""" + return materialization_probe(dict(module.named_parameters()), dict(module.named_buffers())) + + +def restore_tensor_storage( + target: torch.Tensor, + value: torch.Tensor, + *, + device: torch.device | str, +) -> None: + """Detach ``target`` from an offloader-owned backing with a real copy.""" + restored = torch.empty_strided( + value.shape, + value.stride(), + dtype=value.dtype, + device=device, + ) + restored.copy_(value) + set_tensor_storage(target, restored) diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 4fea1ed17..8a352fb7d 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -51,6 +51,7 @@ from vllm_omni.diffusion.models.interface import ( supports_step_execution, ) from vllm_omni.diffusion.offloader import enable_offload_backend +from vllm_omni.diffusion.offloader.config import TEXT_ENCODER_COMPONENT, resolve_offload from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.interface import ( @@ -279,6 +280,26 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): if load_format == "dummy": return + # Resolve environment overrides before model loading. Rank-local + # prompt-cache hits cannot safely skip text-encoder collectives. + enable_pec, pec_size = resolve_prompt_embed_cache_config( + enable=getattr(self.od_config, "enable_prompt_embed_cache", False), + max_size=getattr(self.od_config, "prompt_embed_cache_size", 32), + ) + resolved_offload = resolve_offload(self.od_config) + dp_size = int(getattr(getattr(self.od_config, "parallel_config", None), "data_parallel_size", 1)) + if ( + enable_pec + and dp_size > 1 + and resolved_offload.offloads(TEXT_ENCODER_COMPONENT) + and resolved_offload.uses_allgather(TEXT_ENCODER_COMPONENT) + ): + raise ValueError( + "Prompt embedding cache cannot be combined with text_encoder " + "AllGather across data-parallel ranks; disable the cache or use " + "rank-local text_encoder transfer." + ) + current_omni_platform.init_diffusion_model_runner_runtime( vllm_config=self.vllm_config, od_config=self.od_config, @@ -398,10 +419,6 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): # Install prompt-embedding cache (transparent wrapper around # ``pipeline.encode_prompt``). Enabled via config or env var; a no-op # when the pipeline does not expose ``encode_prompt``. - enable_pec, pec_size = resolve_prompt_embed_cache_config( - enable=getattr(self.od_config, "enable_prompt_embed_cache", False), - max_size=getattr(self.od_config, "prompt_embed_cache_size", 32), - ) if enable_pec: self.prompt_embed_cache = install_prompt_embed_cache( self.pipeline, diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 22790fbd4..812f4c3e7 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -572,6 +572,9 @@ class OrchestratorArgs: vae_use_tiling: bool = False enable_multithread_weight_load: bool = True num_weight_load_threads: int = 4 + diffusion_offload_config: dict[str, Any] | None = None + # Compatibility aliases for existing callers and model-specific stage + # lifecycles that are broader than the compact dit/text_encoder selector. enable_cpu_offload: bool = False enable_layerwise_offload: bool = False enable_distributed_layerwise_offload: bool = False diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 8501d8730..608d605e9 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1137,6 +1137,7 @@ class AsyncOmniEngine: "cache_backend": cache_backend, "cache_config": cache_config, "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False), + "diffusion_offload_config": kwargs.get("diffusion_offload_config", None), "enable_cpu_offload": kwargs.get("enable_cpu_offload", False), "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False), "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False), diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index f4c890816..ea56acb1a 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -703,30 +703,41 @@ class OmniServeCommand(CLISubcommand): ) # diffusion model offload parameters + omni_config_group.add_argument( + "--diffusion-offload-config", + type=json.loads, + default=None, + help="Diffusion CPU-offload config as JSON. " + "Set mode to module or layer, list dit and/or text_encoder in " + "components, and put layer-only tuning under layer_options. " + "Layer settings are weight_transfer (rank-local or allgather) and " + "resident_layers (DiT only).", + ) omni_config_group.add_argument( "--enable-cpu-offload", action="store_true", - help="Enable CPU offloading for diffusion models.", + help="Compatibility alias for model-level CPU offload. New integrations should use " + "--diffusion-offload-config with mode=module and explicit components.", ) omni_config_group.add_argument( "--enable-layerwise-offload", action="store_true", - help="Enable layerwise (blockwise) offloading on DiT modules.", + help="Compatibility alias for layerwise CPU offload. New integrations should use " + "--diffusion-offload-config with mode=layer and explicit components.", ) omni_config_group.add_argument( "--enable-distributed-layerwise-offload", action="store_true", - help="Enable distributed layerwise offloading with H2D + AllGather overlap. " - "Shards weights across DP ranks, stores only 1/DP_size on each host, " - "and overlaps H2D transfers and AllGather with computation. " - "DP size is automatically derived from the parallel configuration.", + help="Compatibility alias for distributed layerwise CPU offload. " + "New integrations should use mode=layer and configure weight transfer per component.", ) omni_config_group.add_argument( "--dlo-use-allgather", dest="dlo_use_allgather", action="store_true", default=True, - help="Use shard + AllGather for weight reconstruction (default: True). " + help="Compatibility option; use component weight_transfer=allgather in new configurations. " + "Use shard + AllGather for weight reconstruction (default: True). " "When disabled (--dlo-no-use-allgather), each rank streams the " "standard loader's rank-local tensors via H2D only — no additional " "DP sharding, no AllGather, and no concurrent-request requirement.", @@ -736,6 +747,7 @@ class OmniServeCommand(CLISubcommand): dest="dlo_use_allgather", action="store_false", help=( + "Compatibility option; use component weight_transfer=rank-local in new configurations. " "Disable AllGather and stream standard-loader rank-local weights " "independently (including existing TP shards)." ), @@ -744,8 +756,7 @@ class OmniServeCommand(CLISubcommand): "--dlo-resident-layers", type=int, default=0, - help="Keep this many leading main-DiT blocks resident on the device " - "while distributed layerwise offload streams the remaining blocks.", + help="Compatibility option; use layer_options.dit.resident_layers in new configurations.", ) omni_config_group.add_argument( "--host-weight-runtime-mode",