397 lines
15 KiB
Python
397 lines
15 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Loader-owned host-weight plans shared with diffusion offload backends."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import glob
|
|
import json
|
|
import os
|
|
from collections.abc import Callable, Sequence
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING
|
|
|
|
import torch
|
|
from safetensors import safe_open
|
|
from torch import nn
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
|
|
get_direct_mmap_adapter,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm_omni.host_weight_runtime import HostWeightLeaseCarrier
|
|
|
|
TensorTransform = Callable[[torch.Tensor], torch.Tensor]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TensorBinding:
|
|
"""One runtime tensor backed by one safetensors entry."""
|
|
|
|
checkpoint_key: str
|
|
file_path: str
|
|
transform: TensorTransform | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class HostWeightPlan:
|
|
"""Complete, prevalidated host backing consumed by an offload backend."""
|
|
|
|
backing_kind: str
|
|
bindings: dict[str, TensorBinding]
|
|
planned_source_prefixes: frozenset[str] = frozenset()
|
|
lease_carrier: HostWeightLeaseCarrier | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class HostWeightPlanResult:
|
|
"""A complete plan or a fail-closed reason to use the ordinary loader."""
|
|
|
|
plan: HostWeightPlan | None
|
|
fallback_reason: str | None = None
|
|
|
|
|
|
class _PlanIncompatibleError(RuntimeError):
|
|
pass
|
|
|
|
|
|
_SAFETENSORS_DTYPES: dict[str, torch.dtype] = {
|
|
"BOOL": torch.bool,
|
|
"U8": torch.uint8,
|
|
"I8": torch.int8,
|
|
"I16": torch.int16,
|
|
"I32": torch.int32,
|
|
"I64": torch.int64,
|
|
"F16": torch.float16,
|
|
"BF16": torch.bfloat16,
|
|
"F32": torch.float32,
|
|
"F64": torch.float64,
|
|
}
|
|
for _safe_name, _torch_name in (
|
|
("F8_E4M3", "float8_e4m3fn"),
|
|
("F8_E5M2", "float8_e5m2"),
|
|
):
|
|
if (_dtype := getattr(torch, _torch_name, None)) is not None:
|
|
_SAFETENSORS_DTYPES[_safe_name] = _dtype
|
|
|
|
|
|
def has_online_quantization(model: nn.Module) -> bool:
|
|
"""Whether any module defers online-quantized weights on ``meta``."""
|
|
return any(getattr(getattr(module, "quant_method", None), "uses_meta_device", False) for module in model.modules())
|
|
|
|
|
|
def _resolve_source_root(source: object) -> str:
|
|
source_root = os.fspath(getattr(source, "model_or_path"))
|
|
subfolder = getattr(source, "subfolder", None)
|
|
if not os.path.isdir(source_root):
|
|
from vllm.model_executor.model_loader.weight_utils import (
|
|
download_weights_from_hf,
|
|
)
|
|
|
|
source_root = download_weights_from_hf(
|
|
model_name_or_path=source_root,
|
|
cache_dir=None,
|
|
allow_patterns=["*.safetensors", "*.safetensors.index.json"],
|
|
revision=getattr(source, "revision", None),
|
|
subfolder=subfolder,
|
|
)
|
|
return os.path.join(source_root, subfolder) if subfolder is not None else source_root
|
|
|
|
|
|
def _record_binding(
|
|
model_to_ckpt: dict[str, tuple[str, str]],
|
|
model_name: str,
|
|
checkpoint_key: str,
|
|
file_path: str,
|
|
) -> None:
|
|
existing = model_to_ckpt.get(model_name)
|
|
candidate = (checkpoint_key, file_path)
|
|
if existing is not None and existing != candidate:
|
|
raise _PlanIncompatibleError(
|
|
f"multiple checkpoint entries map to runtime tensor {model_name!r}: {existing!r} and {candidate!r}"
|
|
)
|
|
model_to_ckpt[model_name] = candidate
|
|
|
|
|
|
def _add_safetensors_source(
|
|
checkpoint_dir: str,
|
|
prefix: str,
|
|
remap_fn: Callable[[str], str | None] | None,
|
|
model_to_ckpt: dict[str, tuple[str, str]],
|
|
) -> int:
|
|
indexed_keys = 0
|
|
index_files = sorted(glob.glob(os.path.join(checkpoint_dir, "*.safetensors.index.json")))
|
|
if index_files:
|
|
for index_file in index_files:
|
|
with open(index_file, encoding="utf-8") as handle:
|
|
index = json.load(handle)
|
|
for checkpoint_key, filename in index.get("weight_map", {}).items():
|
|
source_name = prefix + checkpoint_key
|
|
model_name = remap_fn(source_name) if remap_fn is not None else source_name
|
|
if model_name is not None:
|
|
_record_binding(
|
|
model_to_ckpt,
|
|
model_name,
|
|
checkpoint_key,
|
|
os.path.join(checkpoint_dir, filename),
|
|
)
|
|
indexed_keys += 1
|
|
return indexed_keys
|
|
|
|
for filename in sorted(glob.glob(os.path.join(checkpoint_dir, "*.safetensors"))):
|
|
with safe_open(filename, framework="pt", device="cpu") as handle:
|
|
for checkpoint_key in handle.keys():
|
|
source_name = prefix + checkpoint_key
|
|
model_name = remap_fn(source_name) if remap_fn is not None else source_name
|
|
if model_name is not None:
|
|
_record_binding(model_to_ckpt, model_name, checkpoint_key, filename)
|
|
indexed_keys += 1
|
|
return indexed_keys
|
|
|
|
|
|
def _build_source_map(
|
|
*,
|
|
sources: Sequence[object],
|
|
model_path: str | None,
|
|
remap_fn: Callable[[str], str | None] | None,
|
|
) -> dict[str, tuple[str, str]]:
|
|
model_to_ckpt: dict[str, tuple[str, str]] = {}
|
|
indexed_keys = 0
|
|
|
|
for source in sources:
|
|
indexed_keys += _add_safetensors_source(
|
|
_resolve_source_root(source),
|
|
getattr(source, "prefix", ""),
|
|
remap_fn,
|
|
model_to_ckpt,
|
|
)
|
|
|
|
if not sources:
|
|
if remap_fn is None:
|
|
raise _PlanIncompatibleError(
|
|
"the loader has no component weight sources and the pipeline has no checkpoint-key remapper"
|
|
)
|
|
if not model_path:
|
|
raise _PlanIncompatibleError("the loader has no model path for legacy checkpoint remapping")
|
|
|
|
resolved_model_path = model_path
|
|
if not os.path.isdir(resolved_model_path):
|
|
from vllm.model_executor.model_loader.weight_utils import (
|
|
download_weights_from_hf,
|
|
)
|
|
|
|
resolved_model_path = download_weights_from_hf(
|
|
model_name_or_path=resolved_model_path,
|
|
cache_dir=None,
|
|
allow_patterns=["*.safetensors", "*.safetensors.index.json"],
|
|
)
|
|
checkpoint_dirs = {
|
|
os.path.dirname(path)
|
|
for pattern in ("*.safetensors.index.json", "*.safetensors")
|
|
for path in glob.glob(os.path.join(resolved_model_path, "**", pattern), recursive=True)
|
|
}
|
|
for checkpoint_dir in sorted(checkpoint_dirs):
|
|
indexed_keys += _add_safetensors_source(
|
|
checkpoint_dir,
|
|
"",
|
|
remap_fn,
|
|
model_to_ckpt,
|
|
)
|
|
|
|
if not model_to_ckpt:
|
|
raise _PlanIncompatibleError("no compatible safetensors entries were found in the loader's model sources")
|
|
logger.info(
|
|
"Indexed %d runtime tensor names from %d checkpoint keys for DLO preflight",
|
|
len(model_to_ckpt),
|
|
indexed_keys,
|
|
)
|
|
return model_to_ckpt
|
|
|
|
|
|
def _is_persistent_buffer(module: nn.Module, local_name: str) -> bool:
|
|
parts = local_name.split(".")
|
|
owner = module.get_submodule(".".join(parts[:-1])) if len(parts) > 1 else module
|
|
return parts[-1] not in owner._non_persistent_buffers_set
|
|
|
|
|
|
def _collect_required_dit_tensors(
|
|
dit_modules: Sequence[tuple[str, nn.Module]],
|
|
) -> dict[str, torch.Tensor]:
|
|
required: dict[str, torch.Tensor] = {}
|
|
for dit_name, dit_module in dit_modules:
|
|
for local_name, tensor in dit_module.named_parameters():
|
|
required[f"{dit_name}.{local_name}"] = tensor
|
|
for local_name, tensor in dit_module.named_buffers():
|
|
if _is_persistent_buffer(dit_module, local_name):
|
|
required[f"{dit_name}.{local_name}"] = tensor
|
|
return required
|
|
|
|
|
|
def _planned_source_prefixes(
|
|
sources: Sequence[object],
|
|
dit_modules: Sequence[tuple[str, nn.Module]],
|
|
) -> frozenset[str]:
|
|
"""Identify component sources that the DiT-only plan owns completely."""
|
|
if not sources:
|
|
return frozenset()
|
|
|
|
expected = frozenset(f"{dit_name}." for dit_name, _ in dit_modules)
|
|
available = {getattr(source, "prefix", "") for source in sources}
|
|
missing = sorted(expected - available)
|
|
if missing:
|
|
raise _PlanIncompatibleError(
|
|
f"direct mmap requires a dedicated component weight source for each DiT; missing prefixes: {missing}"
|
|
)
|
|
return expected
|
|
|
|
|
|
def _validate_source_metadata(
|
|
required: dict[str, torch.Tensor],
|
|
bindings: dict[str, TensorBinding],
|
|
) -> None:
|
|
by_file: dict[str, list[tuple[str, torch.Tensor, TensorBinding]]] = {}
|
|
for runtime_name, target in required.items():
|
|
by_file.setdefault(bindings[runtime_name].file_path, []).append((runtime_name, target, bindings[runtime_name]))
|
|
|
|
for file_path, entries in by_file.items():
|
|
if not os.path.isfile(file_path):
|
|
raise _PlanIncompatibleError(f"checkpoint file does not exist: {file_path}")
|
|
with safe_open(file_path, framework="pt", device="cpu") as handle:
|
|
available = set(handle.keys())
|
|
for runtime_name, target, binding in entries:
|
|
if binding.checkpoint_key not in available:
|
|
raise _PlanIncompatibleError(
|
|
f"checkpoint key {binding.checkpoint_key!r} for {runtime_name!r} is missing from {file_path}"
|
|
)
|
|
source = handle.get_slice(binding.checkpoint_key)
|
|
source_shape = tuple(source.get_shape())
|
|
source_dtype = _SAFETENSORS_DTYPES.get(source.get_dtype())
|
|
if source_dtype is None:
|
|
raise _PlanIncompatibleError(
|
|
f"dtype mismatch for {runtime_name!r}: checkpoint={source.get_dtype()}, runtime={target.dtype}"
|
|
)
|
|
|
|
runtime_shape = source_shape
|
|
runtime_dtype = source_dtype
|
|
if binding.transform is not None:
|
|
# Adapter transforms can map full checkpoint tensors to a
|
|
# rank-local runtime shape. Validate that contract without
|
|
# materializing the source tensor during preflight.
|
|
try:
|
|
transformed = binding.transform(
|
|
torch.empty(
|
|
source_shape,
|
|
dtype=source_dtype,
|
|
device="meta",
|
|
)
|
|
)
|
|
except (NotImplementedError, RuntimeError, TypeError, ValueError) as exc:
|
|
raise _PlanIncompatibleError(
|
|
f"checkpoint transform for {runtime_name!r} cannot be validated on tensor metadata: {exc}"
|
|
) from exc
|
|
if not isinstance(transformed, torch.Tensor):
|
|
raise _PlanIncompatibleError(
|
|
f"checkpoint transform for {runtime_name!r} returned {type(transformed).__name__}, "
|
|
"expected torch.Tensor"
|
|
)
|
|
runtime_shape = tuple(transformed.shape)
|
|
runtime_dtype = transformed.dtype
|
|
|
|
if runtime_shape != tuple(target.shape):
|
|
raise _PlanIncompatibleError(
|
|
f"shape mismatch for {runtime_name!r}: checkpoint={source_shape}, "
|
|
f"transformed={runtime_shape}, runtime={tuple(target.shape)}"
|
|
)
|
|
if runtime_dtype != target.dtype:
|
|
raise _PlanIncompatibleError(
|
|
f"dtype mismatch for {runtime_name!r}: checkpoint={source.get_dtype()}, "
|
|
f"transformed={runtime_dtype}, runtime={target.dtype}"
|
|
)
|
|
|
|
|
|
def build_checkpoint_mmap_plan(
|
|
pipeline: nn.Module,
|
|
*,
|
|
dit_modules: Sequence[tuple[str, nn.Module]],
|
|
sources: Sequence[object],
|
|
model_path: str | None,
|
|
tensor_parallel_size: int,
|
|
use_hsdp: bool,
|
|
online_quantization: bool,
|
|
has_distilled_lora: bool = False,
|
|
) -> HostWeightPlanResult:
|
|
"""Build a complete direct-checkpoint plan or return a fallback reason."""
|
|
if tensor_parallel_size != 1:
|
|
return HostWeightPlanResult(None, f"TP={tensor_parallel_size} requires the ordinary loader")
|
|
if use_hsdp:
|
|
return HostWeightPlanResult(None, "HSDP requires the ordinary loader")
|
|
if online_quantization:
|
|
return HostWeightPlanResult(None, "online quantization requires the ordinary loader")
|
|
if has_distilled_lora:
|
|
return HostWeightPlanResult(None, "distilled LoRA requires the ordinary loader")
|
|
|
|
remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None)
|
|
if not callable(remap_fn):
|
|
remap_fn = None
|
|
adapter = get_direct_mmap_adapter(pipeline)
|
|
|
|
try:
|
|
planned_source_prefixes = _planned_source_prefixes(sources, dit_modules)
|
|
planned_sources = tuple(
|
|
source for source in sources if getattr(source, "prefix", "") in planned_source_prefixes
|
|
)
|
|
model_to_ckpt = _build_source_map(
|
|
sources=planned_sources,
|
|
model_path=model_path,
|
|
remap_fn=remap_fn,
|
|
)
|
|
required = _collect_required_dit_tensors(dit_modules)
|
|
if not required:
|
|
raise _PlanIncompatibleError("no DiT parameters or persistent buffers were discovered")
|
|
|
|
missing = sorted(set(required) - set(model_to_ckpt))
|
|
if missing:
|
|
raise _PlanIncompatibleError(
|
|
f"{len(missing)} required DiT tensors have no checkpoint binding (first 5: {missing[:5]})"
|
|
)
|
|
|
|
bindings: dict[str, TensorBinding] = {}
|
|
for runtime_name, target in required.items():
|
|
policy = adapter.policy_for(runtime_name, target) if adapter is not None else None
|
|
has_custom_loader = callable(getattr(target, "weight_loader", None))
|
|
if has_custom_loader and not (policy and policy.allow_custom_loader):
|
|
raise _PlanIncompatibleError(
|
|
f"{runtime_name!r} requires a custom weight_loader with no direct-mmap adapter"
|
|
)
|
|
checkpoint_key, file_path = model_to_ckpt[runtime_name]
|
|
bindings[runtime_name] = TensorBinding(
|
|
checkpoint_key=checkpoint_key,
|
|
file_path=file_path,
|
|
transform=policy.transform if policy is not None else None,
|
|
)
|
|
|
|
_validate_source_metadata(required, bindings)
|
|
except (OSError, ValueError, _PlanIncompatibleError) as exc:
|
|
return HostWeightPlanResult(None, str(exc))
|
|
|
|
return HostWeightPlanResult(
|
|
HostWeightPlan(
|
|
backing_kind="checkpoint_mmap",
|
|
bindings=bindings,
|
|
planned_source_prefixes=planned_source_prefixes,
|
|
)
|
|
)
|
|
|
|
|
|
__all__ = [
|
|
"HostWeightPlan",
|
|
"HostWeightPlanResult",
|
|
"TensorBinding",
|
|
"build_checkpoint_mmap_plan",
|
|
"has_online_quantization",
|
|
]
|