wehub-resource-sync
153 lines
5.2 KiB
Python
153 lines
5.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Focused tests for Diffusers backend quantization conversion helpers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from vllm_omni.diffusion.models.diffusers_adapter import quantization_utils
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
|
|
|
|
|
class _FakePipelineQuantizationConfig:
|
|
def __init__(self, *, quant_mapping=None, quant_backend=None, quant_kwargs=None, components_to_quantize=None):
|
|
self.quant_mapping = quant_mapping
|
|
self.quant_backend = quant_backend
|
|
self.quant_kwargs = quant_kwargs
|
|
self.components_to_quantize = components_to_quantize
|
|
|
|
|
|
class _FakeTorchAoConfig:
|
|
def __init__(self, quant_type):
|
|
self.quant_type = quant_type
|
|
|
|
@classmethod
|
|
def from_dict(cls, config):
|
|
return cls(quant_type=config["quant_type"])
|
|
|
|
|
|
class _FakeFloat8DynamicActivationFloat8WeightConfig:
|
|
pass
|
|
|
|
|
|
class _FakeInt8DynamicActivationInt8WeightConfig:
|
|
pass
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def patch_quantization_backends(monkeypatch):
|
|
monkeypatch.setattr(
|
|
quantization_utils,
|
|
"PipelineQuantizationConfig",
|
|
_FakePipelineQuantizationConfig,
|
|
)
|
|
monkeypatch.setattr(quantization_utils, "TorchAoConfig", _FakeTorchAoConfig)
|
|
|
|
def fake_get_torchao_quant_type_cls(class_name: str):
|
|
return {
|
|
"Float8DynamicActivationFloat8WeightConfig": _FakeFloat8DynamicActivationFloat8WeightConfig,
|
|
"Int8DynamicActivationInt8WeightConfig": _FakeInt8DynamicActivationInt8WeightConfig,
|
|
}[class_name]
|
|
|
|
monkeypatch.setattr(quantization_utils, "_get_torchao_quant_type_cls", fake_get_torchao_quant_type_cls)
|
|
monkeypatch.setattr(quantization_utils, "_get_diffusers_quantization_config_cls", lambda _: _FakeTorchAoConfig)
|
|
|
|
|
|
def _quant_config(method: str, **kwargs):
|
|
return SimpleNamespace(get_name=lambda: method, **kwargs)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("method", "expected_type"),
|
|
[
|
|
("fp8", _FakeFloat8DynamicActivationFloat8WeightConfig),
|
|
("int8", _FakeInt8DynamicActivationInt8WeightConfig),
|
|
],
|
|
)
|
|
def test_apply_injects_converted_quantization_config(method, expected_type):
|
|
od_config = SimpleNamespace(quantization_config=_quant_config(method))
|
|
load_kwargs = {}
|
|
|
|
injected = quantization_utils.apply_diffusers_quantization_config(
|
|
od_config,
|
|
load_kwargs,
|
|
{"transformer": ["diffusers", "Transformer2DModel"]},
|
|
)
|
|
|
|
assert injected is True
|
|
torchao_config = load_kwargs["quantization_config"].quant_mapping["transformer"]
|
|
assert isinstance(torchao_config.quant_type, expected_type)
|
|
|
|
|
|
def test_apply_injects_converted_quantization_config_for_transformer_2():
|
|
od_config = SimpleNamespace(quantization_config=_quant_config("int8"))
|
|
load_kwargs = {}
|
|
|
|
quantization_utils.apply_diffusers_quantization_config(
|
|
od_config,
|
|
load_kwargs,
|
|
{
|
|
"transformer": ["diffusers", "Transformer2DModel"],
|
|
"transformer_2": ["diffusers", "Transformer2DModel"],
|
|
},
|
|
)
|
|
|
|
assert sorted(load_kwargs["quantization_config"].quant_mapping) == ["transformer", "transformer_2"]
|
|
|
|
|
|
@pytest.mark.parametrize("method", ["gguf", "modelopt", "mxfp4", "mxfp8", "inc"])
|
|
def test_unsupported_methods_fail_explicitly(method):
|
|
with pytest.raises(NotImplementedError, match=method):
|
|
quantization_utils.ensure_supported_diffusers_quantization(_quant_config(method))
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"quant_config",
|
|
[
|
|
_quant_config("fp8", activation_scheme="static"),
|
|
_quant_config("int8", is_checkpoint_int8_serialized=True),
|
|
_quant_config("fp8", weight_block_size=[128, 128]),
|
|
_quant_config("int8", ignored_layers=["transformer.proj_out"]),
|
|
],
|
|
)
|
|
def test_ambiguous_mappings_fail_explicitly(quant_config):
|
|
with pytest.raises(NotImplementedError):
|
|
quantization_utils.ensure_supported_diffusers_quantization(quant_config)
|
|
|
|
|
|
def test_apply_preserves_diffusers_load_kwargs_quantization_config(mocker):
|
|
od_config = SimpleNamespace(quantization_config=_quant_config("fp8"))
|
|
existing = object()
|
|
load_kwargs = {"quantization_config": existing}
|
|
mock_warning = mocker.patch.object(quantization_utils.logger, "warning")
|
|
|
|
injected = quantization_utils.apply_diffusers_quantization_config(od_config, load_kwargs, {})
|
|
|
|
assert injected is False
|
|
assert load_kwargs["quantization_config"] is existing
|
|
assert "Using the Diffusers-native quantization_config" in mock_warning.call_args.args[0]
|
|
|
|
|
|
def test_convert_diffusers_quantization_config_from_dict():
|
|
load_kwargs = {
|
|
"quantization_config": {
|
|
"quant_mapping": {
|
|
"transformer": {
|
|
"quant_method": "torchao",
|
|
"quant_type": "fake-int8-config",
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
quantization_utils.convert_diffusers_quantization_config(load_kwargs)
|
|
|
|
quant_config = load_kwargs["quantization_config"]
|
|
assert isinstance(quant_config, _FakePipelineQuantizationConfig)
|
|
assert isinstance(quant_config.quant_mapping["transformer"], _FakeTorchAoConfig)
|
|
assert quant_config.quant_mapping["transformer"].quant_type == "fake-int8-config"
|