A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 14:45 UTC
[Bugfix] Fix bugs of running Quark quantized models (#16236)
Signed-off-by: chaow <chaow@amd.com>
chaow-amd committed
Apr 11, 2025 at 22:18 UTC
9e90c9f73f94d16c382b8c55fc99ab24a9baf09b
3 files changed
+67
-22
tests/quantization/test_quark.py
+37
-8
@@ -4,17 +4,28 @@
4
Run `pytest tests/quantization/test_quark.py`.
5
"""
6
7
-import torch
7
+import pytest
8
9
from vllm.model_executor.layers.quantization.quark.quark import ( # noqa: E501
10
- QuarkLinearMethod, QuarkW8A8Fp8)
10
+ QuarkLinearMethod, QuarkW8A8Fp8, QuarkW8A8Int8)
11
+from vllm.platforms import current_platform
12
13
13
-def test_quark_fp8(vllm_runner, monkeypatch):
14
- # vllm_runner.apply_model() relies on V0 internals.
15
- monkeypatch.setenv("VLLM_USE_V1", "0")
14
+@pytest.fixture(scope="function", autouse=True)
15
+def use_v0_only(monkeypatch):
16
+ """
17
+ This module relies on V0 internals, so set VLLM_USE_V1=0.
18
+ """
19
+ monkeypatch.setenv('VLLM_USE_V1', '0')
20
+
21
+
22
+@pytest.mark.parametrize('kv_cache_dtype', ['auto', 'fp8'])
23
+@pytest.mark.parametrize('tp', [1])
24
+def test_quark_fp8_w_per_tensor_a_per_tensor(vllm_runner, kv_cache_dtype, tp):
25
model_path = "amd/Llama-3.1-8B-Instruct-FP8-KV-Quark-test"
17
- with vllm_runner(model_path) as llm:
26
+ with vllm_runner(model_path,
27
+ kv_cache_dtype=kv_cache_dtype,
28
+ tensor_parallel_size=tp) as llm:
29
30
def check_model(model):
31
layer = model.model.layers[0]
@@ -26,11 +37,29 @@ def test_quark_fp8(vllm_runner, monkeypatch):
37
38
if isinstance(qkv_proj.scheme, QuarkW8A8Fp8):
39
assert len(qkv_proj.input_scale.shape) == 0
29
- assert qkv_proj.weight.dtype is torch.float8_e4m3fn
30
- #assert qkv_proj.weight.dtype is torch.float8_e4m3fnuz
40
+ assert qkv_proj.weight.dtype is current_platform.fp8_dtype()
41
assert len(qkv_proj.weight_scale.shape) == 0
42
43
llm.apply_model(check_model)
44
45
output = llm.generate_greedy("Hello my name is", max_tokens=20)
46
assert output
47
+
48
+
49
+@pytest.mark.parametrize('tp', [1])
50
+def test_quark_int8_w_per_tensor_a_per_tensor(vllm_runner, tp):
51
+ model_path = "amd/Llama-3.1-8B-Instruct-w-int8-a-int8-sym-test"
52
+ with vllm_runner(model_path, tensor_parallel_size=tp) as llm:
53
+
54
+ def check_model(model):
55
+ layer = model.model.layers[0]
56
+
57
+ qkv_proj = layer.self_attn.qkv_proj
58
+
59
+ assert isinstance(qkv_proj.quant_method, QuarkLinearMethod)
60
+ assert isinstance(qkv_proj.scheme, QuarkW8A8Int8)
61
+
62
+ llm.apply_model(check_model)
63
+
64
+ output = llm.generate_greedy("Hello my name is", max_tokens=20)
65
+ assert output
vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py
+6
-4
@@ -21,7 +21,7 @@ class QuarkW8A8Fp8(QuarkScheme):
21
def __init__(self, qscheme: str, is_static_input_scheme: Optional[bool]):
22
self.qscheme = qscheme
23
self.is_static_input_scheme = is_static_input_scheme
24
- self.fp8_linear = Fp8LinearOp(use_per_token_if_dynamic=True)
24
+ self.fp8_linear = Fp8LinearOp(use_per_token_if_dynamic=False)
25
self.out_dtype = torch.get_default_dtype()
26
27
@classmethod
@@ -41,10 +41,11 @@ class QuarkW8A8Fp8(QuarkScheme):
41
)
42
43
if current_platform.is_fp8_fnuz():
44
+ input_scale = getattr(layer, 'input_scale', None)
45
weight, max_w_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz(
46
weight=weight,
47
weight_scale=max_w_scale,
47
- input_scale=layer.input_scale)
48
+ input_scale=input_scale)
49
if input_scale is not None:
50
layer.input_scale = Parameter(input_scale,
51
requires_grad=False)
@@ -57,11 +58,12 @@ class QuarkW8A8Fp8(QuarkScheme):
58
weight = layer.weight
59
60
if current_platform.is_fp8_fnuz():
61
+ input_scale = getattr(layer, 'input_scale', None)
62
weight, weight_scale, input_scale = \
63
normalize_e4m3fn_to_e4m3fnuz(
64
weight=weight,
65
weight_scale=layer.weight_scale,
64
- input_scale=layer.input_scale)
66
+ input_scale=input_scale)
67
if input_scale is not None:
68
layer.input_scale = Parameter(input_scale,
69
requires_grad=False)
@@ -105,7 +107,7 @@ class QuarkW8A8Fp8(QuarkScheme):
107
# the newly added parameters
108
if self.qscheme == "per_channel":
109
weight_scale = ChannelQuantScaleParameter(
108
- data=torch.empty((sum(output_partition_sizes), 1),
110
+ data=torch.empty((sum(output_partition_sizes)),
111
dtype=torch.float32),
112
output_dim=0,
113
weight_loader=weight_loader)
vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py
+24
-10
@@ -35,7 +35,7 @@ class QuarkW8A8Int8(QuarkScheme):
35
input_size_per_partition: int,
36
params_dtype: torch.dtype, weight_loader: Callable,
37
**kwargs):
38
- self.logical_widths = output_partition_sizes
38
+ layer.logical_widths = output_partition_sizes
39
40
scaled_mm_linear_kernel_config = ScaledMMLinearLayerConfig(
41
is_channelwise=(self.qscheme == "per_channel"),
@@ -63,16 +63,28 @@ class QuarkW8A8Int8(QuarkScheme):
63
# WEIGHT SCALE
64
if self.qscheme == "per_channel":
65
weight_scale = ChannelQuantScaleParameter(
66
- data=torch.empty((sum(output_partition_sizes), 1),
66
+ data=torch.empty((sum(output_partition_sizes)),
67
dtype=torch.float32),
68
output_dim=0,
69
weight_loader=weight_loader)
70
+ ChannelQuantZPParameter = ChannelQuantScaleParameter
71
+ weight_zero_point = ChannelQuantZPParameter(
72
+ data=torch.empty((sum(output_partition_sizes)),
73
+ dtype=torch.int8),
74
+ output_dim=0,
75
+ weight_loader=weight_loader)
76
else:
77
assert self.qscheme == "per_tensor"
78
weight_scale = PerTensorScaleParameter(data=torch.empty(
79
len(output_partition_sizes), dtype=torch.float32),
80
weight_loader=weight_loader)
81
+ PerTensorZPParameter = PerTensorScaleParameter
82
+ weight_zero_point = PerTensorZPParameter(
83
+ data=torch.empty(len(output_partition_sizes),
84
+ dtype=torch.int8),
85
+ weight_loader=weight_loader)
86
layer.register_parameter("weight_scale", weight_scale)
87
+ layer.register_parameter("weight_zero_point", weight_zero_point)
88
89
# INPUT SCALE
90
if self.is_static_input_scheme:
@@ -81,14 +93,10 @@ class QuarkW8A8Int8(QuarkScheme):
93
weight_loader=weight_loader)
94
layer.register_parameter("input_scale", input_scale)
95
84
- if not self.input_symmetric:
85
- # Note: quark stores the zp using the same dtype
86
- # as the weights
87
- # AZP loaded as int8 but used as int32
88
- input_zero_point = BasevLLMParameter(
89
- data=torch.empty(1, dtype=torch.int8),
90
- weight_loader=weight_loader)
91
- layer.register_parameter("input_zero_point", input_zero_point)
96
+ input_zero_point = BasevLLMParameter(data=torch.empty(
97
+ 1, dtype=torch.int8),
98
+ weight_loader=weight_loader)
99
+ layer.register_parameter("input_zero_point", input_zero_point)
100
101
self.kernel = kernel_type(c=scaled_mm_linear_kernel_config,
102
w_q_param_name="weight",
@@ -100,6 +108,12 @@ class QuarkW8A8Int8(QuarkScheme):
108
# Checkpoints are serialized in quark format, which is
109
# different from the format the kernel may want. Handle repacking here.
110
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
111
+ layer.register_parameter("weight_zero_point", None)
112
+ delattr(layer, 'weight_zero_point')
113
+ if self.input_symmetric:
114
+ layer.register_parameter("input_zero_point", None)
115
+ delattr(layer, 'input_zero_point')
116
+
117
self.kernel.process_weights_after_loading(layer)
118
119
def apply_weights(self, layer: torch.nn.Module, x: torch.Tensor,