| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests register custom quantization config. |
| 4 | |
| 5 | See https://github.com/vllm-project/vllm/issues/11926 for more details. |
| 6 | |
| 7 | Run `pytest tests/quantization/test_register_quantization_config.py`. |
| 8 | """ |
| 9 | |
| 10 | import logging |
| 11 | from typing import Any |
| 12 | |
| 13 | import pytest |
| 14 | import torch |
| 15 | import torch.nn.functional as F |
| 16 | |
| 17 | from tests.quantization.utils import load_model_without_vllm_runner |
| 18 | from vllm.config import set_current_vllm_config |
| 19 | from vllm.forward_context import set_forward_context |
| 20 | from vllm.model_executor.layers.attention import Attention |
| 21 | from vllm.model_executor.layers.linear import ( |
| 22 | LinearBase, # noqa: E501 |
| 23 | UnquantizedLinearMethod, |
| 24 | ) |
| 25 | from vllm.model_executor.layers.quantization import ( |
| 26 | QuantizationMethods, |
| 27 | get_quantization_config, |
| 28 | register_quantization_config, |
| 29 | ) |
| 30 | from vllm.model_executor.layers.quantization.base_config import ( |
| 31 | QuantizationConfig, # noqa: E501 |
| 32 | ) |
| 33 | from vllm.platforms import current_platform |
| 34 | |
| 35 | |
| 36 | class FakeQuantLinearMethod(UnquantizedLinearMethod): |
| 37 | """Fake quantization linear method for per-token dynamic quantization.""" |
| 38 | |
| 39 | def __init__(self, num_bits: int = 8) -> None: |
| 40 | """Initialize the quantization method.""" |
| 41 | super().__init__() |
| 42 | self.num_bits = num_bits |
| 43 | |
| 44 | def apply( |
| 45 | self, |
| 46 | layer: torch.nn.Module, |
| 47 | x: torch.Tensor, |
| 48 | bias: torch.Tensor | None = None, |
| 49 | ) -> torch.Tensor: |
| 50 | """Perform fake quantization before the linear layer.""" |
| 51 | # Calculate the scales dynamically |
| 52 | max_val = torch.amax(x, dim=(0, -1), keepdims=True) |
| 53 | min_val = torch.amin(x, dim=(0, -1), keepdims=True) |
| 54 | scales = (max_val - min_val) / (2**self.num_bits - 1) |
| 55 | |
| 56 | # Fake quantize the input |
| 57 | quant_x = torch.clamp( |
| 58 | torch.round(x / scales), |
| 59 | -(2 ** (self.num_bits - 1)), |
| 60 | 2 ** (self.num_bits - 1) - 1, |
| 61 | ) |
| 62 | dequant_x = quant_x * scales |
| 63 | |
| 64 | return F.linear(dequant_x, layer.weight, bias) |
| 65 | |
| 66 | |
| 67 | @register_quantization_config("custom_quant") |
| 68 | class CustomQuantConfig(QuantizationConfig): |
| 69 | """Custom quantization config for per-token dynamic fake quantization.""" |
| 70 | |
| 71 | def __init__(self, num_bits: int = 8) -> None: |
| 72 | """Initialize the quantization config.""" |
| 73 | super().__init__() |
| 74 | self.num_bits = num_bits |
| 75 | |
| 76 | def get_name(self) -> QuantizationMethods: |
| 77 | """Name of the quantization method.""" |
| 78 | return "custom_quant" # type: ignore[return-value] # Runtime registration extends built-in names. |
| 79 | |
| 80 | def get_supported_act_dtypes(self) -> list[torch.dtype]: |
| 81 | """List of supported activation dtypes.""" |
| 82 | return [torch.float16, torch.bfloat16] |
| 83 | |
| 84 | @classmethod |
| 85 | def get_min_capability(cls) -> int: |
| 86 | """Minimum GPU capability to support the quantization method.""" |
| 87 | return -1 |
| 88 | |
| 89 | @staticmethod |
| 90 | def get_config_filenames() -> list[str]: |
| 91 | """List of filenames to search for in the model directory.""" |
| 92 | return [] |
| 93 | |
| 94 | @classmethod |
| 95 | def from_config(cls, config: dict[str, Any]) -> "CustomQuantConfig": |
| 96 | """Create a config class from the model's quantization config.""" |
| 97 | return CustomQuantConfig(num_bits=config.get("num_bits", 8)) |
| 98 | |
| 99 | def get_quant_method( |
| 100 | self, layer: torch.nn.Module, prefix: str |
| 101 | ) -> FakeQuantLinearMethod | None: |
| 102 | """Get the quantize method to use for the quantized layer.""" |
| 103 | if isinstance(layer, LinearBase): |
| 104 | return FakeQuantLinearMethod(num_bits=self.num_bits) |
| 105 | return None |
| 106 | |
| 107 | |
| 108 | def test_register_quantization_config(caplog_vllm): |
| 109 | """Test register custom quantization config.""" |
| 110 | # The quantization method `custom_quant` should be registered. |
| 111 | assert get_quantization_config("custom_quant") == CustomQuantConfig |
| 112 | |
| 113 | # The quantization method `custom_quant` is already exists, |
| 114 | # should raise a debug message when re-registering it. |
| 115 | with caplog_vllm.at_level(logging.DEBUG, logger="vllm"): |
| 116 | register_quantization_config("custom_quant")(CustomQuantConfig) |
| 117 | |
| 118 | assert any( |
| 119 | "The quantization method 'custom_quant' already exists" in message |
| 120 | for message in caplog_vllm.messages |
| 121 | ), "Expected a debug message when re-registering custom_quant" |
| 122 | |
| 123 | |
| 124 | @pytest.mark.parametrize( |
| 125 | argnames="model", |
| 126 | argvalues=[ |
| 127 | "meta-llama/Llama-3.2-1B-Instruct", |
| 128 | ], |
| 129 | ) |
| 130 | def test_custom_quant(model, monkeypatch, dist_init, workspace_init): |
| 131 | """Test infer with the custom quantization method.""" |
| 132 | vllm_model, vllm_config = load_model_without_vllm_runner( |
| 133 | model, |
| 134 | quantization="custom_quant", |
| 135 | ) |
| 136 | layer = vllm_model.model.layers[0] |
| 137 | qkv_proj = layer.self_attn.qkv_proj |
| 138 | |
| 139 | # Check the quantization method is FakeQuantLinearMethod. |
| 140 | assert isinstance(qkv_proj.quant_method, FakeQuantLinearMethod) |
| 141 | |
| 142 | monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous()) |
| 143 | device_type = current_platform.device_type |
| 144 | input_ids = torch.tensor([1, 2, 3, 4], device=device_type) |
| 145 | positions = torch.arange(input_ids.numel(), device=device_type) |
| 146 | with ( |
| 147 | set_current_vllm_config(vllm_config), |
| 148 | set_forward_context(None, vllm_config, num_tokens=input_ids.numel()), |
| 149 | ): |
| 150 | hidden_states = vllm_model(input_ids, positions, None) |
| 151 | logits = vllm_model.compute_logits(hidden_states) |
| 152 | assert torch.isfinite(logits).all() |