430 lines
17 KiB
Python
430 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Public diffusion CPU-offload configuration helpers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Collection, Iterator, Mapping
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
from typing import Any, TypeVar
|
|
|
|
DIT_COMPONENT = "dit"
|
|
TEXT_ENCODER_COMPONENT = "text_encoder"
|
|
OFFLOAD_COMPONENTS = frozenset({DIT_COMPONENT, TEXT_ENCODER_COMPONENT})
|
|
DEFAULT_OFFLOAD_COMPONENTS = frozenset({DIT_COMPONENT})
|
|
|
|
_KeyT = TypeVar("_KeyT")
|
|
_ValueT = TypeVar("_ValueT")
|
|
|
|
|
|
class _FrozenMapping(Mapping[_KeyT, _ValueT]):
|
|
"""Small immutable mapping that remains safe to pickle between workers."""
|
|
|
|
def __init__(self, value: Mapping[_KeyT, _ValueT]) -> None:
|
|
self._items = tuple(value.items())
|
|
|
|
def __getitem__(self, key: _KeyT) -> _ValueT:
|
|
for candidate, value in self._items:
|
|
if candidate == key:
|
|
return value
|
|
raise KeyError(key)
|
|
|
|
def __iter__(self) -> Iterator[_KeyT]:
|
|
return (key for key, _ in self._items)
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._items)
|
|
|
|
def __reduce__(self):
|
|
return type(self), (dict(self._items),)
|
|
|
|
|
|
class OffloadMode(str, Enum):
|
|
"""User-facing offload granularity."""
|
|
|
|
MODULE = "module"
|
|
LAYER = "layer"
|
|
|
|
|
|
class OffloadStrategy(str, Enum):
|
|
"""Resolved internal backend strategy."""
|
|
|
|
NONE = "none"
|
|
MODEL_LEVEL = "model_level"
|
|
LAYER_WISE = "layer_wise"
|
|
DISTRIBUTED_LAYER_WISE = "distributed_layer_wise"
|
|
|
|
|
|
class DLOTransfer(str, Enum):
|
|
"""How one component's next block reaches the device."""
|
|
|
|
ALLGATHER = "allgather"
|
|
RANK_LOCAL = "rank-local"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LayerOffloadOptions:
|
|
"""Layer-mode settings for one selected diffusion component."""
|
|
|
|
weight_transfer: DLOTransfer | None = None
|
|
resident_layers: int = 0
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ParsedDiffusionOffloadConfig:
|
|
"""Validated internal form of the raw public offload mapping."""
|
|
|
|
mode: OffloadMode
|
|
components: frozenset[str]
|
|
layer_options: Mapping[str, LayerOffloadOptions]
|
|
pin_memory: bool | None = None
|
|
|
|
def __post_init__(self) -> None:
|
|
object.__setattr__(self, "layer_options", _FrozenMapping(self.layer_options))
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ResolvedOffload:
|
|
"""Canonical offload policy derived once from public or legacy inputs."""
|
|
|
|
strategy: OffloadStrategy
|
|
public: ParsedDiffusionOffloadConfig | None
|
|
components: frozenset[str]
|
|
transfers: Mapping[str, DLOTransfer]
|
|
pin_memory: bool
|
|
resident_layers: int
|
|
|
|
def __post_init__(self) -> None:
|
|
object.__setattr__(self, "transfers", _FrozenMapping(self.transfers))
|
|
|
|
def offloads(self, component: str) -> bool:
|
|
return component in self.components
|
|
|
|
def uses_allgather(self, component: str = DIT_COMPONENT) -> bool:
|
|
try:
|
|
return self.transfers[component] is DLOTransfer.ALLGATHER
|
|
except KeyError as exc:
|
|
raise ValueError(f"Unknown offload component {component!r}") from exc
|
|
|
|
@property
|
|
def any_allgather(self) -> bool:
|
|
return self.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE and any(
|
|
self.uses_allgather(component) for component in self.components
|
|
)
|
|
|
|
|
|
_LEGACY_STRATEGY_FIELDS = {
|
|
"enable_cpu_offload": OffloadStrategy.MODEL_LEVEL,
|
|
"enable_layerwise_offload": OffloadStrategy.LAYER_WISE,
|
|
"enable_distributed_layerwise_offload": OffloadStrategy.DISTRIBUTED_LAYER_WISE,
|
|
}
|
|
_LEGACY_STRATEGY_PRIORITY = (
|
|
OffloadStrategy.DISTRIBUTED_LAYER_WISE,
|
|
OffloadStrategy.LAYER_WISE,
|
|
OffloadStrategy.MODEL_LEVEL,
|
|
)
|
|
|
|
|
|
_EnumT = TypeVar("_EnumT", bound=Enum)
|
|
|
|
|
|
def _parse_enum(value: Any, enum_type: type[_EnumT], label: str) -> _EnumT:
|
|
if isinstance(value, enum_type):
|
|
return value
|
|
if not isinstance(value, str):
|
|
raise TypeError(f"{label} must be a string, got {type(value).__name__}")
|
|
try:
|
|
return enum_type(value)
|
|
except ValueError as exc:
|
|
choices = ", ".join(item.value for item in enum_type)
|
|
raise ValueError(f"Unknown {label} {value!r}; choose from: {choices}") from exc
|
|
|
|
|
|
def _validate_string_keys(value: Mapping[Any, Any], label: str) -> None:
|
|
if any(not isinstance(key, str) for key in value):
|
|
raise TypeError(f"{label} keys must be strings")
|
|
|
|
|
|
def _parse_layer_options(component: str, value: Any) -> LayerOffloadOptions:
|
|
if not isinstance(value, Mapping):
|
|
raise TypeError(
|
|
f"diffusion_offload_config.layer_options[{component!r}] must be a mapping, got {type(value).__name__}"
|
|
)
|
|
_validate_string_keys(value, f"diffusion_offload_config.layer_options[{component!r}]")
|
|
unknown = sorted(set(value) - {"weight_transfer", "resident_layers"})
|
|
if unknown:
|
|
raise ValueError(
|
|
f"Unknown diffusion offload setting(s) for {component}: {', '.join(unknown)}; "
|
|
"choose from: weight_transfer, resident_layers"
|
|
)
|
|
weight_transfer = (
|
|
_parse_enum(value["weight_transfer"], DLOTransfer, "offload transfer") if "weight_transfer" in value else None
|
|
)
|
|
resident_layers = value.get("resident_layers", 0)
|
|
|
|
if type(resident_layers) is not int or resident_layers < 0:
|
|
raise ValueError(f"resident_layers for {component} must be a non-negative integer")
|
|
parsed = LayerOffloadOptions(
|
|
weight_transfer=weight_transfer,
|
|
resident_layers=resident_layers,
|
|
)
|
|
if component != DIT_COMPONENT and parsed.resident_layers:
|
|
raise ValueError("resident_layers currently supports only the 'dit' component")
|
|
return parsed
|
|
|
|
|
|
def parse_diffusion_offload_config(value: Any) -> ParsedDiffusionOffloadConfig | None:
|
|
"""Validate the compact ``diffusion_offload_config`` public API."""
|
|
if value is None:
|
|
return None
|
|
if not isinstance(value, Mapping):
|
|
raise TypeError(f"diffusion_offload_config must be a mapping, got {type(value).__name__}")
|
|
|
|
_validate_string_keys(value, "diffusion_offload_config")
|
|
unknown = sorted(set(value) - {"mode", "components", "layer_options", "pin_memory"})
|
|
if unknown:
|
|
raise ValueError(
|
|
f"Unknown diffusion_offload_config field(s): {', '.join(unknown)}; "
|
|
"choose from: mode, components, layer_options, pin_memory"
|
|
)
|
|
if "mode" not in value:
|
|
raise ValueError("diffusion_offload_config requires 'mode'")
|
|
if "components" not in value:
|
|
raise ValueError("diffusion_offload_config requires 'components'")
|
|
|
|
mode = _parse_enum(value["mode"], OffloadMode, "diffusion offload mode")
|
|
components = parse_offload_components(value["components"])
|
|
|
|
raw_layer_options = value.get("layer_options", {})
|
|
if not isinstance(raw_layer_options, Mapping):
|
|
raise TypeError("diffusion_offload_config.layer_options must be a mapping")
|
|
_validate_string_keys(raw_layer_options, "diffusion_offload_config.layer_options")
|
|
unselected_options = sorted(set(raw_layer_options) - components)
|
|
if unselected_options:
|
|
raise ValueError(
|
|
"diffusion_offload_config.layer_options requires selecting the same component(s): "
|
|
+ ", ".join(unselected_options)
|
|
)
|
|
layer_options = {
|
|
component: _parse_layer_options(component, raw_layer_options.get(component, {})) for component in components
|
|
}
|
|
|
|
pin_memory = value.get("pin_memory")
|
|
if pin_memory is not None and type(pin_memory) is not bool:
|
|
raise TypeError("diffusion_offload_config.pin_memory must be a bool")
|
|
|
|
if mode is OffloadMode.MODULE and raw_layer_options:
|
|
raise ValueError(
|
|
"diffusion_offload_config.layer_options requires mode='layer'; "
|
|
f"configured for: {', '.join(sorted(raw_layer_options))}"
|
|
)
|
|
dit = layer_options.get(DIT_COMPONENT)
|
|
if dit is not None and dit.resident_layers and dit.weight_transfer is DLOTransfer.ALLGATHER:
|
|
raise ValueError("resident_layers requires dit.weight_transfer='rank-local'")
|
|
|
|
return ParsedDiffusionOffloadConfig(
|
|
mode=mode,
|
|
components=components,
|
|
layer_options=layer_options,
|
|
pin_memory=pin_memory,
|
|
)
|
|
|
|
|
|
def _legacy_strategy(config: Any) -> OffloadStrategy:
|
|
enabled = {strategy for field, strategy in _LEGACY_STRATEGY_FIELDS.items() if bool(getattr(config, field, False))}
|
|
return next((strategy for strategy in _LEGACY_STRATEGY_PRIORITY if strategy in enabled), OffloadStrategy.NONE)
|
|
|
|
|
|
def _public_strategy(public: ParsedDiffusionOffloadConfig) -> OffloadStrategy:
|
|
if public.mode is OffloadMode.MODULE:
|
|
return OffloadStrategy.MODEL_LEVEL
|
|
|
|
needs_distributed_backend = any(
|
|
settings.weight_transfer is DLOTransfer.ALLGATHER or settings.resident_layers
|
|
for settings in public.layer_options.values()
|
|
)
|
|
return OffloadStrategy.DISTRIBUTED_LAYER_WISE if needs_distributed_backend else OffloadStrategy.LAYER_WISE
|
|
|
|
|
|
def _validate_legacy_layer_options(config: Any, public: ParsedDiffusionOffloadConfig | None) -> None:
|
|
"""Reject ambiguous or invalid compatibility DLO tuning before loading."""
|
|
resident_layers = getattr(config, "dlo_resident_layers", 0)
|
|
if type(resident_layers) is not int or resident_layers < 0:
|
|
raise ValueError(f"dlo_resident_layers must be a non-negative integer, got {resident_layers!r}")
|
|
|
|
if public is not None:
|
|
if getattr(config, "host_weight_runtime_mode", "disabled") != "disabled":
|
|
raise ValueError("diffusion_offload_config cannot be combined with Host Weight Runtime")
|
|
conflicting_fields = []
|
|
if getattr(config, "dlo_use_allgather", True) is not True:
|
|
conflicting_fields.append("dlo_use_allgather")
|
|
if resident_layers:
|
|
conflicting_fields.append("dlo_resident_layers")
|
|
if float(getattr(config, "dlo_host_registration_limit_gib", 0.0)) > 0:
|
|
conflicting_fields.append("dlo_host_registration_limit_gib")
|
|
if conflicting_fields:
|
|
raise ValueError(
|
|
"diffusion_offload_config cannot be combined with legacy DLO option(s): "
|
|
+ ", ".join(conflicting_fields)
|
|
)
|
|
return
|
|
|
|
if resident_layers and bool(getattr(config, "dlo_use_allgather", True)):
|
|
raise ValueError(
|
|
"dlo_resident_layers requires the DiT DLO transfer to be rank-local; set dlo_use_allgather=False"
|
|
)
|
|
|
|
|
|
def resolve_offload(config: Any) -> ResolvedOffload:
|
|
"""Resolve and cache the complete offload policy on its owning config."""
|
|
cached = getattr(config, "_resolved_diffusion_offload", None)
|
|
if isinstance(cached, ResolvedOffload):
|
|
return cached
|
|
|
|
public = parse_diffusion_offload_config(getattr(config, "diffusion_offload_config", None))
|
|
_validate_legacy_layer_options(config, public)
|
|
legacy = _legacy_strategy(config)
|
|
if public is None:
|
|
transfers = {
|
|
DIT_COMPONENT: (
|
|
DLOTransfer.ALLGATHER if bool(getattr(config, "dlo_use_allgather", True)) else DLOTransfer.RANK_LOCAL
|
|
),
|
|
TEXT_ENCODER_COMPONENT: DLOTransfer.RANK_LOCAL,
|
|
}
|
|
resolved = ResolvedOffload(
|
|
strategy=legacy,
|
|
public=None,
|
|
components=DEFAULT_OFFLOAD_COMPONENTS,
|
|
transfers=transfers,
|
|
pin_memory=bool(getattr(config, "pin_cpu_memory", True)),
|
|
resident_layers=int(getattr(config, "dlo_resident_layers", 0)),
|
|
)
|
|
else:
|
|
strategy = _public_strategy(public)
|
|
if legacy is not OffloadStrategy.NONE and legacy is not strategy:
|
|
raise ValueError("diffusion_offload_config cannot be combined with legacy enable_*_offload flags")
|
|
transfers = {component: DLOTransfer.RANK_LOCAL for component in OFFLOAD_COMPONENTS}
|
|
for component, options in public.layer_options.items():
|
|
transfers[component] = options.weight_transfer or DLOTransfer.RANK_LOCAL
|
|
dit_options = public.layer_options.get(DIT_COMPONENT)
|
|
resolved = ResolvedOffload(
|
|
strategy=strategy,
|
|
public=public,
|
|
components=public.components,
|
|
transfers=transfers,
|
|
pin_memory=(
|
|
public.pin_memory if public.pin_memory is not None else bool(getattr(config, "pin_cpu_memory", True))
|
|
),
|
|
resident_layers=0 if dit_options is None else dit_options.resident_layers,
|
|
)
|
|
|
|
parallel_config = getattr(config, "parallel_config", None)
|
|
data_parallel_size = int(getattr(parallel_config, "data_parallel_size", 1))
|
|
cache_backend = getattr(config, "cache_backend", None)
|
|
if data_parallel_size > 1 and resolved.any_allgather and cache_backend not in (None, "none"):
|
|
raise ValueError(
|
|
"Diffusion cache acceleration cannot be combined with DLO AllGather "
|
|
"across data-parallel ranks because rank-local cache decisions can skip "
|
|
"different weight collectives. Disable cache_backend or use rank-local transfer."
|
|
)
|
|
if (
|
|
data_parallel_size > 1
|
|
and bool(getattr(config, "enable_prompt_embed_cache", False))
|
|
and resolved.offloads(TEXT_ENCODER_COMPONENT)
|
|
and resolved.uses_allgather(TEXT_ENCODER_COMPONENT)
|
|
):
|
|
raise ValueError(
|
|
"Prompt embedding cache cannot be combined with text_encoder AllGather "
|
|
"across data-parallel ranks: rank-local cache hits would skip different "
|
|
"encoder collectives. Use text_encoder weight_transfer='rank-local' or "
|
|
"disable enable_prompt_embed_cache."
|
|
)
|
|
|
|
# Actual runtime configs are immutable after construction. Keeping the
|
|
# derived object off the public dataclass fields also keeps projections and
|
|
# serialization transport-safe while making hot-path accessors constant-time.
|
|
try:
|
|
setattr(config, "_resolved_diffusion_offload", resolved)
|
|
except (AttributeError, TypeError):
|
|
pass
|
|
return resolved
|
|
|
|
|
|
def resolve_offload_strategy(config: Any) -> OffloadStrategy:
|
|
"""Resolve the compact config or compatibility boolean entry points."""
|
|
return resolve_offload(config).strategy
|
|
|
|
|
|
def materialize_legacy_offload_flags(config: Any) -> OffloadStrategy:
|
|
"""Keep existing strategy readers working after resolving the compact API."""
|
|
resolved = resolve_offload(config)
|
|
for field, field_strategy in _LEGACY_STRATEGY_FIELDS.items():
|
|
setattr(config, field, resolved.strategy is field_strategy)
|
|
setattr(config, "pin_cpu_memory", resolved.pin_memory)
|
|
return resolved.strategy
|
|
|
|
|
|
def parse_offload_components(value: Collection[str]) -> frozenset[str]:
|
|
"""Validate an internal component collection."""
|
|
if isinstance(value, (str, Mapping)) or not isinstance(value, Collection):
|
|
raise TypeError("diffusion_offload_config.components must be a non-empty list of component names")
|
|
if any(not isinstance(item, str) for item in value):
|
|
raise TypeError("offload component entries must be strings")
|
|
components = list(value)
|
|
if not components:
|
|
raise ValueError("offload components must not be empty")
|
|
unknown = sorted(set(components) - OFFLOAD_COMPONENTS)
|
|
if unknown:
|
|
choices = ", ".join(sorted(OFFLOAD_COMPONENTS))
|
|
raise ValueError(f"Unknown diffusion offload component(s): {', '.join(unknown)}; choose from: {choices}")
|
|
return frozenset(components)
|
|
|
|
|
|
def parse_dlo_transfer(value: Mapping[str, str | DLOTransfer]) -> dict[str, DLOTransfer]:
|
|
"""Resolve an internal per-component transfer mapping."""
|
|
if not isinstance(value, Mapping):
|
|
raise TypeError(f"offload transfers must be a mapping, got {type(value).__name__}")
|
|
_validate_string_keys(value, "offload transfers")
|
|
missing = sorted(OFFLOAD_COMPONENTS - set(value))
|
|
if missing:
|
|
raise ValueError("offload transfers require every component; missing: " + ", ".join(missing))
|
|
resolved: dict[str, DLOTransfer] = {}
|
|
for component, raw_transfer in value.items():
|
|
if component not in OFFLOAD_COMPONENTS:
|
|
choices = ", ".join(sorted(OFFLOAD_COMPONENTS))
|
|
raise ValueError(f"Unknown offload transfer component {component!r}; choose from: {choices}")
|
|
resolved[component] = _parse_enum(raw_transfer, DLOTransfer, "offload transfer")
|
|
return resolved
|
|
|
|
|
|
def component_uses_allgather(config: Any, component: str = DIT_COMPONENT) -> bool:
|
|
"""Return whether one selected component uses AllGather transport."""
|
|
resolved = resolve_offload(config)
|
|
if not resolved.offloads(component):
|
|
raise ValueError(f"Offload component {component!r} is not selected")
|
|
return resolved.uses_allgather(component)
|
|
|
|
|
|
def selected_offload_components(config: Any) -> frozenset[str]:
|
|
"""Resolve selected components while preserving legacy topology defaults."""
|
|
return resolve_offload(config).components
|
|
|
|
|
|
def should_offload_component(config: Any, component: str) -> bool:
|
|
"""Return whether an active layer policy selects ``component``."""
|
|
if component not in OFFLOAD_COMPONENTS:
|
|
raise ValueError(f"Unknown offload component: {component}")
|
|
resolved = resolve_offload(config)
|
|
if resolved.strategy not in {
|
|
OffloadStrategy.LAYER_WISE,
|
|
OffloadStrategy.DISTRIBUTED_LAYER_WISE,
|
|
}:
|
|
return False
|
|
return resolved.offloads(component)
|
|
|
|
|
|
def any_selected_component_uses_allgather(config: Any) -> bool:
|
|
"""Return whether an enabled layer backend requires weight collectives."""
|
|
return resolve_offload(config).any_allgather
|