hao-ai-lab-FastVideo/fastvideo/hooks/layerwise_offload.py

from contextlib import contextmanager
from typing import Any
import torch
from torch import nn
from fastvideo.hooks.hooks import ForwardHook, ModuleHookManager
from fastvideo.logger import init_logger

logger = init_logger(__name__)


def _tensor_placeholder(tensor: torch.Tensor, device: torch.device) -> torch.Tensor:
    """Create a rank-preserving empty placeholder on the specified device."""
    shape = (0, ) if tensor.ndim <= 0 else (0, ) * tensor.ndim
    return torch.empty(shape, device=device, dtype=tensor.dtype)


class LayerwiseOffloadState:

    def __init__(
        self,
        async_copy_stream: torch.cuda.Stream,
        device: torch.device,
        next_state: "LayerwiseOffloadState | None" = None,
    ) -> None:
        self.async_copy_stream = async_copy_stream
        self.next_state = next_state
        self.gpu_named_parameters: dict[str, torch.Tensor] = {}
        self.cpu_named_parameters: dict[str, torch.Tensor] = {}
        self.module_ref: nn.Module = None  # type: ignore
        self.device: torch.device = device

    def _will_offload(self, name: str) -> bool:
        return True

    @torch.compiler.disable
    def on_init(self, module: nn.Module):
        self.module_ref = module
        for name, param in self.module_ref.named_parameters():
            if self._will_offload(name):
                self.cpu_named_parameters[name] = (param.data.detach().to("cpu").pin_memory())
                param.data = _tensor_placeholder(param.data, self.device)

    @torch.compiler.disable
    def wait_and_replace_params(self):
        torch.cuda.current_stream().wait_stream(self.async_copy_stream)
        # now gpu_named_parameters are ready
        for name, param in self.module_ref.named_parameters():
            if not self._will_offload(name):
                continue
            if name not in self.gpu_named_parameters:
                # first load with blocking load
                self.gpu_named_parameters[name] = self.cpu_named_parameters[name].to(self.device)
            param.data = self.gpu_named_parameters[name]

    @torch.compiler.disable
    def prefetch_params(self):
        compute_stream = torch.cuda.current_stream()
        with torch.cuda.stream(self.async_copy_stream):
            for name, param in self.module_ref.named_parameters():
                if not self._will_offload(name):
                    continue
                assert name not in self.gpu_named_parameters
                gpu_param = self.cpu_named_parameters[name].to(self.device, non_blocking=True)
                gpu_param.record_stream(compute_stream)  # ensure tensor will not be freed until forward is completed
                self.gpu_named_parameters[name] = gpu_param

    @torch.compiler.disable
    def release_gpu_params(self):
        for name, param in self.module_ref.named_parameters():
            if self._will_offload(name):
                param.data = _tensor_placeholder(param.data, self.device)
                del self.gpu_named_parameters[name]
        assert len(self.gpu_named_parameters) == 0


class LayerwiseOffloadHook(ForwardHook):
    """A hook that enables layerwise CPU offloading during forward pass."""

    def __init__(self, state: LayerwiseOffloadState) -> None:
        self.state = state

    def on_attach(self, module: nn.Module):
        self.state.on_init(module)  # pyright: ignore

    def on_detach(self, module: nn.Module):
        named_parameters = dict(module.named_parameters())
        for name, cpu_tensor in self.state.cpu_named_parameters.items():
            if name not in self.state.gpu_named_parameters:
                if name in named_parameters:
                    named_parameters[name].data = cpu_tensor.to(device=self.state.device)
                else:
                    logger.warning(
                        "Parameter {} not found in module during detachment.",
                        name,
                    )

    @classmethod
    def name(cls) -> str:
        return "LayerwiseOffloadHook"

    # These hook entry points only orchestrate host-side parameter
    # movement (stream waits, pinned H2D/D2H, param swapping) and must
    # always run eager. Without an explicit boundary, torch.compile
    # traces *into* the hook, hits the `@torch.compiler.disable`
    # `LayerwiseOffloadState` methods, and is forced to bail mid-trace
    # — an implicit graph break at the call site, once per layer every
    # step, which fragments the per-layer compiled region (and blocks
    # CUDA-graph capture, which cannot span a break). Marking the entry
    # points disabled makes the hook a clean opaque eager boundary so
    # torch.compile keeps one contiguous region around it. Completes
    # the pattern already applied to the State methods.
    @torch.compiler.disable
    def pre_forward(self, module: nn.Module, *args, **kwargs):
        self.state.wait_and_replace_params()  # pyright: ignore
        if self.state.next_state is not None:
            self.state.next_state.prefetch_params()  # pyright: ignore
        return args, kwargs

    @torch.compiler.disable
    def post_forward(self, module: torch.nn.Module, output: Any):
        self.state.release_gpu_params()  # pyright: ignore
        return output

    @contextmanager
    def mutate_params_scope(self):
        try:
            # load params to GPU and keep them there
            self.state.wait_and_replace_params()  # pyright: ignore
            yield
        finally:
            # instead of releasing, we should overwrite the original params since they have been modified
            self.state.cpu_named_parameters.clear()
            self.state.gpu_named_parameters.clear()
            self.state.on_init(self.state.module_ref)  # pyright: ignore


def enable_layerwise_offload(model: nn.Module, is_replace: bool = False):
    if torch.cuda.is_available():
        device = torch.device("cuda", torch.cuda.current_device())
    else:
        logger.warning("CUDA is not available. Layerwise offloading is disabled.")
        return
    state_list = []
    async_stream = torch.cuda.Stream()
    for name, submodule in model.named_children():
        if isinstance(submodule, nn.ModuleList):
            for idx, module_entry in enumerate(submodule):
                state = LayerwiseOffloadState(async_copy_stream=async_stream, device=device)
                state_list.append(state)
                hook_mgr = ModuleHookManager.get_from_or_default(module_entry)
                hook = LayerwiseOffloadHook(state)
                if is_replace:
                    existing_hook = hook_mgr.forward_hooks.get(hook.name())
                    if existing_hook is not None:
                        hook_mgr.replace_forward_hook(hook.name(), hook)
                    else:
                        raise AssertionError(f"Expect hook exists in {name} for replacement.")
                else:
                    hook_mgr.append_forward_hook(hook)
            break
    if len(state_list) == 0:
        raise ValueError("No nn.ModuleList found in the model for layerwise offloading.")

    # circular linking of states
    for i in range(len(state_list)):
        state_list[i].next_state = state_list[(i + 1) % len(state_list)]