Add normalized online FP8 cache for DLO
Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
This commit is contained in:
Hongsheng Liu committed
1 parent
14bd8f8772
commit
1ae56cfcfb
14 files changed
+714
-11
No files matched your search
@@ -0,0 +1,349 @@
|
||||
# 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",
|
||||
]
|
||||
Reference in new issue
Block a user