vllm-omni/vllm_omni/diffusion/offloader/sequential_backend.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

from collections.abc import Collection, Iterator
from contextlib import contextmanager

import torch
from torch import nn
from torch.distributed._tensor import DTensor  # type: ignore[attr-defined]
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, SupportsModelCpuOffload
from .config import DIT_COMPONENT
from .plan_resolver import resolve_offload_plan

logger = init_logger(__name__)


def _capture_tensor_devices(modules: Collection[nn.Module]) -> list[tuple[torch.Tensor, torch.device]]:
    seen: set[int] = set()
    devices: list[tuple[torch.Tensor, torch.device]] = []
    for module in modules:
        for tensor in (*module.parameters(), *module.buffers()):
            if id(tensor) in seen:
                continue
            seen.add(id(tensor))
            devices.append((tensor, tensor.device))
    return devices


def _restore_tensor_devices(devices: list[tuple[torch.Tensor, torch.device]]) -> None:
    first_error: BaseException | None = None
    for tensor, device in devices:
        try:
            if tensor.device != device:
                tensor.data = tensor.data.to(device)
        except BaseException as exc:
            logger.exception("Failed to restore tensor placement to %s", device)
            first_error = first_error or exc
    if first_error is not None:
        raise RuntimeError("Failed to restore one or more tensor placements") from first_error


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,
        use_hsdp: bool = False,
        offload_after_context: bool = True,
    ):
        # Modules to offload to CPU before this module runs
        self.offload_targets = offload_targets
        self.device = device
        self.pin_memory = pin_memory
        self.use_hsdp = use_hsdp
        self.offload_after_context = offload_after_context

    @staticmethod
    def _move_params(
        module: nn.Module,
        target_device: torch.device,
        *,
        non_blocking: bool = False,
        pin_memory: bool = False,
    ) -> bool:
        """Move module parameters and buffers to device.

        This cls method specifically prevents recursion device movement,
        E.g., Cache-DiT CachedBlocks has attr `transformer` as a ref to original
        transformer blocks, thus `module.to(device)` will fail for recursion calling,
        refer to
        https://github.com/vipshop/cache-dit/blob/v1.2.3/src/cache_dit/caching/cache_blocks/__init__.py#L83
        """
        moved = False
        for p in module.parameters():
            if p.data.device != target_device:
                data = p.data.to(target_device, non_blocking=non_blocking)
                if pin_memory and target_device.type == "cpu" and not isinstance(data, DTensor):
                    data = data.pin_memory()
                p.data = data
                moved = True
        for b in module.buffers():
            if b.device != target_device:
                data = b.data.to(target_device, non_blocking=non_blocking)
                if pin_memory and target_device.type == "cpu" and not isinstance(data, DTensor):
                    data = data.pin_memory()
                b.data = data
                moved = True
        return moved

    def _to_cpu(self, module: nn.Module) -> None:
        # XPU's allocator doesn't respect stream dependencies in empty_cache,
        # so non-blocking copies can race with cache eviction. Use blocking
        # copies on XPU to avoid NULL pointer errors during DMA.
        non_blocking = not self.use_hsdp and not current_omni_platform.is_xpu()
        moved = self._move_params(
            module,
            torch.device("cpu"),
            non_blocking=non_blocking,
            pin_memory=self.pin_memory,
        )
        if moved:
            current_omni_platform.empty_cache()

    def _to_gpu(self, module: nn.Module) -> None:
        self._move_params(module, self.device, non_blocking=False)

    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 -> %s, free memory: %.4f GB",
            [t.__class__.__name__ for t in self.offload_targets],
            module.__class__.__name__,
            f"{self.device.type}:{self.device.index}",
            current_omni_platform.get_free_memory() / 1024 / 1024 / 1024,
        )

        return args, kwargs


