From 8bf849c399f0dd9064e703ac3e07056fd5117dcd Mon Sep 17 00:00:00 2001 From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> Date: Thu, 12 Feb 2026 17:29:55 +0800 Subject: [PATCH] Refactor CPU Offloading Backend Pattern (#1223) Signed-off-by: yuanheng Signed-off-by: Yuanheng Zhao Signed-off-by: Samit <285365963@qq.com> Co-authored-by: Samit <285365963@qq.com> --- .../model/adding_diffusion_model.md | 31 + .../diffusion/cpu_offload_diffusion.md | 154 +++-- .../image_to_image/image_edit.py | 7 - .../image_to_video/image_to_video.py | 7 - .../text_to_image/text_to_image.py | 7 - .../text_to_video/text_to_video.py | 7 - .../test_diffusion_layerwise_offload.py | 20 +- vllm_omni/diffusion/data.py | 3 - vllm_omni/diffusion/offload.py | 535 ------------------ vllm_omni/diffusion/offloader/__init__.py | 77 +++ vllm_omni/diffusion/offloader/base.py | 94 +++ .../diffusion/offloader/layerwise_backend.py | 362 ++++++++++++ .../diffusion/offloader/module_collector.py | 80 +++ .../diffusion/offloader/sequential_backend.py | 227 ++++++++ .../worker/diffusion_model_runner.py | 9 +- vllm_omni/entrypoints/async_omni.py | 1 - vllm_omni/entrypoints/cli/serve.py | 6 - 17 files changed, 980 insertions(+), 647 deletions(-) delete mode 100644 vllm_omni/diffusion/offload.py create mode 100644 vllm_omni/diffusion/offloader/__init__.py create mode 100644 vllm_omni/diffusion/offloader/base.py create mode 100644 vllm_omni/diffusion/offloader/layerwise_backend.py create mode 100644 vllm_omni/diffusion/offloader/module_collector.py create mode 100644 vllm_omni/diffusion/offloader/sequential_backend.py diff --git a/docs/contributing/model/adding_diffusion_model.md b/docs/contributing/model/adding_diffusion_model.md index 529cab752..10755c6c6 100644 --- a/docs/contributing/model/adding_diffusion_model.md +++ b/docs/contributing/model/adding_diffusion_model.md @@ -763,6 +763,37 @@ omni = Omni(model="your-model", ) ``` +### CPU Offload + +See detailed guide: [CPU Offloading for Diffusion Models](../../user_guide/diffusion/cpu_offload_diffusion.md) + +vLLM-Omni provides two offloading strategies to reduce GPU memory usage: + +1. **Model-level offload**: Mutual exclusion between DiT and encoders (only one on GPU at a time) +2. **Layerwise (Blockwise) offload**: Keeps only a single transformer block on GPU at a time with compute-memory overlap + +**Usage:** Enable offload when initializing: +```python +# Model-level offload +omni = Omni(model="your-model", enable_cpu_offload=True) + +# Layerwise offload +omni = Omni(model="your-model", enable_layerwise_offload=True) +``` + +**To support layerwise offloading:** Define the blocks attribute name in your transformer: + +```python +class WanTransformer3DModel(nn.Module): + _layerwise_offload_blocks_attr = "blocks" # Attribute name containing transformer blocks + + def __init__(self): + self.blocks = nn.ModuleList([...]) # Transformer blocks +``` + +**Note:** Layerwise offloading is primarily recommended for large **video generation models** where the compute cost per block is high enough to effectively overlap with memory prefetch operations. + + --- ## Troubleshooting diff --git a/docs/user_guide/diffusion/cpu_offload_diffusion.md b/docs/user_guide/diffusion/cpu_offload_diffusion.md index 1f82fd608..8786ae964 100644 --- a/docs/user_guide/diffusion/cpu_offload_diffusion.md +++ b/docs/user_guide/diffusion/cpu_offload_diffusion.md @@ -1,101 +1,143 @@ -# CPU Offloading for Diffusion Model +# CPU Offloading for Diffusion Models ## Overview -vLLM-Omni provides two offloading strategies to reduce GPU memory usage for diffusion models, allowing you to run larger models on GPUs with limited VRAM: +vLLM-Omni provides two offloading strategies to reduce GPU memory usage for diffusion models: -1. **Model-level (Component) Offloading**: Swaps entire model components (DiT transformer, VAE, encoders) between GPU and CPU. -2. **Layerwise (Blockwise) Offloading**: Keeps only a single or a few transformer blocks on GPU at a time, with compute - memory copy overlap. +1. **Model-level (Sequential) Offloading**: Mutual exclusion between DiT model and encoder - only one is on GPU at a time. +2. **Layerwise (Blockwise) Offloading**: Keeps only one transformer block on GPU at a time with compute-memory overlap. -Both approaches use pinned memory for faster CPU-GPU transfers. For now, the two offloading strategies could not be used at the same time. +Both strategies use pinned memory for faster CPU-GPU transfers. The strategies are **mutually exclusive** for now - if both are enabled, layerwise takes priority. -## Model-level CPU Offloading +## Model-level (Sequential) Offloading -### Implementation +### How It Works -CPU offload lets the diffusion worker move large model components between GPU and CPU memory on demand. It keeps the DiT transformer resident on GPU only while it is actively running, and swaps it out when encoders modules need the device. This reduces peak VRAM usage so bigger checkpoints run on smaller GPUs, or multiple requests can share the same GPU. +Model-level offloading implements mutual exclusion between DiT transformer and encoder modules using pre forward hooks: -**Execution Flow**: -1. Text encoders run on GPU while the DiT transformer is offloaded to CPU. -2. Before denoising, weights are prefetched back to GPU, honoring pinned-memory copies for speed. -3. After the diffusion step, the transformer returns to CPU and the process repeats as needed. +- **When encoders run**: DiT transformer is offloaded to CPU +- **When DiT runs**: Encoders are offloaded to CPU +- **VAE**: Stays resident on GPU -Transfers use pinned host buffers, and the worker coordinates swaps via mutex-style hooks so components never compete for memory. +Before each module's forward pass, the hook automatically moves it to GPU while offloading the other module group to CPU. Transfers use pinned memory for speed. -### Configuration -You can enable CPU offload in two ways: - -1. **Python API**: set `enable_cpu_offload=True`. +### Usage +**Python API:** ```python from vllm_omni import Omni -if __name__ == "__main__": - - m = Omni(model="Qwen/Qwen-Image",enable_cpu_offload=True) +m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True) ``` -2. **CLI**: pass `--enable-cpu-offload` to the diffusion service entrypoint. +**CLI:** +```bash +vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-cpu-offload +``` ### Limitations -- Cold start latency increases for over one minute for some models(e.g., Qwen-Image) +- Cold start latency increases +- Adds overhead from CPU-GPU transfers between encoder and denoising phases +- Support single GPU only for now ## Layerwise (Blockwise) Offloading -### Implementation -Layerwise offload operates at transformer block granularity, keeping a single transformer block, or a specified number of blocks, on GPU while others stay in CPU memory. +### How It Works -Unlike full model-wise CPU offload which swaps entire components like DiT and encoders, layerwise offloading applies a sliding window way of loading and offloading weights between gpu and cpu: while block `i` computes, block `i+1` get prefetched asynchronously via pinned memory. In this way, only partial blocks(s) reside on GPU at any moment during inference, so that greatly decrease the memory occupancy. +Layerwise offloading keeps only one transformer block on GPU at a time. -**Execution Flow**: +As each block completes, the next block is prefetched to GPU while the current block is freed. The pre and forward hooks utilized by layerwise offloading apply a separate CUDA stream (`copy_stream`) to overlap weight transfer with computation, and retain flattened tensors in pinned CPU memory for block parameters re-materialization. Encoders, VAE, and non-block DiT modules (embeddings, norms) always stay on GPU. -1. During model initialization, all components are loaded to CPU first. Then components other than DiT model(s) in the pipeline, such as VAE and encoders, are moved to GPU. The weights of target transformer blocks are collected as contiguous tensors per layer on CPU with pinned memory; and non-block modules (embeddings, norms, etc) in the DiT model are moved to and stay on GPU. -2. The first block(s) are transferred to GPU during initialization of `LayerwiseOffloader`, before the first denoising step of the very first request. -3. As each block executes, the next block prefetches on a separate CUDA stream for compute - memory copy overlap. After execution, the current block is immediately freed from GPU memory. -4. When the last block completes, the first block prefetches for the next denoising step. +**Execution Flow:** +| Block | Pre-forward Hook | Forward | Post-forward Hook | +|-------|------------------|---------|-------------------| +| block-0 | Prefetch block-1 (async) | Compute block-0 | Free block-0 | +| block-1 | Prefetch block-2 (async) | Compute block-1 | Free block-1 | +| ... | ... | ... | ... | +| block-(n-1) | **Prefetch block-0** (async) | Compute block-(n-1) | Free block-(n-1) | -Example of hook executions of a DiT model with n layers, by default keep a single layer on GPU: -| Layer (block) idx | forward pre-hook | forward | forward post-hook | -|-------------------|--------------------------------|------------------|---------------------------| -| layer-0 | prefetch layer 1 (copy stream) | compute layer 0 | free layer-0 gpu weights | -| layer-1 | prefetch layer 2 (copy stream) | compute layer 1 | free layer-1 gpu weights | -| layer-2 | prefetch layer 3 (copy stream) | compute layer 2 | free layer-2 gpu weights | -| ... | ... | ... | ... | -| layer-(n-1) | **prefetch layer 0 (copy stream)** | compute layer (n-1) | free layer (n-1) gpu weights | +Each transformer block has a `LayerwiseOffloadHook` that prefetches the next block before forward and frees the current block after forward. +Layerwise offloading is primarily recommended for large **video generation models** where the compute cost per block is high enough to effectively overlap with memory prefetch operations. For example, Wan2.2 T2V and I2V pipelines. -### Configuration - -1. **Python API**: set `enable_layerwise_offload=True` and optionally `layerwise_num_gpu_layers`. +### Usage +**Python API:** ```python from vllm_omni import Omni -if __name__ == "__main__": - m = Omni( - model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", - enable_layerwise_offload=True, - ... - ) +# Text-to-video +m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_layerwise_offload=True) + +# Or image-to-video +m = Omni(model="Wan-AI/Wan2.2-I2V-A14B-Diffusers", enable_layerwise_offload=True) ``` -2. **CLI**: pass `--enable-layerwise-offload` and `--layerwise-num-gpu-layers` to the diffusion service entrypoint. +**CLI:** +```bash +# Text-to-video +vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-layerwise-offload -### Supported Models +# Or image-to-video +vllm-omni serve diffusion Wan-AI/Wan2.2-I2V-A14B-Diffusers --enable-layerwise-offload +``` -| Architecture | Models | Example HF Models | DiT Model Cls | Blocks Attr Name | -|--------------|--------|-------------------|----------|----------| -| `QwenImagePipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | "transformer_blocks" | -| `Wan22Pipeline` | Wan2.2 | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | "blocks" | +### To Support a Model -NOTE: Models must define `_layerwise_offload_blocks_attr` class attribute so that the layerwise offloader finds the target transformer blocks. +Models must define the blocks attribute name for layerwise offloading: + +```python +class WanTransformer3DModel(nn.Module): + _layerwise_offload_blocks_attr = "blocks" # Attribute name containing transformer blocks + + def __init__(self): + self.blocks = nn.ModuleList([...]) # Transformer blocks +``` ### Limitations - Cold start latency increases because of - 1) components are loaded to CPU first at the very first during initialization, + 1) components are loaded to CPU first at the very first during initialization, 2) weight consolidation and pinning -- Performance depends on CPU <-> GPU interconnection (e.g., PCIe bandwidth). +- Performance depends on compute cost and H2D bandwidth as well - Support single GPU only for now + + +### Implementation Notes + +**Module Discovery** + +The offloader automatically discovers pipeline components: + +- **DiT modules**: `transformer`, `transformer_2`, `dit` +- **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder` +- **VAE**: `vae` + +**Hook System** + +Both strategies use vLLM-Omni's hook registry system (`HookRegistry` and `ModelHook`) to register pre/post forward callbacks on modules, enabling automatic swapping without modifying model code. + +**Backend Architecture** + +``` +OffloadBackend (base class) +├── ModelLevelOffloadBackend → uses SequentialOffloadHook +└── LayerWiseOffloadBackend → uses LayerwiseOffloadHook +``` + +Factory function `get_offload_backend()` selects the appropriate backend based on configuration. + + +## Supported Models + +| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Blocks Attr (Layerwise specific) | +|--------------|----------------|-----------|---------------------|-------------------|-------------| +| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` | +| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` | +| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | `"transformer_blocks"` | + +**Notes:** +- Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally +- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attr` pointing to transformer blocks diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index 6ff7b561b..14cff2726 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -300,12 +300,6 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable layerwise (blockwise) offloading on DiT modules.", ) - parser.add_argument( - "--layerwise-num-gpu-layers", - type=int, - default=1, - help="Number of ready layers (blocks) to keep on GPU during generation.", - ) return parser.parse_args() @@ -362,7 +356,6 @@ def main(): omni = Omni( model=args.model, enable_layerwise_offload=args.enable_layerwise_offload, - layerwise_num_gpu_layers=args.layerwise_num_gpu_layers, vae_use_slicing=args.vae_use_slicing, vae_use_tiling=args.vae_use_tiling, cache_backend=args.cache_backend, diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index e57f54819..5a91c4c73 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -80,12 +80,6 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable layerwise (blockwise) offloading on DiT modules.", ) - parser.add_argument( - "--layerwise-num-gpu-layers", - type=int, - default=1, - help="Number of ready layers (blocks) to keep on GPU during generation.", - ) parser.add_argument( "--cfg-parallel-size", type=int, @@ -139,7 +133,6 @@ def main(): omni = Omni( model=args.model, enable_layerwise_offload=args.enable_layerwise_offload, - layerwise_num_gpu_layers=args.layerwise_num_gpu_layers, vae_use_slicing=args.vae_use_slicing, vae_use_tiling=args.vae_use_tiling, boundary_ratio=args.boundary_ratio, diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 5da0c1879..db87abf00 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -113,12 +113,6 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable layerwise (blockwise) offloading on DiT modules.", ) - parser.add_argument( - "--layerwise-num-gpu-layers", - type=int, - default=1, - help="Number of ready layers (blocks) to keep on GPU during generation.", - ) parser.add_argument( "--quantization", type=str, @@ -222,7 +216,6 @@ def main(): omni = Omni( model=args.model, enable_layerwise_offload=args.enable_layerwise_offload, - layerwise_num_gpu_layers=args.layerwise_num_gpu_layers, vae_use_slicing=args.vae_use_slicing, vae_use_tiling=args.vae_use_tiling, cache_backend=args.cache_backend, 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 f645a47c0..40fafa100 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -84,12 +84,6 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable layerwise (blockwise) offloading on DiT modules.", ) - parser.add_argument( - "--layerwise-num-gpu-layers", - type=int, - default=1, - help="Number of ready layers (blocks) to keep on GPU during generation.", - ) parser.add_argument( "--ulysses-degree", type=int, @@ -155,7 +149,6 @@ def main(): omni = Omni( model=args.model, enable_layerwise_offload=args.enable_layerwise_offload, - layerwise_num_gpu_layers=args.layerwise_num_gpu_layers, vae_use_slicing=args.vae_use_slicing, vae_use_tiling=args.vae_use_tiling, boundary_ratio=args.boundary_ratio, diff --git a/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py b/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py index 87a9e0a9e..1a8b78969 100644 --- a/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py +++ b/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py @@ -17,13 +17,15 @@ if str(REPO_ROOT) not in sys.path: from vllm_omni import Omni # Models to test and expected saved memory in MB, correspondingly -MODELS_SAVED_MEMORY_MB = {"riverclouds/qwen_image_random": 4500} +MODELS_SAVED_MEMORY_MB = { + "riverclouds/qwen_image_random": 4500, + # "Wan-AI/Wan2.2-T2V-A14B-Diffusers": 45000, +} def run_inference( model_name: str, layerwise_offload: bool = False, - num_gpu_layers: int = 1, num_inference_steps: int = 3, ) -> float: # For now, only support on GPU, so apply torch.cuda operations here @@ -36,7 +38,6 @@ def run_inference( m = Omni( model=model_name, enable_layerwise_offload=layerwise_offload, - layerwise_num_gpu_layers=num_gpu_layers, boundary_ratio=0.875, flow_shift=5.0, ) @@ -83,16 +84,12 @@ def test_layerwise_offload_diffusion_model(model_name: str): cleanup_dist_env_and_memory() # Run with layerwise offloading (1 layer on device) - layerwise_offload_peak_memory = run_inference(model_name, layerwise_offload=True, num_gpu_layers=1) + layerwise_offload_peak_memory = run_inference(model_name, layerwise_offload=True) cleanup_dist_env_and_memory() - - # Run with 2 layers on device - layerwise_offload_two_layers_peak = run_inference(model_name, layerwise_offload=True, num_gpu_layers=2) except Exception: pytest.fail("Inference failed") print(f"Layerwise offload peak memory (1 GPU layer): {layerwise_offload_peak_memory} MB") - print(f"Layerwise offload peak memory (2 GPU layers): {layerwise_offload_two_layers_peak} MB") print(f"No offload peak memory: {no_offload_peak_memory} MB") # Verify that layerwise offloading significantly reduces memory usage @@ -101,10 +98,3 @@ def test_layerwise_offload_diffusion_model(model_name: str): f"Layerwise offload peak memory {layerwise_offload_peak_memory} MB " f"should be significantly less than no offload peak memory {no_offload_peak_memory} MB" ) - - # Verify that 2 GPU layers uses more memory than 1 GPU layer - # But not excessively more (should be a reasonable increase) - assert layerwise_offload_peak_memory < layerwise_offload_two_layers_peak, ( - f"1 GPU layer peak {layerwise_offload_peak_memory} MB should be < " - f"2 GPU layers peak {layerwise_offload_two_layers_peak} MB" - ) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index e55bbc6f0..a18411dba 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -298,11 +298,8 @@ class OmniDiffusionConfig: # - Text encoders run on GPU while DiT is on CPU # - DiT runs on GPU while encoders are on CPU enable_cpu_offload: bool = False - # Layer-wise offloading (block-level offloading) parameters enable_layerwise_offload: bool = False - # Number of transformer blocks ready for computation to keep on GPU - layerwise_num_gpu_layers: int = 1 use_fsdp_inference: bool = False pin_cpu_memory: bool = True # Use pinned memory for faster transfers when offloading diff --git a/vllm_omni/diffusion/offload.py b/vllm_omni/diffusion/offload.py deleted file mode 100644 index 6f1d8a0db..000000000 --- a/vllm_omni/diffusion/offload.py +++ /dev/null @@ -1,535 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""CPU offloading utilities for diffusion models. - -This module provides mutual-exclusion CPU offloading between DiT and encoders. -When enable_cpu_offload is enabled: -- Text encoders run on GPU while DiT is on CPU -- DiT runs on GPU while encoders are offloaded to CPU - -This allows running large models on limited GPU memory. -""" - -from __future__ import annotations - -from functools import partial -from itertools import chain -from typing import TYPE_CHECKING, Any - -import torch -from torch import nn -from vllm.logger import init_logger - -from vllm_omni.platforms import current_omni_platform - -if TYPE_CHECKING: - from vllm_omni.diffusion.data import OmniDiffusionConfig - -logger = init_logger(__name__) - - -class SequentialOffloader: - """Sequential offloader: DiT and encoders take turns on GPU. - - Uses PyTorch's forward pre-hooks to automatically swap models: - - Before encoder runs: move DiT modules to CPU, move encoder to GPU - - Before DiT runs: move encoders to CPU, move active DiT to GPU - - This ensures only one large model group is on GPU at a time. - """ - - def __init__( - self, - dits: list[nn.Module], - encoders: list[nn.Module], - device: torch.device, - pin_memory: bool = True, - ): - assert all(isinstance(m, nn.Module) for m in dits), "All dits must be nn.Module" - assert all(isinstance(m, nn.Module) for m in encoders), "All encoders must be nn.Module" - self.dits = dits - self.encoders = encoders - self.device = device - self.pin_memory = pin_memory - self._handles: list = [] - - def _to_cpu(self, module: nn.Module) -> None: - """Move module to CPU with optional memory pinning.""" - # Skip if already on CPU - try: - param = next(module.parameters()) - if param.device.type == "cpu": - return - except StopIteration: - return - - previous_device = param.device - module.to("cpu", non_blocking=True) - - # Release allocator blocks when tensors leave the GPU. - if previous_device.type != "cpu": - torch.cuda.empty_cache() - - if self.pin_memory: - for p in module.parameters(): - if p.data.device.type == "cpu" and not p.data.is_pinned(): - p.data = p.data.pin_memory() - - def _to_gpu(self, module: nn.Module) -> None: - """Move module to GPU.""" - # Skip if already on target device - try: - if next(module.parameters()).device == self.device: - return - except StopIteration: - return - - module.to(self.device, non_blocking=True) - - def _dit_pre_hook(self, module: nn.Module, args: tuple) -> None: - """Before DiT forward: offload encoders, load DiT.""" - for enc in self.encoders: - self._to_cpu(enc) - self._to_gpu(module) - - current_omni_platform.synchronize() - - logger.debug("Swapped: encoders -> CPU, DiT -> GPU") - - def _encoder_pre_hook(self, module: nn.Module, args: tuple) -> None: - """Before encoder forward: offload DiT, load encoder.""" - for dit_mod in self.dits: - self._to_cpu(dit_mod) - self._to_gpu(module) - - current_omni_platform.synchronize() - - logger.debug("Swapped: DiT -> CPU, encoder -> GPU") - - def register(self) -> None: - """Register forward pre-hooks on DiT and encoders.""" - # Hook on each DiT-like module - for dit_mod in self.dits: - h = dit_mod.register_forward_pre_hook(self._dit_pre_hook) - self._handles.append(h) - logger.debug("Registered offload hook for %s", dit_mod.__class__.__name__) - - # Hook on each encoder - for enc in self.encoders: - h = enc.register_forward_pre_hook(self._encoder_pre_hook) - self._handles.append(h) - logger.debug("Registered offload hook for %s", enc.__class__.__name__) - - def remove(self) -> None: - """Remove all hooks.""" - for h in self._handles: - h.remove() - self._handles = [] - - -class LayerwiseOffloader: - """Layer-wise CPU offloading for transformer blocks. - - Keeps only a sliding window of layers (blocks), by default a single layer, on GPU, - prefetching the next block while the current block computes to approach compute - memcpy overlap. - Unused blocks are freed on GPU. - - Based on implementations from: - https://github.com/sgl-project/sglang/blob/v0.5.8/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py - """ - - def __init__( - self, - blocks: list[nn.Module], - device: torch.device, - pin_memory: bool = True, - num_gpu_layers: int = 1, - ): - assert all(isinstance(m, nn.Module) for m in blocks), "All transformer blocks must be torch.nn.Module" - assert current_omni_platform.is_cuda(), "Layerwise offloading is only supported on cuda devices for now" - - self.blocks = blocks - self.device = device - self.pin_memory = pin_memory - self.num_gpu_layers = num_gpu_layers - self.num_blocks = len(self.blocks) - if self.num_blocks == 0: - raise ValueError("LayerwiseOffloader requires at least one block, but found 0.") - if not (1 <= self.num_gpu_layers <= self.num_blocks): - raise ValueError(f"Invalid num_gpu_layers {self.num_gpu_layers} with {self.num_blocks} blocks") - - self._pre_hook_handles: list = [] - self._post_hook_handles: list = [] - - self._copy_stream = torch.cuda.Stream() - - # Per-layer synchronization primitive: set after H2D copy completes. - self._prefetch_done: list[torch.cuda.Event | None] = [None] * self.num_blocks - - # Simple state to avoid redundant work. - self._resident: list[bool] = [False] * self.num_blocks - - # Pre-allocate gpu tensors - # layer-id -> {dtype -> flattened aggregated cpu tensor} - self.layer_cpu_weights: list[dict[torch.dtype, torch.Tensor]] = [] - self.layer_metadata: list[dict[torch.dtype, list[dict[str, Any]]]] = [] - - self.block_parameters: dict[int, dict[str, nn.Parameter]] = {} - self.block_buffers: dict[int, dict[str, torch.Tensor]] = {} - for layer_idx, block in enumerate(self.blocks): - self.block_parameters[layer_idx] = dict(block.named_parameters()) - self.block_buffers[layer_idx] = dict(block.named_buffers()) - - dtype_cpu_flattened_weights, dtype_metadata = self._to_cpu( - self.block_parameters[layer_idx], self.block_buffers[layer_idx] - ) - self.layer_cpu_weights.append(dtype_cpu_flattened_weights) - self.layer_metadata.append(dtype_metadata) - - if self.num_blocks != len(self.layer_cpu_weights): - logger.error( - f"Inconsistent block layers happened: # of blocks: {self.num_blocks}; " - f"# of layer cpu weights: {len(self.layer_cpu_weights)}" - ) - - # Register pre and post forward hooks on each of the blocks - self.register_block_hooks() - - # Pre-fetch the first layer - # For subsequent requests, the first layer/block will be pre-fetched - # during the last layer compute of the previous request. - self.prefetch_layer(0, non_blocking=False) - - def _to_cpu( - self, params: dict[str, nn.Parameter], bufs: dict[str, torch.Tensor] - ) -> tuple[dict[torch.dtype, torch.Tensor], dict[torch.dtype, list[dict[str, Any]]]]: - """Move block parameters and buffers to CPU, flattening by dtype. - - Consolidates parameters and buffers into contiguous CPU tensors grouped by dtype - for GPU transfers. Replaces original tensors with empty placeholders. - - Returns: - Tuple of - 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] = {} - # order does matter - 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_weights: - dtype_grouped_weights[dtype] = {} - dtype_grouped_weights[dtype][name] = param_or_buf - - for dtype, name2weights in dtype_grouped_weights.items(): - # total # of parameters + buffers - total_numel = sum(t.numel() for _, t in name2weights.items()) - cpu_tensor = torch.empty(total_numel, dtype=dtype, device="cpu", pin_memory=self.pin_memory) - - current_offset = 0 - for name, param_or_buf in name2weights.items(): - numel = param_or_buf.numel() - cpu_tensor[current_offset : current_offset + numel].copy_(param_or_buf.flatten()) - if dtype not in dtype_metadata: - dtype_metadata[dtype] = [] - dtype_metadata[dtype].append( - { - "name": name, - "offset": current_offset, - "numel": numel, - "shape": param_or_buf.shape, - } - ) - - param_or_buf.data = torch.empty((), device=self.device, dtype=dtype) - current_offset += numel - - dtype_cpu_flattened_weights[dtype] = cpu_tensor - - return dtype_cpu_flattened_weights, dtype_metadata - - def register_block_hooks(self) -> None: - """Register forward hooks on blocks for prefetching and offloading.""" - - def _pre_hook(module: nn.Module, args: tuple, *, layer_idx: int) -> None: - # For the last block / layer, prefetch layer 0 (the first layer) - next_id = (layer_idx + 1) % self.num_blocks - self.prefetch_layer(next_id, non_blocking=True) - - def _post_hook(module: nn.Module, args: tuple, output: tuple, *, layer_idx: int) -> None: - self.offload_layer(layer_idx) - self._resident[layer_idx] = False - self._prefetch_done[layer_idx] = None - - for i, layer in enumerate(self.blocks): - pre_hook_fn = partial(_pre_hook, layer_idx=i) - handle = layer.register_forward_pre_hook(pre_hook_fn) - self._pre_hook_handles.append(handle) - - post_hook_fn = partial(_post_hook, layer_idx=i) - handle = layer.register_forward_hook(post_hook_fn) - self._post_hook_handles.append(handle) - - @torch.compiler.disable - def prefetch_layer(self, layer_idx: int, non_blocking: bool = True) -> None: - """Copy layer weights from CPU -> GPU. - - Pre-fetch target layer in an asynchronous way with compute - memory copy overlap, - with non_blocking set to True. - """ - if layer_idx >= self.num_blocks or layer_idx < 0: - logger.warning(f"Invalid layer id specified: {layer_idx}") - return - - self._copy_stream.wait_stream(torch.cuda.current_stream()) - - layers_to_fetch = [(layer_idx + i) % self.num_blocks for i in range(self.num_gpu_layers)] - - for idx in layers_to_fetch: - if self._resident[idx]: - continue - - layer_params = self.block_parameters[idx] - layer_bufs = self.block_buffers[idx] - - evt = torch.cuda.Event() - gpu_weights: dict[torch.dtype, torch.Tensor] = {} - - with torch.cuda.stream(self._copy_stream): - for dtype, cpu_weight in self.layer_cpu_weights[idx].items(): - gpu_weight = torch.empty(cpu_weight.shape, dtype=dtype, device=self.device) - gpu_weight.copy_(cpu_weight, non_blocking=non_blocking) - gpu_weights[dtype] = gpu_weight - - evt.record(self._copy_stream) - - for dtype in self.layer_metadata[idx]: - ordered_metadata: list[dict[str, Any]] = self.layer_metadata[idx][dtype] - - gpu_weight = gpu_weights[dtype] - - for metadata in ordered_metadata: - target_name = metadata["name"] - target_param_or_buf = ( - layer_params[target_name] if target_name in layer_params else layer_bufs[target_name] - ) - - target_param_or_buf.data = gpu_weight[ - metadata["offset"] : metadata["offset"] + metadata["numel"] - ].view(metadata["shape"]) - - self._prefetch_done[idx] = evt - self._resident[idx] = True - - @torch.compiler.disable - def offload_layer(self, layer_idx: int) -> None: - """Free GPU memory for layer by replacing tensors with empty placeholders.""" - if layer_idx >= self.num_blocks or layer_idx < 0: - logger.warning(f"Invalid layer id specified: {layer_idx}") - return - if not self._resident[layer_idx]: - logger.warning(f"{layer_idx} is not residing on GPU") - return - - evt = self._prefetch_done[layer_idx] - if evt is not None: - torch.cuda.current_stream().wait_event(evt) - - # free GPU residency - for _, param in self.block_parameters[layer_idx].items(): - param.data = torch.empty((), device=self.device, dtype=param.dtype) - for _, buf in self.block_buffers[layer_idx].items(): - buf.data = torch.empty((), device=self.device, dtype=buf.dtype) - - def remove_all_hooks(self) -> None: - """Remove all hooks.""" - for h in self._pre_hook_handles: - h.remove() - for h in self._post_hook_handles: - h.remove() - self._pre_hook_handles.clear() - self._post_hook_handles.clear() - - @staticmethod - def get_blocks_attr_name(model: nn.Module) -> str | None: - """Retrieve blocks attribute name from provided DiT model""" - return getattr(model.__class__, "_layerwise_offload_blocks_attr", None) - - @staticmethod - def get_blocks_from_dit(model: nn.Module) -> list[nn.Module]: - """ - Retrieve a list of blocks from provided DiT model. Blocks attribute name - are found by `_layerwise_offload_blocks_attr` set to DiT models. For example, - - ``` - class WanTransformer3DModel(nn.Module): - _layerwise_offload_blocks_attr = "blocks" - ``` - """ - blocks_attr_name = LayerwiseOffloader.get_blocks_attr_name(model) - if blocks_attr_name is None: - logger.warning( - f"No _layerwise_offload_blocks_attr defined for {model.__class__.__name__}, " - "skipping layerwise offloading" - ) - return [] - - _blocks = getattr(model, blocks_attr_name, None) - if _blocks is None: - logger.warning( - f"Blocks (layers) '{blocks_attr_name}' not found on {model.__class__.__name__}, " - "skipping layerwise offloading" - ) - return [] - - return list(_blocks) - - -def apply_offload_hooks( - model: nn.Module, - od_config: OmniDiffusionConfig, - *, - device: torch.device | None = None, -) -> None: - """Apply mutual-exclusion offload hooks based on config. - - When enable_cpu_offload is enabled, DiT and encoders swap GPU access: - - Encoders (text_encoder, text_encoder_2, text_encoder_3, image_encoder) - run on GPU while DiT is on CPU - - DiT runs on GPU while encoders are on CPU - - Args: - model: Diffusion pipeline model - od_config: OmniDiffusionConfig with offload settings - """ - enable_cpu_offload = getattr(od_config, "enable_cpu_offload", False) - enable_layerwise_offload = getattr(od_config, "enable_layerwise_offload", False) - pin_cpu_memory = getattr(od_config, "pin_cpu_memory", True) - - if not enable_cpu_offload and not enable_layerwise_offload: - return - if enable_cpu_offload and enable_layerwise_offload: - # NOTE: Model-wise and layerwise cpu offloading are not supported together at this moment, - # consider layerwise offloading has higher priority than model-wise offloading - enable_cpu_offload = False - logger.info( - "Model-wise and layer-wise CPU offloading are not supported together at this moment. " - "Automatically disabled model-wise offloading." - ) - # For now, model-wise and layer-wise (block-wise) offloading - # are functioning as expected when cuda device is available - if not current_omni_platform.is_cuda() or current_omni_platform.get_device_count() < 1: - logger.info("CPU Offloading requires cuda devices available. Skipping for now...") - return - - # Find DiT/transformer modules - dit_modules: list[nn.Module] = [] - dit_names: list[str] = [] - candidate_attrs = ["transformer", "transformer_2", "dit"] - for attr in candidate_attrs: - if not hasattr(model, attr): - continue - module_obj = getattr(model, attr) - if module_obj is None: - continue - - assert isinstance(module_obj, nn.Module), f"Expected {attr} to be nn.Module, got {type(module_obj)!r}" - - if module_obj in dit_modules: - continue - - dit_modules.append(module_obj) - dit_names.append(attr) - - if not dit_modules: - logger.warning("enable_cpu_offload enabled but no transformer/dit/unet found") - return - if device is None: - try: - device = next(dit_modules[0].parameters()).device - except StopIteration: - try: - device = current_omni_platform.get_torch_device() - except (NotImplementedError, AttributeError): - logger.error("Fail to get device of pipeline. Skipping applying offloading hooks") - return - - # Collect all encoders - encoders: list[nn.Module] = [] - encoder_names: list[str] = [] - for attr in ["text_encoder", "text_encoder_2", "text_encoder_3", "image_encoder"]: - if hasattr(model, attr) and getattr(model, attr) is not None: - encoders.append(getattr(model, attr)) - encoder_names.append(attr) - if not encoders and enable_cpu_offload: - logger.warning("enable_cpu_offload enabled but no encoders found") - return - for enc in encoders: - enc.to(device) - - # Collect VAE - for name in ["vae"]: - module = getattr(model, name, None) - if module is None: - continue - try: - module.to(device, non_blocking=True) - except Exception as exc: - logger.debug("Failed to move %s to GPU: %s", name, exc) - - if enable_cpu_offload: - # Initial state: keep DiT modules on CPU (encoders typically run first) - for dit_mod in dit_modules: - dit_mod.to("cpu") - - torch.cuda.empty_cache() - - if pin_cpu_memory: - for dit_mod in dit_modules: - for p in dit_mod.parameters(): - if p.data.device.type == "cpu" and not p.data.is_pinned(): - p.data = p.data.pin_memory() - - # Register sequential offload hooks - SequentialOffloader(dit_modules, encoders, device, pin_cpu_memory).register() - logger.info( - "CPU offload enabled: %s <-> %s (mutual exclusion)", - ", ".join(dit_names), - ", ".join(encoder_names), - ) - elif enable_layerwise_offload: - logger.info(f"Applying offloading hooks on {dit_names}") - - for i, dit_module in enumerate(dit_modules): - logger.info(f"Applying hook on {dit_names[i]} ({dit_module.__class__.__name__})") - blocks_attr_name = LayerwiseOffloader.get_blocks_attr_name(dit_module) - blocks = LayerwiseOffloader.get_blocks_from_dit(dit_module) - - if not blocks_attr_name or not blocks: - logger.warning( - "Target layers (blocks) are not found. " - f"Skipping offloading on {dit_names[i]} ({dit_module.__class__.__name__})" - ) - continue - - # move modules other than blocks to gpu and keep them on gpu - for name, m in dit_module.named_children(): - # Skip the blocks module (layers to be offloaded) - if name == blocks_attr_name: - logger.debug(f"Skipped module {name}") - continue - - m.to(device) - logger.debug(f"Moved {name} to device {device}") - - # set to the module (transformer) - offloader = LayerwiseOffloader(blocks, device, pin_cpu_memory, od_config.layerwise_num_gpu_layers) - setattr(dit_module, "_layerwise_offloader", offloader) - - logger.info( - f"Layerwise offloading enabled on {len(blocks)} layers (blocks), " - f"with {od_config.layerwise_num_gpu_layers} kept on device" - ) diff --git a/vllm_omni/diffusion/offloader/__init__.py b/vllm_omni/diffusion/offloader/__init__.py new file mode 100644 index 000000000..f4fd78c7d --- /dev/null +++ b/vllm_omni/diffusion/offloader/__init__.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.platforms import current_omni_platform + +from .base import OffloadBackend, OffloadConfig, OffloadStrategy +from .layerwise_backend import LayerWiseOffloadBackend +from .sequential_backend import ( + ModelLevelOffloadBackend, + apply_sequential_offload, + remove_sequential_offload, +) + +logger = init_logger(__name__) + +__all__ = [ + "OffloadBackend", + "OffloadConfig", + "OffloadStrategy", + "LayerWiseOffloadBackend", + "ModelLevelOffloadBackend", + "apply_sequential_offload", + "remove_sequential_offload", + "get_offload_backend", +] + + +def get_offload_backend( + od_config: OmniDiffusionConfig, + device: torch.device | None = None, +) -> OffloadBackend | None: + """Create appropriate offload backend based on configuration. + + Args: + od_config: OmniDiffusionConfig with offload settings + device: Target device (auto-detected if None) + + Returns: + OffloadBackend instance or None if offloading disabled + + Example: + >>> backend = get_offload_backend(od_config, device=torch.device("cuda:0")) + >>> if backend: + ... backend.enable(pipeline) + """ + # Extract and validate configuration + config = OffloadConfig.from_od_config(od_config) + + # Return None if no offloading requested + if config.strategy == OffloadStrategy.NONE: + return None + + # Validate platform (CUDA required for now) + if not current_omni_platform.is_cuda() or current_omni_platform.get_device_count() < 1: + logger.warning("CPU offloading requires CUDA devices. Skipping offloading.") + return None + + # Detect device if not provided + if device is None: + try: + device = current_omni_platform.get_torch_device() + except (NotImplementedError, AttributeError) as exc: + logger.error("Failed to detect device: %s. Skipping offloading.", exc) + return None + + # Create appropriate backend + if config.strategy == OffloadStrategy.MODEL_LEVEL: + return ModelLevelOffloadBackend(config, device) + elif config.strategy == OffloadStrategy.LAYER_WISE: + return LayerWiseOffloadBackend(config, device) + else: + logger.error("Unknown offload strategy: %s", config.strategy) + return None diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py new file mode 100644 index 000000000..6f838d4dd --- /dev/null +++ b/vllm_omni/diffusion/offloader/base.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from enum import Enum + +import torch +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import OmniDiffusionConfig + +logger = init_logger(__name__) + + +class OffloadStrategy(Enum): + NONE = "none" + MODEL_LEVEL = "model_level" # Sequential offloading between DiT and encoders + LAYER_WISE = "layer_wise" # Block-level + + +@dataclass +class OffloadConfig: + strategy: OffloadStrategy + pin_cpu_memory: bool = True + + @classmethod + def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig": + """Extract and validate offload settings from OmniDiffusionConfig. + + For now, enforces mutual exclusion between model-level and layer-wise offloading. + Layer-wise takes priority if both are enabled. + + Args: + od_config: OmniDiffusionConfig with offload settings + + 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) + pin_cpu_memory = getattr(od_config, "pin_cpu_memory", True) + + # Determine strategy (mutual exclusion, layer-wise takes priority) + if 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 + else: + strategy = OffloadStrategy.NONE + + return cls( + strategy=strategy, + pin_cpu_memory=pin_cpu_memory, + ) + + +class OffloadBackend(ABC): + """Base class for CPU offload backends""" + + def __init__(self, config: OffloadConfig, device: torch.device): + self.config = config + self.device = device + self.enabled = False + + @abstractmethod + def enable(self, pipeline: nn.Module) -> None: + """Enable offloading on the pipeline. + + Discovers modules, moves them to appropriate devices, and + registers forward hooks for swapping/prefetching. + + Args: + pipeline: Diffusion pipeline model (e.g., Wan22Pipeline) + """ + raise NotImplementedError + + @abstractmethod + def disable(self) -> None: + """Disable offloading and cleanup resources. + + Removes all registered hooks. Does NOT move modules back to + original devices (caller responsible for that). + """ + raise NotImplementedError + + def is_enabled(self) -> bool: + return self.enabled diff --git a/vllm_omni/diffusion/offloader/layerwise_backend.py b/vllm_omni/diffusion/offloader/layerwise_backend.py new file mode 100644 index 000000000..a69b13e64 --- /dev/null +++ b/vllm_omni/diffusion/offloader/layerwise_backend.py @@ -0,0 +1,362 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from itertools import chain +from typing import Any + +import torch +from torch import nn +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 .module_collector import ModuleDiscovery + +logger = init_logger(__name__) + + +class LayerwiseOffloadHook(ModelHook): + """Hook for layerwise (transformer-block-wise) CPU offloading. + + The hook instance retains parameters for both the current registered block + module and those for the next block, as well as flattened CPU tensors which + record the parameters of the current block module, so that these parameters + could be re-materialized on device in an overlapping way. + This hook should be registered to each of the transformer blocks in DiT + module(s) of the target pipeline. + + Based on implementations from: + https://github.com/sgl-project/sglang/blob/v0.5.8/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py + """ + + _HOOK_NAME = "layerwise_offload" + + def __init__( + self, + next_block: nn.Module, + device: torch.device, + stream: torch.cuda.Stream | None = None, + pin_memory: bool = True, + ): + assert isinstance(next_block, nn.Module), "transformer block must be type `torch.nn.Module`" + assert current_omni_platform.is_cuda(), "Layerwise offloading is only supported on cuda devices for now" + + self.next_block = next_block + self.device = device + self.copy_stream = stream or torch.cuda.current_stream() + self.pin_memory = pin_memory + + # Per-block synchronization primitive: set after H2D copy completes. + self._prefetch_done: torch.cuda.Event | None = None + + self.next_block_parameters: dict[str, nn.Parameter] = {} + 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]]] = {} + + def initialize_hook(self, module: nn.Module) -> nn.Module: + # This all happen during the hook instance being registered to hook registry; + # the input module is kept intact + module = super().initialize_hook(module) + + self.block_parameters: dict[str, nn.Parameter] = dict(module.named_parameters()) + self.block_buffers: dict[str, torch.Tensor] = dict(module.named_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()) + + # Pre-allocate gpu tensors in a flattened way + self.dtype_cpu_flattened_weights, self.dtype_metadata = LayerwiseOffloadHook._to_cpu( + self.next_block_parameters, self.next_block_buffers, self.device, self.pin_memory + ) + + return module + + @staticmethod + 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. + + Consolidates parameters and buffers into contiguous CPU tensors grouped by dtype + for GPU transfers. Replaces original tensors with empty placeholders. + + Returns: + Tuple of + 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]]] = {} + + 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(): + # total # of parameters + buffers + total_numel = sum(t.numel() for _, t in name2weights.items()) + cpu_tensor = torch.empty(total_numel, dtype=dtype, device="cpu", pin_memory=pin_memory) + + current_offset = 0 + for name, param_or_buf in name2weights.items(): + numel = param_or_buf.numel() + cpu_tensor[current_offset : current_offset + numel].copy_(param_or_buf.flatten()) + if dtype not in dtype_metadata: + dtype_metadata[dtype] = [] + dtype_metadata[dtype].append( + { + "name": name, + "offset": current_offset, + "numel": numel, + "shape": param_or_buf.shape, + } + ) + + param_or_buf.data = torch.empty((), device=device, dtype=dtype) + current_offset += numel + + dtype_cpu_flattened_weights[dtype] = cpu_tensor + + return dtype_cpu_flattened_weights, dtype_metadata + + @torch.compiler.disable + def prefetch_layer(self, non_blocking: bool = True) -> None: + """Copy layer weights from CPU -> GPU. + + Pre-fetch target block in an asynchronous way with compute - memory copy overlap, + with non_blocking set to True. + """ + self.copy_stream.wait_stream(torch.cuda.current_stream()) + + layer_params = self.next_block_parameters + layer_bufs = self.next_block_buffers + + evt = torch.cuda.Event() + gpu_weights: dict[torch.dtype, torch.Tensor] = {} + + with torch.cuda.stream(self.copy_stream): + for dtype, cpu_weight in self.dtype_cpu_flattened_weights.items(): + gpu_weight = torch.empty(cpu_weight.shape, dtype=dtype, device=self.device) + gpu_weight.copy_(cpu_weight, non_blocking=non_blocking) + gpu_weights[dtype] = gpu_weight + + evt.record(self.copy_stream) + + for dtype, ordered_metadata in self.dtype_metadata.items(): + # ordered_metadata: list[dict[str, Any]] + gpu_weight = gpu_weights[dtype] + + for metadata in ordered_metadata: + target_name = metadata["name"] + target_param_or_buf = ( + layer_params[target_name] if target_name in layer_params else layer_bufs[target_name] + ) + + target_param_or_buf.data = gpu_weight[metadata["offset"] : metadata["offset"] + metadata["numel"]].view( + metadata["shape"] + ) + + self._prefetch_done = evt + + @torch.compiler.disable + def offload_layer(self) -> None: + """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: + torch.cuda.current_stream().wait_event(evt) + + self._prefetch_done = None + + # free GPU residency + for _, param in self.block_parameters.items(): + param.data = torch.empty((), device=self.device, dtype=param.dtype) + for _, buf in self.block_buffers.items(): + buf.data = torch.empty((), device=self.device, dtype=buf.dtype) + + def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]: + self.prefetch_layer(non_blocking=True) + + return args, kwargs + + def post_forward(self, module: nn.Module, output: Any) -> Any: + self.offload_layer() + + return output + + +def apply_block_hook( + module: nn.Module, + next_block: nn.Module, + device: torch.device, + stream: torch.cuda.Stream | None = None, + pin_memory: bool = True, +) -> LayerwiseOffloadHook: + registry = HookRegistry.get_or_create(module) + hook = LayerwiseOffloadHook(next_block, device, stream, pin_memory) + registry.register_hook(LayerwiseOffloadHook._HOOK_NAME, hook) + + return hook + + +def remove_block_hook(module: nn.Module) -> None: + registry: HookRegistry | None = getattr(module, "_hook_registry", None) + if registry is not None: + registry.remove_hook(LayerwiseOffloadHook._HOOK_NAME) + logger.debug("Removed offload hook from %s", module.__class__.__name__) + + +class LayerWiseOffloadBackend(OffloadBackend): + """Layer-wise (block-level) offloading backend. + + Implements sliding window offloading where only a small number of transformer + blocks reside on GPU at a time. Blocks are prefetched asynchronously while + previous blocks compute, and freed after use. + """ + + def __init__(self, config: OffloadConfig, device: torch.device): + super().__init__(config, device) + + self.copy_stream = torch.cuda.Stream() + self._blocks: list[list[nn.Module]] = [] + + 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") + return + + # Move encoders to GPU (they stay resident) + for enc in modules.encoders: + enc.to(self.device) + + # Move VAE to GPU if available + if modules.vae is not None: + try: + modules.vae.to(self.device, non_blocking=True) + except Exception as exc: + logger.debug("Failed to move VAE to GPU: %s", exc) + + 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_name = LayerWiseOffloadBackend.get_blocks_attr_name(dit_module) + blocks = LayerWiseOffloadBackend.get_blocks_from_dit(dit_module) + + if not blocks_attr_name or 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 + + num_blocks = len(blocks) + if num_blocks <= 1: + logger.warning( + "#Target layers (blocks) <= 1. Skipping offloading on %s (%s)", + dit_name, + dit_module.__class__.__name__, + ) + dit_module.to(self.device) + continue + + # Move non-block modules to GPU (they stay resident) + for name, m in dit_module.named_children(): + if name == blocks_attr_name: + logger.debug(f"Skipped blocks module {name}") + continue + m.to(self.device) + logger.debug(f"Moved {name} to device {self.device}") + + # 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] + hook = apply_block_hook(last_block, first_block, self.device, self.copy_stream, self.config.pin_cpu_memory) + hook.prefetch_layer(non_blocking=False) + + # Register hook for each of blocks + for i, block in enumerate(blocks[:-1]): + next_block = blocks[(i + 1) % num_blocks] + apply_block_hook(block, next_block, self.device, self.copy_stream, self.config.pin_cpu_memory) + + 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 + + def disable(self) -> None: + if not self.enabled: + return + + for blocks in self._blocks: + for block in blocks: + remove_block_hook(block) + + self._blocks.clear() + self.enabled = False + logger.info("Layer-wise offloading disabled") + + @staticmethod + def get_blocks_attr_name(model: nn.Module) -> str | None: + """Retrieve blocks attribute name from provided DiT model""" + return getattr(model.__class__, "_layerwise_offload_blocks_attr", None) + + @staticmethod + def set_blocks_attr_name(model: nn.Module, name: str) -> None: + if not hasattr(model.__class__, "_layerwise_offload_blocks_attr"): + setattr(model.__class__, "_layerwise_offload_blocks_attr", name) + + @staticmethod + def get_blocks_from_dit(model: nn.Module) -> list[nn.Module]: + """ + Retrieve a list of blocks from provided DiT model. Blocks attribute name + are found by `_layerwise_offload_blocks_attr` set to DiT models. For example, + + ``` + class WanTransformer3DModel(nn.Module): + _layerwise_offload_blocks_attr = "blocks" + ``` + """ + blocks_attr_name = LayerWiseOffloadBackend.get_blocks_attr_name(model) + if blocks_attr_name is None: + logger.warning( + f"No _layerwise_offload_blocks_attr defined for {model.__class__.__name__}, " + "skipping layerwise offloading" + ) + return [] + + _blocks = getattr(model, blocks_attr_name, None) + if _blocks is None: + logger.warning( + f"Blocks (layers) '{blocks_attr_name}' not found on {model.__class__.__name__}, " + "skipping layerwise offloading" + ) + return [] + + return list(_blocks) diff --git a/vllm_omni/diffusion/offloader/module_collector.py b/vllm_omni/diffusion/offloader/module_collector.py new file mode 100644 index 000000000..307ca53a8 --- /dev/null +++ b/vllm_omni/diffusion/offloader/module_collector.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass + +from torch import nn +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +@dataclass +class PipelineModules: + dits: list[nn.Module] + dit_names: list[str] + encoders: list[nn.Module] + encoder_names: list[str] + vae: nn.Module | None = None + + +class ModuleDiscovery: + """Discovers pipeline components for offloading""" + + DIT_ATTRS = ["transformer", "transformer_2", "dit", "language_model", "transformer_blocks"] + ENCODER_ATTRS = ["text_encoder", "text_encoder_2", "text_encoder_3", "image_encoder"] + VAE_ATTRS = ["vae"] + + @staticmethod + def discover(pipeline: nn.Module) -> PipelineModules: + """Discover DiT, encoder, and VAE modules from pipeline. + + Args: + pipeline: Diffusion pipeline model + + Returns: + PipelineModules with lists of discovered modules and names + """ + # Collect DiT/transformer modules + dit_modules: list[nn.Module] = [] + dit_names: list[str] = [] + for attr in ModuleDiscovery.DIT_ATTRS: + if not hasattr(pipeline, attr): + continue + module_obj = getattr(pipeline, attr) + if module_obj is None: + continue + + if not isinstance(module_obj, nn.Module): + logger.warning(f"Expected {attr} to be nn.Module, got {type(module_obj)!r}") + continue + + if module_obj in dit_modules: + continue + + dit_modules.append(module_obj) + dit_names.append(attr) + + # Collect all encoders + encoders: list[nn.Module] = [] + encoder_names: list[str] = [] + for attr in ModuleDiscovery.ENCODER_ATTRS: + if hasattr(pipeline, attr) and getattr(pipeline, attr) is not None: + encoders.append(getattr(pipeline, attr)) + encoder_names.append(attr) + + # Collect VAE + vae = None + for attr in ModuleDiscovery.VAE_ATTRS: + module = getattr(pipeline, attr, None) + if module is not None: + vae = module + break + + return PipelineModules( + dits=dit_modules, + dit_names=dit_names, + encoders=encoders, + encoder_names=encoder_names, + vae=vae, + ) diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py new file mode 100644 index 000000000..30edc8256 --- /dev/null +++ b/vllm_omni/diffusion/offloader/sequential_backend.py @@ -0,0 +1,227 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch import nn +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 .module_collector import ModuleDiscovery + +logger = init_logger(__name__) + + +class SequentialOffloadHook(ModelHook): + """Hook for sequential offloading with mutual exclusion on encoder and DiT modules. + + To be used as a model-level (or "component-level") of CPU offloading method; + When a module's forward is called, this hook offloads target modules to CPU + and loads the current module to GPU. + """ + + _HOOK_NAME = "sequential_offload" + + def __init__( + self, + offload_targets: list[nn.Module], + device: torch.device, + pin_memory: bool = True, + ): + # Modules to offload to CPU before this module runs + self.offload_targets = offload_targets + self.device = device + self.pin_memory = pin_memory + + def _to_cpu(self, module: nn.Module) -> None: + """Move module to CPU.""" + try: + param = next(module.parameters()) + except StopIteration: + return + + previous_device = param.device + # Skip if already on CPU + if previous_device.type == "cpu": + return + + module.to("cpu") + torch.cuda.empty_cache() + + if self.pin_memory: + for p in module.parameters(): + if p.data.device.type == "cpu" and not p.data.is_pinned(): + p.data = p.data.pin_memory() + + def _to_gpu(self, module: nn.Module) -> None: + """Move module to GPU.""" + try: + # Skip if already on target device + if next(module.parameters()).device == self.device: + return + except StopIteration: + return + + module.to(self.device) + + def pre_forward(self, module: nn.Module, *args, **kwargs) -> tuple[tuple, dict]: + # Offload target modules to CPU + for target in self.offload_targets: + self._to_cpu(target) + + # Load current module to GPU + self._to_gpu(module) + current_omni_platform.synchronize() + + logger.debug( + "Swapped: %s -> CPU, %s -> GPU", + [t.__class__.__name__ for t in self.offload_targets], + module.__class__.__name__, + ) + + return args, kwargs + + +def apply_sequential_offload( + dit_modules: list[nn.Module], + encoder_modules: list[nn.Module], + device: torch.device, + pin_memory: bool = True, +) -> None: + """Apply sequential offloading hooks to DiT and encoder modules. + + Registers hooks on modules to implement mutual-exclusion GPU allocation. + - Before DiT runs, encoders are offloaded to CPU. + - Before encoders run, DiT is offloaded to CPU. + + Args: + dit_modules: DiT/transformer modules to register hooks on + encoder_modules: Encoder modules to register hooks on + device: Target GPU device for loading + pin_memory: Whether to pin CPU memory for faster transfers + + Example: + >>> apply_sequential_offload( + ... dit_modules=[pipeline.transformer], + ... encoder_modules=[pipeline.text_encoder, pipeline.vae], + ... device=torch.device("cuda:0"), + ... ) + >>> # Modules of pipeline now automatically swap between CPU and GPU + """ + # Register hooks on DiT modules (offload encoders when DiT runs) + for dit_mod in dit_modules: + registry = HookRegistry.get_or_create(dit_mod) + hook = SequentialOffloadHook( + offload_targets=encoder_modules, + device=device, + pin_memory=pin_memory, + ) + registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook) + logger.debug("Registered offload hook for %s", dit_mod.__class__.__name__) + + # 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, + ) + registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook) + logger.debug("Registered offload hook for %s", enc.__class__.__name__) + + +def remove_sequential_offload(modules: list[nn.Module]) -> None: + """Remove sequential offloading hooks from modules. + + Args: + modules: Modules to remove hooks from + + Example: + >>> all_modules = [*dit_modules, *encoder_modules] + >>> remove_sequential_offload(all_modules) + """ + 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__) + + +class ModelLevelOffloadBackend(OffloadBackend): + """Model-level (sequential) offloading backend. + + Uses SequentialOffloadHook registered via HookRegistry for automatic module swapping. + """ + + def __init__(self, config: OffloadConfig, device: torch.device): + super().__init__(config, device) + self._offload_modules: list[nn.Module] = [] # Track modules with hooks + + def enable(self, pipeline: nn.Module) -> None: + if self.enabled: + logger.warning("ModelLevelOffloadBackend already enabled") + return + + modules = ModuleDiscovery.discover(pipeline) + if not modules.dits: + logger.warning("No DiT/transformer modules found, skipping model-level offloading") + return + if not modules.encoders: + logger.warning("No encoder modules found, skipping model-level offloading") + return + + # Move encoders to GPU + for enc in modules.encoders: + enc.to(self.device) + + # Move VAE to GPU if available + if modules.vae is not None: + try: + modules.vae.to(self.device, non_blocking=True) + except Exception as exc: + logger.debug("Failed to move VAE to GPU: %s", exc) + + # Initial state: keep DiT modules on CPU (encoders typically run first) + # TODO: This part seems to be unnecessary, remove it after testing + # for dit_mod in modules.dits: + # dit_mod.to("cpu") + + # torch.cuda.empty_cache() + + # if self.config.pin_cpu_memory: + # for dit_mod in modules.dits: + # for p in dit_mod.parameters(): + # if p.data.device.type == "cpu" and not p.data.is_pinned(): + # p.data = p.data.pin_memory() + + # 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, + ) + + # Track modules for cleanup + self._offload_modules = [*modules.dits, *modules.encoders] + + self.enabled = True + + logger.info( + "Model-level offloading enabled: %s <-> %s (mutual exclusion)", + ", ".join(modules.dit_names), + ", ".join(modules.encoder_names), + ) + + def disable(self) -> None: + if not self.enabled: + return + + remove_sequential_offload(self._offload_modules) + + self._offload_modules.clear() + self.enabled = False + logger.info("Model-level offloading disabled") diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 1e00e4b0c..8362eef18 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -26,7 +26,7 @@ from vllm_omni.diffusion.compile import regionally_compile from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader -from vllm_omni.diffusion.offload import apply_offload_hooks +from vllm_omni.diffusion.offloader import get_offload_backend from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager from vllm_omni.platforms import current_omni_platform @@ -62,6 +62,7 @@ class DiffusionModelRunner: self.device = device self.pipeline = None self.cache_backend = None + self.offload_backend = None # Initialize KV cache manager for connector management self.kv_transfer_manager = OmniKVTransferManager.from_od_config(od_config) @@ -107,8 +108,10 @@ class DiffusionModelRunner: logger.info("Model runner: Model loaded successfully.") # Apply CPU offloading - if self.od_config.enable_cpu_offload or self.od_config.enable_layerwise_offload: - apply_offload_hooks(self.pipeline, self.od_config, device=self.device) + self.offload_backend = get_offload_backend(self.od_config, device=self.device) + if self.offload_backend is not None: + logger.info(f" Enabling offloader backend: {self.offload_backend.__class__.__name__}") + self.offload_backend.enable(self.pipeline) # Apply torch.compile if not in eager mode if not self.od_config.enforce_eager: diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index ad55c048d..0be77bf73 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -162,7 +162,6 @@ class AsyncOmni(OmniBase): "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False), "enable_cpu_offload": kwargs.get("enable_cpu_offload", False), "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False), - "layerwise_num_gpu_layers": kwargs.get("layerwise_num_gpu_layers", False), "enforce_eager": kwargs.get("enforce_eager", False), }, "final_output": True, diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 1be3b02f1..a59cc5673 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -204,12 +204,6 @@ class OmniServeCommand(CLISubcommand): action="store_true", help="Enable layerwise (blockwise) offloading on DiT modules.", ) - omni_config_group.add_argument( - "--layerwise-num-gpu-layers", - type=int, - default=1, - help="Number of layers (blocks) to keep on GPU during generation.", - ) # Video model parameters (e.g., Wan2.2) - engine-level omni_config_group.add_argument(