197 lines
8.4 KiB
Python
197 lines
8.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Generic validation-first restorer for exact final-layout tensor artifacts."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
from vllm_omni.host_weight_runtime import HostWeightLease, TensorKind, WeightRestorePlan
|
|
|
|
from .identity_adapter import (
|
|
FinalLayoutIdentityContext,
|
|
validate_final_layout_identity,
|
|
)
|
|
from .tensor_layout import collect_final_layout_targets, validate_final_layout_model_contract
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _Replacement:
|
|
name: str
|
|
parent: nn.Module
|
|
leaf_name: str
|
|
target: torch.Tensor
|
|
source: torch.Tensor
|
|
kind: TensorKind
|
|
|
|
def current_target(self) -> torch.Tensor | None:
|
|
if self.kind is TensorKind.PARAMETER:
|
|
return self.parent._parameters.get(self.leaf_name)
|
|
return self.parent._buffers.get(self.leaf_name)
|
|
|
|
|
|
class FinalLayoutTensorRestorePlan(WeightRestorePlan):
|
|
"""One-shot exact tensor rebinding validated before the first mutation."""
|
|
|
|
def __init__(
|
|
self,
|
|
lease: HostWeightLease,
|
|
replacements: tuple[_Replacement, ...],
|
|
validators: tuple[Callable[[], None], ...],
|
|
source_guard: Callable[[], None],
|
|
) -> None:
|
|
self._lease = lease
|
|
self._replacements = replacements
|
|
self._validators = validators
|
|
self._source_guard = source_guard
|
|
self._committed = False
|
|
|
|
def commit(self) -> None:
|
|
if self._committed:
|
|
raise RuntimeError("final-layout tensor restore plan was already committed")
|
|
if self._lease.closed:
|
|
raise RuntimeError("cannot commit a final-layout tensor restore from a closed lease")
|
|
self._source_guard()
|
|
|
|
for replacement in self._replacements:
|
|
current = replacement.current_target()
|
|
if current is not replacement.target:
|
|
raise RuntimeError(f"restore target {replacement.name!r} changed after preflight")
|
|
if (
|
|
tuple(current.shape) != tuple(replacement.source.shape)
|
|
or current.dtype != replacement.source.dtype
|
|
or tuple(current.stride()) != tuple(replacement.source.stride())
|
|
or not replacement.source.is_contiguous()
|
|
):
|
|
raise RuntimeError(f"restore target {replacement.name!r} changed layout after preflight")
|
|
|
|
# Mark first: assignment or model-validator failure makes the plan
|
|
# permanently non-retryable and the partially restored model disposable.
|
|
self._committed = True
|
|
for replacement in self._replacements:
|
|
if replacement.target.is_meta:
|
|
if replacement.kind is TensorKind.PARAMETER:
|
|
replacement.parent._parameters[replacement.leaf_name] = nn.Parameter(
|
|
replacement.source,
|
|
requires_grad=replacement.target.requires_grad,
|
|
)
|
|
else:
|
|
replacement.parent._buffers[replacement.leaf_name] = replacement.source
|
|
else:
|
|
replacement.target.data = replacement.source
|
|
for validator in self._validators:
|
|
validator()
|
|
|
|
|
|
class FinalLayoutTensorRestorer:
|
|
"""Validate and plan restoration for one exact policy-defined lease."""
|
|
|
|
def __init__(self, context: FinalLayoutIdentityContext) -> None:
|
|
if not context.dit_names or any(not name for name in context.dit_names):
|
|
raise ValueError("final-layout restorer requires DiT component names")
|
|
self._context = context
|
|
|
|
@property
|
|
def schema(self) -> str:
|
|
return self._context.spec.producer.restorer_schema
|
|
|
|
def plan_restore(self, model: object, lease: HostWeightLease) -> FinalLayoutTensorRestorePlan:
|
|
if not isinstance(model, nn.Module):
|
|
raise TypeError("final-layout restoration requires an nn.Module pipeline")
|
|
if type(model) is not self._context.pipeline_type:
|
|
raise ValueError("restore model implementation differs from the exact identity context")
|
|
if lease.closed:
|
|
raise ValueError("cannot restore from a closed HostWeightLease")
|
|
if lease.identity.canonical_bytes != self._context.identity.canonical_bytes:
|
|
raise ValueError("lease semantic identity differs from the exact restore request")
|
|
if lease.manifest.identity.canonical_bytes != self._context.identity.canonical_bytes:
|
|
raise ValueError("lease manifest identity differs from the exact restore request")
|
|
if lease.identity.producer.restorer_schema != self.schema or lease.manifest.restorer_schema != self.schema:
|
|
raise ValueError("lease restorer schema is incompatible with the exact restore request")
|
|
self._context.ensure_sources_unchanged()
|
|
|
|
try:
|
|
dit_modules = tuple((name, model.get_submodule(name)) for name in self._context.dit_names)
|
|
except AttributeError as exc:
|
|
raise ValueError("one or more lease-owned DiT modules do not exist in the pipeline") from exc
|
|
if tuple(type(module) for _, module in dit_modules) != self._context.dit_types:
|
|
raise ValueError("restore DiT implementation differs from the exact identity context")
|
|
bindings = validate_final_layout_model_contract(
|
|
dit_modules,
|
|
expected_schema=self._context.spec.model_contract_schema,
|
|
)
|
|
targets = collect_final_layout_targets(
|
|
model,
|
|
dit_modules,
|
|
policy=self._context.policy,
|
|
require_materialized=False,
|
|
)
|
|
contract_digest = validate_final_layout_identity(self._context, targets)
|
|
|
|
manifest_entries = {entry.name: entry for entry in lease.manifest.tensors}
|
|
target_names = {target.name for target in targets}
|
|
if set(manifest_entries) != target_names or target_names != self._context.tensor_names:
|
|
missing = sorted(target_names - set(manifest_entries))
|
|
unexpected = sorted(set(manifest_entries) - target_names)
|
|
raise ValueError(
|
|
"lease tensor coverage differs from the structurally initialized DiT: "
|
|
f"missing={missing[:5]}, unexpected={unexpected[:5]}"
|
|
)
|
|
self._context.policy.validate_format_metadata(
|
|
lease.manifest.format_metadata,
|
|
component_names=self._context.dit_names,
|
|
tensor_contract_digest=contract_digest,
|
|
tensor_count=len(manifest_entries),
|
|
)
|
|
|
|
replacements: list[_Replacement] = []
|
|
source_storages: dict[tuple[int, int], str] = {}
|
|
for target in targets:
|
|
entry = manifest_entries[target.name]
|
|
source = lease.tensors[target.name]
|
|
if entry.kind is not target.kind or entry.role != target.role:
|
|
raise ValueError(f"lease tensor kind or role differs for {target.name!r}")
|
|
if (
|
|
tuple(source.shape) != tuple(target.tensor.shape)
|
|
or source.dtype != target.tensor.dtype
|
|
or tuple(source.stride()) != tuple(target.tensor.stride())
|
|
or not source.is_contiguous()
|
|
or source.device.type != "cpu"
|
|
):
|
|
raise ValueError(f"lease tensor layout differs for {target.name!r}")
|
|
if source.numel():
|
|
storage = source.untyped_storage()
|
|
storage_id = (storage.data_ptr(), storage.nbytes())
|
|
owner = source_storages.setdefault(storage_id, target.name)
|
|
if owner != target.name:
|
|
raise ValueError(f"lease tensor {target.name!r} aliases {owner!r}")
|
|
parent_path, _, leaf_name = target.name.rpartition(".")
|
|
replacements.append(
|
|
_Replacement(
|
|
name=target.name,
|
|
parent=model.get_submodule(parent_path),
|
|
leaf_name=leaf_name,
|
|
target=target.tensor,
|
|
source=source,
|
|
kind=target.kind,
|
|
)
|
|
)
|
|
|
|
self._context.ensure_sources_unchanged()
|
|
return FinalLayoutTensorRestorePlan(
|
|
lease,
|
|
tuple(replacements),
|
|
tuple(binding.validator for binding in bindings),
|
|
self._context.ensure_sources_unchanged,
|
|
)
|
|
|
|
|
|
__all__ = [
|
|
"FinalLayoutTensorRestorePlan",
|
|
"FinalLayoutTensorRestorer",
|
|
]
|