2923d5a91e
Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com> Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
536 lines
21 KiB
Python
536 lines
21 KiB
Python
# 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"
|
|
)
|