def apply_sequential_offload(
    dit_modules: list[nn.Module],
    encoder_modules: list[nn.Module],
    device: torch.device,
    pin_memory: bool = True,
    use_hsdp: bool = False,
    offload_initial_dits: bool = False,
    offload_dit_modules: Collection[nn.Module] | None = None,
    offload_encoder_modules: Collection[nn.Module] | None = None,
) -> 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
        use_hsdp: Whether HSDP is enabled (affects non_blocking behavior)
        offload_initial_dits: Whether to begin with all DiT modules on CPU.
        offload_dit_modules: DiT modules allowed to move to CPU. None selects
            every DiT for backward compatibility.
        offload_encoder_modules: Encoder/stage modules allowed to move to CPU.
            None selects every supplied module for backward compatibility.

    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
    """
    selected_dit_ids = {id(module) for module in (dit_modules if offload_dit_modules is None else offload_dit_modules)}
    selected_encoder_ids = {
        id(module) for module in (encoder_modules if offload_encoder_modules is None else offload_encoder_modules)
    }

    all_modules = [*dit_modules, *encoder_modules]
    initial_devices = _capture_tensor_devices(dit_modules) if offload_initial_dits else []
    try:
        selected_encoders = [encoder for encoder in encoder_modules if id(encoder) in selected_encoder_ids]
        # Register hooks on DiT modules (offload selected encoders and other selected DiTs).
        for i, dit_mod in enumerate(dit_modules):
            other_dits = [d for j, d in enumerate(dit_modules) if j != i and id(d) in selected_dit_ids]
            registry = HookRegistry.get_or_create(dit_mod)
            hook = SequentialOffloadHook(
                offload_targets=selected_encoders + other_dits,
                device=device,
                pin_memory=pin_memory,
                use_hsdp=use_hsdp,
                offload_after_context=id(dit_mod) in selected_dit_ids,
            )
            registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook)
            logger.debug("Registered offload hook for %s", dit_mod.__class__.__name__)

        # Register hooks on all execution stages so unselected resident modules can
        # still evict selected DiTs before they run.
        selected_dits = [dit for dit in dit_modules if id(dit) in selected_dit_ids]
        for enc in encoder_modules:
            registry = HookRegistry.get_or_create(enc)
            hook = SequentialOffloadHook(
                offload_targets=selected_dits,
                device=device,
                pin_memory=pin_memory,
                use_hsdp=use_hsdp,
                offload_after_context=id(enc) in selected_encoder_ids,
            )
            registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook)
            logger.debug("Registered offload hook for %s", enc.__class__.__name__)

        if offload_initial_dits:
            for dit_mod in dit_modules:
                if id(dit_mod) not in selected_dit_ids:
                    continue
                _get_sequential_offload_hook(dit_mod)._to_cpu(dit_mod)
    except BaseException:
        try:
            remove_sequential_offload(all_modules)
        except BaseException:
            logger.exception("Failed to remove every sequential hook during rollback")
        try:
            _restore_tensor_devices(initial_devices)
        except BaseException:
            logger.exception("Failed to restore initial DiT placement during rollback")
        raise


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)
    """
    first_error: BaseException | None = None
    for module in modules:
        try:
            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__)
        except BaseException as exc:
            logger.exception("Failed to remove offload hook from %s", module.__class__.__name__)
            first_error = first_error or exc
    if first_error is not None:
        raise RuntimeError("Failed to remove one or more sequential offload hooks") from first_error


def _get_sequential_offload_hook(module: nn.Module) -> SequentialOffloadHook:
    registry: HookRegistry | None = getattr(module, "_hook_registry", None)
    hook = registry.get_hook(SequentialOffloadHook._HOOK_NAME) if registry is not None else None
    if not isinstance(hook, SequentialOffloadHook):
        raise RuntimeError(f"{module.__class__.__name__} has no sequential offload hook")
    return hook


