From 1ae56cfcfb0ba340bc1b6312b5b03281350a9684 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Wed, 19 Aug 2026 11:55:29 +0800 Subject: [PATCH] Add normalized online FP8 cache for DLO Signed-off-by: hsliuustc0106 --- .../distributed_layerwise_offload.md | 33 +- .../distributed_layerwise_offload.md | 32 +- .../diffusion/model_loader/test_fp8_cache.py | 153 ++++++++ .../test_async_omni_diffusion_config.py | 4 + vllm_omni/config/omni_config.py | 1 + vllm_omni/config/stage_config.py | 1 + vllm_omni/diffusion/data.py | 4 + .../checkpoint_adapters/direct_mmap.py | 2 + .../model_loader/diffusers_loader.py | 87 ++++- vllm_omni/diffusion/model_loader/fp8_cache.py | 349 ++++++++++++++++++ .../distributed_layerwise_backend.py | 49 +++ vllm_omni/engine/arg_utils.py | 1 + vllm_omni/engine/async_omni_engine.py | 1 + vllm_omni/entrypoints/cli/serve.py | 8 + 14 files changed, 714 insertions(+), 11 deletions(-) create mode 100644 tests/diffusion/model_loader/test_fp8_cache.py create mode 100644 vllm_omni/diffusion/model_loader/fp8_cache.py diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index e254e3a43..8ad5df0b2 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -105,6 +105,31 @@ flags such as `_supports_mmap_loading` or parameter attributes for mmap-only transforms. Model-specific direct-layout knowledge, when required, lives in a checkpoint adapter beside the ordinary loader. +### Phase-I normalized online-FP8 cache + +The direct-checkpoint plan above cannot quantize a BF16 source online. Phase I +adds an explicit, opt-in normalized cache for the narrower case of one TP1 +transformer with per-tensor online FP8 and DLO AllGather: + +```text +cache miss: BF16 source -> ordinary online FP8 loader -> atomic FP8 cache +cache hit: FP8 cache -> direct mmap plan -> DLO host shard -> AllGather +``` + +The cache stores the serialized runtime representation, including finalized FP8 +weights, generated scales, and persistent transformer buffers. Its manifest +records the schema, component, quantization contract, tensor count, source +files, and a source fingerprint. A missing, stale, incomplete, incompatible, or +invalid artifact is a cache miss; cache publication is best-effort and falls +back to the ordinary loader on failure. + +Phase I keeps the cache binding deliberately simple: one source per transformer +component, with existing per-tensor transforms applied before publication. The +cache is not selected for DLO-disabled, no-AllGather, TP>1, or HSDP deployments, +and it normalizes the transformer only. Cache sharing and compatibility for +other components, transformed TP coordinates, and broader DP/SP layouts remain +in [RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231). + ### AllGather path With the default `dlo_use_allgather=True`, each rank stores approximately @@ -211,10 +236,10 @@ reconstructs those tensors with their recorded layouts. Other online methods must use `--dlo-no-use-allgather` or disable online quantization until their runtime layouts are validated. -A normalized runtime mmap cache, built through the ordinary loader, is the -proposed general mechanism for sharing transformed TP or quantized layouts. -That cache and its publication/lifecycle protocol are intentionally outside -this phase; see [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195). +The Phase-I normalized online-FP8 cache is the supported TP1 transformer case. +The general cache-compatibility contract for additional components, transformed +TP coordinates, and cross-DP/SP reuse remains in +[RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231). ## Validation coverage diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md index 49139e695..f07ae594d 100644 --- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md +++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md @@ -42,6 +42,12 @@ vllm serve /path/to/model --omni \ vllm serve /path/to/model --omni \ --enable-distributed-layerwise-offload \ --usp 4 + +# Opt-in normalized online-FP8 cache (TP1 + DLO AllGather) +vllm serve /path/to/model --omni \ + --quantization fp8 \ + --enable-distributed-layerwise-offload \ + --dlo-fp8-cache-dir /path/to/node-local/fp8-cache ``` ```python @@ -63,6 +69,7 @@ omni = Omni( | `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` | | `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` | | `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` | +| `--dlo-fp8-cache-dir PATH` | Opt into the Phase-I normalized online-FP8 cache for the transformer | unset | ## Host-weight loading @@ -104,6 +111,24 @@ process. When the effective DLO group size is one, `dlo_use_allgather=True` does not perform a collective and uses the same rank-local transfer behavior. +## Normalized online-FP8 cache + +`--dlo-fp8-cache-dir` is an opt-in Phase-I cache for per-tensor online FP8 with +TP1 and DLO AllGather. On a cache miss, the ordinary loader quantizes the +transformer once and publishes finalized FP8 weights and scales as an atomic, +serialized safetensors artifact. On a valid cache hit, the loader selects that +serialized source before DLO planning, so the transformer can be loaded through +the direct checkpoint-mmap path without re-materializing the full BF16 source. + +The cache is keyed to the source files through the manifest's schema, component, +quantization contract, and source fingerprint. Incomplete, incompatible, or +stale artifacts are treated as cache misses. Cache creation is best-effort: a +publication failure logs a warning and keeps the ordinary online-FP8 path. + +Phase I normalizes the transformer only. Other pipeline components keep their +ordinary loading and quantization behavior. The cache is not selected when DLO +is disabled, AllGather is disabled, TP is greater than one, or HSDP is enabled. + ## Declarative topology Models may declare an `OffloadPlan` instead of embedding offload logic: @@ -147,10 +172,9 @@ must enter each collective. `OffloadPlan` that declares eligible `resident_dit_paths`. - DP concurrency requires an explicit, identical inference-step count. -Sharing transformed TP or quantized runtime layouts through a normalized mmap -cache is a follow-up design in -[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of -the direct-checkpoint path. +Sharing transformed TP or quantized runtime layouts beyond this TP1 transformer +cache remains a follow-up design in +[RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231). See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) for an end-to-end example. diff --git a/tests/diffusion/model_loader/test_fp8_cache.py b/tests/diffusion/model_loader/test_fp8_cache.py new file mode 100644 index 000000000..d08ce063f --- /dev/null +++ b/tests/diffusion/model_loader/test_fp8_cache.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for the Phase-I normalized online-FP8 cache contract.""" + +import json +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import load_file, save_file +from torch import nn + +from vllm_omni.diffusion.model_loader.fp8_cache import ( + CACHE_COMPLETE_NAME, + CACHE_COMPONENT_SUBFOLDER, + CACHE_MANIFEST_NAME, + _source_fingerprint, + _transformer_quant_config, + build_normalized_fp8_cache, + is_online_fp8_config, + load_fp8_cache_manifest, +) + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + + +def _write_valid_cache(tmp_path): + source = tmp_path / "source.safetensors" + save_file({"weight": torch.ones(2, 2)}, str(source)) + + cache = tmp_path / "cache" + component = cache / CACHE_COMPONENT_SUBFOLDER + component.mkdir(parents=True) + shard = component / "part-00000.safetensors" + save_file({"weight": torch.ones(2, 2, dtype=torch.float8_e4m3fn)}, str(shard)) + (component / "model.safetensors.index.json").write_text( + json.dumps( + { + "metadata": {"total_size": 4}, + "weight_map": {"weight": shard.name}, + } + ) + ) + + manifest = { + "schema_version": 1, + "cache_kind": "normalized_fp8_runtime", + "component": "transformer", + "component_subfolder": CACHE_COMPONENT_SUBFOLDER, + "source_prefix": "transformer.", + "quantization": {"method": "fp8", "serialized": True, "activation_scheme": "dynamic"}, + "tensor_count": 1, + "source_files": [str(source)], + "source_fingerprint": _source_fingerprint({str(source)}), + } + cache.mkdir(exist_ok=True) + (cache / CACHE_MANIFEST_NAME).write_text(json.dumps(manifest)) + (cache / CACHE_COMPLETE_NAME).write_text("phase1\n") + return cache, source, manifest + + +def test_online_and_serialized_fp8_configs_are_classified(): + from vllm.model_executor.layers.quantization.fp8 import Fp8Config + + assert is_online_fp8_config(Fp8Config()) + assert not is_online_fp8_config(Fp8Config(is_checkpoint_fp8_serialized=True)) + + cache_config = _transformer_quant_config() + assert not is_online_fp8_config(cache_config) + assert cache_config.resolve("transformer").is_checkpoint_fp8_serialized + assert cache_config.resolve("text_encoder") is None + + +def test_load_fp8_cache_manifest_validates_artifact_and_sources(tmp_path): + cache, source, expected = _write_valid_cache(tmp_path) + + assert load_fp8_cache_manifest(cache) == expected + + source.write_bytes(source.read_bytes() + b"changed") + assert load_fp8_cache_manifest(cache) is None + + +def test_load_fp8_cache_manifest_rejects_wrong_serialization_contract(tmp_path): + cache, _, manifest = _write_valid_cache(tmp_path) + + manifest["quantization"]["serialized"] = False + (cache / CACHE_MANIFEST_NAME).write_text(json.dumps(manifest)) + + assert load_fp8_cache_manifest(cache) is None + + +def test_load_fp8_cache_manifest_rejects_path_traversal(tmp_path): + cache, _, manifest = _write_valid_cache(tmp_path) + index_path = cache / CACHE_COMPONENT_SUBFOLDER / "model.safetensors.index.json" + index_path.write_text(json.dumps({"weight_map": {"weight": "../outside.safetensors"}})) + + assert load_fp8_cache_manifest(cache) is None + + +def test_load_fp8_cache_manifest_rejects_non_object_json(tmp_path): + cache, _, _ = _write_valid_cache(tmp_path) + (cache / CACHE_MANIFEST_NAME).write_text("[]") + + assert load_fp8_cache_manifest(cache) is None + + +def test_build_normalized_fp8_cache_serializes_weight_and_scale(tmp_path): + from vllm.model_executor.layers.quantization.online.fp8 import ( + Fp8PerTensorOnlineLinearMethod, + ) + + source_dir = tmp_path / "source" + source_dir.mkdir() + source_file = source_dir / "model.safetensors" + save_file({"block.weight": torch.ones((3, 2), dtype=torch.bfloat16)}, str(source_file)) + + model = nn.Module() + model._dit_modules = ("transformer",) + model.transformer = nn.Module() + model.transformer.block = nn.Module() + model.transformer.block.weight = nn.Parameter(torch.ones((2, 3), dtype=torch.float8_e4m3fn), requires_grad=False) + model.transformer.block.weight_scale = nn.Parameter(torch.tensor([0.5]), requires_grad=False) + model.transformer.block.quant_method = object.__new__(Fp8PerTensorOnlineLinearMethod) + + source = SimpleNamespace( + model_or_path=str(source_dir), + subfolder=None, + revision=None, + prefix="transformer.", + ) + cache_dir = tmp_path / "cache" + + manifest = build_normalized_fp8_cache( + model, + sources=(source,), + cache_dir=cache_dir, + ) + + assert manifest["cache_kind"] == "normalized_fp8_runtime" + tensors = load_file(str(cache_dir / CACHE_COMPONENT_SUBFOLDER / "part-00000.safetensors")) + scale_tensors = load_file(str(cache_dir / CACHE_COMPONENT_SUBFOLDER / "part-00001.safetensors")) + assert tensors["block.weight"].dtype == torch.float8_e4m3fn + assert tuple(tensors["block.weight"].shape) == (3, 2) + assert torch.equal(scale_tensors["block.weight_scale"], torch.tensor([0.5])) + + +@pytest.mark.parametrize("missing", [CACHE_COMPLETE_NAME, CACHE_MANIFEST_NAME]) +def test_load_fp8_cache_manifest_requires_completion_marker(tmp_path, missing): + cache, _, _ = _write_valid_cache(tmp_path) + (cache / missing).unlink() + + assert load_fp8_cache_manifest(cache) is None diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 1d5d58011..22ccb7b2f 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -347,6 +347,8 @@ def test_serve_cli_forwards_distributed_offload_residency(): "--dlo-no-use-allgather", "--dlo-resident-layers", "20", + "--dlo-fp8-cache-dir", + "/tmp/dlo-fp8-cache", ] ) @@ -357,9 +359,11 @@ def test_serve_cli_forwards_distributed_offload_residency(): assert args.enable_distributed_layerwise_offload is True assert args.dlo_use_allgather is False assert args.dlo_resident_layers == 20 + assert args.dlo_fp8_cache_dir == "/tmp/dlo-fp8-cache" assert engine_args["enable_distributed_layerwise_offload"] is True assert engine_args["dlo_use_allgather"] is False assert engine_args["dlo_resident_layers"] == 20 + assert engine_args["dlo_fp8_cache_dir"] == "/tmp/dlo-fp8-cache" def test_serve_cli_accepts_diffusion_compile_controls(): diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 8dd9c3db6..04641fc74 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -693,6 +693,7 @@ class _DiffusionConfigProjection: enable_distributed_layerwise_offload: bool = False dlo_use_allgather: bool = True dlo_resident_layers: int = Field(default=0, ge=0) + dlo_fp8_cache_dir: str | None = None pin_cpu_memory: bool = True diffusion_compile_granularity: Literal["regional", "full"] = "regional" diffusion_compile_dynamic: bool = Field(default=True, strict=True) diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index cdfd35ee2..c9bc4555b 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -383,6 +383,7 @@ class StageDeployConfig: enable_distributed_layerwise_offload: bool | None = None dlo_use_allgather: bool | None = None dlo_resident_layers: int | None = None + dlo_fp8_cache_dir: str | None = None # Diffusion-specific debug and observability knobs. enable_diffusion_pipeline_profiler: bool | None = None diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 2abd478c7..a6fe98ba2 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -752,6 +752,10 @@ class OmniDiffusionConfig: dlo_use_allgather: bool = True # Leading main-DiT blocks kept resident by distributed layerwise offload. dlo_resident_layers: int = 0 + # Optional Phase-I normalized online-FP8 cache. A cache miss builds the + # artifact after the ordinary online loader; a valid cache hit uses the + # serialized FP8 checkpoint-mmap path. + dlo_fp8_cache_dir: str | None = None pin_cpu_memory: bool = True # Use pinned memory for faster transfers when offloading diff --git a/vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py b/vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py index 86b23429b..b2aa2f9cd 100644 --- a/vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py +++ b/vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py @@ -41,6 +41,8 @@ class _MiniMaxH3DirectMmapAdapter: target: torch.Tensor, ) -> DirectMmapTensorPolicy: del target + if getattr(self.pipeline, "_dlo_fp8_cache_normalized", False): + return DirectMmapTensorPolicy(allow_custom_loader=True) transform = None suffix = ".qkv_proj.weight" if runtime_name.endswith(suffix): diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 62bc0525f..36fe60a83 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -8,7 +8,7 @@ import re import time from collections.abc import Generator, Iterable, Sequence from pathlib import Path -from typing import cast +from typing import Any, cast import torch from torch import nn @@ -34,6 +34,12 @@ from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp from vllm_omni.diffusion.model_loader.checkpoint_adapters import ( get_checkpoint_adapter, ) +from vllm_omni.diffusion.model_loader.fp8_cache import ( + _transformer_quant_config, + build_normalized_fp8_cache, + is_online_fp8_config, + load_fp8_cache_manifest, +) from vllm_omni.diffusion.model_loader.host_weight_plan import ( HostWeightPlan, build_checkpoint_mmap_plan, @@ -141,6 +147,10 @@ class DiffusersPipelineLoader: self.quant_config = od_config.quantization_config self.parallel_config = od_config.parallel_config self.host_weight_plan: HostWeightPlan | None = None + self._fp8_cache_dir: Path | None = None + self._fp8_cache_requested = False + self._fp8_cache_hit = False + self._fp8_cache_manifest: dict[str, Any] | None = None def take_host_weight_plan(self) -> HostWeightPlan | None: """Transfer the loader-produced plan to the offload backend.""" @@ -367,12 +377,61 @@ class DiffusersPipelineLoader: ) def _get_weight_sources(self, model: nn.Module) -> tuple["ComponentSource", ...]: - return tuple( + sources = tuple( cast( Iterable[DiffusersPipelineLoader.ComponentSource], getattr(model, "weights_sources", ()), ) ) + if not self._fp8_cache_hit or self._fp8_cache_dir is None or self._fp8_cache_manifest is None: + return sources + + cache_prefix = str(self._fp8_cache_manifest["source_prefix"]) + cache_subfolder = str(self._fp8_cache_manifest.get("component_subfolder", "transformer")) + return tuple( + dataclasses.replace( + source, + model_or_path=str(self._fp8_cache_dir), + subfolder=cache_subfolder, + revision=None, + ) + if source.prefix == cache_prefix + else source + for source in sources + ) + + def _prepare_fp8_cache(self) -> None: + """Select the Phase-I cache path before constructing the pipeline.""" + cache_dir = getattr(self.od_config, "dlo_fp8_cache_dir", None) + if not cache_dir or not is_online_fp8_config(self.quant_config): + return + if not getattr(self.od_config, "enable_distributed_layerwise_offload", False): + logger.warning("Ignoring --dlo-fp8-cache-dir because distributed layerwise offload is disabled") + return + if not getattr(self.od_config, "dlo_use_allgather", True): + logger.warning("Ignoring --dlo-fp8-cache-dir because DLO AllGather is disabled") + return + if int(getattr(self.parallel_config, "tensor_parallel_size", 1)) != 1: + logger.warning("Ignoring --dlo-fp8-cache-dir because Phase-I cache requires TP=1") + return + if bool(getattr(self.parallel_config, "use_hsdp", False)): + logger.warning("Ignoring --dlo-fp8-cache-dir because Phase-I cache does not support HSDP") + return + + self._fp8_cache_dir = Path(cache_dir) + self._fp8_cache_requested = True + self._fp8_cache_manifest = load_fp8_cache_manifest(self._fp8_cache_dir) + if self._fp8_cache_manifest is None: + logger.info("DLO normalized FP8 cache miss: %s", self._fp8_cache_dir) + return + + self._fp8_cache_hit = True + self.quant_config = _transformer_quant_config() + self.od_config.quantization_config = self.quant_config + logger.info( + "DLO normalized FP8 cache hit: %s; selecting serialized FP8 checkpoint path", + self._fp8_cache_dir, + ) def _get_expected_parameter_names(self, model: nn.Module) -> set[str]: """Return parameter names that should be covered by strict load checks.""" @@ -403,6 +462,7 @@ class DiffusersPipelineLoader: self.host_weight_plan = None if load_format is None: load_format = "default" + self._prepare_fp8_cache() # CPU offload + quantization: for offline-quantized models (e.g., AutoRound MXFP8), # weights are already quantized in the checkpoint — load directly on CPU. # For online quantization, load on device so quantization can run on accelerator, @@ -413,7 +473,9 @@ class DiffusersPipelineLoader: is_offline = getattr(quant_cfg, "data_type", None) == "mx_fp" or getattr( quant_cfg, "is_checkpoint_quantized", False ) - if not is_offline: + if self._fp8_cache_hit: + logger.info("Normalized FP8 cache is already serialized; keeping model initialization on CPU") + elif not is_offline: load_device = device.type offload_after_quant = True logger.info( @@ -431,6 +493,11 @@ class DiffusersPipelineLoader: ) else: model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False) + if self._fp8_cache_hit: + # The direct-mmap adapter must see this before plan + # construction so normalized QKV tensors bypass the raw + # checkpoint grouped-QKV transform. + setattr(model, "_dlo_fp8_cache_normalized", True) _dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False) _use_ag = getattr(self.od_config, "dlo_use_allgather", True) @@ -520,6 +587,20 @@ class DiffusersPipelineLoader: model.to("cpu") logger.info("Quantization complete, offloaded model back to CPU") + if self._fp8_cache_requested and self._fp8_cache_dir is not None and not self._fp8_cache_hit: + try: + build_normalized_fp8_cache( + model, + sources=self._get_weight_sources(model), + cache_dir=self._fp8_cache_dir, + ) + except Exception as exc: + logger.warning( + "DLO normalized FP8 cache build failed; retaining ordinary-loader path: %s", + exc, + exc_info=True, + ) + self._apply_skip_softmax_calibration(model) return model.eval() diff --git a/vllm_omni/diffusion/model_loader/fp8_cache.py b/vllm_omni/diffusion/model_loader/fp8_cache.py new file mode 100644 index 000000000..53ec3495d --- /dev/null +++ b/vllm_omni/diffusion/model_loader/fp8_cache.py @@ -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", +] diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index d184e71ea..c53f8962f 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -1034,6 +1034,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): logger.info("Realized %d loader-planned tensors as mmap views", len(loaded_names)) + if getattr(pipeline, "_dlo_fp8_cache_normalized", False): + self._process_normalized_fp8_cache(pipeline) + # Keep file handles open — _shard_and_pin will read from the mmap views. # They will be released after _shard_and_pin completes (when params are # replaced with offload placeholders). @@ -1103,6 +1106,52 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): f"(first 5: {remaining_meta[:5]})." ) + @staticmethod + def _process_normalized_fp8_cache(pipeline: nn.Module) -> None: + """Finalize serialized FP8 parameters after mmap replacement. + + Phase-I cache files use the serialized FP8 checkpoint layout: the + weight is stored before the quantization method's final transpose, and + ``weight_scale`` is stored as an explicit cache tensor. The ordinary + loader normally runs this method before DLO owns the plan; mmap loading + must run the same finalizer after replacing meta parameters. + """ + from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod + from vllm.model_executor.utils import replace_parameter + + for module in pipeline.modules(): + quant_method = getattr(module, "quant_method", None) + if not isinstance(quant_method, Fp8LinearMethod): + continue + if getattr(module, "_already_called_process_weights_after_loading", False): + continue + + # ``Fp8LinearMethod.process_weights_after_loading`` also handles + # non-serialized checkpoints. Its common path calls + # ``process_fp8_weight_tensor_strategy`` to requantize fused + # tensors, which is a CUDA-only operation. A normalized cache + # already contains the serialized FP8 values and the generated + # scales, so only perform the layout transition and the kernel's + # own preprocessing here. This keeps the warm path CPU/mmap + # compatible and preserves the cache bytes exactly. + if quant_method.block_quant: + raise ValueError("Phase-I normalized FP8 cache does not support block-quantized weights") + replace_parameter(module, "weight", module.weight.t().data) + # The online quantizer records one scalar. A serialized FP8 + # checkpoint may expose one repeated slot per logical shard in a + # fused module, but those slots must agree before they can be + # collapsed to the canonical per-tensor ``[1, 1]`` form. + weight_scale = module.weight_scale.reshape(-1) + if weight_scale.numel() > 1 and not torch.all(weight_scale == weight_scale[0]): + raise ValueError("Phase-I normalized FP8 cache contains non-uniform fused weight scales") + weight_scale = weight_scale[:1].reshape(1, 1) + replace_parameter(module, "weight_scale", weight_scale.data) + module.input_scale = None + if quant_method.use_marlin and hasattr(quant_method.fp8_linear, "marlin_input_dtype"): + quant_method.fp8_linear.marlin_input_dtype = quant_method.marlin_input_dtype + quant_method.fp8_linear.process_weights_after_loading(module) + module._already_called_process_weights_after_loading = True + def _init_dp_group(self) -> None: """Reuse the process group initialized by parallel_state. diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index fcd11b173..868e69939 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -546,6 +546,7 @@ class OrchestratorArgs: enable_distributed_layerwise_offload: bool = False dlo_use_allgather: bool = True dlo_resident_layers: int = 0 + dlo_fp8_cache_dir: str | None = None boundary_ratio: float | None = None flow_shift: float | None = None diffusion_kv_cache_dtype: str | None = None diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index d30db03cb..ca7a8b314 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1046,6 +1046,7 @@ class AsyncOmniEngine: "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False), "dlo_use_allgather": kwargs.get("dlo_use_allgather", True), "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0), + "dlo_fp8_cache_dir": kwargs.get("dlo_fp8_cache_dir", None), "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"), "diffusion_compile_granularity": ( "regional" diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 62ee10d94..7c04692d7 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -729,6 +729,14 @@ class OmniServeCommand(CLISubcommand): help="Keep this many leading main-DiT blocks resident on the device " "while distributed layerwise offload streams the remaining blocks.", ) + omni_config_group.add_argument( + "--dlo-fp8-cache-dir", + type=str, + default=None, + help="Optional Phase-I normalized online-FP8 cache directory. " + "A cache miss builds the cache after the ordinary loader; a valid " + "cache hit uses serialized FP8 checkpoint mmap with DLO.", + ) # Video model parameters (e.g., Wan2.2) - engine-level omni_config_group.add_argument( "--boundary-ratio",