350 lines
14 KiB
Python
350 lines
14 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Phase-I normalized FP8 cache support for DLO checkpoint mmap."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import shutil
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import torch
|
|
from safetensors import safe_open
|
|
from safetensors.torch import save_file
|
|
from torch import nn
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.model_loader.host_weight_plan import (
|
|
_build_source_map,
|
|
_collect_required_dit_tensors,
|
|
)
|
|
from vllm_omni.quantization.component_config import ComponentQuantizationConfig
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
CACHE_SCHEMA_VERSION = 1
|
|
CACHE_KIND = "normalized_fp8_runtime"
|
|
CACHE_MANIFEST_NAME = "manifest.json"
|
|
CACHE_COMPLETE_NAME = "COMPLETE"
|
|
CACHE_COMPONENT_SUBFOLDER = "transformer"
|
|
|
|
|
|
def _transformer_quant_config() -> ComponentQuantizationConfig:
|
|
"""Build the serialized FP8 config used by a normalized cache."""
|
|
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
|
|
|
|
return ComponentQuantizationConfig(
|
|
{
|
|
"transformer": Fp8Config(
|
|
is_checkpoint_fp8_serialized=True,
|
|
activation_scheme="dynamic",
|
|
)
|
|
}
|
|
)
|
|
|
|
|
|
def is_online_fp8_config(quant_config: Any) -> bool:
|
|
"""Return whether the transformer component requests online FP8."""
|
|
if quant_config is None:
|
|
return False
|
|
if hasattr(quant_config, "resolve"):
|
|
quant_config = quant_config.resolve("transformer")
|
|
get_name = getattr(quant_config, "get_name", None)
|
|
if not callable(get_name) or str(get_name()).lower() != "fp8":
|
|
return False
|
|
return not bool(getattr(quant_config, "is_checkpoint_fp8_serialized", False))
|
|
|
|
|
|
def _source_fingerprint(files: set[str]) -> str:
|
|
entries: list[dict[str, int | str]] = []
|
|
for filename in sorted(files):
|
|
stat = os.stat(filename)
|
|
entries.append(
|
|
{
|
|
"path": os.path.abspath(filename),
|
|
"size": stat.st_size,
|
|
"mtime_ns": stat.st_mtime_ns,
|
|
}
|
|
)
|
|
payload = json.dumps(entries, sort_keys=True, separators=(",", ":")).encode()
|
|
return hashlib.sha256(payload).hexdigest()
|
|
|
|
|
|
def load_fp8_cache_manifest(cache_dir: str | os.PathLike[str] | None) -> dict[str, Any] | None:
|
|
"""Return a validated Phase-I manifest, or ``None`` for a cache miss."""
|
|
if cache_dir is None:
|
|
return None
|
|
root = Path(cache_dir)
|
|
manifest_path = root / CACHE_MANIFEST_NAME
|
|
if not (root / CACHE_COMPLETE_NAME).is_file() or not manifest_path.is_file():
|
|
return None
|
|
try:
|
|
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
|
|
except (OSError, ValueError):
|
|
return None
|
|
if not isinstance(manifest, dict):
|
|
return None
|
|
if manifest.get("schema_version") != CACHE_SCHEMA_VERSION:
|
|
return None
|
|
if manifest.get("cache_kind") != CACHE_KIND:
|
|
return None
|
|
if manifest.get("component") != CACHE_COMPONENT_SUBFOLDER:
|
|
return None
|
|
quantization = manifest.get("quantization")
|
|
if not isinstance(quantization, dict) or quantization != {
|
|
"method": "fp8",
|
|
"serialized": True,
|
|
"activation_scheme": "dynamic",
|
|
}:
|
|
return None
|
|
if manifest.get("component_subfolder") != CACHE_COMPONENT_SUBFOLDER:
|
|
return None
|
|
if manifest.get("source_prefix") != f"{CACHE_COMPONENT_SUBFOLDER}.":
|
|
return None
|
|
|
|
index_path = root / CACHE_COMPONENT_SUBFOLDER / "model.safetensors.index.json"
|
|
if not index_path.is_file():
|
|
return None
|
|
try:
|
|
index = json.loads(index_path.read_text(encoding="utf-8"))
|
|
except (OSError, ValueError):
|
|
return None
|
|
if not isinstance(index, dict):
|
|
return None
|
|
weight_map = index.get("weight_map")
|
|
if not isinstance(weight_map, dict) or not weight_map or not all(isinstance(key, str) for key in weight_map):
|
|
return None
|
|
if manifest.get("tensor_count") != len(weight_map):
|
|
return None
|
|
for filename in weight_map.values():
|
|
if not isinstance(filename, str):
|
|
return None
|
|
relative_path = Path(filename)
|
|
if (
|
|
relative_path.is_absolute()
|
|
or relative_path.name != filename
|
|
or relative_path.suffix != ".safetensors"
|
|
or not (root / CACHE_COMPONENT_SUBFOLDER / relative_path).is_file()
|
|
):
|
|
return None
|
|
|
|
source_files = manifest.get("source_files")
|
|
source_fingerprint = manifest.get("source_fingerprint")
|
|
if not isinstance(source_files, list) or not source_files or not isinstance(source_fingerprint, str):
|
|
return None
|
|
try:
|
|
if _source_fingerprint({str(filename) for filename in source_files}) != source_fingerprint:
|
|
return None
|
|
except (OSError, TypeError):
|
|
return None
|
|
return manifest
|
|
|
|
|
|
def _cache_key_for_runtime_name(runtime_name: str, prefix: str) -> str:
|
|
if not runtime_name.startswith(prefix):
|
|
raise ValueError(f"runtime tensor {runtime_name!r} is outside cache source prefix {prefix!r}")
|
|
return runtime_name[len(prefix) :]
|
|
|
|
|
|
def build_normalized_fp8_cache(
|
|
model: nn.Module,
|
|
*,
|
|
sources: tuple[object, ...],
|
|
cache_dir: str | os.PathLike[str],
|
|
) -> dict[str, Any]:
|
|
"""Write a normalized FP8 cache from a completed online-FP8 model.
|
|
|
|
The builder writes tensors one at a time into individual safetensors shards,
|
|
so it does not create a second full-model state dict. The first cache build
|
|
still follows the existing online loader; warm launches consume the cache
|
|
through the direct DLO mmap path.
|
|
"""
|
|
cache_root = Path(cache_dir)
|
|
existing_manifest = load_fp8_cache_manifest(cache_root)
|
|
if existing_manifest is not None:
|
|
return existing_manifest
|
|
if cache_root.exists() and any(cache_root.iterdir()):
|
|
raise FileExistsError(f"FP8 cache directory is not empty and is not valid: {cache_root}")
|
|
cache_root.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
dit_names = tuple(getattr(model, "_dit_modules", ()))
|
|
dit_names = tuple(name for name in dit_names if hasattr(model, name))
|
|
if len(dit_names) != 1:
|
|
raise ValueError(f"Phase-I FP8 cache requires exactly one DiT module, got {dit_names}")
|
|
dit_name = dit_names[0]
|
|
if dit_name != CACHE_COMPONENT_SUBFOLDER:
|
|
raise ValueError(
|
|
f"Phase-I FP8 cache requires a top-level {CACHE_COMPONENT_SUBFOLDER!r} module, got {dit_name!r}"
|
|
)
|
|
source_prefix = f"{dit_name}."
|
|
dit_module = model.get_submodule(dit_name)
|
|
planned_sources = tuple(source for source in sources if getattr(source, "prefix", "") == source_prefix)
|
|
if len(planned_sources) != 1:
|
|
raise ValueError(f"Phase-I FP8 cache requires one source for {source_prefix!r}, got {len(planned_sources)}")
|
|
|
|
source_map = _build_source_map(
|
|
sources=planned_sources,
|
|
model_path=None,
|
|
remap_fn=getattr(type(model), "_remap_ckpt_key", None),
|
|
)
|
|
required = _collect_required_dit_tensors(((dit_name, dit_module),))
|
|
quantized_modules: dict[str, nn.Module] = {}
|
|
try:
|
|
from vllm.model_executor.layers.quantization.online.fp8 import Fp8PerTensorOnlineLinearMethod
|
|
except ImportError as exc: # pragma: no cover - vLLM 0.27 provides this class
|
|
raise RuntimeError("Phase-I FP8 cache requires vLLM's online FP8 implementation") from exc
|
|
|
|
for module_name, module in model.named_modules():
|
|
if not module_name.startswith(source_prefix):
|
|
continue
|
|
if isinstance(getattr(module, "quant_method", None), Fp8PerTensorOnlineLinearMethod):
|
|
if not hasattr(module, "weight_scale"):
|
|
raise ValueError(f"online FP8 module {module_name!r} has no generated weight_scale")
|
|
quantized_modules[module_name] = module
|
|
|
|
tmp_root = cache_root.parent / f".{cache_root.name}.tmp-{uuid.uuid4().hex}"
|
|
component_root = tmp_root / CACHE_COMPONENT_SUBFOLDER
|
|
component_root.mkdir(parents=True, exist_ok=False)
|
|
weight_map: dict[str, str] = {}
|
|
source_files: set[str] = set()
|
|
total_size = 0
|
|
part_index = 0
|
|
|
|
file_cache: dict[str, Any] = {}
|
|
|
|
def write_tensor(key: str, tensor: torch.Tensor) -> None:
|
|
nonlocal part_index, total_size
|
|
if key in weight_map:
|
|
raise ValueError(f"normalized FP8 cache has duplicate tensor key {key!r}")
|
|
tensor = tensor.detach().to("cpu").contiguous()
|
|
filename = f"part-{part_index:05d}.safetensors"
|
|
part_index += 1
|
|
save_file({key: tensor}, str(component_root / filename))
|
|
weight_map[key] = filename
|
|
total_size += tensor.numel() * tensor.element_size()
|
|
|
|
try:
|
|
for runtime_name, target in required.items():
|
|
binding_info = source_map.get(runtime_name)
|
|
module_name = runtime_name.removesuffix(".weight")
|
|
quant_module = quantized_modules.get(module_name) if runtime_name.endswith(".weight") else None
|
|
if runtime_name.endswith(".weight_scale"):
|
|
owner_name = runtime_name.removesuffix(".weight_scale")
|
|
if owner_name in quantized_modules:
|
|
# Generated online scales are published in the dedicated
|
|
# pass below; they have no raw checkpoint binding.
|
|
continue
|
|
if quant_module is not None:
|
|
if binding_info is None:
|
|
raise ValueError(f"no source binding for quantized weight {runtime_name!r}")
|
|
qweight = quant_module.weight.detach()
|
|
if qweight.ndim != 2:
|
|
raise ValueError(f"Phase-I FP8 cache only supports 2D weights, got {runtime_name!r}")
|
|
# Online FP8 stores the final runtime weight as qweight.T. The
|
|
# serialized FP8 method expects the pre-process qweight shape.
|
|
qweight = qweight.t().contiguous()
|
|
# Normalized caches contain the post-checkpoint-adapter QKV
|
|
# order. The cache-hit adapter bypasses the raw-checkpoint
|
|
# grouped-QKV transform, then the serialized FP8 finalizer
|
|
# performs only the standard weight transpose.
|
|
write_tensor(binding_info[0], qweight)
|
|
source_files.add(binding_info[1])
|
|
continue
|
|
|
|
if binding_info is not None:
|
|
checkpoint_key, file_path = binding_info
|
|
source_files.add(file_path)
|
|
handle = file_cache.get(file_path)
|
|
if handle is None:
|
|
handle = safe_open(file_path, framework="pt", device="cpu")
|
|
file_cache[file_path] = handle
|
|
write_tensor(checkpoint_key, handle.get_tensor(checkpoint_key))
|
|
else:
|
|
# Constructor-derived persistent buffers have no checkpoint
|
|
# entry; preserve them as canonical cache tensors.
|
|
write_tensor(_cache_key_for_runtime_name(runtime_name, source_prefix), target)
|
|
|
|
for module_name, module in quantized_modules.items():
|
|
scale = module.weight_scale.detach()
|
|
logical_widths = getattr(module, "logical_widths", None)
|
|
if scale.numel() == 1 and logical_widths is not None and len(logical_widths) > 1:
|
|
# Serialized Fp8LinearMethod represents a fused per-tensor
|
|
# weight with one scale slot per logical shard. Online FP8
|
|
# produces one scalar for the fused tensor, so replicate that
|
|
# scalar to the offline runtime shape without changing its
|
|
# value.
|
|
scale = scale.reshape(1).expand(len(logical_widths)).contiguous()
|
|
write_tensor(
|
|
_cache_key_for_runtime_name(f"{module_name}.weight_scale", source_prefix),
|
|
scale,
|
|
)
|
|
|
|
index = {
|
|
"metadata": {"total_size": total_size},
|
|
"weight_map": dict(sorted(weight_map.items())),
|
|
}
|
|
(component_root / "model.safetensors.index.json").write_text(
|
|
json.dumps(index, indent=2, sort_keys=True) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
manifest = {
|
|
"schema_version": CACHE_SCHEMA_VERSION,
|
|
"cache_kind": CACHE_KIND,
|
|
"component": "transformer",
|
|
"component_subfolder": CACHE_COMPONENT_SUBFOLDER,
|
|
"source_prefix": source_prefix,
|
|
"quantization": {
|
|
"method": "fp8",
|
|
"serialized": True,
|
|
"activation_scheme": "dynamic",
|
|
},
|
|
"tensor_count": len(weight_map),
|
|
"source_fingerprint": _source_fingerprint(source_files),
|
|
"source_files": sorted(source_files),
|
|
}
|
|
(tmp_root / CACHE_MANIFEST_NAME).write_text(
|
|
json.dumps(manifest, indent=2, sort_keys=True) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
(tmp_root / CACHE_COMPLETE_NAME).write_text("phase1\n", encoding="utf-8")
|
|
for handle in file_cache.values():
|
|
close = getattr(handle, "__exit__", None)
|
|
if callable(close):
|
|
close(None, None, None)
|
|
try:
|
|
os.replace(tmp_root, cache_root)
|
|
except FileExistsError:
|
|
# Multiple DLO workers may build the same cache concurrently. If
|
|
# another worker published a complete artifact first, treat this
|
|
# publisher as a successful cache race rather than a build error.
|
|
published = load_fp8_cache_manifest(cache_root)
|
|
if published is None:
|
|
raise
|
|
shutil.rmtree(tmp_root, ignore_errors=True)
|
|
logger.info("Normalized FP8 cache was published concurrently at %s", cache_root)
|
|
return published
|
|
logger.info("Published normalized FP8 cache at %s (%d tensors)", cache_root, len(weight_map))
|
|
return manifest
|
|
except BaseException:
|
|
for handle in file_cache.values():
|
|
close = getattr(handle, "__exit__", None)
|
|
if callable(close):
|
|
close(None, None, None)
|
|
shutil.rmtree(tmp_root, ignore_errors=True)
|
|
raise
|
|
|
|
|
|
__all__ = [
|
|
"CACHE_SCHEMA_VERSION",
|
|
"CACHE_KIND",
|
|
"CACHE_COMPONENT_SUBFOLDER",
|
|
"build_normalized_fp8_cache",
|
|
"is_online_fp8_config",
|
|
"load_fp8_cache_manifest",
|
|
"_transformer_quant_config",
|
|
]
|