59 lines
2.1 KiB
Python
59 lines
2.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.quantization import build_quant_config
|
|
from vllm_omni.quantization.component_config import ComponentQuantizationConfig
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"transformer_spec",
|
|
[
|
|
{
|
|
"method": "torchao",
|
|
"quant_type": {
|
|
"default": {
|
|
"_type": "Float8WeightOnlyConfig",
|
|
"_version": 2,
|
|
"_data": {
|
|
"weight_dtype": {
|
|
"_type": "torch.dtype",
|
|
"_data": "float8_e4m3fn",
|
|
},
|
|
"set_inductor_config": False,
|
|
},
|
|
}
|
|
},
|
|
},
|
|
{"method": "torchao_float8_weight_only"},
|
|
],
|
|
ids=["serialized-json", "checkpoint-shorthand"],
|
|
)
|
|
def test_build_quant_config_torchao_checkpoint(transformer_spec):
|
|
quantization = pytest.importorskip("torchao.quantization")
|
|
result = build_quant_config({"transformer": transformer_spec})
|
|
|
|
assert isinstance(result, ComponentQuantizationConfig)
|
|
transformer = result.resolve("transformer")
|
|
assert transformer.get_name() == "torchao"
|
|
assert transformer.is_checkpoint_torchao_serialized is True
|
|
assert isinstance(transformer.torchao_config, quantization.Float8WeightOnlyConfig)
|
|
assert transformer.torchao_config.version == 2
|
|
assert transformer.torchao_config.weight_dtype is torch.float8_e4m3fn
|
|
assert transformer.torchao_config.set_inductor_config is False
|
|
|
|
|
|
def test_build_quant_config_torchao_runtime():
|
|
from vllm.model_executor.layers.quantization.torchao import TorchAOConfig
|
|
|
|
torchao_config = object()
|
|
result = build_quant_config("torchao", torchao_config=torchao_config)
|
|
|
|
assert isinstance(result, TorchAOConfig)
|
|
assert result.torchao_config is torchao_config
|
|
assert result.is_checkpoint_torchao_serialized is False
|