@contextmanager
def sequential_offload_component(module: nn.Module) -> Iterator[None]:
    """Activate and release a hooked component called outside ``forward``."""
    hook = _get_sequential_offload_hook(module)
    try:
        hook.pre_forward(module)
        yield
    except BaseException:
        try:
            if hook.offload_after_context:
                hook._to_cpu(module)
        except Exception:
            logger.exception("Failed to release %s after component failure", module.__class__.__name__)
        raise
    else:
        if hook.offload_after_context:
            hook._to_cpu(module)


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
        self._custom_pipeline: SupportsModelCpuOffload | None = None

    def enable(self, pipeline: nn.Module) -> None:
        if self.enabled:
            logger.warning("ModelLevelOffloadBackend already enabled")
            return

        # Pipelines with non-forward component entry points own their complete
        # mutual-exclusion lifecycle. Delegate through the explicit protocol.
        if isinstance(pipeline, SupportsModelCpuOffload):
            try:
                if self.config.components is not None:
                    pipeline.enable_omni_model_cpu_offload(
                        device=self.device,
                        pin_memory=self.config.pin_cpu_memory,
                        use_hsdp=self.config.use_hsdp,
                        offload_components=self.config.components,
                    )
                else:
                    pipeline.enable_omni_model_cpu_offload(
                        device=self.device,
                        pin_memory=self.config.pin_cpu_memory,
                        use_hsdp=self.config.use_hsdp,
                    )
            except BaseException:
                try:
                    pipeline.disable_omni_model_cpu_offload()
                except BaseException:
                    logger.exception("Model-level cleanup failed while handling an enable failure")
                raise
            self._custom_pipeline = pipeline
            self.enabled = True
            logger.info(
                "Model-level offloading enabled through %s.enable_omni_model_cpu_offload",
                pipeline.__class__.__name__,
            )
            return

        resolved = resolve_offload_plan(pipeline, self.config)
        dits = [component.module for component in resolved.dits]
        encoders = [component.module for component in resolved.encoders]
        vaes = [component.module for component in resolved.vaes]
        residents = [component.module for component in resolved.residents]
        selected_encoders = [component.module for component in resolved.encoders if component.selected]

        all_modules = [*dits, *encoders, *vaes, *residents]
        initial_devices = _capture_tensor_devices(all_modules)
        try:
            for encoder in encoders:
                encoder.to(self.device)
            for vae in vaes:
                vae.to(self.device, non_blocking=True)
            for resident in residents:
                resident.to(self.device)

            if not dits:
                logger.warning("No DiT/transformer modules found, skipping model-level offloading")
                return
            if not encoders:
                for dit in dits:
                    dit.to(self.device)
                logger.warning("No encoder modules found, skipping model-level offloading")
                return

            apply_sequential_offload(
                dit_modules=dits,
                encoder_modules=encoders,
                device=self.device,
                pin_memory=self.config.pin_cpu_memory,
                use_hsdp=self.config.use_hsdp,
                offload_dit_modules=(dits if self.config.offloads(DIT_COMPONENT) else ()),
                offload_encoder_modules=selected_encoders,
            )
        except BaseException:
            try:
                remove_sequential_offload([*dits, *encoders])
            except BaseException:
                logger.exception("Failed to remove every model-level hook during rollback")
            try:
                _restore_tensor_devices(initial_devices)
            except BaseException:
                logger.exception("Failed to restore model placement during rollback")
            raise

        # Track modules for cleanup
        self._offload_modules = [*dits, *encoders]

        self.enabled = True

        logger.info(
            "Model-level offloading enabled: %s <-> %s (mutual exclusion)%s",
            ", ".join(component.path for component in resolved.dits),
            ", ".join(component.path for component in resolved.encoders),
            (
                f"; resident on GPU: {', '.join(component.path for component in resolved.residents)}"
                if resolved.residents
                else ""
            ),
        )

    def disable(self) -> None:
        if not self.enabled:
            return

        if self._custom_pipeline is not None:
            self._custom_pipeline.disable_omni_model_cpu_offload()
            self._custom_pipeline = None
            self.enabled = False
            logger.info("Model-level offloading disabled")
            return

        remove_sequential_offload(self._offload_modules)

        self._offload_modules.clear()
        self.enabled = False
        logger.info("Model-level offloading disabled")