222 lines
8.5 KiB
Python
222 lines
8.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Shared component-plan helpers for diffusion offload backends."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable, Iterator
|
|
from itertools import chain
|
|
from operator import attrgetter
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import torch
|
|
from torch import nn
|
|
from torch.distributed.tensor import DTensor
|
|
from vllm.logger import init_logger
|
|
|
|
from .block_discovery import get_blocks_from_dit
|
|
from .config import DIT_COMPONENT, TEXT_ENCODER_COMPONENT
|
|
from .offload_plan import OffloadPlan
|
|
from .tensor_utils import set_tensor_storage
|
|
|
|
if TYPE_CHECKING:
|
|
from .base import OffloadConfig
|
|
from .module_collector import PipelineModules
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def encoder_component_type(name: str, plan: OffloadPlan | None) -> str | None:
|
|
"""Map a discovered encoder path to its public offload component."""
|
|
declared = None if plan is None else plan.encoder_component_types.get(name)
|
|
if declared is not None:
|
|
if declared != TEXT_ENCODER_COMPONENT:
|
|
raise ValueError(f"OffloadPlan maps encoder {name!r} to unknown component {declared!r}")
|
|
return declared
|
|
|
|
leaf_name = name.rsplit(".", 1)[-1]
|
|
if leaf_name.startswith(TEXT_ENCODER_COMPONENT) or leaf_name.endswith(TEXT_ENCODER_COMPONENT):
|
|
return TEXT_ENCODER_COMPONENT
|
|
return None
|
|
|
|
|
|
def get_encoder_block_groups(
|
|
module: nn.Module,
|
|
name: str,
|
|
plan: OffloadPlan | None,
|
|
*,
|
|
strict: bool = False,
|
|
) -> list[nn.ModuleList]:
|
|
"""Resolve the streamable block lists declared for one encoder."""
|
|
if plan is None:
|
|
return []
|
|
|
|
# Some distributed models expose an unloaded stub on ranks where the
|
|
# component never executes. That is a valid local no-op even when the
|
|
# component was selected explicitly.
|
|
if getattr(module, "is_loaded", True) is False:
|
|
return []
|
|
|
|
groups: list[nn.ModuleList] = []
|
|
for block_path in plan.encoder_block_attrs.get(name, ()):
|
|
try:
|
|
blocks = attrgetter(block_path)(module)
|
|
except AttributeError:
|
|
if strict:
|
|
raise ValueError(f"Encoder offload path {name}.{block_path} was not found") from None
|
|
logger.warning("Encoder offload path %s.%s was not found", name, block_path)
|
|
continue
|
|
if not isinstance(blocks, nn.ModuleList) or len(blocks) <= 1:
|
|
if strict:
|
|
raise ValueError(f"Encoder offload path {name}.{block_path} is not a streamable block list")
|
|
logger.warning("Encoder offload path %s.%s is not a streamable block list", name, block_path)
|
|
continue
|
|
groups.append(blocks)
|
|
return groups
|
|
|
|
|
|
def iter_streamable_dits(
|
|
modules: PipelineModules,
|
|
config: OffloadConfig,
|
|
device: torch.device,
|
|
plan: OffloadPlan | None,
|
|
) -> Iterator[tuple[str, nn.Module, list[str], list[nn.Module]]]:
|
|
"""Yield selected DiTs whose block metadata resolves successfully."""
|
|
if not config.offloads(DIT_COMPONENT):
|
|
return
|
|
for name, module in zip(modules.dit_names, modules.dits):
|
|
logger.info("Applying hooks on %s (%s)", name, module.__class__.__name__)
|
|
planned_attrs = None if plan is None else plan.block_attrs.get(name)
|
|
block_attrs, blocks = get_blocks_from_dit(module, planned_attrs)
|
|
if blocks:
|
|
yield name, module, block_attrs, blocks
|
|
continue
|
|
if config.components is not None:
|
|
raise ValueError(f"Selected DiT {name!r} has no streamable layerwise-offload blocks")
|
|
logger.warning("Target layers (blocks) not found. Skipping offloading on %s (%s)", name, type(module).__name__)
|
|
module.to(device)
|
|
|
|
|
|
def move_non_block_state_to_device(
|
|
module: nn.Module,
|
|
block_groups: list[nn.ModuleList],
|
|
device: torch.device,
|
|
) -> None:
|
|
"""Keep component state outside streamed block lists resident on device."""
|
|
block_tensors = {
|
|
id(tensor)
|
|
for blocks in block_groups
|
|
for block in blocks
|
|
for tensor in chain(block.parameters(), block.buffers())
|
|
}
|
|
for tensor in chain(module.parameters(), module.buffers()):
|
|
if id(tensor) in block_tensors:
|
|
continue
|
|
local = tensor.to_local() if isinstance(tensor, DTensor) else tensor
|
|
if local.device != device:
|
|
set_tensor_storage(tensor, local.to(device, non_blocking=True))
|
|
|
|
|
|
def set_encoder_layerwise_state(
|
|
module: nn.Module,
|
|
hooks: list[Any],
|
|
block_groups: list[nn.ModuleList],
|
|
) -> None:
|
|
"""Publish the backend-neutral state used by encoder stage lifecycles."""
|
|
module._omni_layerwise_hooks = hooks
|
|
module._omni_layerwise_block_groups = block_groups
|
|
module._omni_layerwise_enabled = True
|
|
|
|
|
|
def clear_encoder_layerwise_state(module: nn.Module) -> None:
|
|
"""Clear encoder layerwise state after its backend hooks are removed."""
|
|
module._omni_layerwise_hooks = []
|
|
module._omni_layerwise_block_groups = []
|
|
module._omni_layerwise_enabled = False
|
|
|
|
|
|
def validate_on_demand_component(module: nn.Module, name: str) -> None:
|
|
"""Require the explicit lifecycle used by pipeline-managed components."""
|
|
if not callable(getattr(module, "load_to_device", None)) or not callable(getattr(module, "offload_to_cpu", None)):
|
|
raise ValueError(
|
|
f"Component {name!r} declares on-demand offload but must implement load_to_device() and offload_to_cpu()"
|
|
)
|
|
|
|
|
|
def prepare_component(
|
|
module: nn.Module,
|
|
name: str,
|
|
*,
|
|
device: torch.device,
|
|
stage_on_demand: bool,
|
|
blockwise: bool,
|
|
staged_components: list[nn.Module],
|
|
) -> None:
|
|
"""Stage a selected component or keep its non-streamed form resident."""
|
|
if stage_on_demand:
|
|
validate_on_demand_component(module, name)
|
|
getattr(module, "offload_to_cpu")()
|
|
staged_components.append(module)
|
|
logger.info("Prepared %s for pipeline-managed staged offload", name)
|
|
elif not blockwise:
|
|
module.to(device)
|
|
|
|
|
|
def prepare_pipeline_components(
|
|
modules: PipelineModules,
|
|
config: OffloadConfig,
|
|
plan: OffloadPlan | None,
|
|
*,
|
|
device: torch.device,
|
|
staged_components: list[nn.Module],
|
|
enable_encoder_blocks: Callable[[nn.Module, str, OffloadPlan | None, bool], bool],
|
|
) -> None:
|
|
"""Apply the shared encoder/VAE/resident placement policy."""
|
|
if config.components is not None and config.offloads(TEXT_ENCODER_COMPONENT):
|
|
selected_encoder_names = [name for name in modules.encoder_names if config.offloads_encoder(name, plan)]
|
|
if not selected_encoder_names:
|
|
raise ValueError("No text encoder modules found for selected text_encoder offload")
|
|
|
|
if plan is not None:
|
|
for encoder, name in zip(modules.encoders, modules.encoder_names):
|
|
if config.should_offload_encoder(name, plan) and name in plan.on_demand_component_paths:
|
|
validate_on_demand_component(encoder, name)
|
|
|
|
for encoder, name in zip(modules.encoders, modules.encoder_names):
|
|
selected = config.should_offload_encoder(name, plan)
|
|
stage_on_demand = bool(selected and plan is not None and name in plan.on_demand_component_paths)
|
|
blockwise = selected and enable_encoder_blocks(encoder, name, plan, stage_on_demand)
|
|
if stage_on_demand and not blockwise and config.uses_allgather(TEXT_ENCODER_COMPONENT):
|
|
raise ValueError(
|
|
f"Text encoder {name!r} cannot use AllGather without a model-declared streamable block plan"
|
|
)
|
|
if selected and config.components is not None and not (blockwise or stage_on_demand):
|
|
raise ValueError(f"Selected text encoder {name!r} requires a model-declared streamable or on-demand plan")
|
|
prepare_component(
|
|
encoder,
|
|
name,
|
|
device=device,
|
|
stage_on_demand=stage_on_demand,
|
|
blockwise=blockwise,
|
|
staged_components=staged_components,
|
|
)
|
|
|
|
for vae, name in zip(modules.vaes, modules.vae_names):
|
|
legacy_staged = config.components is None and plan is not None and name in plan.on_demand_component_paths
|
|
prepare_component(
|
|
vae,
|
|
name,
|
|
device=device,
|
|
stage_on_demand=legacy_staged,
|
|
blockwise=False,
|
|
staged_components=staged_components,
|
|
)
|
|
|
|
for name, module in zip(modules.resident_names, modules.resident_modules):
|
|
module.to(device)
|
|
logger.debug("Moved resident module %s to %s", name, device)
|
|
|
|
if not config.offloads(DIT_COMPONENT):
|
|
for dit in modules.dits:
|
|
dit.to(device)
|