From 779f5987d9806530a0ea35a900261fe441a282ce Mon Sep 17 00:00:00 2001 From: Prajwal A <34590600+LawJarp-A@users.noreply.github.com> Date: Mon, 12 Jan 2026 16:11:18 +0530 Subject: [PATCH] [feature] cpu offloading support for diffusion (#497) Signed-off-by: zjy0516 Signed-off-by: Prajwal A Co-authored-by: zjy0516 --- .buildkite/pipeline.yml | 17 ++ docs/.nav.yml | 1 + docs/features/cpu_offload_diffusion.md | 29 +++ .../test_diffusion_cpu_offload.py | 90 ++++++++ vllm_omni/diffusion/data.py | 10 +- vllm_omni/diffusion/offload.py | 202 ++++++++++++++++++ vllm_omni/diffusion/registry.py | 1 + vllm_omni/diffusion/worker/gpu_worker.py | 18 +- vllm_omni/entrypoints/async_omni.py | 1 + vllm_omni/entrypoints/cli/serve.py | 7 + 10 files changed, 370 insertions(+), 6 deletions(-) create mode 100644 docs/features/cpu_offload_diffusion.md create mode 100644 tests/e2e/offline_inference/test_diffusion_cpu_offload.py create mode 100644 vllm_omni/diffusion/offload.py diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 63fea9199..f18aef617 100644 --- a/.buildkite/pipeline.yml +++ b/.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 diff --git a/docs/.nav.yml b/docs/.nav.yml index ad8545524..abe4865af 100644 --- a/docs/.nav.yml +++ b/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 diff --git a/docs/features/cpu_offload_diffusion.md b/docs/features/cpu_offload_diffusion.md new file mode 100644 index 000000000..aaa4243a3 --- /dev/null +++ b/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) diff --git a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py new file mode 100644 index 000000000..0066d49b1 --- /dev/null +++ b/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" + ) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index b2faf067c..49a49829b 100644 --- a/vllm_omni/diffusion/data.py +++ b/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 diff --git a/vllm_omni/diffusion/offload.py b/vllm_omni/diffusion/offload.py new file mode 100644 index 000000000..4220ee9ff --- /dev/null +++ b/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), + ) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 7c57ac3a8..c49ba0a3c 100644 --- a/vllm_omni/diffusion/registry.py +++ b/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.") diff --git a/vllm_omni/diffusion/worker/gpu_worker.py b/vllm_omni/diffusion/worker/gpu_worker.py index 7ae02b16c..99a718e38 100644 --- a/vllm_omni/diffusion/worker/gpu_worker.py +++ b/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( diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index b3c54736e..1cd7f5063 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/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", diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 4c585953b..be1a1bf14 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/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",