[feature] cpu offloading support for diffusion (#497)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Signed-off-by: Prajwal A <prajwalanagani@gmail.com>
Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>

.buildkite/pipeline.yml

@@ -42,6 +42,23 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
+  - label: "Diffusion Model CPU offloading Test"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
   - label: "Audio Generation Model Test"
     timeout_in_minutes: 20
     depends_on: image-build

docs/.nav.yml

@@ -36,6 +36,7 @@ nav:
     - models/supported_models.md
   - Features:
     - Sleep Mode: features/sleep_mode.md
+    - CPU Offloading for Diffusion Model: features/cpu_offload_diffusion.md
 - Developer Guide:
   - General:
     - contributing/README.md

docs/features/cpu_offload_diffusion.md

@@ -0,0 +1,29 @@
+# CPU Offloading for Diffusion Model
+
+## Overview
+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.
+
+## Execution Model
+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.
+
+Transfers use pinned host buffers, and the worker coordinates swaps via mutex-style hooks so components never compete for memory.
+
+## Configuration
+You can enable CPU offload in two ways:
+
+- **Python API**: set `enable_cpu_offload=True`.
+
+```python
+from vllm_omni import Omni
+
+if __name__ == "__main__":
+
+    m = Omni(model="Qwen/Qwen-Image",enable_cpu_offload=True)
+```
+
+- **CLI**: pass `--dit-cpu-offload` to the diffusion service entrypoint.
+
+## Known Limitations
+- Cold start latency increases for over one minute for some models(e.g., Qwen-Image)

tests/e2e/offline_inference/test_diffusion_cpu_offload.py

@@ -0,0 +1,90 @@
+import sys
+import threading
+import time
+from pathlib import Path
+
+import pytest
+import torch
+
+from vllm_omni.utils.platform_utils import is_npu, is_rocm
+
+# ruff: noqa: E402
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from vllm_omni import Omni
+
+
+class GPUMemoryMonitor:
+    """Poll global device memory usage via CUDA APIs."""
+
+    def __init__(self, device_index: int, interval: float = 0.05):
+        self.device_index = device_index
+        self.interval = interval
+        self.peak_used_mb = 0.0
+        self._stop_event = threading.Event()
+        self._thread: threading.Thread | None = None
+
+    def start(self) -> None:
+        def monitor_loop() -> None:
+            while not self._stop_event.is_set():
+                try:
+                    with torch.cuda.device(self.device_index):
+                        free_bytes, total_bytes = torch.cuda.mem_get_info()
+                    used_mb = (total_bytes - free_bytes) / (1024**2)
+                    self.peak_used_mb = max(self.peak_used_mb, used_mb)
+                except Exception:
+                    pass
+                time.sleep(self.interval)
+
+        self._thread = threading.Thread(target=monitor_loop, daemon=True)
+        self._thread.start()
+
+    def stop(self) -> None:
+        if self._thread is None:
+            return
+        self._stop_event.set()
+        self._thread.join(timeout=2.0)
+
+
+models = ["riverclouds/qwen_image_random"]
+
+
+@pytest.mark.skipif(is_npu() or is_rocm(), reason="Hardware not supported")
+@pytest.mark.parametrize("model_name", models)
+def test_cpu_offload_diffusion_model(model_name: str):
+    def inference(offload: bool = True):
+        torch.cuda.empty_cache()
+        device_index = torch.cuda.current_device()
+        monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
+        monitor.start()
+        m = Omni(model=model_name, enable_cpu_offload=offload)
+        torch.cuda.reset_peak_memory_stats(device=device_index)
+        height = 256
+        width = 256
+
+        m.generate(
+            "a photo of a cat sitting on a laptop keyboard",
+            height=height,
+            width=width,
+            num_inference_steps=9,
+            guidance_scale=0.0,
+            generator=torch.Generator("cuda").manual_seed(42),
+        )
+
+        monitor.stop()
+        torch.cuda.synchronize(device_index)
+        fallback_alloc = torch.cuda.max_memory_allocated(device=device_index) / (1024**2)
+        fallback_reserved = torch.cuda.max_memory_reserved(device=device_index) / (1024**2)
+        peak_memory_mb = max(monitor.peak_used_mb, fallback_alloc, fallback_reserved)
+
+        return peak_memory_mb
+
+    offload_peak_memory = inference(offload=True)
+    no_offload_peak_memory = inference(offload=False)
+    print(f"Offload peak memory: {offload_peak_memory} MB")
+    print(f"No offload peak memory: {no_offload_peak_memory} MB")
+    assert offload_peak_memory + 2500 < no_offload_peak_memory, (
+        f"Offload peak memory {offload_peak_memory} MB should be less than no offload peak memory {no_offload_peak_memory} MB"
+    )

vllm_omni/diffusion/data.py

@@ -286,12 +286,12 @@ class OmniDiffusionConfig:
     output_type: str = "pil"
 
     # CPU offload parameters
-    dit_cpu_offload: bool = True
+    # 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
+    enable_cpu_offload: bool = False
     use_fsdp_inference: bool = False
-    text_encoder_cpu_offload: bool = True
-    image_encoder_cpu_offload: bool = True
-    vae_cpu_offload: bool = True
-    pin_cpu_memory: bool = True
+    pin_cpu_memory: bool = True  # Use pinned memory for faster transfers when offloading
 
     # VAE memory optimization parameters
     vae_use_slicing: bool = False

vllm_omni/diffusion/offload.py

@@ -0,0 +1,202 @@
+# 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 typing import TYPE_CHECKING
+
+import torch
+from torch import nn
+from vllm.logger import init_logger
+
+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 == "cuda" and torch.cuda.is_available():
+            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)
+        torch.cuda.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)
+        torch.cuda.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 = []
+
+
+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
+    """
+    if not getattr(od_config, "enable_cpu_offload", False):
+        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:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # 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:
+        logger.warning("enable_cpu_offload enabled but no encoders found")
+        return
+
+    # Initial state: keep DiT modules on CPU (encoders typically run first)
+    pin = getattr(od_config, "pin_cpu_memory", True)
+    for dit_mod in dit_modules:
+        dit_mod.to("cpu")
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    if pin and torch.cuda.is_available():
+        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).register()
+
+    logger.info(
+        "CPU offload enabled: %s <-> %s (mutual exclusion)",
+        ", ".join(dit_names),
+        ", ".join(encoder_names),
+    )

vllm_omni/diffusion/registry.py

@@ -99,6 +99,7 @@ def initialize_model(
             model.vae.use_slicing = od_config.vae_use_slicing
         if hasattr(model.vae, "use_tiling"):
             model.vae.use_tiling = od_config.vae_use_tiling
+
         return model
     else:
         raise ValueError(f"Model class {od_config.model_class_name} not found in diffusion model registry.")

vllm_omni/diffusion/worker/gpu_worker.py

@@ -26,6 +26,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 )
 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.request import OmniDiffusionRequest
 
 logger = init_logger(__name__)
@@ -69,6 +70,8 @@ class GPUWorker:
         vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size
         vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size
         self.vllm_config = vllm_config
+        load_device = "cpu" if self.od_config.enable_cpu_offload else str(self.device)
+
         with set_forward_context(vllm_config=vllm_config, omni_diffusion_config=self.od_config):
             init_distributed_environment(world_size=world_size, rank=rank)
             logger.info(f"Worker {self.rank}: Initialized device and distributed environment.")
@@ -90,7 +93,7 @@ class GPUWorker:
                 with DeviceMemoryProfiler() as m:
                     self.pipeline = model_loader.load_model(
                         od_config=self.od_config,
-                        load_device=str(self.device),
+                        load_device=load_device,
                     )
             time_after_load = time.perf_counter()
 
@@ -101,6 +104,19 @@ class GPUWorker:
         )
         logger.info(f"Worker {self.rank}: Model loaded successfully.")
 
+        # Apply CPU offloading (DiT <-> encoders mutual exclusion)
+        if self.od_config.enable_cpu_offload:
+            for name in ["vae"]:
+                module = getattr(self.pipeline, name, None)
+                if module is None:
+                    continue
+                try:
+                    module.to(self.device, non_blocking=True)
+                except Exception as exc:
+                    logger.debug("Failed to move %s to GPU: %s", name, exc)
+
+            apply_offload_hooks(self.pipeline, self.od_config, device=self.device)
+
         if not self.od_config.enforce_eager:
             try:
                 self.pipeline.transformer = regionally_compile(

vllm_omni/entrypoints/async_omni.py

@@ -160,6 +160,7 @@ class AsyncOmni(OmniBase):
                     "vae_use_tiling": kwargs.get("vae_use_tiling", False),
                     "cache_backend": cache_backend,
                     "cache_config": cache_config,
+                    "enable_cpu_offload": kwargs.get("enable_cpu_offload", False),
                 },
                 "final_output": True,
                 "final_output_type": "image",

vllm_omni/entrypoints/cli/serve.py

@@ -181,6 +181,13 @@ class OmniServeCommand(CLISubcommand):
             help="Enable VAE tiling for memory optimization (useful for mitigating OOM issues).",
         )
 
+        # diffusion model offload parameters
+        serve_parser.add_argument(
+            "--enable-cpu-offload",
+            action="store_true",
+            help="Enable CPU offloading for diffusion models.",
+        )
+
         # Video model parameters (e.g., Wan2.2) - engine-level
         serve_parser.add_argument(
             "--boundary-ratio",