154 lines
5.4 KiB
Python
154 lines
5.4 KiB
Python
# 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
|