Files
vllm-omni/tests/quantization/test_quant_factory.py
T
admin 321e8bb47f fix(quant): load offline-quantized compressed-tensors checkpoints directly on CPU under offload
The diffusion loader treated any quant config without the MXFP8/
is_checkpoint_quantized markers as online quantization: it loaded every
component on the accelerator before offloading, so serving the int8
weight-only compressed-tensors artifact OOM'd at model construction
(encoder 17 GiB + VAE + packed DiT skeleton > 24 GiB). compressed-tensors
is always a serialized format with no on-load quantization pass — detect
it (through component wrappers too) and load on the host, round-tripping
each quantized module through the accelerator only for the post-load
Marlin repack, which has no CPU path.
2026-10-08 07:36:01 -04:00

94 lines
3.7 KiB
Python

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for the quantization config factory's method-name and disk-dict resolution."""
from __future__ import annotations
import pytest
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import (
CompressedTensorsConfig,
)
from vllm_omni.quantization import build_quant_config, resolve_component_quant_config
from vllm_omni.quantization.component_config import (
ComponentQuantizationConfig,
is_prequantized_config,
resolve_encoder_quant_config,
)
from vllm_omni.quantization.factory import SUPPORTED_QUANTIZATION_METHODS, _build_single
def _int8_weight_only_disk_config() -> dict:
"""Serialized compressed-tensors int8 weight-only metadata as it appears on disk."""
weights = {
"num_bits": 8,
"type": "int",
"symmetric": True,
"strategy": "channel",
"dynamic": False,
"observer": "minmax",
}
return {
"quant_method": "compressed-tensors",
"format": "pack-quantized",
"config_groups": {"group_0": {"targets": ["Linear"], "weights": weights}},
"ignore": ["img_in"],
"quantization_status": "frozen",
}
def test_hyphenated_registry_method_builds_from_disk_schema() -> None:
# A bare str spec cannot construct compressed-tensors (a serialized-format
# method with no parameterless default); the serialized schema must build.
with pytest.raises(TypeError, match="Cannot instantiate CompressedTensorsConfig"):
_build_single("compressed-tensors")
disk = _int8_weight_only_disk_config()
disk.pop("quant_method")
config = _build_single("compressed-tensors", **disk)
assert isinstance(config, CompressedTensorsConfig)
@pytest.mark.parametrize("spelling", ["compressed-tensors", "compressed_tensors"])
def test_disk_dict_builds_via_from_config(spelling: str) -> None:
disk = _int8_weight_only_disk_config()
disk["quant_method"] = spelling
config = build_quant_config(disk)
assert isinstance(config, CompressedTensorsConfig)
def test_per_component_dict_routes_transformer_only() -> None:
disk = {"transformer": _int8_weight_only_disk_config()}
config = build_quant_config(disk)
assert isinstance(config, ComponentQuantizationConfig)
assert isinstance(
resolve_component_quant_config(config, "transformer"), CompressedTensorsConfig
)
encoder = resolve_encoder_quant_config(
resolve_component_quant_config(config, "text_encoder")
)
assert encoder is None
def test_unknown_method_still_raises() -> None:
with pytest.raises(ValueError, match="Unknown quantization method"):
_build_single("not-a-real-method")
def test_supported_list_advertises_hyphenated_registry_names() -> None:
assert "compressed-tensors" in SUPPORTED_QUANTIZATION_METHODS
def test_is_prequantized_config_recognizes_serialized_formats() -> None:
# compressed-tensors is always a serialized format: its linear methods
# consume packed checkpoint tensors and there is no on-load quantization.
assert is_prequantized_config(build_quant_config(_int8_weight_only_disk_config()))
component = build_quant_config({"transformer": _int8_weight_only_disk_config()})
assert isinstance(component, ComponentQuantizationConfig)
assert is_prequantized_config(component)
assert not is_prequantized_config(None)
# Online (non-serialized) methods stay unrecognized so the loader keeps
# the accelerator-side load-quantize-offload path for them.
from vllm_omni.quantization.fp8_config import DiffusionFp8Config
assert not is_prequantized_config(DiffusionFp8Config())