[Core] Unified quantization framework (#1764)
This commit is contained in:
SYLAR authored and GitHub committed
1 parent
2a630313ce
commit
6b93459826
40 files changed
+3306
-605
No files matched your search
@@ -10,11 +10,11 @@ from pytest_mock import MockerFixture
|
||||
from torch.nn import Module, Parameter
|
||||
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
|
||||
|
||||
from vllm_omni.diffusion.quantization import (
|
||||
get_diffusion_quant_config,
|
||||
get_vllm_quant_config_for_layers,
|
||||
)
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.factory import SUPPORTED_QUANTIZATION_METHODS
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
||||
|
||||
npu_available = pytest.mark.skipif(not current_omni_platform.is_npu(), reason="NPU platform not available.")
|
||||
|
||||
@@ -23,36 +23,26 @@ cuda_available = pytest.mark.skipif(not current_omni_platform.is_cuda(), reason=
|
||||
|
||||
def test_int8_config_creation():
|
||||
"""Test that Int8 config can be created."""
|
||||
config = get_diffusion_quant_config("int8")
|
||||
config = build_quant_config("int8")
|
||||
assert config is not None
|
||||
assert config.get_name() == "int8"
|
||||
|
||||
|
||||
def test_vllm_config_extraction():
|
||||
"""Test that vLLM config can be extracted from diffusion config."""
|
||||
diff_config = get_diffusion_quant_config("int8")
|
||||
vllm_config = get_vllm_quant_config_for_layers(diff_config)
|
||||
assert vllm_config is not None
|
||||
assert vllm_config.activation_scheme == "dynamic"
|
||||
|
||||
|
||||
def test_none_quantization():
|
||||
"""Test that None quantization returns None config."""
|
||||
config = get_diffusion_quant_config(None)
|
||||
config = build_quant_config(None)
|
||||
assert config is None
|
||||
vllm_config = get_vllm_quant_config_for_layers(config)
|
||||
assert vllm_config is None
|
||||
|
||||
|
||||
def test_invalid_quantization():
|
||||
"""Test that invalid quantization method raises error."""
|
||||
with pytest.raises(ValueError, match="Unknown quantization method"):
|
||||
get_diffusion_quant_config("invalid_method")
|
||||
build_quant_config("invalid_method")
|
||||
|
||||
|
||||
def test_int8_config_with_custom_params():
|
||||
"""Test Int8 config with custom parameters."""
|
||||
config = get_diffusion_quant_config(
|
||||
config = build_quant_config(
|
||||
"int8",
|
||||
activation_scheme="dynamic",
|
||||
ignored_layers=["proj_out"],
|
||||
@@ -64,8 +54,6 @@ def test_int8_config_with_custom_params():
|
||||
|
||||
def test_supported_methods():
|
||||
"""Test that supported methods list is correct."""
|
||||
from vllm_omni.diffusion.quantization import SUPPORTED_QUANTIZATION_METHODS
|
||||
|
||||
assert "int8" in SUPPORTED_QUANTIZATION_METHODS
|
||||
|
||||
|
||||
@@ -73,8 +61,8 @@ def test_quantization_integration():
|
||||
"""Test end-to-end quantization flow through OmniDiffusionConfig."""
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
|
||||
# Test with quantization string only
|
||||
config = OmniDiffusionConfig(model="test", quantization="int8")
|
||||
# Test with quantization_config string
|
||||
config = OmniDiffusionConfig(model="test", quantization_config="int8")
|
||||
assert config.quantization_config is not None
|
||||
assert config.quantization_config.get_name() == "int8"
|
||||
|
||||
@@ -87,10 +75,6 @@ def test_quantization_integration():
|
||||
assert config2.quantization_config.get_name() == "int8"
|
||||
assert config2.quantization_config.activation_scheme == "dynamic"
|
||||
|
||||
# Test that vLLM config can be extracted
|
||||
vllm_config = config.quantization_config.get_vllm_quant_config()
|
||||
assert vllm_config is not None
|
||||
|
||||
|
||||
def test_quantization_dict_not_mutated():
|
||||
"""Test that passing a dict to quantization_config doesn't mutate it."""
|
||||
@@ -105,28 +89,24 @@ def test_quantization_dict_not_mutated():
|
||||
assert original_dict == dict_copy
|
||||
|
||||
|
||||
def test_quantization_conflicting_methods_warning(caplog):
|
||||
"""Test warning when quantization and quantization_config['method'] conflict."""
|
||||
import logging
|
||||
|
||||
def test_quantization_config_string_and_dict_equivalent():
|
||||
"""Test that string and dict forms produce equivalent configs."""
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
config = OmniDiffusionConfig(
|
||||
model="test",
|
||||
quantization="int8", # This should be overridden
|
||||
quantization_config={"method": "int8", "activation_scheme": "dynamic"},
|
||||
)
|
||||
# No warning when methods match
|
||||
assert config.quantization_config is not None
|
||||
config_str = OmniDiffusionConfig(model="test", quantization_config="int8")
|
||||
config_dict = OmniDiffusionConfig(
|
||||
model="test",
|
||||
quantization_config={"method": "int8", "activation_scheme": "dynamic"},
|
||||
)
|
||||
assert config_str.quantization_config.get_name() == config_dict.quantization_config.get_name()
|
||||
assert config_str.quantization_config.activation_scheme == config_dict.quantization_config.activation_scheme
|
||||
|
||||
|
||||
def test_get_quant_method(mocker: MockerFixture):
|
||||
"""Test for get_quant_method method for GPU"""
|
||||
from vllm_omni.diffusion.quantization.int8 import Int8OnlineLinearMethod
|
||||
from vllm_omni.quantization.int8_config import Int8OnlineLinearMethod
|
||||
|
||||
diff_config = get_diffusion_quant_config("int8")
|
||||
vllm_config = get_vllm_quant_config_for_layers(diff_config)
|
||||
config = build_quant_config("int8")
|
||||
|
||||
def _fake_init(self, quant_config):
|
||||
pass
|
||||
@@ -141,21 +121,20 @@ def test_get_quant_method(mocker: MockerFixture):
|
||||
patch("vllm_omni.platforms.current_omni_platform.is_cuda", return_value=True),
|
||||
patch("vllm_omni.platforms.current_omni_platform.is_npu", return_value=False),
|
||||
):
|
||||
method = vllm_config.get_quant_method(layer, prefix)
|
||||
method = config.get_quant_method(layer, prefix)
|
||||
assert isinstance(method, Int8OnlineLinearMethod)
|
||||
|
||||
# Test skipping quantization for a layer
|
||||
vllm_config.ignored_layers = [prefix]
|
||||
method = vllm_config.get_quant_method(layer, prefix)
|
||||
config.ignored_layers = [prefix]
|
||||
method = config.get_quant_method(layer, prefix)
|
||||
assert isinstance(method, UnquantizedLinearMethod)
|
||||
|
||||
|
||||
def test_get_npu_quant_method():
|
||||
"""Test for get_quant_method method for NPU"""
|
||||
from vllm_omni.diffusion.quantization.int8 import NPUInt8OnlineLinearMethod
|
||||
from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod
|
||||
|
||||
diff_config = get_diffusion_quant_config("int8")
|
||||
vllm_config = get_vllm_quant_config_for_layers(diff_config)
|
||||
config = build_quant_config("int8")
|
||||
|
||||
layer = MagicMock(spec=LinearBase)
|
||||
prefix = "test_layer"
|
||||
@@ -165,12 +144,12 @@ def test_get_npu_quant_method():
|
||||
patch("vllm_omni.platforms.current_omni_platform.is_cuda", return_value=False),
|
||||
patch("vllm_omni.platforms.current_omni_platform.is_npu", return_value=True),
|
||||
):
|
||||
method = vllm_config.get_quant_method(layer, prefix)
|
||||
method = config.get_quant_method(layer, prefix)
|
||||
assert isinstance(method, NPUInt8OnlineLinearMethod)
|
||||
|
||||
# Test skipping quantization for a layer
|
||||
vllm_config.ignored_layers = [prefix]
|
||||
method = vllm_config.get_quant_method(layer, prefix)
|
||||
config.ignored_layers = [prefix]
|
||||
method = config.get_quant_method(layer, prefix)
|
||||
assert isinstance(method, UnquantizedLinearMethod)
|
||||
|
||||
|
||||
@@ -189,12 +168,12 @@ class TestInt8LinearMethod:
|
||||
@pytest.fixture
|
||||
def patch_deps(self, mocker, mock_kernel):
|
||||
# mock init_int8_linear_kernel
|
||||
mocker.patch("vllm_omni.diffusion.quantization.int8.init_int8_linear_kernel", return_value=mock_kernel)
|
||||
mocker.patch("vllm_omni.quantization.int8_config.init_int8_linear_kernel", return_value=mock_kernel)
|
||||
return mock_kernel
|
||||
|
||||
def test_init(self, patch_deps, mock_quant_config):
|
||||
# test for Int8LinearMethod init
|
||||
from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod, init_int8_linear_kernel
|
||||
from vllm_omni.quantization.int8_config import Int8LinearMethod, init_int8_linear_kernel
|
||||
|
||||
method = Int8LinearMethod(mock_quant_config)
|
||||
|
||||
@@ -205,7 +184,7 @@ class TestInt8LinearMethod:
|
||||
assert method.int8_linear == patch_deps
|
||||
|
||||
def test_process_weights_after_loading(self, patch_deps, mock_quant_config):
|
||||
from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod
|
||||
from vllm_omni.quantization.int8_config import Int8LinearMethod
|
||||
|
||||
method = Int8LinearMethod(mock_quant_config)
|
||||
layer = Module()
|
||||
@@ -214,7 +193,7 @@ class TestInt8LinearMethod:
|
||||
patch_deps.process_weights_after_loading.assert_called_once_with(layer)
|
||||
|
||||
def test_apply(self, patch_deps, mock_quant_config):
|
||||
from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod
|
||||
from vllm_omni.quantization.int8_config import Int8LinearMethod
|
||||
|
||||
method = Int8LinearMethod(mock_quant_config)
|
||||
layer = Module()
|
||||
@@ -237,19 +216,19 @@ class TestInt8OnlineLinearMethod:
|
||||
# mock kernel
|
||||
mock_kernel = mocker.Mock()
|
||||
mock_kernel.layer_param_names = ("weight", "weight_scale", "input_scale", "input_zero_point", "azp_adj")
|
||||
mocker.patch("vllm_omni.diffusion.quantization.int8.init_int8_linear_kernel", return_value=mock_kernel)
|
||||
mocker.patch("vllm_omni.diffusion.quantization.int8.replace_parameter")
|
||||
mocker.patch("vllm_omni.quantization.int8_config.init_int8_linear_kernel", return_value=mock_kernel)
|
||||
mocker.patch("vllm_omni.quantization.int8_config.replace_parameter")
|
||||
|
||||
# mock scaled_int8_quant return value
|
||||
mock_qweight = torch.ones((128, 64), dtype=torch.int8)
|
||||
mock_scale = torch.randn(128)
|
||||
mock_quant = mocker.patch(
|
||||
"vllm_omni.diffusion.quantization.int8.ops.scaled_int8_quant", return_value=(mock_qweight, mock_scale, None)
|
||||
"vllm_omni.quantization.int8_config.ops.scaled_int8_quant", return_value=(mock_qweight, mock_scale, None)
|
||||
)
|
||||
return {"kernel": mock_kernel, "quant": mock_quant, "mock_qweight": mock_qweight, "mock_scale": mock_scale}
|
||||
|
||||
def test_process_weights_after_loading(self, mock_deps, mock_quant_config):
|
||||
from vllm_omni.diffusion.quantization.int8 import Int8OnlineLinearMethod
|
||||
from vllm_omni.quantization.int8_config import Int8OnlineLinearMethod
|
||||
|
||||
method = Int8OnlineLinearMethod(mock_quant_config)
|
||||
layer = Module()
|
||||
@@ -268,12 +247,12 @@ class TestNPUInt8LinearMethod:
|
||||
def mock_torch_npu(self, mocker):
|
||||
torch_npu = MagicMock()
|
||||
|
||||
mocker.patch("vllm_omni.diffusion.quantization.int8.torch_npu", return_value=torch_npu)
|
||||
mocker.patch("vllm_omni.quantization.int8_config.torch_npu", return_value=torch_npu)
|
||||
mocker.patch(
|
||||
"vllm_omni.diffusion.quantization.int8.torch_npu.npu_dynamic_quant",
|
||||
"vllm_omni.quantization.int8_config.torch_npu.npu_dynamic_quant",
|
||||
return_value=(self.qweight_mock, self.scale_mock),
|
||||
)
|
||||
mocker.patch("vllm_omni.diffusion.quantization.int8.torch_npu.npu_quant_matmul", return_value=self.out_mock)
|
||||
mocker.patch("vllm_omni.quantization.int8_config.torch_npu.npu_quant_matmul", return_value=self.out_mock)
|
||||
return torch_npu
|
||||
|
||||
@pytest.fixture
|
||||
@@ -288,7 +267,7 @@ class TestNPUInt8LinearMethod:
|
||||
return layer
|
||||
|
||||
def test_npu_int8_process_weights_after_loading(self, mock_layer, mock_quant_config, mock_torch_npu):
|
||||
from vllm_omni.diffusion.quantization.int8 import NPUInt8LinearMethod
|
||||
from vllm_omni.quantization.int8_config import NPUInt8LinearMethod
|
||||
|
||||
method = NPUInt8LinearMethod(mock_quant_config)
|
||||
ori_weight_shape = mock_layer.weight.shape
|
||||
@@ -299,7 +278,7 @@ class TestNPUInt8LinearMethod:
|
||||
assert mock_layer.weight.is_contiguous()
|
||||
|
||||
def test_npu_int8_apply(self, mock_layer, mock_quant_config, mock_torch_npu):
|
||||
from vllm_omni.diffusion.quantization.int8 import NPUInt8LinearMethod
|
||||
from vllm_omni.quantization.int8_config import NPUInt8LinearMethod
|
||||
|
||||
method = NPUInt8LinearMethod(mock_quant_config)
|
||||
x = torch.randn(1, 16, 64)
|
||||
@@ -308,7 +287,7 @@ class TestNPUInt8LinearMethod:
|
||||
assert output.shape == (1, 16, 128)
|
||||
|
||||
def test_npu_int8_online_process_weights(self, mock_layer, mock_quant_config, mock_torch_npu):
|
||||
from vllm_omni.diffusion.quantization.int8 import NPUInt8OnlineLinearMethod
|
||||
from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod
|
||||
|
||||
method = NPUInt8OnlineLinearMethod(mock_quant_config)
|
||||
method.process_weights_after_loading(mock_layer)
|
||||
@@ -320,9 +299,9 @@ class TestNPUInt8LinearMethod:
|
||||
@pytest.fixture
|
||||
def quant_config():
|
||||
"""Shared quant config fixture for smoke tests."""
|
||||
from vllm_omni.diffusion.quantization.int8 import Int8Config
|
||||
from vllm_omni.quantization.int8_config import DiffusionInt8Config
|
||||
|
||||
return Int8Config(
|
||||
return DiffusionInt8Config(
|
||||
is_checkpoint_int8_serialized=False,
|
||||
activation_scheme="dynamic",
|
||||
)
|
||||
@@ -360,7 +339,7 @@ class TestNPUInt8Smoke:
|
||||
|
||||
def test_real_npu_online_process_weights_after_loading(self, quant_config, real_layer):
|
||||
"""Smoke test: full process_weights_after_loading with real torch_npu."""
|
||||
from vllm_omni.diffusion.quantization.int8 import NPUInt8OnlineLinearMethod
|
||||
from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod
|
||||
|
||||
method = NPUInt8OnlineLinearMethod(quant_config)
|
||||
|
||||
@@ -375,7 +354,7 @@ class TestNPUInt8Smoke:
|
||||
"""Smoke test: forward pass with real npu_quant_matmul."""
|
||||
import torch_npu
|
||||
|
||||
from vllm_omni.diffusion.quantization.int8 import NPUInt8LinearMethod
|
||||
from vllm_omni.quantization.int8_config import NPUInt8LinearMethod
|
||||
|
||||
method = NPUInt8LinearMethod(quant_config)
|
||||
|
||||
@@ -425,7 +404,7 @@ class TestCudaInt8Smoke:
|
||||
|
||||
def test_real_cuda_online_process_weights_after_loading(self, quant_config, real_layer):
|
||||
"""Smoke test: full process_weights_after_loading with real CUDA ops."""
|
||||
from vllm_omni.diffusion.quantization.int8 import Int8OnlineLinearMethod
|
||||
from vllm_omni.quantization.int8_config import Int8OnlineLinearMethod
|
||||
|
||||
method = Int8OnlineLinearMethod(quant_config)
|
||||
|
||||
@@ -439,7 +418,7 @@ class TestCudaInt8Smoke:
|
||||
"""Smoke test: forward pass with real CUDA int8 kernel."""
|
||||
from vllm import _custom_ops as ops
|
||||
|
||||
from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod
|
||||
from vllm_omni.quantization.int8_config import Int8LinearMethod
|
||||
|
||||
method = Int8LinearMethod(quant_config)
|
||||
|
||||
|
||||
Reference in new issue
Block a user