508 lines
22 KiB
Python
508 lines
22 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Host Weight Runtime integration for diffusion model loaders.
|
|
|
|
This mixin keeps HWR identity, restore, publication, and startup-recovery
|
|
policy out of the ordinary diffusers loader. The host loader supplies the
|
|
canonical source-preparation and weight-loading hooks; this module owns only
|
|
the optional final-layout transaction.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import dataclasses
|
|
import hashlib
|
|
import os
|
|
from collections.abc import Iterable, Sequence
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, cast
|
|
|
|
import torch
|
|
from torch import nn
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.model_loader.host_weight_plan import HostWeightPlan, TensorBinding
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
from vllm_omni.diffusion.model_loader.host_weights.identity_adapter import FinalLayoutIdentityContext
|
|
from vllm_omni.diffusion.model_loader.host_weights.source_identity import (
|
|
NodeSourceDigestCache,
|
|
PreparedWeightSource,
|
|
)
|
|
|
|
|
|
class _HWRCommitError(RuntimeError):
|
|
"""A committed warm restore made the current model disposable."""
|
|
|
|
|
|
class HWRLoaderMixin:
|
|
"""Optional final-layout HWR behavior shared by diffusion loaders."""
|
|
|
|
od_config: OmniDiffusionConfig
|
|
|
|
@staticmethod
|
|
def _identity_value(value: object) -> object:
|
|
"""Convert config objects into deterministic identity metadata."""
|
|
if value is None or isinstance(value, (bool, int, float, str)):
|
|
return value
|
|
if isinstance(value, torch.dtype):
|
|
return str(value)
|
|
if isinstance(value, dict):
|
|
return {str(key): HWRLoaderMixin._identity_value(item) for key, item in value.items()}
|
|
if isinstance(value, (list, tuple)):
|
|
return [HWRLoaderMixin._identity_value(item) for item in value]
|
|
to_dict = getattr(value, "to_dict", None)
|
|
if callable(to_dict):
|
|
try:
|
|
return HWRLoaderMixin._identity_value(to_dict())
|
|
except TypeError:
|
|
pass
|
|
if dataclasses.is_dataclass(value):
|
|
return HWRLoaderMixin._identity_value(dataclasses.asdict(value))
|
|
return f"{type(value).__module__}.{type(value).__qualname__}:{value!r}"
|
|
|
|
@staticmethod
|
|
def _identity_fingerprint(value: object) -> str:
|
|
from vllm_omni.host_weight_runtime import CanonicalJson
|
|
|
|
return hashlib.sha256(CanonicalJson.from_value(value).encoded).hexdigest()
|
|
|
|
@staticmethod
|
|
def _snapshot_final_layout_tensors(model: nn.Module, names: Iterable[str]) -> dict[str, tuple[int, str]]:
|
|
snapshot: dict[str, tuple[int, str]] = {}
|
|
for name in names:
|
|
parent_path, _, leaf_name = name.rpartition(".")
|
|
parent = model.get_submodule(parent_path)
|
|
tensor = parent._parameters.get(leaf_name)
|
|
if tensor is None:
|
|
tensor = parent._buffers.get(leaf_name)
|
|
if tensor is None or tensor.is_meta:
|
|
raise RuntimeError(f"cannot snapshot missing or meta final-layout tensor {name!r}")
|
|
value = tensor.detach()
|
|
if value.device.type != "cpu":
|
|
value = value.cpu()
|
|
value = value.contiguous()
|
|
digest = hashlib.sha256(memoryview(value.view(torch.uint8).numpy())).hexdigest()
|
|
snapshot[name] = (tensor.untyped_storage().data_ptr(), digest)
|
|
return snapshot
|
|
|
|
@classmethod
|
|
def _assert_final_layout_tensors_unchanged(
|
|
cls,
|
|
model: nn.Module,
|
|
snapshot: dict[str, tuple[int, str]],
|
|
) -> None:
|
|
current = cls._snapshot_final_layout_tensors(model, snapshot)
|
|
changed = [name for name in snapshot if current[name] != snapshot[name]]
|
|
if changed:
|
|
raise RuntimeError(
|
|
"shared warm finalization changed restored final-layout tensor bytes or backing pointers: "
|
|
f"{changed[:5]}"
|
|
)
|
|
|
|
def _hwr_eligibility_mode(
|
|
self,
|
|
model: nn.Module,
|
|
modules: object,
|
|
*,
|
|
dist_offload: bool,
|
|
use_allgather: bool,
|
|
load_format: str,
|
|
) -> object | None:
|
|
"""Return the enabled HWR mode only after all zero-interaction gates."""
|
|
from vllm_omni.host_weight_runtime import RuntimeMode
|
|
|
|
raw_mode = getattr(self.od_config, "host_weight_runtime_mode", "disabled")
|
|
try:
|
|
mode = RuntimeMode(raw_mode)
|
|
except ValueError as exc:
|
|
raise ValueError("host_weight_runtime_mode must be disabled, preferred, or required") from exc
|
|
|
|
# These gates intentionally precede HWR imports, source preparation,
|
|
# identity construction, and store creation.
|
|
if mode is RuntimeMode.DISABLED or not dist_offload or use_allgather:
|
|
return None
|
|
|
|
parallel = self.parallel_config
|
|
reasons: list[str] = []
|
|
if load_format != "default":
|
|
reasons.append("load_format must be 'default'")
|
|
if bool(getattr(parallel, "use_hsdp", False)):
|
|
reasons.append("HSDP layouts are not supported by the final-layout BF16 consumer")
|
|
if self.quant_config is not None:
|
|
reasons.append("quantized layouts require a representation-specific HWR producer")
|
|
if getattr(self.od_config, "lora_path", None):
|
|
reasons.append("adapted weights are not reusable base-model artifacts")
|
|
|
|
dit_modules = tuple(zip(getattr(modules, "dit_names", ()), getattr(modules, "dits", ())))
|
|
if not dit_modules:
|
|
reasons.append("no DiT modules were discovered")
|
|
elif any(
|
|
getattr(dit, "host_weight_restore_contract", None) is None
|
|
or not callable(getattr(dit, "validate_restored_host_weights", None))
|
|
for _, dit in dit_modules
|
|
):
|
|
reasons.append("every owned DiT must declare the final-layout restore contract")
|
|
|
|
if reasons:
|
|
message = "; ".join(reasons)
|
|
if mode is RuntimeMode.REQUIRED:
|
|
raise ValueError(f"required Host Weight Runtime path is ineligible: {message}")
|
|
logger.info("Host Weight Runtime is ineligible; using the canonical DLO path: %s", message)
|
|
return None
|
|
return mode
|
|
|
|
def _prepare_hwr_sources(
|
|
self,
|
|
model: nn.Module,
|
|
modules: object,
|
|
sources: Sequence[object],
|
|
) -> tuple[PreparedWeightSource, ...]:
|
|
from vllm_omni.diffusion.model_loader.host_weights import (
|
|
ImplementationIdentity,
|
|
PreparedWeightSource,
|
|
WeightSourceKind,
|
|
)
|
|
|
|
dit_prefixes = tuple(f"{name}." for name in getattr(modules, "dit_names", ()))
|
|
selected_sources = tuple(
|
|
source
|
|
for source in sources
|
|
if not source.prefix or any(prefix.startswith(source.prefix) for prefix in dit_prefixes)
|
|
)
|
|
if not selected_sources:
|
|
raise ValueError("final-layout HWR requires canonical weight sources covering every DiT")
|
|
|
|
prepared: list[PreparedWeightSource] = []
|
|
for source in selected_sources:
|
|
resolved_root, weight_files, use_safetensors = self._prepare_weights(
|
|
source.model_or_path,
|
|
source.subfolder,
|
|
source.revision,
|
|
source.fall_back_to_pt,
|
|
source.allow_patterns_overrides,
|
|
)
|
|
adapter = self._get_checkpoint_adapter(model, source, use_safetensors)
|
|
adapter_identity = None
|
|
if adapter is not None:
|
|
adapter_name = f"{type(adapter).__module__}.{type(adapter).__qualname__}"
|
|
adapter_identity = ImplementationIdentity(
|
|
implementation_id=adapter_name,
|
|
version="1",
|
|
fingerprint=self._identity_fingerprint({"adapter": adapter_name}),
|
|
)
|
|
source_kind = (
|
|
WeightSourceKind.LOCAL_PATH
|
|
if os.path.isdir(os.fspath(source.model_or_path))
|
|
else WeightSourceKind.HUGGING_FACE_HUB
|
|
)
|
|
prepared.append(
|
|
PreparedWeightSource(
|
|
model_or_path=os.fspath(source.model_or_path),
|
|
subfolder=source.subfolder,
|
|
requested_revision=source.revision,
|
|
prefix=source.prefix,
|
|
resolved_root=Path(os.fspath(resolved_root)),
|
|
weight_files=tuple(Path(os.fspath(path)) for path in weight_files),
|
|
use_safetensors=use_safetensors,
|
|
checkpoint_adapter=adapter_identity,
|
|
source_kind=source_kind,
|
|
)
|
|
)
|
|
return tuple(prepared)
|
|
|
|
def _build_hwr_context(
|
|
self,
|
|
model: nn.Module,
|
|
modules: object,
|
|
*,
|
|
load_format: str,
|
|
sources: Sequence[object],
|
|
source_digest_cache: NodeSourceDigestCache | None = None,
|
|
) -> FinalLayoutIdentityContext:
|
|
from vllm.distributed.parallel_state import get_tensor_model_parallel_rank
|
|
|
|
from vllm_omni.diffusion.model_loader.host_weights import (
|
|
FINAL_LAYOUT_BF16_POLICY,
|
|
FinalLayoutLoaderIdentity,
|
|
FinalLayoutParallelIdentity,
|
|
FinalLayoutRequest,
|
|
ImplementationIdentity,
|
|
build_final_layout_identity,
|
|
)
|
|
|
|
parallel = self.parallel_config
|
|
tp_size = int(getattr(parallel, "tensor_parallel_size", 1))
|
|
try:
|
|
tp_rank = int(get_tensor_model_parallel_rank()) if tp_size > 1 else 0
|
|
except Exception:
|
|
tp_rank = int(getattr(parallel, "tensor_parallel_rank", 0))
|
|
sp_size = int(getattr(parallel, "sequence_parallel_size", 1) or 1)
|
|
loader_config = {
|
|
"dtype": str(self.od_config.dtype),
|
|
"model_class": f"{type(model).__module__}.{type(model).__qualname__}",
|
|
"model_config": self._identity_value(getattr(self.od_config, "tf_model_config", None)),
|
|
"load_format": load_format,
|
|
"quantization": self._identity_value(self.quant_config),
|
|
}
|
|
contracts = [
|
|
self._identity_value(getattr(dit, "host_weight_restore_contract")) for dit in getattr(modules, "dits", ())
|
|
]
|
|
loader_identity = FinalLayoutLoaderIdentity(
|
|
implementation=ImplementationIdentity(
|
|
implementation_id="vllm-omni.diffusion.diffusers-loader",
|
|
version="final-layout-v1",
|
|
fingerprint=self._identity_fingerprint(
|
|
{
|
|
"loader": "diffusers-loader-final-layout-v1",
|
|
"pipeline": f"{type(model).__module__}.{type(model).__qualname__}",
|
|
}
|
|
),
|
|
),
|
|
model_config_fingerprint=self._identity_fingerprint(loader_config),
|
|
weight_transform_fingerprint=self._identity_fingerprint(
|
|
{
|
|
"contracts": contracts,
|
|
"transform": "diffusion-final-layout-loader-transforms-v1",
|
|
}
|
|
),
|
|
)
|
|
semantic_parallel = FinalLayoutParallelIdentity(
|
|
tensor_parallel_size=tp_size,
|
|
tensor_parallel_rank=tp_rank,
|
|
sequence_parallel_size=sp_size,
|
|
ulysses_degree=int(getattr(parallel, "ulysses_degree", 1)),
|
|
ring_degree=int(getattr(parallel, "ring_degree", 1)),
|
|
allgather_degree=int(getattr(parallel, "allgather_degree", 1)),
|
|
ulysses_mode=str(getattr(parallel, "ulysses_mode", "strict")),
|
|
pipeline_parallel_size=int(getattr(parallel, "pipeline_parallel_size", 1)),
|
|
cfg_parallel_size=int(getattr(parallel, "cfg_parallel_size", 1)),
|
|
use_hsdp=bool(getattr(parallel, "use_hsdp", False)),
|
|
enable_expert_parallel=bool(getattr(parallel, "enable_expert_parallel", False)),
|
|
)
|
|
model_id = str(getattr(self.od_config, "model", "") or "")
|
|
if not model_id:
|
|
raise ValueError("final-layout HWR requires a canonical model identifier")
|
|
request = FinalLayoutRequest(
|
|
model_id=model_id,
|
|
loader=loader_identity,
|
|
parallel=semantic_parallel,
|
|
load_format=load_format,
|
|
)
|
|
prepared_sources = self._prepare_hwr_sources(model, modules, sources)
|
|
dit_modules = tuple(zip(getattr(modules, "dit_names", ()), getattr(modules, "dits", ())))
|
|
return build_final_layout_identity(
|
|
model,
|
|
dit_modules=dit_modules,
|
|
prepared_sources=prepared_sources,
|
|
request=request,
|
|
policy=FINAL_LAYOUT_BF16_POLICY,
|
|
source_digest_cache=source_digest_cache,
|
|
)
|
|
|
|
def _resolve_hwr(
|
|
self,
|
|
model: nn.Module,
|
|
modules: object,
|
|
*,
|
|
dist_offload: bool,
|
|
use_allgather: bool,
|
|
load_format: str,
|
|
sources: Sequence[object],
|
|
) -> dict[str, object] | None:
|
|
"""Resolve an eligible no-AllGather final-layout HWR transaction."""
|
|
from vllm_omni.diffusion.model_loader.host_weights import (
|
|
FinalLayoutTensorRestorer,
|
|
NodeSourceDigestCache,
|
|
)
|
|
from vllm_omni.host_weight_runtime import (
|
|
HostWeightLeaseCarrier,
|
|
HostWeightRuntime,
|
|
HostWeightRuntimeConfig,
|
|
IntegrityPolicy,
|
|
ProductionPolicy,
|
|
ResolutionOutcome,
|
|
RuntimeMode,
|
|
StorageDomainPolicy,
|
|
ValidationLevel,
|
|
)
|
|
from vllm_omni.host_weight_runtime.filesystem import detect_storage_class
|
|
|
|
mode = self._hwr_eligibility_mode(
|
|
model,
|
|
modules,
|
|
dist_offload=dist_offload,
|
|
use_allgather=use_allgather,
|
|
load_format=load_format,
|
|
)
|
|
if mode is None:
|
|
return None
|
|
assert isinstance(mode, RuntimeMode)
|
|
expected_prefixes = frozenset(f"{name}." for name in getattr(modules, "dit_names", ()))
|
|
available_prefixes = frozenset(getattr(source, "prefix", "") for source in sources)
|
|
if not expected_prefixes <= available_prefixes:
|
|
message = "final-layout HWR requires one dedicated source prefix per owned DiT"
|
|
if mode is RuntimeMode.REQUIRED:
|
|
raise ValueError(f"required Host Weight Runtime path is ineligible: {message}")
|
|
logger.info("Host Weight Runtime is ineligible; using the canonical DLO path: %s", message)
|
|
return None
|
|
overlapping_prefixes = frozenset(
|
|
prefix
|
|
for prefix in available_prefixes
|
|
if any(dit_prefix.startswith(prefix) for dit_prefix in expected_prefixes)
|
|
)
|
|
if not overlapping_prefixes <= expected_prefixes:
|
|
message = "final-layout HWR requires dedicated DiT sources and rejects mixed component sources"
|
|
if mode is RuntimeMode.REQUIRED:
|
|
raise ValueError(f"required Host Weight Runtime path is ineligible: {message}")
|
|
logger.info("Host Weight Runtime is ineligible; using the canonical DLO path: %s", message)
|
|
return None
|
|
root_value = getattr(self.od_config, "host_weight_runtime_root", None)
|
|
if not isinstance(root_value, str) or not root_value.strip():
|
|
raise ValueError("enabled Host Weight Runtime requires host_weight_runtime_root")
|
|
root = Path(root_value).expanduser()
|
|
runtime = HostWeightRuntime.from_config(
|
|
HostWeightRuntimeConfig(
|
|
mode=mode,
|
|
domain=StorageDomainPolicy(root=root, storage_class=detect_storage_class(root)),
|
|
integrity=IntegrityPolicy(
|
|
local_lookup=ValidationLevel(
|
|
getattr(self.od_config, "host_weight_runtime_validation", "manifest_and_metadata")
|
|
),
|
|
),
|
|
production=ProductionPolicy(
|
|
allow_local_build=False,
|
|
allow_post_load_publish=True,
|
|
),
|
|
)
|
|
)
|
|
source_digest_cache = NodeSourceDigestCache(
|
|
root,
|
|
timeout_seconds=runtime.config.wait.coordination_timeout_seconds,
|
|
)
|
|
context = self._build_hwr_context(
|
|
model,
|
|
modules,
|
|
load_format=load_format,
|
|
sources=sources,
|
|
source_digest_cache=source_digest_cache,
|
|
)
|
|
resolution = runtime.resolve(context.identity)
|
|
state: dict[str, object] = {
|
|
"mode": mode,
|
|
"context": context,
|
|
"runtime": runtime,
|
|
"outcome": resolution.report.outcome,
|
|
}
|
|
if resolution.report.outcome is ResolutionOutcome.FAILED:
|
|
failure = next(
|
|
(attempt.failure for attempt in reversed(resolution.report.attempts) if attempt.failure is not None),
|
|
None,
|
|
)
|
|
detail = failure.message if failure is not None else "resolution failed without a typed detail"
|
|
raise RuntimeError(f"Host Weight Runtime resolution failed: {detail}")
|
|
if resolution.report.outcome is not ResolutionOutcome.LOCAL_HIT:
|
|
return state
|
|
|
|
lease = resolution.lease
|
|
if lease is None:
|
|
raise RuntimeError("Host Weight Runtime returned LOCAL_HIT without a lease")
|
|
restorer = FinalLayoutTensorRestorer(context)
|
|
try:
|
|
restore_plan = restorer.plan_restore(model, lease)
|
|
except Exception:
|
|
lease.close()
|
|
if mode is RuntimeMode.REQUIRED:
|
|
raise
|
|
logger.warning("HWR warm restore planning failed; falling back to canonical loading", exc_info=True)
|
|
state["outcome"] = ResolutionOutcome.CANONICAL_FALLBACK
|
|
return state
|
|
|
|
try:
|
|
restore_plan.commit()
|
|
except Exception as exc:
|
|
lease.close()
|
|
raise _HWRCommitError(
|
|
"Host Weight Runtime restore commit failed; the partially restored model must be discarded"
|
|
) from exc
|
|
|
|
try:
|
|
carrier = HostWeightLeaseCarrier(lease)
|
|
warm_snapshot = self._snapshot_final_layout_tensors(model, context.tensor_names)
|
|
except Exception as exc:
|
|
lease.close()
|
|
raise _HWRCommitError(
|
|
"Host Weight Runtime committed restore could not establish its startup ownership boundary"
|
|
) from exc
|
|
source_metadata = context.identity.source.metadata.to_value()
|
|
target_bindings = source_metadata.get("target_bindings") if isinstance(source_metadata, dict) else None
|
|
planned_prefixes = frozenset(
|
|
binding["source_prefix"]
|
|
for binding in (target_bindings or ())
|
|
if isinstance(binding, dict) and isinstance(binding.get("source_prefix"), str)
|
|
)
|
|
if not planned_prefixes:
|
|
lease.close()
|
|
raise _HWRCommitError("Host Weight Runtime identity did not record owned canonical source prefixes")
|
|
state["plan"] = HostWeightPlan(
|
|
backing_kind="host_weight_runtime",
|
|
bindings={name: TensorBinding(name, "") for name in context.tensor_names},
|
|
planned_source_prefixes=planned_prefixes,
|
|
lease_carrier=carrier,
|
|
)
|
|
state["warm_snapshot"] = warm_snapshot
|
|
return state
|
|
|
|
def load_fresh_canonical_model(self) -> nn.Module:
|
|
"""Reload a disposable startup model without HWR or checkpoint mmap."""
|
|
if self._last_load_request is None:
|
|
raise RuntimeError("cannot construct a fresh canonical model before the initial load")
|
|
request = dict(self._last_load_request)
|
|
self._force_canonical_load = True
|
|
try:
|
|
return self.load_model(**cast(dict[str, object], request))
|
|
finally:
|
|
self._force_canonical_load = False
|
|
|
|
def _publish_hwr_after_load(self, model: nn.Module, modules: object, state: dict[str, object] | None) -> None:
|
|
from vllm_omni.diffusion.model_loader.host_weights import FinalLayoutBF16Producer
|
|
from vllm_omni.host_weight_runtime import PostLoadPublicationOutcome, ResolutionOutcome
|
|
|
|
if state is None or state.get("outcome") is not ResolutionOutcome.CANONICAL_FALLBACK:
|
|
return
|
|
context = state["context"]
|
|
runtime = state["runtime"]
|
|
dit_modules = tuple(zip(getattr(modules, "dit_names", ()), getattr(modules, "dits", ())))
|
|
try:
|
|
producer = FinalLayoutBF16Producer(context, model, dit_modules)
|
|
report = runtime.publish_after_load(context.identity, producer=producer)
|
|
except Exception:
|
|
logger.warning("Host Weight Runtime post-load publication failed", exc_info=True)
|
|
return
|
|
if report.outcome is PostLoadPublicationOutcome.FAILED:
|
|
logger.warning("Host Weight Runtime post-load publication failed: %s", report.failure)
|
|
|
|
def _attach_offload_startup_state(self, model: nn.Module) -> None:
|
|
"""Publish generic offload startup state on the loaded pipeline."""
|
|
if self.host_weight_plan is None:
|
|
return
|
|
from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state
|
|
|
|
hwr_state = self._hwr_state
|
|
allow_retry = False
|
|
if hwr_state is not None and isinstance(hwr_state.get("plan"), HostWeightPlan):
|
|
from vllm_omni.host_weight_runtime import RuntimeMode
|
|
|
|
allow_retry = hwr_state.get("mode") is RuntimeMode.PREFERRED
|
|
attach_offload_startup_state(
|
|
model,
|
|
OffloadStartupState(
|
|
host_weight_plan=self.host_weight_plan,
|
|
fresh_model_loader=self.load_fresh_canonical_model if allow_retry else None,
|
|
allow_fresh_retry=allow_retry,
|
|
),
|
|
)
|