[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>

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),
+    )