Anjie HouandClaude Opus 5
85da943391
Signed-off-by: specture724 <specture724@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
524 lines
20 KiB
Python
524 lines
20 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
from itertools import chain
|
|
from typing import Any
|
|
|
|
import torch
|
|
from torch import nn
|
|
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, run_cleanup_steps
|
|
from .component_utils import (
|
|
clear_encoder_layerwise_state,
|
|
iter_streamable_dits,
|
|
move_non_block_state_to_device,
|
|
prepare_pipeline_components,
|
|
set_encoder_layerwise_state,
|
|
)
|
|
from .plan_resolver import ResolvedComponent, resolve_offload_plan
|
|
from .tensor_utils import (
|
|
clear_block_storage,
|
|
clear_tensor_storage,
|
|
describe_tensor_storage,
|
|
flatten_physical_storage,
|
|
group_named_tensors_by_dtype,
|
|
is_materialized_tensor,
|
|
materialization_probe,
|
|
module_materialization_probe,
|
|
restore_tensor_storage,
|
|
set_tensor_storage,
|
|
tensor_storage_metadata,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class LayerwiseOffloadHook(ModelHook):
|
|
"""Hook for layerwise (transformer-block-wise) CPU offloading.
|
|
|
|
The hook instance retains parameters for both the current registered block
|
|
module and those for the next block, as well as flattened CPU tensors which
|
|
record the parameters of the current block module, so that these parameters
|
|
could be re-materialized on device in an overlapping way.
|
|
This hook should be registered to each of the transformer blocks in DiT
|
|
module(s) of the target pipeline.
|
|
|
|
Based on implementations from:
|
|
https://github.com/sgl-project/sglang/blob/v0.5.8/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py
|
|
"""
|
|
|
|
_HOOK_NAME = "layerwise_offload"
|
|
|
|
def __init__(
|
|
self,
|
|
next_block: nn.Module,
|
|
device: torch.device,
|
|
stream: current_omni_platform.Stream | None = None,
|
|
pin_memory: bool = True,
|
|
materialization_probe_tensor: torch.Tensor | None = None,
|
|
):
|
|
assert isinstance(next_block, nn.Module), "transformer block must be type `torch.nn.Module`"
|
|
|
|
self.next_block = next_block
|
|
self.device = device
|
|
self.copy_stream = stream or current_omni_platform.current_stream()
|
|
self.pin_memory = pin_memory
|
|
|
|
# Per-block synchronization primitive: set after H2D copy completes.
|
|
self._prefetch_done: current_omni_platform.Event | None = None
|
|
|
|
# Backward link to the hook that is responsible for prefetching *this* block's weights
|
|
self._prev_hook: LayerwiseOffloadHook | None = None
|
|
|
|
self.next_block_parameters: dict[str, nn.Parameter] = {}
|
|
self.next_block_buffers: dict[str, torch.Tensor] = {}
|
|
self.dtype_cpu_flattened_weights: dict[torch.dtype, torch.Tensor] = {}
|
|
self.dtype_metadata: dict[torch.dtype, list[dict[str, Any]]] = {}
|
|
self._materialization_probe = materialization_probe_tensor
|
|
|
|
def initialize_hook(self, module: nn.Module) -> nn.Module:
|
|
# This all happen during the hook instance being registered to hook registry;
|
|
# the input module is kept intact
|
|
module = super().initialize_hook(module)
|
|
|
|
self.block_parameters: dict[str, nn.Parameter] = dict(module.named_parameters())
|
|
self.block_buffers: dict[str, torch.Tensor] = dict(module.named_buffers())
|
|
if self._materialization_probe is None:
|
|
self._materialization_probe = materialization_probe(self.block_parameters, self.block_buffers)
|
|
|
|
self.next_block_parameters: dict[str, nn.Parameter] = dict(self.next_block.named_parameters())
|
|
self.next_block_buffers: dict[str, torch.Tensor] = dict(self.next_block.named_buffers())
|
|
|
|
# Pre-allocate gpu tensors in a flattened way
|
|
self.dtype_cpu_flattened_weights, self.dtype_metadata = LayerwiseOffloadHook._to_cpu(
|
|
self.next_block_parameters,
|
|
self.next_block_buffers,
|
|
self.pin_memory,
|
|
)
|
|
|
|
return module
|
|
|
|
@staticmethod
|
|
def _to_cpu(
|
|
params: dict[str, nn.Parameter],
|
|
bufs: dict[str, torch.Tensor],
|
|
pin_memory: bool = True,
|
|
) -> tuple[dict[torch.dtype, torch.Tensor], dict[torch.dtype, list[dict[str, Any]]]]:
|
|
"""Helper method to 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_cpu_flattened_weights: dict[torch.dtype, torch.Tensor] = {}
|
|
# NOTE: order does matter
|
|
dtype_metadata: dict[torch.dtype, list[dict[str, Any]]] = {}
|
|
targets_to_offload: list[torch.Tensor] = []
|
|
|
|
for dtype, named_weights in group_named_tensors_by_dtype(params, bufs).items():
|
|
# total # of parameters + buffers
|
|
specs = describe_tensor_storage(named_weights)
|
|
total_numel = sum(spec.storage_numel for spec in specs)
|
|
cpu_tensor = torch.empty(total_numel, dtype=dtype, device="cpu", pin_memory=pin_memory)
|
|
|
|
current_offset = 0
|
|
for spec in specs:
|
|
flat_storage = flatten_physical_storage(spec.value, spec.storage_numel)
|
|
cpu_tensor[current_offset : current_offset + spec.storage_numel].copy_(flat_storage)
|
|
dtype_metadata.setdefault(dtype, []).append(
|
|
tensor_storage_metadata(spec, current_offset, include_device=True)
|
|
)
|
|
targets_to_offload.append(spec.target)
|
|
current_offset += spec.storage_numel
|
|
|
|
dtype_cpu_flattened_weights[dtype] = cpu_tensor
|
|
|
|
# Do not mutate the module until every host master has been built.
|
|
# Allocation/copy failures therefore leave the input module intact.
|
|
clear_tensor_storage(targets_to_offload)
|
|
|
|
return dtype_cpu_flattened_weights, dtype_metadata
|
|
|
|
@property
|
|
def is_materialized(self) -> bool:
|
|
"""Check whether this block's parameters hold real data on device."""
|
|
return self._materialization_probe is None or is_materialized_tensor(self._materialization_probe)
|
|
|
|
@torch.compiler.disable
|
|
def prefetch_layer(self, non_blocking: bool = True) -> None:
|
|
"""Copy layer weights from CPU -> GPU.
|
|
|
|
Pre-fetch target block in an asynchronous way with compute - memory copy overlap,
|
|
with non_blocking set to True.
|
|
"""
|
|
self.copy_stream.wait_stream(current_omni_platform.current_stream())
|
|
|
|
layer_params = self.next_block_parameters
|
|
layer_bufs = self.next_block_buffers
|
|
|
|
evt = current_omni_platform.Event()
|
|
gpu_weights: dict[torch.dtype, torch.Tensor] = {}
|
|
|
|
with current_omni_platform.stream(self.copy_stream):
|
|
for dtype, cpu_weight in self.dtype_cpu_flattened_weights.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, ordered_metadata in self.dtype_metadata.items():
|
|
# ordered_metadata: list[dict[str, Any]]
|
|
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]
|
|
)
|
|
|
|
set_tensor_storage(
|
|
target_param_or_buf,
|
|
torch.as_strided(
|
|
gpu_weight[metadata["offset"] : metadata["offset"] + metadata["numel"]],
|
|
size=metadata["shape"],
|
|
stride=metadata["stride"],
|
|
),
|
|
)
|
|
|
|
self._prefetch_done = evt
|
|
|
|
@torch.compiler.disable
|
|
def offload_layer(self) -> None:
|
|
"""Free GPU memory for layer by replacing tensors with empty placeholders.
|
|
This function does not actually offload weights from GPU back to CPU.
|
|
"""
|
|
clear_block_storage(self.block_parameters, self.block_buffers, self._prefetch_done)
|
|
self._prefetch_done = None
|
|
|
|
@torch.compiler.disable
|
|
def restore_next_block(self) -> None:
|
|
"""Detach the next block from this hook's host backing store."""
|
|
for dtype, ordered_metadata in self.dtype_metadata.items():
|
|
flat = self.dtype_cpu_flattened_weights[dtype]
|
|
for metadata in ordered_metadata:
|
|
value = torch.as_strided(
|
|
flat[metadata["offset"] : metadata["offset"] + metadata["numel"]],
|
|
size=metadata["shape"],
|
|
stride=metadata["stride"],
|
|
)
|
|
target_name = metadata["name"]
|
|
target = (
|
|
self.next_block_parameters[target_name]
|
|
if target_name in self.next_block_parameters
|
|
else self.next_block_buffers[target_name]
|
|
)
|
|
restore_tensor_storage(target, value, device=metadata["device"])
|
|
|
|
def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]:
|
|
# if the previous hook was skipped and the weights are not on device,
|
|
# (e.g. by cache-dit block caching), ask the previous hook to
|
|
# synchronously prefetch *this* block's weights before computation
|
|
if not self.is_materialized and self._prev_hook is not None:
|
|
self._prev_hook.prefetch_layer(non_blocking=False)
|
|
|
|
self.prefetch_layer(non_blocking=True)
|
|
|
|
return args, kwargs
|
|
|
|
def post_forward(self, module: nn.Module, output: Any) -> Any:
|
|
self.offload_layer()
|
|
|
|
return output
|
|
|
|
|
|
def apply_block_hook(
|
|
module: nn.Module,
|
|
next_block: nn.Module,
|
|
device: torch.device,
|
|
stream: current_omni_platform.Stream | None = None,
|
|
pin_memory: bool = True,
|
|
*,
|
|
materialization_probe_tensor: torch.Tensor | None = None,
|
|
) -> LayerwiseOffloadHook:
|
|
registry = HookRegistry.get_or_create(module)
|
|
hook = LayerwiseOffloadHook(
|
|
next_block,
|
|
device,
|
|
stream,
|
|
pin_memory,
|
|
materialization_probe_tensor,
|
|
)
|
|
registry.register_hook(LayerwiseOffloadHook._HOOK_NAME, hook)
|
|
|
|
return hook
|
|
|
|
|
|
def remove_block_hook(module: nn.Module) -> None:
|
|
registry: HookRegistry | None = getattr(module, "_hook_registry", None)
|
|
if registry is not None:
|
|
registry.remove_hook(LayerwiseOffloadHook._HOOK_NAME)
|
|
logger.debug("Removed offload hook from %s", module.__class__.__name__)
|
|
|
|
|
|
def _install_layerwise_hook_group(
|
|
blocks: Sequence[nn.Module] | nn.ModuleList,
|
|
device: torch.device,
|
|
stream: Any,
|
|
pin_memory: bool,
|
|
) -> list[LayerwiseOffloadHook]:
|
|
"""Install one circular hook ring and roll it back transactionally."""
|
|
block_list = list(blocks)
|
|
if len(block_list) <= 1:
|
|
raise ValueError("A layerwise hook group requires at least two blocks")
|
|
|
|
hooks: list[LayerwiseOffloadHook] = []
|
|
hooked_blocks: list[nn.Module] = []
|
|
probes = {id(block): module_materialization_probe(block) for block in block_list}
|
|
try:
|
|
for block, next_block in zip(
|
|
chain((block_list[-1],), block_list[:-1]),
|
|
block_list,
|
|
strict=True,
|
|
):
|
|
hooks.append(
|
|
apply_block_hook(
|
|
block,
|
|
next_block,
|
|
device,
|
|
stream,
|
|
pin_memory,
|
|
materialization_probe_tensor=probes[id(block)],
|
|
)
|
|
)
|
|
hooked_blocks.append(block)
|
|
except BaseException:
|
|
run_cleanup_steps(
|
|
[
|
|
*(("restoring a partially installed layerwise block", hook.restore_next_block) for hook in hooks),
|
|
*(
|
|
("removing a partially installed layerwise hook", lambda block=block: remove_block_hook(block))
|
|
for block in hooked_blocks
|
|
),
|
|
]
|
|
)
|
|
raise
|
|
|
|
for index, hook in enumerate(hooks):
|
|
hook._prev_hook = hooks[index - 1]
|
|
return hooks
|
|
|
|
|
|
def enable_plan_encoder_layerwise_offload(
|
|
component: ResolvedComponent,
|
|
*,
|
|
device: torch.device,
|
|
stream: current_omni_platform.Stream,
|
|
pin_memory: bool,
|
|
) -> bool:
|
|
"""Apply rank-local layerwise hooks to the encoder's resolved stacks."""
|
|
module = component.module
|
|
if getattr(module, "_omni_layerwise_enabled", False):
|
|
return True
|
|
|
|
block_groups = [stack.blocks for stack in component.stacks]
|
|
if not block_groups:
|
|
return False
|
|
|
|
hooks: list[LayerwiseOffloadHook] = []
|
|
hooked_blocks: list[nn.Module] = []
|
|
try:
|
|
for blocks in block_groups:
|
|
group_hooks = _install_layerwise_hook_group(blocks, device, stream, pin_memory)
|
|
hooks.extend(group_hooks)
|
|
hooked_blocks.extend(blocks)
|
|
if not component.on_demand:
|
|
move_non_block_state_to_device(module, block_groups, device)
|
|
except BaseException:
|
|
run_cleanup_steps(
|
|
[
|
|
*(("restoring a partially installed encoder block", hook.restore_next_block) for hook in hooks),
|
|
*(
|
|
("removing a partially installed encoder hook", lambda block=block: remove_block_hook(block))
|
|
for block in hooked_blocks
|
|
),
|
|
]
|
|
)
|
|
raise
|
|
set_encoder_layerwise_state(
|
|
module,
|
|
hooks,
|
|
block_groups,
|
|
)
|
|
logger.info(
|
|
"Enabled rank-local layerwise offload for encoder %s (%d blocks across %d stacks)",
|
|
component.path,
|
|
sum(len(blocks) for blocks in block_groups),
|
|
len(block_groups),
|
|
)
|
|
return True
|
|
|
|
|
|
def disable_plan_encoder_layerwise_offload(
|
|
module: nn.Module,
|
|
*,
|
|
restore_weights: bool = True,
|
|
) -> None:
|
|
"""Remove hooks installed by :func:`enable_plan_encoder_layerwise_offload`."""
|
|
if not getattr(module, "_omni_layerwise_enabled", False):
|
|
return
|
|
hooks = getattr(module, "_omni_layerwise_hooks", [])
|
|
block_groups = getattr(module, "_omni_layerwise_block_groups", [])
|
|
steps = []
|
|
if restore_weights:
|
|
steps.extend(("restoring an encoder block", hook.restore_next_block) for hook in hooks)
|
|
steps.extend(
|
|
("removing an encoder hook", lambda block=block: remove_block_hook(block))
|
|
for blocks in block_groups
|
|
for block in blocks
|
|
)
|
|
cleanup_error = run_cleanup_steps(steps)
|
|
if cleanup_error is not None:
|
|
raise RuntimeError("Failed to fully disable encoder layerwise offload") from cleanup_error
|
|
clear_encoder_layerwise_state(module)
|
|
|
|
|
|
class LayerWiseOffloadBackend(OffloadBackend):
|
|
"""Layer-wise (block-level) offloading backend.
|
|
|
|
Implements sliding window offloading where only a small number of transformer
|
|
blocks reside on GPU at a time. Blocks are prefetched asynchronously while
|
|
previous blocks compute, and freed after use.
|
|
"""
|
|
|
|
def __init__(self, config: OffloadConfig, device: torch.device):
|
|
super().__init__(config, device)
|
|
|
|
self.copy_stream = current_omni_platform.Stream()
|
|
self._blocks: list[list[nn.Module]] = []
|
|
self._dit_hooks: list[LayerwiseOffloadHook] = []
|
|
self._hooked_dit_blocks: list[nn.Module] = []
|
|
self._encoder_modules: list[nn.Module] = []
|
|
self._staged_components: list[nn.Module] = []
|
|
|
|
def enable(self, pipeline: nn.Module) -> None:
|
|
try:
|
|
self._enable(pipeline)
|
|
except BaseException:
|
|
try:
|
|
self.disable()
|
|
except BaseException:
|
|
logger.exception("LayerWiseOffloadBackend cleanup failed while handling an enable failure")
|
|
raise
|
|
|
|
def _enable(self, pipeline: nn.Module) -> None:
|
|
if self.enabled:
|
|
logger.warning("LayerWiseOffloadBackend already enabled")
|
|
return
|
|
|
|
resolved = resolve_offload_plan(pipeline, self.config)
|
|
if resolved.skip_reason is not None:
|
|
logger.warning("%s", resolved.skip_reason)
|
|
return
|
|
|
|
def enable_encoder_blocks(component: ResolvedComponent) -> bool:
|
|
enabled = enable_plan_encoder_layerwise_offload(
|
|
component,
|
|
device=self.device,
|
|
stream=self.copy_stream,
|
|
pin_memory=self.config.pin_cpu_memory,
|
|
)
|
|
if enabled:
|
|
# Record each successful installation immediately so a later
|
|
# component failure can remove these hooks transactionally.
|
|
self._encoder_modules.append(component.module)
|
|
return enabled
|
|
|
|
prepare_pipeline_components(
|
|
resolved,
|
|
device=self.device,
|
|
staged_components=self._staged_components,
|
|
enable_encoder_blocks=enable_encoder_blocks,
|
|
)
|
|
|
|
logger.info("Applying layer-wise offloading on %s", [component.path for component in resolved.dits])
|
|
|
|
# Apply block-wise offloading hook for each of the blocks in DiT model(s)
|
|
# Note that there might exist multiple DiT models in specific pipelines
|
|
for component, stack in iter_streamable_dits(resolved, self.device):
|
|
dit_module = component.module
|
|
blocks = list(stack.blocks)
|
|
|
|
# Place the remainder by resolved block tensor identity, just as
|
|
# for encoders. Attribute aliases must not move streamed weights.
|
|
move_non_block_state_to_device(dit_module, (stack.blocks,), self.device)
|
|
|
|
block_hooks = _install_layerwise_hook_group(
|
|
blocks,
|
|
self.device,
|
|
self.copy_stream,
|
|
self.config.pin_cpu_memory,
|
|
)
|
|
self._dit_hooks.extend(block_hooks)
|
|
self._hooked_dit_blocks.extend(blocks)
|
|
|
|
# The last block owns block zero's host backing. Materialize block
|
|
# zero once; later denoising iterations prefetch it from the ring.
|
|
block_hooks[0].prefetch_layer(non_blocking=False)
|
|
|
|
logger.info(f"Layer-wise offloading enabled on {len(blocks)} layers (blocks)")
|
|
|
|
# Track hooked blocks for cleanup
|
|
self._blocks.append(blocks)
|
|
|
|
self.enabled = bool(self._blocks or self._encoder_modules or self._staged_components)
|
|
|
|
def _disable(self, *, restore_weights: bool) -> None:
|
|
if not self.enabled and not (
|
|
self._dit_hooks or self._hooked_dit_blocks or self._encoder_modules or self._staged_components
|
|
):
|
|
return
|
|
|
|
steps = []
|
|
if restore_weights:
|
|
steps.extend(("restoring a DiT block", hook.restore_next_block) for hook in self._dit_hooks)
|
|
steps.extend(
|
|
("removing a DiT block hook", lambda block=block: remove_block_hook(block))
|
|
for block in self._hooked_dit_blocks
|
|
)
|
|
steps.extend(
|
|
(
|
|
"disabling encoder layerwise offload",
|
|
lambda module=module: disable_plan_encoder_layerwise_offload(
|
|
module,
|
|
restore_weights=restore_weights,
|
|
),
|
|
)
|
|
for module in self._encoder_modules
|
|
)
|
|
cleanup_error = run_cleanup_steps(steps)
|
|
if cleanup_error is not None:
|
|
# Keep every host-master reference so a transient teardown error
|
|
# can be retried without losing weights behind placeholders.
|
|
raise RuntimeError("Failed to fully disable layerwise offload") from cleanup_error
|
|
self._blocks.clear()
|
|
self._dit_hooks.clear()
|
|
self._hooked_dit_blocks.clear()
|
|
self._encoder_modules.clear()
|
|
self._staged_components.clear()
|
|
self.enabled = False
|
|
logger.info("Layer-wise offloading disabled")
|
|
|
|
def disable(self) -> None:
|
|
self._disable(restore_weights=True)
|