0
tests/ut/quantization/__init__.py
Normal file
0
tests/ut/quantization/__init__.py
Normal file
219
tests/ut/quantization/conftest_quantization.py
Normal file
219
tests/ut/quantization/conftest_quantization.py
Normal file
@@ -0,0 +1,219 @@
|
||||
from unittest.mock import Mock
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
FAKQUANT_CONFIG = {
|
||||
"version": "1.0.0",
|
||||
"model_quant_type": "W8A8_DYNAMIC",
|
||||
"fa_quant_type": "FAKQuant",
|
||||
"model.embed_tokens.weight": "FLOAT",
|
||||
"model.layers.3.self_attn.fa_q.scale": "FAQuant",
|
||||
"model.layers.3.self_attn.fa_k.scale": "FAQuant",
|
||||
"model.layers.3.self_attn.fa_v.scale": "FAQuant",
|
||||
"model.layers.3.self_attn.fa_q.offset": "FAQuant",
|
||||
"model.layers.3.self_attn.fa_k.offset": "FAQuant",
|
||||
"model.layers.3.self_attn.fa_v.offset": "FAQuant",
|
||||
}
|
||||
|
||||
W8A8_CONFIG = {
|
||||
"version": "1.0.0",
|
||||
"model_quant_type": "W8A8_DYNAMIC",
|
||||
"model.embed_tokens.weight": "FLOAT",
|
||||
"model.layers.0.self_attn.q_a_proj.weight": "W8A8",
|
||||
"model.layers.0.mlp.gate_proj.weight": "W8A8_DYNAMIC",
|
||||
"model.layers.0.mlp.up_proj.weight": "W8A8_DYNAMIC",
|
||||
"model.layers.0.mlp.down_proj.weight": "W8A8_DYNAMIC",
|
||||
"model.layers.3.mlp.experts.0.gate_proj.weight": "W8A8_DYNAMIC",
|
||||
"model.layers.3.mlp.experts.0.up_proj.weight": "W8A8_DYNAMIC",
|
||||
"model.layers.3.mlp.experts.0.down_proj.weight": "W8A8_DYNAMIC",
|
||||
"model.layers.3.mlp.experts.1.gate_proj.weight": "W8A8_DYNAMIC",
|
||||
"model.layers.3.mlp.experts.1.up_proj.weight": "W8A8_DYNAMIC",
|
||||
"model.layers.3.mlp.experts.1.down_proj.weight": "W8A8_DYNAMIC",
|
||||
}
|
||||
|
||||
COMPRESSED_TENSORS_W8A8_CONFIG = {
|
||||
"config_groups": {
|
||||
"group_0": {
|
||||
"format": "int-quantized",
|
||||
"input_activations": {
|
||||
"actorder": None,
|
||||
"block_structure": None,
|
||||
"dynamic": True,
|
||||
"group_size": None,
|
||||
"num_bits": 8,
|
||||
"observer": None,
|
||||
"observer_kwargs": {},
|
||||
"strategy": "token",
|
||||
"symmetric": True,
|
||||
"type": "int",
|
||||
},
|
||||
"output_activations": None,
|
||||
"targets": ["Linear"],
|
||||
"weights": {
|
||||
"actorder": None,
|
||||
"block_structure": None,
|
||||
"dynamic": False,
|
||||
"group_size": None,
|
||||
"num_bits": 8,
|
||||
"observer": "minmax",
|
||||
"observer_kwargs": {},
|
||||
"strategy": "channel",
|
||||
"symmetric": True,
|
||||
"type": "int",
|
||||
},
|
||||
}
|
||||
},
|
||||
"format": "int-quantized",
|
||||
"global_compression_ratio": None,
|
||||
"ignore": ["lm_head"],
|
||||
"kv_cache_scheme": None,
|
||||
"quant_method": "compressed-tensors",
|
||||
"quantization_status": "compressed",
|
||||
}
|
||||
|
||||
|
||||
def identity(*args):
|
||||
return args[0]
|
||||
|
||||
|
||||
def create_mock_vllm_config(
|
||||
quant_description=None,
|
||||
model_dtype=torch.bfloat16,
|
||||
scheduler_config=None,
|
||||
compilation_mode=None,
|
||||
enforce_eager=True,
|
||||
kv_transfer_config=None,
|
||||
parallel_config=None,
|
||||
):
|
||||
if quant_description is None:
|
||||
quant_description = {"group_size": 32}
|
||||
|
||||
mock_config = Mock()
|
||||
mock_config.quant_config = Mock(quant_description=quant_description)
|
||||
mock_config.model_config = Mock(
|
||||
dtype=model_dtype,
|
||||
hf_config=Mock(model_type=None),
|
||||
enforce_eager=enforce_eager,
|
||||
)
|
||||
|
||||
if scheduler_config is None:
|
||||
mock_config.scheduler_config = Mock(
|
||||
max_num_batched_tokens=2048,
|
||||
max_model_len=2048,
|
||||
enable_chunked_prefill=False,
|
||||
)
|
||||
else:
|
||||
mock_config.scheduler_config = scheduler_config
|
||||
|
||||
if compilation_mode is not None:
|
||||
mock_config.compilation_config = Mock(mode=compilation_mode)
|
||||
else:
|
||||
mock_config.compilation_config = Mock()
|
||||
|
||||
mock_config.kv_transfer_config = kv_transfer_config
|
||||
|
||||
if parallel_config is None:
|
||||
mock_config.parallel_config = Mock(enable_expert_parallel=True)
|
||||
else:
|
||||
mock_config.parallel_config = parallel_config
|
||||
|
||||
return mock_config
|
||||
|
||||
|
||||
def create_mock_ascend_config(
|
||||
multistream_overlap_gate=False,
|
||||
dynamic_eplb=False,
|
||||
flashcomm2_oproj_tensor_parallel_size=1,
|
||||
):
|
||||
mock_config = Mock()
|
||||
mock_config.multistream_overlap_gate = multistream_overlap_gate
|
||||
mock_config.eplb_config = Mock(dynamic_eplb=dynamic_eplb)
|
||||
mock_config.flashcomm2_oproj_tensor_parallel_size = flashcomm2_oproj_tensor_parallel_size
|
||||
return mock_config
|
||||
|
||||
|
||||
def create_moe_layer(
|
||||
num_experts=8,
|
||||
hidden_size=128,
|
||||
intermediate_size=128,
|
||||
weight_dtype=torch.int8,
|
||||
params_dtype=torch.bfloat16,
|
||||
):
|
||||
layer = nn.Module()
|
||||
layer.w13_weight = nn.Parameter(
|
||||
torch.randint(-8, 8, (num_experts, 2 * intermediate_size, hidden_size), dtype=weight_dtype),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w2_weight = nn.Parameter(
|
||||
torch.randint(-8, 8, (num_experts, hidden_size, intermediate_size), dtype=weight_dtype),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w13_weight_scale = nn.Parameter(
|
||||
torch.ones((num_experts, 2 * intermediate_size, 1), dtype=params_dtype), requires_grad=False
|
||||
)
|
||||
layer.w13_weight_offset = nn.Parameter(
|
||||
torch.zeros((num_experts, 2 * intermediate_size, 1), dtype=params_dtype), requires_grad=False
|
||||
)
|
||||
layer.w2_weight_scale = nn.Parameter(
|
||||
torch.ones((num_experts, hidden_size, 1), dtype=params_dtype), requires_grad=False
|
||||
)
|
||||
layer.w2_weight_offset = nn.Parameter(
|
||||
torch.zeros((num_experts, hidden_size, 1), dtype=params_dtype), requires_grad=False
|
||||
)
|
||||
return layer
|
||||
|
||||
|
||||
def create_mxfp_moe_layer(
|
||||
num_experts=8,
|
||||
hidden_size=128,
|
||||
intermediate_size=128,
|
||||
group_size=32,
|
||||
weight_dtype=torch.float8_e4m3fn,
|
||||
scale_dtype=torch.uint8,
|
||||
):
|
||||
layer = nn.Module()
|
||||
layer.w13_weight = nn.Parameter(
|
||||
torch.randn(num_experts, 2 * intermediate_size, hidden_size).to(weight_dtype), requires_grad=False
|
||||
)
|
||||
layer.w2_weight = nn.Parameter(
|
||||
torch.randn(num_experts, hidden_size, intermediate_size).to(weight_dtype), requires_grad=False
|
||||
)
|
||||
layer.w13_weight_scale = nn.Parameter(
|
||||
torch.randint(0, 255, (num_experts, 2 * intermediate_size, hidden_size // group_size), dtype=scale_dtype),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w2_weight_scale = nn.Parameter(
|
||||
torch.randint(0, 255, (num_experts, hidden_size, intermediate_size // group_size), dtype=scale_dtype),
|
||||
requires_grad=False,
|
||||
)
|
||||
return layer
|
||||
|
||||
|
||||
def create_linear_layer(
|
||||
quant_method,
|
||||
input_size=128,
|
||||
output_size=256,
|
||||
params_dtype=torch.bfloat16,
|
||||
):
|
||||
layer = nn.Module()
|
||||
weight_dict = quant_method.get_weight(input_size, output_size, params_dtype)
|
||||
for weight_name, weight_param in weight_dict.items():
|
||||
param = torch.nn.Parameter(weight_param.npu(), requires_grad=False)
|
||||
layer.register_parameter(weight_name, param)
|
||||
|
||||
pertensor_dict = quant_method.get_pertensor_param(params_dtype)
|
||||
for pertensor_name, pertensor_param in pertensor_dict.items():
|
||||
param = torch.nn.Parameter(pertensor_param.npu(), requires_grad=False)
|
||||
layer.register_parameter(pertensor_name, param)
|
||||
|
||||
perchannel_dict = quant_method.get_perchannel_param(output_size, params_dtype)
|
||||
for perchannel_name, perchannel_param in perchannel_dict.items():
|
||||
param = torch.nn.Parameter(perchannel_param.npu(), requires_grad=False)
|
||||
layer.register_parameter(perchannel_name, param)
|
||||
|
||||
pergroup_dict = quant_method.get_pergroup_param(input_size, output_size, params_dtype, layer_type="row")
|
||||
for pergroup_name, pergroup_param in pergroup_dict.items():
|
||||
param = torch.nn.Parameter(pergroup_param.npu(), requires_grad=False)
|
||||
layer.register_parameter(pergroup_name, param)
|
||||
return layer
|
||||
0
tests/ut/quantization/methods/__init__.py
Normal file
0
tests/ut/quantization/methods/__init__.py
Normal file
0
tests/ut/quantization/methods/a2/__init__.py
Normal file
0
tests/ut/quantization/methods/a2/__init__.py
Normal file
334
tests/ut/quantization/methods/a2/test_w4a16.py
Normal file
334
tests/ut/quantization/methods/a2/test_w4a16.py
Normal file
@@ -0,0 +1,334 @@
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import regex as re
|
||||
import torch
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
from vllm_ascend.ascend_forward_context import MoECommType
|
||||
from vllm_ascend.quantization.methods.w4a16 import AscendW4A16FusedMoEMethod, pack_to_int32, unpack_from_int32
|
||||
|
||||
|
||||
class TestUnpackFromInt32(TestBase):
|
||||
def test_unpack_from_int32_restores_values_and_crops_padding(self):
|
||||
weight = torch.tensor([[0x76543210]], dtype=torch.int32)
|
||||
shape = torch.Size([1, 6])
|
||||
|
||||
result = unpack_from_int32(weight, shape, num_bits=4, packed_dim=1)
|
||||
|
||||
self.assertEqual(result.dtype, torch.int8)
|
||||
self.assertEqual(result.shape, shape)
|
||||
self.assertTrue(torch.equal(result, torch.tensor([[-8, -7, -6, -5, -4, -3]], dtype=torch.int8)))
|
||||
|
||||
def test_unpack_from_int32_packed_dim_1(self):
|
||||
weight = torch.tensor([[305419896, -1420531520]], dtype=torch.int32)
|
||||
shape = torch.Size([1, 8])
|
||||
num_bits = 4
|
||||
|
||||
result = unpack_from_int32(weight, shape, num_bits, packed_dim=1)
|
||||
|
||||
self.assertEqual(result.dtype, torch.int8)
|
||||
self.assertEqual(result.shape, shape)
|
||||
|
||||
def test_unpack_from_int32_packed_dim_0(self):
|
||||
weight = torch.tensor([[305419896], [-1420531520]], dtype=torch.int32)
|
||||
shape = torch.Size([8, 1])
|
||||
num_bits = 4
|
||||
|
||||
result = unpack_from_int32(weight, shape, num_bits, packed_dim=0)
|
||||
|
||||
self.assertEqual(result.dtype, torch.int8)
|
||||
self.assertEqual(result.shape, shape)
|
||||
|
||||
def test_unpack_from_int32_packed_dim_0_restores_values(self):
|
||||
weight = torch.tensor([[0x76543210]], dtype=torch.int32)
|
||||
shape = torch.Size([6, 1])
|
||||
|
||||
result = unpack_from_int32(weight, shape, num_bits=4, packed_dim=0)
|
||||
|
||||
self.assertEqual(result.dtype, torch.int8)
|
||||
self.assertEqual(result.shape, shape)
|
||||
expected = torch.tensor([[-8], [-7], [-6], [-5], [-4], [-3]], dtype=torch.int8)
|
||||
self.assertTrue(torch.equal(result, expected))
|
||||
|
||||
def test_unpack_from_int32_assertion_dtype_message(self):
|
||||
weight = torch.tensor([[1, 2]], dtype=torch.int64)
|
||||
message = "Expecting `weight.dtype` is torch.int32 but got torch.int64."
|
||||
|
||||
with self.assertRaisesRegex(AssertionError, re.escape(message)):
|
||||
unpack_from_int32(weight, torch.Size([8, 1]), 4)
|
||||
|
||||
def test_unpack_from_int32_assertion_num_bits_positive_message(self):
|
||||
weight = torch.tensor([[1, 2]], dtype=torch.int32)
|
||||
message = "Expecting `num_bits` should be positive but got 0."
|
||||
|
||||
with self.assertRaisesRegex(AssertionError, re.escape(message)):
|
||||
unpack_from_int32(weight, torch.Size([8, 1]), 0)
|
||||
|
||||
def test_unpack_from_int32_assertion_num_bits_upper_bound_message(self):
|
||||
weight = torch.tensor([[1, 2]], dtype=torch.int32)
|
||||
message = "Expecting `num_bits` should not be larger than 8 but got 16."
|
||||
|
||||
with self.assertRaisesRegex(AssertionError, re.escape(message)):
|
||||
unpack_from_int32(weight, torch.Size([8, 1]), 16)
|
||||
|
||||
def test_unpack_from_int32_assertion_num_bits_divides_int32_message(self):
|
||||
weight = torch.tensor([[1, 2]], dtype=torch.int32)
|
||||
message = "Expecting `num_bits` 3 to divide 32 exactly."
|
||||
|
||||
with self.assertRaisesRegex(AssertionError, re.escape(message)):
|
||||
unpack_from_int32(weight, torch.Size([8, 1]), 3)
|
||||
|
||||
def test_unpack_from_int32_assertion_packed_dim_message(self):
|
||||
weight = torch.tensor([[1, 2]], dtype=torch.int32)
|
||||
message = "Expecting `packed_dim` is 0 or 1 but got 2."
|
||||
|
||||
with self.assertRaisesRegex(AssertionError, re.escape(message)):
|
||||
unpack_from_int32(weight, torch.Size([8, 1]), 4, packed_dim=2)
|
||||
|
||||
|
||||
class TestPackToInt32(TestBase):
|
||||
@patch("vllm_ascend.quantization.methods.w4a16.torch_npu.npu_convert_weight_to_int4pack")
|
||||
def test_pack_to_int32_int8(self, mock_npu_convert_weight_to_int4pack):
|
||||
mock_npu_convert_weight_to_int4pack.return_value = torch.zeros((2, 4), dtype=torch.int32)
|
||||
|
||||
weight = torch.zeros((2, 8, 16), dtype=torch.int8)
|
||||
result = pack_to_int32(weight)
|
||||
|
||||
self.assertEqual(result.dtype, torch.int32)
|
||||
mock_npu_convert_weight_to_int4pack.assert_not_called()
|
||||
|
||||
self.assertEqual(result.shape, torch.Size([2, 8, 4]))
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a16.torch_npu.npu_convert_weight_to_int4pack")
|
||||
def test_pack_to_int32_int32(self, mock_npu_convert_weight_to_int4pack):
|
||||
def mock_convert_weight(weight):
|
||||
return weight
|
||||
|
||||
mock_npu_convert_weight_to_int4pack.side_effect = mock_convert_weight
|
||||
weight = torch.zeros((2, 8, 8), dtype=torch.int32)
|
||||
result = pack_to_int32(weight)
|
||||
|
||||
self.assertEqual(result.dtype, torch.int32)
|
||||
self.assertEqual(result.shape, weight.shape)
|
||||
|
||||
def test_pack_to_int32_assertion_dim(self):
|
||||
weight = torch.zeros((8, 8), dtype=torch.int8)
|
||||
message = (
|
||||
"Expecting `weight.dim()` is 3 ([expert, output_channel, input_channel] or "
|
||||
"[expert, input_channel, output_channel]) but got 2."
|
||||
)
|
||||
|
||||
with self.assertRaisesRegex(AssertionError, re.escape(message)):
|
||||
pack_to_int32(weight)
|
||||
|
||||
def test_pack_to_int32_assertion_dtype(self):
|
||||
weight = torch.zeros((2, 8, 8), dtype=torch.float32)
|
||||
message = "Expecting `weight.dtype` is torch.int8 or torch.int32 but got torch.float32."
|
||||
|
||||
with self.assertRaisesRegex(AssertionError, re.escape(message)):
|
||||
pack_to_int32(weight)
|
||||
|
||||
def test_pack_to_int32_assertion_int32_divisible_message(self):
|
||||
weight = torch.zeros((2, 8, 7), dtype=torch.int32)
|
||||
message = "the last dim of weight needs to be divided by 8."
|
||||
|
||||
with self.assertRaisesRegex(AssertionError, re.escape(message)):
|
||||
pack_to_int32(weight)
|
||||
|
||||
def test_pack_to_int32_assertion_int8_divisible_message(self):
|
||||
weight = torch.zeros((2, 8, 7), dtype=torch.int8)
|
||||
message = "the last dim of weight needs to be divided by 4."
|
||||
|
||||
with self.assertRaisesRegex(AssertionError, re.escape(message)):
|
||||
pack_to_int32(weight)
|
||||
|
||||
|
||||
class TestAscendW4A16FusedMoEMethod(TestBase):
|
||||
experts = 8
|
||||
input_size = 32
|
||||
output_size = 128
|
||||
group_size = 32
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a16.get_ascend_config")
|
||||
@patch("vllm_ascend.quantization.methods.w4a16.get_current_vllm_config")
|
||||
def setUp(self, mock_get_current_vllm_config, mock_get_ascend_config):
|
||||
mock_ascend_config = Mock()
|
||||
mock_ascend_config.eplb_config.dynamic_eplb = False
|
||||
mock_ascend_config.eplb_config.expert_map_record_path = None
|
||||
mock_get_ascend_config.return_value = mock_ascend_config
|
||||
|
||||
mock_vllm_config = Mock()
|
||||
mock_vllm_config.quant_config = Mock(
|
||||
quant_description={
|
||||
"group_size": self.group_size,
|
||||
}
|
||||
)
|
||||
mock_get_current_vllm_config.return_value = mock_vllm_config
|
||||
|
||||
self.quant_method = AscendW4A16FusedMoEMethod()
|
||||
|
||||
def test_get_weight(self):
|
||||
param_dict = self.quant_method.get_weight(self.experts, self.input_size, self.output_size, torch.bfloat16)
|
||||
|
||||
self.assertEqual(param_dict["w13_weight_packed"].dtype, torch.int32)
|
||||
expected_w13_shape = (self.experts, 2 * self.input_size, self.output_size // self.quant_method.pack_factor)
|
||||
self.assertEqual(param_dict["w13_weight_packed"].shape, expected_w13_shape)
|
||||
|
||||
self.assertEqual(param_dict["w2_weight_packed"].dtype, torch.int32)
|
||||
expected_w2_shape = (self.experts, self.output_size, self.input_size // self.quant_method.pack_factor)
|
||||
self.assertEqual(param_dict["w2_weight_packed"].shape, expected_w2_shape)
|
||||
|
||||
def test_get_weight_assertion_intermediate_size_message(self):
|
||||
message = "Expecting `intermediate_size_per_partition` 33 can be divided by `pack_factor` 8"
|
||||
|
||||
with self.assertRaisesRegex(AssertionError, re.escape(message)):
|
||||
self.quant_method.get_weight(self.experts, self.input_size + 1, self.output_size, torch.bfloat16)
|
||||
|
||||
def test_get_weight_assertion_hidden_sizes_message(self):
|
||||
message = "Expecting `hidden_sizes` 129 can be divided by `pack_factor` 8"
|
||||
|
||||
with self.assertRaisesRegex(AssertionError, re.escape(message)):
|
||||
self.quant_method.get_weight(self.experts, self.input_size, self.output_size + 1, torch.bfloat16)
|
||||
|
||||
def test_get_dynamic_quant_param(self):
|
||||
param_dict = self.quant_method.get_dynamic_quant_param(
|
||||
self.experts, self.input_size, self.output_size, torch.bfloat16
|
||||
)
|
||||
|
||||
self.assertEqual(param_dict["w13_weight_scale"].dtype, torch.bfloat16)
|
||||
expected_w13_scale_shape = (self.experts, 2 * self.input_size, self.output_size // self.group_size)
|
||||
self.assertEqual(param_dict["w13_weight_scale"].shape, expected_w13_scale_shape)
|
||||
|
||||
self.assertEqual(param_dict["w2_weight_shape"].dtype, torch.int32)
|
||||
self.assertEqual(param_dict["w2_weight_shape"].shape, (self.experts, 2))
|
||||
|
||||
self.assertEqual(param_dict["w13_weight_offset"].dtype, torch.bfloat16)
|
||||
self.assertEqual(param_dict["w13_weight_offset"].shape, expected_w13_scale_shape)
|
||||
|
||||
def test_get_dynamic_quant_param_assertion_intermediate_size_message(self):
|
||||
message = "Expecting `intermediate_size_per_partition` 33 can be divided by `group_size` 32"
|
||||
|
||||
with self.assertRaisesRegex(AssertionError, re.escape(message)):
|
||||
self.quant_method.get_dynamic_quant_param(
|
||||
self.experts, self.input_size + 1, self.output_size, torch.bfloat16
|
||||
)
|
||||
|
||||
def test_get_dynamic_quant_param_assertion_hidden_sizes_message(self):
|
||||
message = "Expecting `hidden_sizes` 129 can be divided by `group_size` 32"
|
||||
|
||||
with self.assertRaisesRegex(AssertionError, re.escape(message)):
|
||||
self.quant_method.get_dynamic_quant_param(
|
||||
self.experts, self.input_size, self.output_size + 1, torch.bfloat16
|
||||
)
|
||||
|
||||
def build_layer(self):
|
||||
"""Build a mock layer for testing"""
|
||||
layer = torch.nn.Module()
|
||||
|
||||
w13_shape = (self.experts, 2 * self.input_size, self.output_size // self.quant_method.pack_factor)
|
||||
w2_shape = (self.experts, self.output_size, self.input_size // self.quant_method.pack_factor)
|
||||
|
||||
layer.w13_weight_packed = torch.nn.Parameter(
|
||||
torch.randint(-100, 100, w13_shape, dtype=torch.int32), requires_grad=False
|
||||
)
|
||||
layer.w2_weight_packed = torch.nn.Parameter(
|
||||
torch.randint(-100, 100, w2_shape, dtype=torch.int32), requires_grad=False
|
||||
)
|
||||
|
||||
w13_scale_shape = (self.experts, 2 * self.input_size, self.output_size // self.group_size)
|
||||
w2_scale_shape = (self.experts, self.output_size, self.input_size // self.group_size)
|
||||
|
||||
layer.w13_weight_scale = torch.nn.Parameter(
|
||||
torch.ones(w13_scale_shape, dtype=torch.bfloat16), requires_grad=False
|
||||
)
|
||||
layer.w2_weight_scale = torch.nn.Parameter(
|
||||
torch.ones(w2_scale_shape, dtype=torch.bfloat16), requires_grad=False
|
||||
)
|
||||
|
||||
layer.w13_weight_offset = torch.nn.Parameter(
|
||||
torch.zeros(w13_scale_shape, dtype=torch.bfloat16), requires_grad=False
|
||||
)
|
||||
layer.w2_weight_offset = torch.nn.Parameter(
|
||||
torch.zeros(w2_scale_shape, dtype=torch.bfloat16), requires_grad=False
|
||||
)
|
||||
|
||||
layer.w13_weight_shape = torch.nn.Parameter(
|
||||
torch.tensor([[2 * self.input_size, self.output_size]] * self.experts, dtype=torch.int32),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w2_weight_shape = torch.nn.Parameter(
|
||||
torch.tensor([[self.output_size, self.input_size]] * self.experts, dtype=torch.int32), requires_grad=False
|
||||
)
|
||||
|
||||
return layer
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a16.torch_npu.npu_convert_weight_to_int4pack")
|
||||
def test_process_weights_after_loading_with_transpose(self, mock_npu_convert_weight_to_int4pack):
|
||||
def mock_convert_weight(weight):
|
||||
new_shape = list(weight.shape)
|
||||
new_shape[-1] = new_shape[-1] // 8
|
||||
return torch.zeros(new_shape, dtype=torch.int32)
|
||||
|
||||
mock_npu_convert_weight_to_int4pack.side_effect = mock_convert_weight
|
||||
layer = self.build_layer()
|
||||
self.quant_method.process_weights_after_loading(layer)
|
||||
|
||||
self.assertEqual(layer.w13_weight_packed.data.shape, torch.Size([8, 128, 8]))
|
||||
self.assertEqual(layer.w2_weight_packed.data.shape, torch.Size([8, 32, 16]))
|
||||
self.assertEqual(layer.w13_weight_scale.data.shape, torch.Size([8, 4, 64]))
|
||||
self.assertEqual(layer.w2_weight_offset.data.shape, torch.Size([8, 1, 128]))
|
||||
self.assertTrue(layer.w13_weight_scale.data.is_contiguous())
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a16._EXTRA_CTX")
|
||||
@patch("vllm_ascend.quantization.methods.w4a16.select_experts")
|
||||
def test_apply_uses_explicit_dispatch_and_mlp_args(self, mock_select_experts, mock_extra_ctx):
|
||||
tokens = 3
|
||||
hidden_size = self.output_size
|
||||
layer = self.build_layer()
|
||||
x = torch.randn(tokens, hidden_size, dtype=torch.float32)
|
||||
router_logits = torch.randn(tokens, self.experts, dtype=torch.float32)
|
||||
topk_weights = torch.randn(tokens, 2, dtype=torch.float32)
|
||||
topk_ids = torch.randint(0, self.experts, (tokens, 2), dtype=torch.int64)
|
||||
mc2_mask = torch.tensor([1, 0, 1], dtype=torch.bool)
|
||||
pertoken_scale = torch.randn(tokens, dtype=torch.float32)
|
||||
layer.swiglu_limit = 1000000
|
||||
|
||||
mock_select_experts.return_value = (topk_weights, topk_ids)
|
||||
mock_comm = Mock()
|
||||
mock_comm.fused_experts.return_value = torch.randn(tokens, hidden_size, dtype=torch.float32)
|
||||
mock_extra_ctx.moe_comm_method = mock_comm
|
||||
mock_extra_ctx.moe_comm_type = MoECommType.ALLGATHER
|
||||
|
||||
self.quant_method.apply(
|
||||
layer=layer,
|
||||
x=x,
|
||||
router_logits=router_logits,
|
||||
top_k=2,
|
||||
renormalize=True,
|
||||
num_experts=self.experts,
|
||||
activation="gelu",
|
||||
apply_router_weight_on_input=True,
|
||||
mc2_mask=mc2_mask,
|
||||
pertoken_scale=pertoken_scale,
|
||||
)
|
||||
|
||||
mock_select_experts.assert_called_once()
|
||||
fused_experts_input = mock_comm.fused_experts.call_args.kwargs["fused_experts_input"]
|
||||
self.assertEqual(fused_experts_input.activation, "gelu")
|
||||
self.assertTrue(fused_experts_input.routing.apply_router_weight_on_input)
|
||||
self.assertIs(fused_experts_input.routing.mc2_mask, mc2_mask)
|
||||
self.assertIs(fused_experts_input.routing.pertoken_scale, pertoken_scale)
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a16._EXTRA_CTX")
|
||||
@patch("vllm_ascend.quantization.methods.w4a16.select_experts")
|
||||
def test_apply_router_logits_mismatch_raises(self, mock_select, mock_ctx):
|
||||
layer = self.build_layer()
|
||||
x = torch.randn(4, self.output_size, dtype=torch.float32)
|
||||
router_logits = torch.randn(4, self.experts + 1, dtype=torch.float32)
|
||||
message = (
|
||||
"Number of global experts mismatch (excluding redundancy): router_logits.shape[1]=9, num_logical_experts=8"
|
||||
)
|
||||
|
||||
with self.assertRaisesRegex(AssertionError, re.escape(message)):
|
||||
self.quant_method.apply(layer, x, router_logits, top_k=2, renormalize=True, num_experts=self.experts)
|
||||
mock_select.assert_not_called()
|
||||
247
tests/ut/quantization/methods/a2/test_w4a4_flatquant.py
Normal file
247
tests/ut/quantization/methods/a2/test_w4a4_flatquant.py
Normal file
@@ -0,0 +1,247 @@
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from tests.ut.quantization.conftest_quantization import create_linear_layer
|
||||
from vllm_ascend.quantization.methods.w4a4_flatquant import (
|
||||
KRONECKER_QUANT_MAX_BATCH_SIZE,
|
||||
AscendW4A4FlatQuantDynamicLinearMethod,
|
||||
batched_kronecker_quant,
|
||||
get_decompose_dim,
|
||||
pack_int4_weights,
|
||||
)
|
||||
|
||||
|
||||
class TestW4A4FlatQuantDynamic(unittest.TestCase):
|
||||
"""
|
||||
Unit test suite for AscendW4A4FlatQuantDynamicLinearMethod and its helper functions.
|
||||
"""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up the test environment before each test."""
|
||||
self.method = AscendW4A4FlatQuantDynamicLinearMethod()
|
||||
self.output_size = 64
|
||||
self.input_size = 768 # 768 = 24 * 32, divisible by 8
|
||||
self.params_dtype = torch.float16
|
||||
|
||||
## Test Helper Functions
|
||||
## --------------------
|
||||
|
||||
def test_get_decompose_dim(self):
|
||||
"""
|
||||
Tests the get_decompose_dim function with various inputs.
|
||||
"""
|
||||
self.assertEqual(get_decompose_dim(1024), (32, 32))
|
||||
self.assertEqual(get_decompose_dim(768), (24, 32))
|
||||
self.assertEqual(get_decompose_dim(100), (10, 10))
|
||||
self.assertEqual(get_decompose_dim(99), (9, 11))
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a4_flatquant.torch_npu")
|
||||
def test_pack_int4_weights_npu_success(self, mock_torch_npu):
|
||||
"""
|
||||
Tests weight packing using the mocked NPU kernel.
|
||||
"""
|
||||
weight_tensor = torch.randn(self.output_size, self.input_size)
|
||||
mock_packed_tensor = torch.randint(0, 100, (self.output_size, self.input_size // 8), dtype=torch.int32)
|
||||
mock_npu_tensor = MagicMock()
|
||||
mock_npu_tensor.to.return_value = mock_packed_tensor
|
||||
mock_torch_npu.npu_convert_weight_to_int4pack.return_value = mock_npu_tensor
|
||||
with patch("torch.Tensor.npu", return_value=weight_tensor):
|
||||
result = pack_int4_weights(weight_tensor)
|
||||
|
||||
mock_torch_npu.npu_convert_weight_to_int4pack.assert_called_once()
|
||||
self.assertTrue(torch.equal(result, mock_packed_tensor))
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a4_flatquant.torch_npu")
|
||||
def test_large_batch_multiple_calls(self, mock_npu):
|
||||
batch_size = 50000
|
||||
x = torch.randn(batch_size, 24, 32)
|
||||
left_trans = torch.randn(24, 24)
|
||||
right_trans = torch.randn(32, 32)
|
||||
num_chunks = batch_size // KRONECKER_QUANT_MAX_BATCH_SIZE + 1
|
||||
mock_returns = [
|
||||
(
|
||||
torch.randint(0, 255, (KRONECKER_QUANT_MAX_BATCH_SIZE, 24, 4), dtype=torch.int32),
|
||||
torch.randn(KRONECKER_QUANT_MAX_BATCH_SIZE),
|
||||
)
|
||||
for _ in range(num_chunks - 1)
|
||||
]
|
||||
last_chunk_size = batch_size - (num_chunks - 1) * KRONECKER_QUANT_MAX_BATCH_SIZE
|
||||
mock_returns.append(
|
||||
(torch.randint(0, 255, (last_chunk_size, 24, 4), dtype=torch.int32), torch.randn(last_chunk_size))
|
||||
)
|
||||
mock_npu.npu_kronecker_quant.side_effect = mock_returns
|
||||
batched_kronecker_quant(x, left_trans, right_trans, 0.95)
|
||||
self.assertEqual(mock_npu.npu_kronecker_quant.call_count, num_chunks)
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a4_flatquant.torch_npu")
|
||||
def test_exact_max_batch_size(self, mock_npu):
|
||||
batch_size = KRONECKER_QUANT_MAX_BATCH_SIZE
|
||||
x = torch.randn(batch_size, 24, 32)
|
||||
left_trans = torch.randn(24, 24)
|
||||
right_trans = torch.randn(32, 32)
|
||||
mock_npu.npu_kronecker_quant.return_value = (
|
||||
torch.randint(0, 255, (batch_size, 24, 4), dtype=torch.int32),
|
||||
torch.randn(batch_size, dtype=torch.float32),
|
||||
)
|
||||
batched_kronecker_quant(x, left_trans, right_trans, 0.95)
|
||||
mock_npu.npu_kronecker_quant.assert_called_once()
|
||||
|
||||
## Test AscendW4A4FlatQuantDynamicLinearMethod Class
|
||||
## --------------------------------------------------
|
||||
|
||||
def test_get_weight(self):
|
||||
"""Tests the get_weight static method for correct output."""
|
||||
params = self.method.get_weight(self.input_size, self.output_size, self.params_dtype)
|
||||
self.assertIn("weight", params)
|
||||
self.assertEqual(params["weight"].shape, (self.output_size, self.input_size))
|
||||
self.assertEqual(params["weight"].dtype, torch.int8)
|
||||
self.assertEqual(AscendW4A4FlatQuantDynamicLinearMethod.input_size, self.input_size)
|
||||
|
||||
def test_get_weight_value_error(self):
|
||||
"""Tests that get_weight raises ValueError for invalid input_size."""
|
||||
with self.assertRaisesRegex(ValueError, "must be divisible by 8"):
|
||||
self.method.get_weight(127, self.output_size, self.params_dtype)
|
||||
|
||||
def test_get_pertensor_param(self):
|
||||
"""Tests the get_pertensor_param static method."""
|
||||
self.method.get_weight(self.input_size, self.output_size, self.params_dtype)
|
||||
params = self.method.get_pertensor_param(self.params_dtype)
|
||||
left_dim, right_dim = get_decompose_dim(self.input_size)
|
||||
self.assertIn("left_trans", params)
|
||||
self.assertIn("right_trans", params)
|
||||
self.assertIn("clip_ratio", params)
|
||||
self.assertEqual(params["left_trans"].shape, (left_dim, left_dim))
|
||||
self.assertEqual(params["right_trans"].shape, (right_dim, right_dim))
|
||||
self.assertEqual(params["clip_ratio"].shape, (1,))
|
||||
self.assertEqual(params["left_trans"].dtype, self.params_dtype)
|
||||
self.assertEqual(params["clip_ratio"].dtype, torch.float32)
|
||||
|
||||
def test_get_perchannel_param(self):
|
||||
"""Tests the get_perchannel_param static method."""
|
||||
params = self.method.get_perchannel_param(self.output_size, self.params_dtype)
|
||||
self.assertIn("weight_scale", params)
|
||||
self.assertIn("weight_offset", params)
|
||||
self.assertEqual(params["weight_scale"].shape, (self.output_size, 1))
|
||||
self.assertEqual(params["weight_offset"].shape, (self.output_size, 1))
|
||||
self.assertEqual(params["weight_scale"].dtype, torch.float32)
|
||||
self.assertEqual(params["weight_offset"].dtype, torch.float32)
|
||||
|
||||
def test_get_pergroup_param(self):
|
||||
"""Tests the get_pergroup_param method."""
|
||||
params = self.method.get_pergroup_param(self.input_size, self.output_size, self.params_dtype)
|
||||
self.assertEqual(params, {})
|
||||
|
||||
def _prepare_apply_mocks_and_layer(self, batch_size):
|
||||
"""Helper to create a mock layer and input tensor for apply tests."""
|
||||
layer = nn.Module()
|
||||
m, n = get_decompose_dim(self.input_size)
|
||||
layer.left_trans = torch.randn(m, m, dtype=self.params_dtype)
|
||||
layer.right_trans = torch.randn(n, n, dtype=self.params_dtype)
|
||||
layer.aclnn_clip_ratio = 0.95
|
||||
layer.weight_packed = torch.randint(-8, 7, (self.output_size, self.input_size // 8), dtype=torch.int32)
|
||||
layer.weight_scale = torch.randn(self.output_size, 1, dtype=torch.float32)
|
||||
x = torch.randn(batch_size, self.input_size, dtype=self.params_dtype)
|
||||
return layer, x, m, n
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a4_flatquant.torch_npu")
|
||||
def test_apply_small_batch(self, mock_torch_npu):
|
||||
"""Tests the apply method with a batch size smaller than MAX_BATCH_SIZE."""
|
||||
batch_size = 128
|
||||
layer, x, m, n = self._prepare_apply_mocks_and_layer(batch_size)
|
||||
mock_quant_x = torch.randint(0, 255, (batch_size, self.input_size // 8), dtype=torch.int32)
|
||||
mock_act_scale = torch.randn(batch_size, 1, dtype=torch.float32)
|
||||
mock_torch_npu.npu_kronecker_quant.return_value = (mock_quant_x.view(batch_size, m, n // 8), mock_act_scale)
|
||||
mock_output = torch.randn(batch_size, self.output_size, dtype=self.params_dtype)
|
||||
mock_torch_npu.npu_quant_matmul.return_value = mock_output
|
||||
bias = torch.randn(self.output_size, dtype=self.params_dtype)
|
||||
output = self.method.apply(layer, x, bias=bias)
|
||||
mock_torch_npu.npu_kronecker_quant.assert_called_once()
|
||||
mock_torch_npu.npu_quant_matmul.assert_called_once()
|
||||
self.assertTrue(torch.allclose(output, mock_output + bias.to(self.params_dtype)))
|
||||
self.assertEqual(output.shape, (batch_size, self.output_size))
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a4_flatquant.KRONECKER_QUANT_MAX_BATCH_SIZE", 10)
|
||||
@patch("vllm_ascend.quantization.methods.w4a4_flatquant.torch_npu")
|
||||
def test_apply_large_batch(self, mock_torch_npu):
|
||||
"""Tests the apply method with a batch size larger than MAX_BATCH_SIZE."""
|
||||
batch_size = 25
|
||||
layer, x, m, n = self._prepare_apply_mocks_and_layer(batch_size)
|
||||
mock_quant_x = torch.randint(0, 255, (batch_size, self.input_size // 8), dtype=torch.int32)
|
||||
mock_act_scale = torch.randn(batch_size, 1, dtype=torch.float32)
|
||||
mock_torch_npu.npu_kronecker_quant.side_effect = [
|
||||
(mock_quant_x[:10].view(10, m, n // 8), mock_act_scale[:10]),
|
||||
(mock_quant_x[10:20].view(10, m, n // 8), mock_act_scale[10:20]),
|
||||
(mock_quant_x[20:].view(5, m, n // 8), mock_act_scale[20:]),
|
||||
]
|
||||
mock_output = torch.randn(batch_size, self.output_size, dtype=self.params_dtype)
|
||||
mock_torch_npu.npu_quant_matmul.return_value = mock_output
|
||||
output = self.method.apply(layer, x, bias=None)
|
||||
self.assertEqual(mock_torch_npu.npu_kronecker_quant.call_count, 3)
|
||||
mock_torch_npu.npu_quant_matmul.assert_called_once()
|
||||
self.assertTrue(torch.equal(output, mock_output))
|
||||
self.assertEqual(output.shape, (batch_size, self.output_size))
|
||||
|
||||
def test_apply_dimension_mismatch_error(self):
|
||||
"""Tests that apply raises ValueError on transform matrix dimension mismatch."""
|
||||
layer, x, _, _ = self._prepare_apply_mocks_and_layer(16)
|
||||
layer.left_trans = torch.randn(20, 20)
|
||||
layer.right_trans = torch.randn(30, 30) # 20 * 30 != 768
|
||||
with self.assertRaisesRegex(ValueError, "FlatQuant transform matrices dimension mismatch"):
|
||||
self.method.apply(layer, x)
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a4_flatquant.pack_int4_weights")
|
||||
def test_process_weights_after_loading(self, mock_pack_weights):
|
||||
"""Tests weight processing after loading, without transpose."""
|
||||
layer = nn.Module()
|
||||
layer.weight = torch.randint(-8, 7, (self.output_size, self.input_size), dtype=torch.int8)
|
||||
layer.weight_scale = torch.randn(self.output_size, 1, dtype=torch.bfloat16)
|
||||
layer.weight_offset = torch.randn(self.output_size, 1, dtype=torch.bfloat16)
|
||||
layer.left_trans = torch.randn(24, 24)
|
||||
layer.right_trans = torch.randn(32, 32)
|
||||
layer.clip_ratio = torch.tensor([0.9])
|
||||
mock_packed = torch.randint(0, 100, (self.output_size, self.input_size // 8), dtype=torch.int32)
|
||||
mock_pack_weights.return_value = mock_packed
|
||||
self.method.process_weights_after_loading(layer)
|
||||
mock_pack_weights.assert_called_once()
|
||||
self.assertFalse(hasattr(layer, "weight"))
|
||||
self.assertTrue(hasattr(layer, "weight_packed"))
|
||||
self.assertTrue(torch.equal(layer.weight_packed.data, mock_packed))
|
||||
self.assertEqual(layer.weight_scale.dtype, torch.float32)
|
||||
self.assertEqual(layer.weight_offset.dtype, torch.float32)
|
||||
self.assertEqual(layer.clip_ratio.dtype, torch.float32)
|
||||
self.assertTrue(layer.aclnn_clip_ratio - 0.9 < 0.01)
|
||||
self.assertEqual(layer.left_trans.shape, (24, 24))
|
||||
self.assertTrue(layer.left_trans.is_contiguous())
|
||||
|
||||
|
||||
class TestW4A4FlatQuantDynamicWithNpu(unittest.TestCase):
|
||||
"""
|
||||
Unit test suite for AscendW4A4FlatQuantDynamicLinearMethod and its helper functions.
|
||||
"""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up the test environment before each test."""
|
||||
self.method = AscendW4A4FlatQuantDynamicLinearMethod()
|
||||
self.output_size = 64
|
||||
self.input_size = 768 # 768 = 24 * 32, divisible by 8
|
||||
self.params_dtype = torch.bfloat16
|
||||
|
||||
def test_apply_with_npu(self):
|
||||
"""Tests the apply method with NPU."""
|
||||
batch_size = 16
|
||||
layer = create_linear_layer(self.method, self.input_size, self.output_size, self.params_dtype)
|
||||
layer.clip_ratio = nn.Parameter(torch.tensor([0.95], dtype=torch.float32).npu(), requires_grad=False)
|
||||
self.method.process_weights_after_loading(layer)
|
||||
|
||||
x = torch.randn(batch_size, self.input_size, dtype=self.params_dtype).npu()
|
||||
output = self.method.apply(layer, x)
|
||||
|
||||
self.assertEqual(output.shape, (batch_size, self.output_size))
|
||||
self.assertEqual(output.dtype, self.params_dtype)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(argv=["first-arg-is-ignored"], exit=False)
|
||||
78
tests/ut/quantization/methods/a2/test_w4a4_laos_dynamic.py
Normal file
78
tests/ut/quantization/methods/a2/test_w4a4_laos_dynamic.py
Normal file
@@ -0,0 +1,78 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
from tests.ut.quantization.conftest_quantization import create_linear_layer
|
||||
from vllm_ascend.quantization.methods.w4a4_laos_dynamic import AscendW4A4LaosDynamicLinearMethod
|
||||
|
||||
|
||||
class TestAscendW4A4LaosDynamicLinearMethod(TestBase):
|
||||
def setUp(self):
|
||||
self.method = AscendW4A4LaosDynamicLinearMethod()
|
||||
|
||||
def test_get_weight_various_sizes(self):
|
||||
sizes = [(64, 128), (256, 512), (1024, 2048)]
|
||||
for input_size, output_size in sizes:
|
||||
result = self.method.get_weight(input_size, output_size, torch.bfloat16)
|
||||
self.assertEqual(result["weight"].shape, (output_size, input_size))
|
||||
self.assertEqual(result["weight"].dtype, torch.int8)
|
||||
|
||||
def test_get_perchannel_param_various_output_sizes(self):
|
||||
output_sizes = [1, 64, 128, 512]
|
||||
for output_size in output_sizes:
|
||||
result = self.method.get_perchannel_param(output_size, torch.bfloat16)
|
||||
self.assertEqual(result["weight_scale"].shape, (output_size, 1))
|
||||
self.assertEqual(result["weight_offset"].shape, (output_size, 1))
|
||||
self.assertEqual(result["weight_scale"].dtype, torch.float32)
|
||||
self.assertEqual(result["weight_offset"].dtype, torch.float32)
|
||||
|
||||
@patch("torch_npu.npu_quant_matmul")
|
||||
@patch("torch_npu.npu_dynamic_quant")
|
||||
def test_apply_with_bias(self, mock_dyn_quant, mock_matmul):
|
||||
mock_dyn_quant.return_value = (
|
||||
torch.randint(0, 15, (32, 128), dtype=torch.int32),
|
||||
torch.randn(32, dtype=torch.float32),
|
||||
)
|
||||
expected_output = torch.randn(32, 256, dtype=torch.bfloat16)
|
||||
mock_matmul.return_value = expected_output
|
||||
layer = MagicMock()
|
||||
layer.weight = MagicMock(data=torch.randint(-8, 7, (256, 128), dtype=torch.int8))
|
||||
layer.weight_scale = MagicMock(data=torch.randn(256, dtype=torch.float32))
|
||||
x = torch.randn(32, 128, dtype=torch.bfloat16)
|
||||
bias = torch.randn(256, dtype=torch.bfloat16)
|
||||
output = self.method.apply(layer, x, bias)
|
||||
expected_output = expected_output + bias
|
||||
self.assertTrue(torch.equal(output, expected_output))
|
||||
|
||||
@patch("torch_npu.npu_convert_weight_to_int4pack")
|
||||
def test_process_weights_various_input_sizes(self, mock_convert):
|
||||
for input_size, output_size in [(64, 128), (256, 512)]:
|
||||
mock_convert.return_value = torch.randint(0, 15, (output_size, input_size // 8), dtype=torch.int32)
|
||||
layer = nn.Module()
|
||||
layer.weight = nn.Parameter(
|
||||
torch.randint(-8, 7, (output_size, input_size), dtype=torch.int8), requires_grad=False
|
||||
)
|
||||
layer.weight_scale = nn.Parameter(torch.randn(output_size, 1, dtype=torch.float32), requires_grad=False)
|
||||
self.method.process_weights_after_loading(layer)
|
||||
mock_convert.assert_called()
|
||||
self.assertEqual(layer.weight_scale.data.dtype, torch.float32)
|
||||
self.assertEqual(layer.weight.shape, (input_size // 8, output_size))
|
||||
|
||||
|
||||
class TestAscendW4A4LaosDynamicLinearMethodWithNpu(TestBase):
|
||||
def setUp(self):
|
||||
self.method = AscendW4A4LaosDynamicLinearMethod()
|
||||
|
||||
def test_apply_with_npu(self):
|
||||
token_num = 32
|
||||
input_size, output_size = 128, 256
|
||||
params_dtype = torch.bfloat16
|
||||
layer = create_linear_layer(self.method, input_size, output_size, params_dtype)
|
||||
self.method.process_weights_after_loading(layer)
|
||||
|
||||
x = torch.randn(token_num, input_size, dtype=params_dtype).npu()
|
||||
bias = torch.randn(output_size, dtype=params_dtype).npu()
|
||||
output = self.method.apply(layer, x, bias)
|
||||
self.assertEqual(output.shape, (token_num, output_size))
|
||||
456
tests/ut/quantization/methods/a2/test_w4a8.py
Normal file
456
tests/ut/quantization/methods/a2/test_w4a8.py
Normal file
@@ -0,0 +1,456 @@
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import regex as re
|
||||
import torch
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
from tests.ut.quantization.conftest_quantization import identity
|
||||
from vllm_ascend.quantization.methods.w4a8 import AscendW4A8DynamicFusedMoEMethod, AscendW4A8DynamicLinearMethod
|
||||
from vllm_ascend.utils import COMPRESSED_TENSORS_METHOD
|
||||
|
||||
|
||||
class TestAscendW4A8DynamicLinearMethod(TestBase):
|
||||
@patch("vllm_ascend.quantization.methods.w4a8.get_tensor_model_parallel_world_size")
|
||||
@patch("vllm_ascend.quantization.methods.w4a8.get_current_vllm_config")
|
||||
def setUp(self, mock_get_current_vllm_config, mock_get_tp_world_size):
|
||||
mock_get_tp_world_size.return_value = 1
|
||||
mock_vllm_config = Mock()
|
||||
mock_vllm_config.quant_config = Mock(quant_description={"group_size": 256})
|
||||
mock_vllm_config.scheduler_config = Mock(
|
||||
max_num_batched_tokens=2048, max_model_len=2048, enable_chunked_prefill=False
|
||||
)
|
||||
mock_get_current_vllm_config.return_value = mock_vllm_config
|
||||
self.method = AscendW4A8DynamicLinearMethod()
|
||||
self.method.group_size = 8
|
||||
|
||||
def test_get_weight(self):
|
||||
weight = self.method.get_weight(8, 32, torch.bfloat16)
|
||||
self.assertEqual(weight["weight"].dtype, torch.int8)
|
||||
self.assertEqual(weight["weight"].shape, (32, 8))
|
||||
# new quant version weight
|
||||
self.method.new_quant_version = True
|
||||
weight = self.method.get_weight(8, 32, torch.bfloat16)
|
||||
self.assertEqual(weight["weight"].dtype, torch.int8)
|
||||
self.assertEqual(weight["weight"].shape, (16, 8))
|
||||
self.assertEqual(weight["_packed_dim"], 0)
|
||||
self.assertEqual(weight["_packed_factor"], 2)
|
||||
|
||||
def test_get_pergroup_param(self):
|
||||
params = self.method.get_pergroup_param(8, 32, torch.bfloat16)
|
||||
self.assertEqual(params["weight_scale"].dtype, torch.bfloat16)
|
||||
self.assertEqual(params["weight_scale"].shape, (32, 1))
|
||||
self.assertEqual(params["weight_offset"].dtype, torch.bfloat16)
|
||||
self.assertEqual(params["weight_offset"].shape, (32, 1))
|
||||
self.assertEqual(params["weight_scale_second"].dtype, torch.bfloat16)
|
||||
self.assertEqual(params["weight_scale_second"].shape, (32, 1))
|
||||
self.assertEqual(params["weight_offset_second"].dtype, torch.bfloat16)
|
||||
self.assertEqual(params["weight_offset_second"].shape, (32, 1))
|
||||
# new quant version weight
|
||||
self.method.new_quant_version = True
|
||||
params = self.method.get_pergroup_param(8, 32, torch.bfloat16, layer_type="column")
|
||||
self.assertEqual(params["scale_bias"].dtype, torch.float32)
|
||||
self.assertEqual(params["scale_bias"].shape, (32, 1))
|
||||
params = self.method.get_pergroup_param(8, 32, torch.bfloat16, layer_type="row")
|
||||
self.assertEqual(params["scale_bias"].dtype, torch.float32)
|
||||
self.assertEqual(params["scale_bias"].shape, (32, 16))
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a8.maybe_trans_nz")
|
||||
@patch("torch_npu.npu_convert_weight_to_int4pack")
|
||||
@patch("torch.Tensor.npu")
|
||||
@patch("torch_npu.npu_format_cast")
|
||||
def test_process_weights_after_loading(
|
||||
self, mock_format_cast, mock_npu, mock_npu_convert_weight, mock_maybe_trans_nz
|
||||
):
|
||||
mock_npu.side_effect = lambda: torch.zeros((1, 32), dtype=torch.float32)
|
||||
mock_npu_convert_weight.return_value = torch.zeros((32, 4), dtype=torch.int32)
|
||||
mock_maybe_trans_nz.side_effect = identity
|
||||
# old quant version weight
|
||||
layer = torch.nn.Module()
|
||||
layer.weight = torch.nn.Parameter(torch.zeros((32, 8), dtype=torch.int8), requires_grad=False)
|
||||
layer.weight_scale = torch.nn.Parameter(torch.ones((32, 1), dtype=torch.float32), requires_grad=False)
|
||||
layer.weight_offset = torch.nn.Parameter(torch.empty_like(layer.weight_scale.data), requires_grad=False)
|
||||
layer.weight_scale_second = torch.nn.Parameter(torch.ones((32, 1), dtype=torch.float32), requires_grad=False)
|
||||
layer.weight_offset_second = torch.nn.Parameter(
|
||||
torch.empty_like(layer.weight_scale_second.data), requires_grad=False
|
||||
)
|
||||
mock_format_cast.return_value = layer.weight.data.transpose(0, 1).contiguous()
|
||||
self.method.process_weights_after_loading(layer)
|
||||
self.assertTrue(hasattr(layer, "weight_scale_bias"))
|
||||
self.assertEqual(layer.weight_scale_bias.data.shape, (32,))
|
||||
self.assertEqual(layer.weight_scale_bias.data.dtype, torch.float32)
|
||||
# new quant version weight
|
||||
self.method.new_quant_version = True
|
||||
new_layer = torch.nn.Module()
|
||||
new_layer.weight = torch.nn.Parameter(torch.zeros((16, 8), dtype=torch.int8), requires_grad=False)
|
||||
new_layer.weight_scale = torch.nn.Parameter(torch.ones((32, 1), dtype=torch.float32), requires_grad=False)
|
||||
new_layer.weight_offset = torch.nn.Parameter(torch.empty_like(new_layer.weight_scale.data), requires_grad=False)
|
||||
new_layer.weight_scale_second = torch.nn.Parameter(
|
||||
torch.ones((32, 1), dtype=torch.float32), requires_grad=False
|
||||
)
|
||||
new_layer.weight_offset_second = torch.nn.Parameter(
|
||||
torch.empty_like(new_layer.weight_scale_second.data), requires_grad=False
|
||||
)
|
||||
new_layer.scale_bias = torch.nn.Parameter(torch.zeros((32, 1), dtype=torch.float32), requires_grad=False)
|
||||
mock_format_cast.return_value = new_layer.weight.data.transpose(0, 1).contiguous()
|
||||
self.method.process_weights_after_loading(new_layer)
|
||||
self.assertEqual(new_layer.scale_bias.data.shape, (32,))
|
||||
self.assertTrue(hasattr(new_layer, "weight_scale_second"))
|
||||
self.assertEqual(new_layer.weight_scale_second.data.shape, (1, 32))
|
||||
|
||||
@patch("torch_npu.npu_weight_quant_batchmatmul")
|
||||
def test_apply_basic(self, mock_matmul):
|
||||
layer = MagicMock()
|
||||
layer.weight = MagicMock(data=torch.randint(-8, 8, (256, 512), dtype=torch.int8))
|
||||
layer.weight_scale_second = MagicMock(data=torch.randn(1, 512, dtype=torch.float32))
|
||||
mock_matmul.return_value = torch.randn(32, 512)
|
||||
x = torch.randn(32, 256)
|
||||
self.method.apply(layer, x)
|
||||
mock_matmul.assert_called_once()
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a8.maybe_trans_nz")
|
||||
def test_process_weights_after_loading_asserts_new_quant_packed_dim(self, mock_maybe_trans_nz):
|
||||
self.method.new_quant_version = True
|
||||
mock_maybe_trans_nz.side_effect = identity
|
||||
layer = torch.nn.Module()
|
||||
layer.weight = torch.nn.Parameter(torch.zeros((10, 16), dtype=torch.int8), requires_grad=False)
|
||||
layer.weight_scale = torch.nn.Parameter(torch.ones((20, 1), dtype=torch.float32), requires_grad=False)
|
||||
layer.weight_offset = torch.nn.Parameter(torch.empty_like(layer.weight_scale.data), requires_grad=False)
|
||||
layer.weight_scale_second = torch.nn.Parameter(torch.ones((20, 2), dtype=torch.float32), requires_grad=False)
|
||||
layer.weight_offset_second = torch.nn.Parameter(
|
||||
torch.empty_like(layer.weight_scale_second.data), requires_grad=False
|
||||
)
|
||||
layer.scale_bias = torch.nn.Parameter(torch.zeros((20, 1), dtype=torch.float32), requires_grad=False)
|
||||
expected_message = "the last dim of weight needs to be divided by 4 but got shape torch.Size([16, 10])"
|
||||
|
||||
with (
|
||||
patch.object(self.method, "process_scale_second", return_value=(torch.ones((2, 20)), None)),
|
||||
self.assertRaisesRegex(AssertionError, re.escape(expected_message)),
|
||||
):
|
||||
self.method.process_weights_after_loading(layer)
|
||||
|
||||
|
||||
class TestAscendW4A8DynamicLinearMethodWithNpu(TestBase):
|
||||
@patch("vllm_ascend.quantization.methods.w4a8.get_tensor_model_parallel_world_size")
|
||||
@patch("vllm_ascend.quantization.methods.w4a8.get_current_vllm_config")
|
||||
def setUp(self, mock_get_current_vllm_config, mock_get_tp_world_size):
|
||||
mock_get_tp_world_size.return_value = 1
|
||||
mock_vllm_config = Mock()
|
||||
mock_vllm_config.quant_config = Mock(quant_description={"group_size": 64})
|
||||
mock_get_current_vllm_config.return_value = mock_vllm_config
|
||||
self.method = AscendW4A8DynamicLinearMethod()
|
||||
|
||||
def test_apply_with_npu(self):
|
||||
layer = torch.nn.Module()
|
||||
layer.weight = torch.nn.Parameter(
|
||||
torch.randint(-128, 127, (128, 32), dtype=torch.int32).npu(), requires_grad=False
|
||||
)
|
||||
layer.weight_scale_second = torch.nn.Parameter(
|
||||
torch.randn(2, 256, dtype=torch.bfloat16).npu(), requires_grad=False
|
||||
)
|
||||
|
||||
x = torch.randn(32, 128, dtype=torch.bfloat16).npu()
|
||||
output = self.method.apply(layer, x)
|
||||
self.assertEqual(output.shape, (32, 256))
|
||||
|
||||
|
||||
class TestAscendW4A8DynamicFusedMoEMethod(TestBase):
|
||||
experts = 8
|
||||
input_size = 16
|
||||
output_size = 56
|
||||
group_size = 2
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a8.get_ascend_config")
|
||||
@patch("vllm_ascend.quantization.methods.w4a8.get_current_vllm_config")
|
||||
@patch("vllm_ascend.quantization.methods.w4a8.get_mc2_group")
|
||||
@patch("torch.distributed.get_rank", return_value=0)
|
||||
def setUp(self, mock_get_rank, mock_get_mc2_group, get_current_vllm_config, mock_get_ascend_config):
|
||||
# Mock ascend config
|
||||
mock_ascend_config = Mock()
|
||||
mock_ascend_config.eplb_config.dynamic_eplb = False
|
||||
mock_get_ascend_config.return_value = mock_ascend_config
|
||||
|
||||
mock_vllm_config = Mock()
|
||||
mock_vllm_config.quant_config = Mock(quant_description={"group_size": self.group_size, "version": "0.0.0"})
|
||||
mock_vllm_config.parallel_config = Mock(enable_expert_parallel=True)
|
||||
mock_vllm_config.scheduler_config = Mock(
|
||||
max_num_batched_tokens=2048, max_model_len=2048, enable_chunked_prefill=False
|
||||
)
|
||||
get_current_vllm_config.return_value = mock_vllm_config
|
||||
self.quant_method = AscendW4A8DynamicFusedMoEMethod()
|
||||
|
||||
def test_get_weight(self):
|
||||
# old quant version w4a8 weight
|
||||
param_dict = self.quant_method.get_weight(self.experts, self.input_size, self.output_size, torch.bfloat16)
|
||||
self.assertEqual(param_dict["w13_weight"].dtype, torch.int8)
|
||||
self.assertEqual(param_dict["w13_weight"].shape, (self.experts, 2 * self.input_size, self.output_size))
|
||||
# new quant version weight
|
||||
self.quant_method.new_quant_version = True
|
||||
param_dict = self.quant_method.get_weight(self.experts, self.input_size, self.output_size, torch.bfloat16)
|
||||
self.assertEqual(param_dict["w13_weight"].dtype, torch.int8)
|
||||
self.assertEqual(param_dict["w13_weight"].shape, (self.experts, self.input_size, self.output_size))
|
||||
|
||||
def test_get_dynamic_quant_param(self):
|
||||
# old quant version weight
|
||||
param_dict = self.quant_method.get_dynamic_quant_param(
|
||||
self.experts, self.input_size, self.output_size, torch.bfloat16
|
||||
)
|
||||
self.assertEqual(param_dict["w13_weight_scale"].dtype, torch.float32)
|
||||
self.assertEqual(param_dict["w13_weight_scale"].shape, (self.experts, 2 * self.input_size, 1))
|
||||
self.assertEqual(param_dict["w13_weight_scale_second"].dtype, torch.float32)
|
||||
self.assertEqual(
|
||||
param_dict["w13_weight_scale_second"].shape,
|
||||
(self.experts, 2 * self.input_size, self.output_size // self.group_size),
|
||||
)
|
||||
self.assertEqual(param_dict["w2_weight_scale"].dtype, torch.float32)
|
||||
self.assertEqual(param_dict["w2_weight_scale"].shape, (self.experts, self.output_size, 1))
|
||||
self.assertEqual(param_dict["w2_weight_scale_second"].dtype, torch.float32)
|
||||
self.assertEqual(
|
||||
param_dict["w2_weight_scale_second"].shape,
|
||||
(self.experts, self.output_size, self.input_size // self.group_size),
|
||||
)
|
||||
# new quant version weight
|
||||
self.quant_method.new_quant_version = True
|
||||
param_dict = self.quant_method.get_dynamic_quant_param(
|
||||
self.experts, self.input_size, self.output_size, torch.bfloat16
|
||||
)
|
||||
self.assertEqual(param_dict["w2_scale_bias"].dtype, torch.float32)
|
||||
self.assertEqual(
|
||||
param_dict["w2_scale_bias"].shape, (self.experts, self.output_size, 16 // self.quant_method.tp_size)
|
||||
)
|
||||
# per-channel weight
|
||||
self.quant_method.is_per_channel_weight = True
|
||||
param_dict = self.quant_method.get_dynamic_quant_param(
|
||||
self.experts, self.input_size, self.output_size, torch.bfloat16
|
||||
)
|
||||
pergroup_param = [
|
||||
"w13_weight_scale_second",
|
||||
"w13_weight_offset_second",
|
||||
"w2_weight_scale_second",
|
||||
"w2_weight_offset_second",
|
||||
]
|
||||
is_contains = any(key in param_dict for key in pergroup_param)
|
||||
self.assertFalse(is_contains)
|
||||
|
||||
def build_layer(self, is_new_quant_version=True, is_per_channel_weight=False):
|
||||
layer = torch.nn.Module()
|
||||
if is_new_quant_version:
|
||||
layer.w13_weight = torch.nn.Parameter(
|
||||
torch.zeros((self.experts, self.input_size, self.output_size), dtype=torch.int8), requires_grad=False
|
||||
)
|
||||
layer.w2_weight = torch.nn.Parameter(
|
||||
torch.zeros((self.experts, self.output_size // 2, self.input_size), dtype=torch.int8),
|
||||
requires_grad=False,
|
||||
)
|
||||
w13_scale_bias = torch.zeros((self.experts, 2 * self.input_size, 1), dtype=torch.float32)
|
||||
layer.w13_scale_bias = torch.nn.Parameter(w13_scale_bias, requires_grad=False)
|
||||
w2_scale_bias = torch.zeros(
|
||||
(self.experts, self.output_size, 16 // self.quant_method.tp_size), dtype=torch.float32
|
||||
)
|
||||
layer.w2_scale_bias = torch.nn.Parameter(w2_scale_bias, requires_grad=False)
|
||||
else:
|
||||
layer.w13_weight = torch.nn.Parameter(
|
||||
torch.zeros((self.experts, 2 * self.input_size, self.output_size), dtype=torch.int8),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w2_weight = torch.nn.Parameter(
|
||||
torch.zeros((self.experts, self.output_size, self.input_size), dtype=torch.int8), requires_grad=False
|
||||
)
|
||||
layer.w13_weight_scale = torch.nn.Parameter(
|
||||
torch.ones((self.experts, 2 * self.input_size, 1), dtype=torch.float32), requires_grad=False
|
||||
)
|
||||
layer.w2_weight_scale = torch.nn.Parameter(
|
||||
torch.ones((self.experts, self.output_size, 1), dtype=torch.float32), requires_grad=False
|
||||
)
|
||||
if not is_per_channel_weight:
|
||||
layer.w13_weight_scale_second = torch.nn.Parameter(
|
||||
torch.ones(
|
||||
(self.experts, 2 * self.input_size, self.output_size // self.group_size), dtype=torch.float32
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w13_weight_offset_second = torch.nn.Parameter(
|
||||
torch.empty_like(layer.w13_weight_scale_second.data), requires_grad=False
|
||||
)
|
||||
layer.w2_weight_scale_second = torch.nn.Parameter(
|
||||
torch.ones((self.experts, self.output_size, self.input_size // self.group_size), dtype=torch.float32),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w2_weight_offset_second = torch.nn.Parameter(
|
||||
torch.empty_like(layer.w2_weight_scale_second.data), requires_grad=False
|
||||
)
|
||||
return layer
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a8.maybe_trans_nz")
|
||||
@patch("torch_npu.npu_format_cast")
|
||||
@patch("torch_npu.npu_quantize")
|
||||
@patch("torch.Tensor.npu", new=lambda self: self)
|
||||
def test_process_weights_after_loading(self, mock_npu_quantize, mock_npu_format_cast, mock_maybe_trans_nz):
|
||||
mock_npu_quantize.return_value = torch.Tensor()
|
||||
mock_npu_format_cast.side_effect = identity
|
||||
mock_maybe_trans_nz.side_effect = identity
|
||||
# old quant version weight
|
||||
layer = self.build_layer(is_new_quant_version=False)
|
||||
self.quant_method.process_weights_after_loading(layer)
|
||||
self.assertTrue(hasattr(layer, "w13_scale_bias"))
|
||||
self.assertEqual(layer.w13_scale_bias.data.shape, (self.experts, 2 * self.input_size))
|
||||
self.assertEqual(layer.w13_scale_bias.data.dtype, torch.float32)
|
||||
self.assertTrue(hasattr(layer, "w2_scale_bias"))
|
||||
self.assertEqual(layer.w2_scale_bias.data.shape, (self.experts, self.output_size))
|
||||
self.assertEqual(layer.w2_scale_bias.data.dtype, torch.float32)
|
||||
# new quant version weight
|
||||
self.quant_method.new_quant_version = True
|
||||
new_layer = self.build_layer(is_new_quant_version=True)
|
||||
self.quant_method.process_weights_after_loading(new_layer)
|
||||
self.assertEqual(new_layer.w13_scale_bias.data.shape, (self.experts, 2 * self.input_size))
|
||||
self.assertEqual(new_layer.w2_scale_bias.data.shape, (self.experts, self.output_size))
|
||||
self.assertFalse(hasattr(new_layer, "w13_weight_scale_second"))
|
||||
# per-channel weight
|
||||
self.quant_method.is_per_channel_weight = True
|
||||
per_channel_layer = self.build_layer(is_new_quant_version=True, is_per_channel_weight=True)
|
||||
self.quant_method.process_weights_after_loading(per_channel_layer)
|
||||
self.assertEqual(new_layer.w13_scale_bias.data.shape, (self.experts, 2 * self.input_size))
|
||||
self.assertEqual(per_channel_layer.w13_weight_scale.data.shape, (self.experts, 2 * self.input_size))
|
||||
|
||||
def test_pack_to_int32_asserts_new_quant_packed_dim(self):
|
||||
self.quant_method.new_quant_version = True
|
||||
weight = torch.zeros((self.experts, self.output_size, 10), dtype=torch.int8)
|
||||
expected_message = f"the last dim of weight needs to be divided by 4 but got shape {weight.shape}"
|
||||
|
||||
with self.assertRaisesRegex(AssertionError, re.escape(expected_message)):
|
||||
self.quant_method.pack_to_int32(weight)
|
||||
|
||||
def test_get_weight_compressed_tensors(self):
|
||||
self.quant_method.quant_method = COMPRESSED_TENSORS_METHOD
|
||||
result = self.quant_method.get_weight(self.experts, self.input_size, self.output_size, torch.bfloat16)
|
||||
self.assertEqual(result["w13_weight"].dtype, torch.int8)
|
||||
|
||||
def test_get_dynamic_quant_param_compressed_tensors(self):
|
||||
self.quant_method.quant_method = COMPRESSED_TENSORS_METHOD
|
||||
result = self.quant_method.get_dynamic_quant_param(
|
||||
self.experts, self.input_size, self.output_size, torch.bfloat16
|
||||
)
|
||||
self.assertIn("w13_weight_scale", result)
|
||||
self.assertIn("w2_weight_scale", result)
|
||||
self.assertEqual(result["w13_weight_scale"].dtype, torch.bfloat16)
|
||||
self.assertEqual(result["w2_weight_scale"].dtype, torch.bfloat16)
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a8.maybe_trans_nz")
|
||||
@patch("torch_npu.npu_format_cast")
|
||||
@patch("torch_npu.npu_quantize")
|
||||
@patch("torch.Tensor.npu", new=lambda self: self)
|
||||
def test_process_weights_after_loading_compressed_tensors(
|
||||
self, mock_npu_quantize, mock_npu_format_cast, mock_maybe_trans_nz
|
||||
):
|
||||
mock_npu_quantize.return_value = torch.Tensor()
|
||||
mock_npu_format_cast.side_effect = identity
|
||||
mock_maybe_trans_nz.side_effect = identity
|
||||
|
||||
layer = self.build_layer(is_new_quant_version=False)
|
||||
self.quant_method.quant_method = COMPRESSED_TENSORS_METHOD
|
||||
self.quant_method.weight_strategy = "group"
|
||||
self.quant_method.process_weights_after_loading(layer)
|
||||
self.assertTrue(hasattr(layer, "w13_scale_bias"))
|
||||
self.assertEqual(layer.w13_scale_bias.data.shape, (self.experts, 2 * self.input_size))
|
||||
self.assertEqual(layer.w13_scale_bias.data.dtype, torch.float32)
|
||||
|
||||
self.quant_method.is_per_channel_weight = True
|
||||
self.quant_method.weight_strategy = "channel"
|
||||
per_channel_layer = self.build_layer(is_new_quant_version=False)
|
||||
self.quant_method.process_weights_after_loading(per_channel_layer)
|
||||
self.assertEqual(per_channel_layer.w13_weight_scale.data.shape, (self.experts, 2 * self.input_size))
|
||||
self.assertEqual(per_channel_layer.w2_weight_scale.data.shape, (self.experts, 1, self.output_size))
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a8._EXTRA_CTX")
|
||||
@patch("vllm_ascend.quantization.methods.w4a8.select_experts")
|
||||
@patch("vllm_ascend.quantization.methods.w4a8.build_fused_experts_input")
|
||||
def test_apply_comprehensive(self, mock_build_input, mock_select, mock_ctx):
|
||||
tokens = 4
|
||||
num_experts = self.experts
|
||||
hidden_size = self.output_size
|
||||
top_k = 2
|
||||
|
||||
layer = self.build_layer(is_new_quant_version=True, is_per_channel_weight=True)
|
||||
self.quant_method.is_per_channel_weight = True
|
||||
layer.swiglu_limit = 1000000
|
||||
x = torch.randn(tokens, hidden_size, dtype=torch.bfloat16)
|
||||
router_logits = torch.randn(tokens, num_experts, dtype=torch.float32)
|
||||
topk_weights = torch.randn(tokens, top_k, dtype=torch.float32)
|
||||
topk_ids = torch.randint(0, num_experts, (tokens, top_k), dtype=torch.int64)
|
||||
expert_map = torch.randint(0, num_experts, (num_experts,), dtype=torch.int64)
|
||||
mc2_mask = torch.tensor([1, 0, 1, 0], dtype=torch.bool)
|
||||
pertoken_scale = torch.randn(tokens, dtype=torch.float32)
|
||||
log2phy = torch.randint(0, num_experts, (num_experts,), dtype=torch.int64)
|
||||
e_score_correction_bias = torch.randn(num_experts, dtype=torch.float32)
|
||||
|
||||
mock_select.return_value = (topk_weights, topk_ids)
|
||||
|
||||
mock_fused_input = Mock()
|
||||
mock_fused_input.hidden_states = x
|
||||
mock_fused_input.topk_weights = topk_weights
|
||||
mock_fused_input.topk_ids = topk_ids
|
||||
mock_fused_input.activation = "silu"
|
||||
mock_build_input.return_value = mock_fused_input
|
||||
|
||||
mock_comm = Mock()
|
||||
expected_output = torch.randn(tokens, hidden_size, dtype=torch.bfloat16)
|
||||
mock_comm.fused_experts.return_value = expected_output
|
||||
mock_ctx.moe_comm_method = mock_comm
|
||||
|
||||
output = self.quant_method.apply(
|
||||
layer=layer,
|
||||
x=x,
|
||||
router_logits=router_logits,
|
||||
top_k=top_k,
|
||||
renormalize=True,
|
||||
use_grouped_topk=False,
|
||||
num_experts=num_experts,
|
||||
expert_map=expert_map,
|
||||
scoring_func="softmax",
|
||||
routed_scaling_factor=1.0,
|
||||
e_score_correction_bias=e_score_correction_bias,
|
||||
is_prefill=True,
|
||||
enable_force_load_balance=False,
|
||||
log2phy=log2phy,
|
||||
global_redundant_expert_num=0,
|
||||
pertoken_scale=pertoken_scale,
|
||||
activation="silu",
|
||||
apply_router_weight_on_input=False,
|
||||
mc2_mask=mc2_mask,
|
||||
)
|
||||
|
||||
mock_select.assert_called_once()
|
||||
select_call_args = mock_select.call_args
|
||||
self.assertTrue(torch.equal(select_call_args.kwargs["hidden_states"], x))
|
||||
self.assertEqual(select_call_args.kwargs["top_k"], top_k)
|
||||
self.assertEqual(select_call_args.kwargs["num_experts"], num_experts)
|
||||
|
||||
mock_build_input.assert_called_once()
|
||||
build_kwargs = mock_build_input.call_args.kwargs
|
||||
self.assertTrue(torch.equal(build_kwargs["hidden_states"], x))
|
||||
self.assertEqual(build_kwargs["quant_type"], self.quant_method.quant_type)
|
||||
self.assertTrue(build_kwargs["is_per_channel_weight"])
|
||||
self.assertEqual(build_kwargs["activation"], "silu")
|
||||
self.assertEqual(build_kwargs["apply_router_weight_on_input"], False)
|
||||
|
||||
mock_comm.fused_experts.assert_called_once()
|
||||
self.assertEqual(mock_comm.fused_experts.call_args.kwargs["fused_experts_input"], mock_fused_input)
|
||||
self.assertTrue(torch.equal(output, expected_output))
|
||||
|
||||
def test_apply_asserts_router_logits_expert_mismatch(self):
|
||||
layer = self.build_layer(is_new_quant_version=True, is_per_channel_weight=True)
|
||||
x = torch.randn(4, self.output_size, dtype=torch.bfloat16)
|
||||
router_logits = torch.randn(4, self.experts - 1, dtype=torch.float32)
|
||||
expected_message = (
|
||||
"Number of global experts mismatch (excluding redundancy): "
|
||||
f"router_logits.shape[1]={self.experts - 1}, num_logical_experts={self.experts}"
|
||||
)
|
||||
|
||||
with self.assertRaisesRegex(AssertionError, re.escape(expected_message)):
|
||||
self.quant_method.apply(
|
||||
layer=layer,
|
||||
x=x,
|
||||
router_logits=router_logits,
|
||||
top_k=2,
|
||||
renormalize=True,
|
||||
num_experts=self.experts,
|
||||
)
|
||||
94
tests/ut/quantization/methods/a2/test_w8a16.py
Normal file
94
tests/ut/quantization/methods/a2/test_w8a16.py
Normal file
@@ -0,0 +1,94 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
from tests.ut.quantization.conftest_quantization import create_linear_layer, identity
|
||||
from vllm_ascend.quantization.methods.w8a16 import AscendW8A16LinearMethod
|
||||
|
||||
|
||||
class TestAscendW8A16LinearMethod(TestBase):
|
||||
def setUp(self):
|
||||
self.method = AscendW8A16LinearMethod()
|
||||
|
||||
def test_get_weight(self):
|
||||
sizes = [(64, 128), (256, 512), (1024, 2048), (1, 1)]
|
||||
for input_size, output_size in sizes:
|
||||
weight = self.method.get_weight(input_size, output_size)
|
||||
self.assertEqual(weight["weight"].dtype, torch.int8)
|
||||
self.assertEqual(weight["weight"].shape, (output_size, input_size))
|
||||
self.assertEqual(len(weight), 1)
|
||||
|
||||
weight = self.method.get_weight(256, 128, torch.float16)
|
||||
self.assertEqual(weight["weight"].dtype, torch.int8)
|
||||
|
||||
def test_get_per_channel_param(self):
|
||||
for output_size, dtype in [(128, torch.bfloat16), (256, torch.float16)]:
|
||||
per_channel_params = self.method.get_perchannel_param(output_size, dtype)
|
||||
self.assertEqual(per_channel_params["weight_scale"].dtype, dtype)
|
||||
self.assertEqual(per_channel_params["weight_scale"].shape, (output_size, 1))
|
||||
self.assertEqual(per_channel_params["weight_offset"].dtype, dtype)
|
||||
self.assertEqual(per_channel_params["weight_offset"].shape, (output_size, 1))
|
||||
self.assertEqual(len(per_channel_params), 2)
|
||||
|
||||
@patch("torch_npu.npu_weight_quant_batchmatmul")
|
||||
def test_apply_with_x_is_int8(self, mock_npu_weight_quant_batchmatmul):
|
||||
layer = MagicMock()
|
||||
layer.weight.data = torch.randn(128, 256)
|
||||
layer.weight_scale.data = torch.randn(128, 1)
|
||||
layer.weight_offset.data = torch.randn(128, 1)
|
||||
|
||||
x = torch.randn(32, 128)
|
||||
bias = torch.randn(256)
|
||||
|
||||
expected_y_output = torch.randn(32, 256)
|
||||
mock_npu_weight_quant_batchmatmul.return_value = expected_y_output
|
||||
|
||||
output = self.method.apply(layer, x, bias)
|
||||
self.assertTrue(torch.equal(output, expected_y_output))
|
||||
mock_npu_weight_quant_batchmatmul.assert_called_once()
|
||||
|
||||
@patch("vllm_ascend.utils.get_ascend_config")
|
||||
@patch("torch_npu.npu_format_cast")
|
||||
def test_process_weights_after_loading_with_nz1(self, mock_npu_format_cast, mock_get_config):
|
||||
mock_config = MagicMock()
|
||||
mock_config.weight_nz_mode = 1
|
||||
mock_get_config.return_value = mock_config
|
||||
layer = MagicMock()
|
||||
|
||||
layer.weight.data = torch.randint(-128, 127, (128, 256), dtype=torch.int8)
|
||||
layer.weight_scale.data = torch.randn(128, 1)
|
||||
layer.weight_offset.data = torch.randn(128, 1)
|
||||
|
||||
mock_npu_format_cast.side_effect = identity
|
||||
self.method.process_weights_after_loading(layer)
|
||||
|
||||
self.assertEqual(layer.weight.data.shape, (256, 128))
|
||||
self.assertEqual(layer.weight_scale.data.shape, (128,))
|
||||
self.assertEqual(layer.weight_offset.data.shape, (128,))
|
||||
mock_npu_format_cast.assert_called_once()
|
||||
|
||||
|
||||
class TestAscendW8A16LinearMethodWithNpu(TestBase):
|
||||
def setUp(self):
|
||||
self.method = AscendW8A16LinearMethod()
|
||||
self.mock_get_config = patch("vllm_ascend.utils.get_ascend_config")
|
||||
mock_config = self.mock_get_config.start()
|
||||
mock_ascend_config = MagicMock()
|
||||
mock_ascend_config.weight_nz_mode = 0
|
||||
mock_config.return_value = mock_ascend_config
|
||||
|
||||
def tearDown(self):
|
||||
self.mock_get_config.stop()
|
||||
|
||||
def test_apply_with_npu(self):
|
||||
input_size, output_size = 128, 256
|
||||
params_dtype = torch.bfloat16
|
||||
layer = create_linear_layer(self.method, input_size, output_size, params_dtype)
|
||||
self.method.process_weights_after_loading(layer)
|
||||
|
||||
x = torch.randn(32, input_size, dtype=params_dtype).npu()
|
||||
bias = torch.randn(output_size, dtype=torch.float32).npu()
|
||||
|
||||
output = self.method.apply(layer, x, bias)
|
||||
self.assertEqual(output.shape, (32, output_size))
|
||||
205
tests/ut/quantization/methods/a2/test_w8a8_dynamic.py
Normal file
205
tests/ut/quantization/methods/a2/test_w8a8_dynamic.py
Normal file
@@ -0,0 +1,205 @@
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
from tests.ut.quantization.conftest_quantization import (
|
||||
create_linear_layer,
|
||||
create_mock_ascend_config,
|
||||
create_mock_vllm_config,
|
||||
create_moe_layer,
|
||||
)
|
||||
from vllm_ascend.ascend_forward_context import MoECommType
|
||||
from vllm_ascend.quantization.methods.w8a8_dynamic import (
|
||||
AscendW8A8DynamicFusedMoEMethod,
|
||||
AscendW8A8DynamicLinearMethod,
|
||||
)
|
||||
|
||||
|
||||
class TestAscendW8A8DynamicLinearMethod(TestBase):
|
||||
def setUp(self):
|
||||
self.method = AscendW8A8DynamicLinearMethod()
|
||||
|
||||
def test_get_weight_various_sizes(self):
|
||||
sizes = [(64, 128), (256, 512), (1024, 2048)]
|
||||
for input_size, output_size in sizes:
|
||||
weight = self.method.get_weight(input_size, output_size, torch.bfloat16)
|
||||
self.assertEqual(weight["weight"].dtype, torch.int8)
|
||||
self.assertEqual(weight["weight"].shape, (output_size, input_size))
|
||||
|
||||
def test_get_perchannel_param_dtype_variations(self):
|
||||
dtypes = [torch.bfloat16, torch.float16]
|
||||
for dtype in dtypes:
|
||||
params = self.method.get_perchannel_param(128, dtype)
|
||||
self.assertEqual(params["weight_scale"].dtype, dtype)
|
||||
self.assertEqual(params["weight_offset"].dtype, dtype)
|
||||
self.assertEqual(params["weight_scale"].shape, (128, 1))
|
||||
self.assertEqual(params["weight_offset"].shape, (128, 1))
|
||||
|
||||
@patch("torch_npu.npu_quant_matmul")
|
||||
@patch("torch_npu.npu_dynamic_quant")
|
||||
def test_apply_3d_input_with_squeeze(self, mock_dyn_quant, mock_matmul):
|
||||
mock_dyn_quant.return_value = (
|
||||
torch.randint(-128, 127, (32, 1, 128), dtype=torch.int8),
|
||||
torch.randn(32, 1, dtype=torch.float32),
|
||||
)
|
||||
mock_matmul.return_value = torch.randn(32, 1, 256)
|
||||
layer = MagicMock()
|
||||
layer.weight = torch.randint(-128, 127, (128, 256), dtype=torch.int8)
|
||||
layer.weight_scale = torch.randn(256, dtype=torch.float32)
|
||||
x = torch.randn(32, 1, 128, dtype=torch.bfloat16)
|
||||
output = self.method.apply(layer, x)
|
||||
mock_dyn_quant.assert_called_once()
|
||||
mock_matmul.assert_called_once()
|
||||
self.assertEqual(output.shape, (32, 1, 1, 256))
|
||||
|
||||
def test_process_weights_after_loading(self):
|
||||
layer = MagicMock()
|
||||
layer.weight.data = torch.randint(-128, 127, (128, 256), dtype=torch.int8)
|
||||
layer.weight_scale.data = torch.randn(256, 1, dtype=torch.bfloat16)
|
||||
layer.weight_offset.data = torch.randn(256, 1, dtype=torch.bfloat16)
|
||||
with patch("vllm_ascend.quantization.methods.w8a8_dynamic.maybe_trans_nz", side_effect=lambda x: x):
|
||||
self.method.process_weights_after_loading(layer)
|
||||
self.assertEqual(layer.weight_scale_fp32.dtype, torch.float32)
|
||||
self.assertEqual(layer.weight_scale.data.shape, (256,))
|
||||
self.assertEqual(layer.weight_offset.data.shape, (256,))
|
||||
self.assertEqual(layer.weight.data.shape, (256, 128))
|
||||
|
||||
|
||||
class TestAscendW8A8DynamicLinearMethodWithNpu(TestBase):
|
||||
def setUp(self):
|
||||
self.method = AscendW8A8DynamicLinearMethod()
|
||||
self.mock_get_config = patch("vllm_ascend.utils.get_ascend_config")
|
||||
mock_config = self.mock_get_config.start()
|
||||
mock_ascend_config = MagicMock()
|
||||
mock_ascend_config.weight_nz_mode = 0
|
||||
mock_config.return_value = mock_ascend_config
|
||||
|
||||
def tearDown(self):
|
||||
self.mock_get_config.stop()
|
||||
|
||||
def test_apply_with_npu(self):
|
||||
input_size, output_size = 128, 256
|
||||
params_dtype = torch.bfloat16
|
||||
layer = create_linear_layer(self.method, input_size, output_size, params_dtype)
|
||||
self.method.process_weights_after_loading(layer)
|
||||
|
||||
x = torch.randn(32, input_size, dtype=params_dtype).npu()
|
||||
bias = torch.randn(output_size, dtype=torch.float32).npu()
|
||||
|
||||
output = self.method.apply(layer, x, bias)
|
||||
self.assertEqual(output.shape, (32, output_size))
|
||||
|
||||
|
||||
class TestAscendW8A8FusedMoEMethod(TestBase):
|
||||
num_experts = 8
|
||||
hidden_size = 128
|
||||
intermediate_size = 128
|
||||
|
||||
@patch("torch.distributed.get_rank")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_mc2_group")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_ascend_config")
|
||||
def setUp(self, mock_ascend, mock_mc2, mock_rank):
|
||||
with patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_current_vllm_config") as mock_vllm:
|
||||
mock_vllm.return_value = create_mock_vllm_config()
|
||||
mock_ascend.return_value = create_mock_ascend_config()
|
||||
mock_mc2.return_value = MagicMock(
|
||||
device_group=Mock(
|
||||
_get_backend=Mock(return_value=Mock(get_hccl_comm_name=Mock(return_value="test_comm")))
|
||||
)
|
||||
)
|
||||
mock_rank.return_value = 0
|
||||
self.quant_method = AscendW8A8DynamicFusedMoEMethod()
|
||||
|
||||
def test_get_weight_various_expert_counts(self):
|
||||
expert_counts = [4, 8, 16, 32]
|
||||
for num_experts in expert_counts:
|
||||
param_dict = self.quant_method.get_weight(
|
||||
num_experts, self.intermediate_size, self.hidden_size, torch.bfloat16
|
||||
)
|
||||
self.assertEqual(param_dict["w13_weight"].shape[0], num_experts)
|
||||
self.assertEqual(param_dict["w2_weight"].shape[0], num_experts)
|
||||
|
||||
def test_get_dynamic_quant_param_various_sizes(self):
|
||||
param_dict = self.quant_method.get_dynamic_quant_param(
|
||||
self.num_experts, self.intermediate_size, self.hidden_size, torch.bfloat16
|
||||
)
|
||||
self.assertEqual(param_dict["w13_weight_scale"].dtype, torch.bfloat16)
|
||||
self.assertEqual(param_dict["w13_weight_offset"].shape, (self.num_experts, 2 * self.intermediate_size, 1))
|
||||
self.assertEqual(param_dict["w2_weight_scale"].dtype, torch.bfloat16)
|
||||
self.assertEqual(param_dict["w2_weight_offset"].shape, (self.num_experts, self.hidden_size, 1))
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_dynamic._EXTRA_CTX")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_dynamic.select_experts")
|
||||
def test_apply_uses_explicit_dispatch_and_mlp_args(self, mock_select_experts, mock_extra_ctx):
|
||||
tokens = 4
|
||||
hidden_size = self.hidden_size
|
||||
layer = torch.nn.Module()
|
||||
layer.w13_weight = torch.randint(
|
||||
-8,
|
||||
8,
|
||||
(self.num_experts, 2 * self.intermediate_size, hidden_size),
|
||||
dtype=torch.int8,
|
||||
)
|
||||
layer.w2_weight = torch.randint(
|
||||
-8,
|
||||
8,
|
||||
(self.num_experts, hidden_size, self.intermediate_size),
|
||||
dtype=torch.int8,
|
||||
)
|
||||
layer.w13_weight_scale_fp32 = torch.ones(self.num_experts, 2 * self.intermediate_size, dtype=torch.float32)
|
||||
layer.w2_weight_scale = torch.ones(self.num_experts, hidden_size, dtype=torch.float32)
|
||||
layer.swiglu_limit = 1000000
|
||||
|
||||
x = torch.randn(tokens, hidden_size, dtype=torch.float32)
|
||||
router_logits = torch.randn(tokens, self.num_experts, dtype=torch.float32)
|
||||
topk_weights = torch.randn(tokens, 2, dtype=torch.float32)
|
||||
topk_ids = torch.randint(0, self.num_experts, (tokens, 2), dtype=torch.int64)
|
||||
mc2_mask = torch.tensor([1, 0, 1, 0], dtype=torch.bool)
|
||||
pertoken_scale = torch.randn(tokens, dtype=torch.float32)
|
||||
|
||||
mock_select_experts.return_value = (topk_weights, topk_ids)
|
||||
mock_comm = Mock()
|
||||
mock_comm.fused_experts.return_value = torch.randn(tokens, hidden_size, dtype=torch.float32)
|
||||
mock_extra_ctx.moe_comm_method = mock_comm
|
||||
mock_extra_ctx.moe_comm_type = MoECommType.ALLGATHER
|
||||
self.quant_method.multistream_overlap_gate = False
|
||||
self.quant_method.in_dtype = torch.float32
|
||||
|
||||
self.quant_method.apply(
|
||||
layer=layer,
|
||||
x=x,
|
||||
router_logits=router_logits,
|
||||
top_k=2,
|
||||
renormalize=True,
|
||||
num_experts=self.num_experts,
|
||||
activation="gelu",
|
||||
apply_router_weight_on_input=True,
|
||||
mc2_mask=mc2_mask,
|
||||
pertoken_scale=pertoken_scale,
|
||||
)
|
||||
|
||||
fused_experts_input = mock_comm.fused_experts.call_args.kwargs["fused_experts_input"]
|
||||
self.assertEqual(fused_experts_input.activation, "gelu")
|
||||
self.assertTrue(fused_experts_input.routing.apply_router_weight_on_input)
|
||||
self.assertIs(fused_experts_input.routing.mc2_mask, mc2_mask)
|
||||
self.assertIs(fused_experts_input.routing.pertoken_scale, pertoken_scale)
|
||||
self.assertIs(fused_experts_input.topk_weights, topk_weights)
|
||||
self.assertIs(fused_experts_input.topk_ids, topk_ids)
|
||||
|
||||
@patch("torch_npu.npu_format_cast")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_ascend_config")
|
||||
def test_process_weights_after_loading(self, mock_get_config, mock_format_cast):
|
||||
mock_config = MagicMock()
|
||||
mock_config.enable_fused_mc2 = 1
|
||||
mock_get_config.return_value = mock_config
|
||||
self.quant_method.dynamic_eplb = True
|
||||
mock_format_cast.return_value = torch.randint(
|
||||
-8, 8, (self.num_experts, self.hidden_size, 2 * self.intermediate_size), dtype=torch.int8
|
||||
)
|
||||
layer = create_moe_layer(
|
||||
num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size
|
||||
)
|
||||
self.quant_method.process_weights_after_loading(layer)
|
||||
self.assertTrue(hasattr(layer, "w13_weight_list"))
|
||||
self.assertFalse(hasattr(layer, "w13_weight_scale_fp32"))
|
||||
184
tests/ut/quantization/methods/a2/test_w8a8_static.py
Normal file
184
tests/ut/quantization/methods/a2/test_w8a8_static.py
Normal file
@@ -0,0 +1,184 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
from tests.ut.quantization.conftest_quantization import create_linear_layer, identity
|
||||
from vllm_ascend.quantization.methods.w8a8_static import AscendW8A8LinearMethod
|
||||
from vllm_ascend.utils import COMPRESSED_TENSORS_METHOD
|
||||
|
||||
|
||||
class TestAscendW8A8LinearMethod(TestBase):
|
||||
def setUp(self):
|
||||
self.method = AscendW8A8LinearMethod()
|
||||
|
||||
def test_get_weight(self):
|
||||
sizes = [(64, 128), (256, 512), (1024, 2048), (1, 1)]
|
||||
for input_size, output_size in sizes:
|
||||
weight = self.method.get_weight(input_size, output_size)
|
||||
self.assertEqual(weight["weight"].dtype, torch.int8)
|
||||
self.assertEqual(weight["weight"].shape, (output_size, input_size))
|
||||
self.assertEqual(len(weight), 1)
|
||||
|
||||
weight = self.method.get_weight(256, 128, torch.float16)
|
||||
self.assertEqual(weight["weight"].dtype, torch.int8)
|
||||
|
||||
def test_get_pertensor_param(self):
|
||||
dtypes = [torch.bfloat16, torch.float16, torch.float32]
|
||||
for dtype in dtypes:
|
||||
params = self.method.get_pertensor_param(dtype)
|
||||
self.assertEqual(params["input_scale"].dtype, dtype)
|
||||
self.assertEqual(params["input_offset"].dtype, torch.int8)
|
||||
self.assertEqual(params["input_scale"].shape, (1,))
|
||||
self.assertEqual(params["input_offset"].shape, (1,))
|
||||
|
||||
def test_get_perchannel_param(self):
|
||||
for output_size, dtype in [(128, torch.bfloat16), (256, torch.float16)]:
|
||||
params = self.method.get_perchannel_param(output_size, dtype)
|
||||
self.assertEqual(params["quant_bias"].shape, (output_size,))
|
||||
self.assertEqual(params["quant_bias"].dtype, torch.int32)
|
||||
self.assertEqual(params["weight_scale"].shape, (output_size, 1))
|
||||
self.assertEqual(params["weight_scale"].dtype, dtype)
|
||||
self.assertEqual(params["weight_offset"].shape, (output_size, 1))
|
||||
self.assertEqual(params["weight_offset"].dtype, dtype)
|
||||
self.assertEqual(params["deq_scale"].shape, (output_size,))
|
||||
if dtype == torch.bfloat16:
|
||||
self.assertEqual(params["deq_scale"].dtype, torch.float32)
|
||||
elif dtype == torch.float16:
|
||||
self.assertEqual(params["deq_scale"].dtype, torch.int64)
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_static.get_weight_prefetch_method")
|
||||
@patch("torch.ops.vllm.quantize")
|
||||
@patch("torch_npu.npu_quant_matmul")
|
||||
def test_apply_with_x_not_int8(self, mock_npu_quant_matmul, mock_quantize, mock_get_weight_prefetch_method):
|
||||
layer = MagicMock()
|
||||
layer.aclnn_input_scale = 0.1
|
||||
layer.aclnn_input_offset = 0.2
|
||||
layer.weight = torch.randn(128, 256)
|
||||
layer.deq_scale = 0.3
|
||||
quant_bias = torch.zeros(256)
|
||||
layer.quant_bias = quant_bias
|
||||
|
||||
mock_get_weight_prefetch_method.return_value = MagicMock()
|
||||
|
||||
x = torch.randn(32, 128)
|
||||
bias = torch.randn(256)
|
||||
mock_quantize.return_value = torch.randint(-128, 127, x.shape, dtype=torch.int8)
|
||||
|
||||
expected_y_output = torch.randn(32, 256)
|
||||
mock_npu_quant_matmul.return_value = expected_y_output
|
||||
|
||||
output = self.method.apply(layer, x, bias)
|
||||
|
||||
self.assertTrue(torch.equal(output, expected_y_output))
|
||||
mock_quantize.assert_called_once()
|
||||
mock_npu_quant_matmul.assert_called_once()
|
||||
call_kwargs = mock_npu_quant_matmul.call_args.kwargs
|
||||
self.assertTrue(torch.equal(call_kwargs["bias"], quant_bias))
|
||||
|
||||
@patch("torch.ops.vllm.quantize")
|
||||
@patch("torch_npu.npu_quant_matmul")
|
||||
def test_apply_with_x_is_int8(self, mock_npu_quant_matmul, mock_quantize):
|
||||
layer = MagicMock()
|
||||
layer.aclnn_input_scale = 0.1
|
||||
layer.aclnn_input_offset = 0.2
|
||||
layer.weight = torch.randn(128, 256)
|
||||
layer.deq_scale = 0.3
|
||||
layer.ascend_quant_method = COMPRESSED_TENSORS_METHOD
|
||||
|
||||
x = torch.randint(-128, 127, (32, 128), dtype=torch.int8)
|
||||
bias = torch.randn(256)
|
||||
|
||||
expected_y_output = torch.randn(32, 256)
|
||||
mock_npu_quant_matmul.return_value = expected_y_output
|
||||
|
||||
output = self.method.apply(layer, x, bias)
|
||||
self.assertTrue(torch.equal(output, expected_y_output))
|
||||
mock_quantize.assert_not_called()
|
||||
mock_npu_quant_matmul.assert_called_once()
|
||||
call_kwargs = mock_npu_quant_matmul.call_args.kwargs
|
||||
self.assertTrue(torch.equal(call_kwargs["bias"], bias))
|
||||
|
||||
@patch("vllm_ascend.utils.get_ascend_config")
|
||||
@patch("torch_npu.npu_format_cast")
|
||||
def test_process_weights_after_loading_with_nz1(self, mock_npu_format_cast, mock_get_config):
|
||||
mock_config = MagicMock()
|
||||
mock_config.weight_nz_mode = 1
|
||||
mock_get_config.return_value = mock_config
|
||||
layer = MagicMock()
|
||||
|
||||
layer.weight.data = torch.randint(-128, 127, (128, 256), dtype=torch.int8)
|
||||
layer.input_scale.data = torch.tensor([0.1])
|
||||
layer.input_offset.data = torch.tensor([0])
|
||||
layer.weight_scale.data = torch.randn(128, 1)
|
||||
layer.weight_offset.data = torch.randn(128, 1)
|
||||
|
||||
mock_npu_format_cast.side_effect = identity
|
||||
self.method.process_weights_after_loading(layer)
|
||||
|
||||
expected_offset = torch.tensor([0]).repeat(256).to(torch.int8)
|
||||
self.assertTrue(torch.equal(layer.aclnn_input_offset.data, expected_offset))
|
||||
self.assertFalse(layer.aclnn_input_offset.requires_grad)
|
||||
|
||||
self.assertEqual(layer.weight.data.shape, (256, 128))
|
||||
self.assertEqual(layer.weight_scale.data.shape, (128,))
|
||||
self.assertEqual(layer.weight_offset.data.shape, (128,))
|
||||
mock_npu_format_cast.assert_called_once()
|
||||
self.assertTrue(isinstance(layer.deq_scale, MagicMock))
|
||||
|
||||
@patch("vllm_ascend.utils.get_ascend_config")
|
||||
@patch("torch_npu.npu_format_cast")
|
||||
def test_process_weights_after_loading_with_nz2_and_compressed_tensors(self, mock_npu_format_cast, mock_get_config):
|
||||
mock_config = MagicMock()
|
||||
mock_config.weight_nz_mode = 2
|
||||
mock_get_config.return_value = mock_config
|
||||
layer = MagicMock()
|
||||
|
||||
layer.weight.data = torch.randint(-128, 127, (128, 256), dtype=torch.int8)
|
||||
layer.input_scale.data = torch.tensor([0.1])
|
||||
layer.input_offset.data = torch.tensor([0])
|
||||
layer.weight_scale.data = torch.randn(128, 1)
|
||||
layer.weight_offset.data = torch.randn(128, 1)
|
||||
layer.ascend_quant_method = COMPRESSED_TENSORS_METHOD
|
||||
|
||||
mock_npu_format_cast.side_effect = identity
|
||||
self.method.process_weights_after_loading(layer)
|
||||
|
||||
expected_offset = torch.tensor([0]).repeat(256).to(torch.int8)
|
||||
self.assertTrue(torch.equal(layer.aclnn_input_offset.data, expected_offset))
|
||||
self.assertFalse(layer.aclnn_input_offset.requires_grad)
|
||||
|
||||
self.assertEqual(layer.weight.data.shape, (256, 128))
|
||||
self.assertEqual(layer.weight_scale.data.shape, (128,))
|
||||
self.assertEqual(layer.weight_offset.data.shape, (128,))
|
||||
mock_npu_format_cast.assert_called_once()
|
||||
self.assertFalse(isinstance(layer.deq_scale, MagicMock))
|
||||
|
||||
|
||||
class TestAscendW8A8LinearMethodWithNpu(TestBase):
|
||||
def setUp(self):
|
||||
self.method = AscendW8A8LinearMethod()
|
||||
self.mock_get_config = patch("vllm_ascend.utils.get_ascend_config")
|
||||
mock_config = self.mock_get_config.start()
|
||||
mock_ascend_config = MagicMock()
|
||||
mock_ascend_config.weight_nz_mode = 0
|
||||
mock_config.return_value = mock_ascend_config
|
||||
|
||||
def tearDown(self):
|
||||
self.mock_get_config.stop()
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_static.get_weight_prefetch_method")
|
||||
def test_apply_with_npu(self, mock_get_weight_prefetch_method):
|
||||
mock_get_weight_prefetch_method.return_value = MagicMock()
|
||||
|
||||
input_size, output_size = 128, 256
|
||||
params_dtype = torch.bfloat16
|
||||
layer = create_linear_layer(self.method, input_size, output_size, params_dtype)
|
||||
layer.params_dtype = params_dtype
|
||||
self.method.process_weights_after_loading(layer)
|
||||
|
||||
x = torch.randn(32, input_size, dtype=params_dtype).npu()
|
||||
bias = torch.randn(output_size, dtype=torch.float32).npu()
|
||||
|
||||
output = self.method.apply(layer, x, bias)
|
||||
self.assertEqual(output.shape, (32, output_size))
|
||||
692
tests/ut/quantization/methods/test_kv_c8.py
Normal file
692
tests/ut/quantization/methods/test_kv_c8.py
Normal file
@@ -0,0 +1,692 @@
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from vllm.config import KVTransferConfig, VllmConfig
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
|
||||
|
||||
class TestWeightLoader(unittest.TestCase):
|
||||
"""Test cases for weight_loader function in kv_c8.py"""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up test environment before each test"""
|
||||
# Import the module under test
|
||||
from vllm_ascend.quantization.methods.kv_c8 import _fa_quant_weight_loader as weight_loader
|
||||
|
||||
self.weight_loader = weight_loader
|
||||
|
||||
# Mock distributed functions
|
||||
self.tp_rank_patch = patch("vllm_ascend.quantization.methods.kv_c8.get_tensor_model_parallel_rank")
|
||||
self.tp_size_patch = patch("vllm_ascend.quantization.methods.kv_c8.get_tensor_model_parallel_world_size")
|
||||
self.mock_tp_rank = self.tp_rank_patch.start()
|
||||
self.mock_tp_size = self.tp_size_patch.start()
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up after each test"""
|
||||
self.tp_rank_patch.stop()
|
||||
self.tp_size_patch.stop()
|
||||
|
||||
def test_weight_loader_single_element(self):
|
||||
"""Test weight_loader when both tensors contain a single element"""
|
||||
# Create tensors with single element
|
||||
param = torch.tensor([0.0])
|
||||
loaded_weight = torch.tensor([5.0])
|
||||
|
||||
# Call weight_loader
|
||||
self.weight_loader(param, loaded_weight)
|
||||
|
||||
# Verify the value was filled correctly
|
||||
self.assertEqual(param.item(), 5.0)
|
||||
self.assertEqual(param.dtype, torch.float32)
|
||||
|
||||
def test_weight_loader_single_element_int(self):
|
||||
"""Test weight_loader with integer tensors"""
|
||||
param = torch.tensor([0], dtype=torch.int32)
|
||||
loaded_weight = torch.tensor([10], dtype=torch.int32)
|
||||
|
||||
self.weight_loader(param, loaded_weight)
|
||||
|
||||
self.assertEqual(param.item(), 10)
|
||||
|
||||
def test_weight_loader_tp_sharding_first_rank(self):
|
||||
"""Test weight_loader with tensor parallelism sharding for first rank"""
|
||||
# Configure mocks for rank 0 of 4
|
||||
self.mock_tp_rank.return_value = 0
|
||||
self.mock_tp_size.return_value = 4
|
||||
|
||||
# Create test tensors
|
||||
param = torch.zeros(2, 5) # Target param shape (2,5)
|
||||
loaded_weight = torch.ones(8, 5) # Full weight (8,5)
|
||||
|
||||
# Mock narrow to track the call
|
||||
with patch.object(loaded_weight, "narrow", wraps=loaded_weight.narrow) as mock_narrow:
|
||||
self.weight_loader(param, loaded_weight)
|
||||
|
||||
# Verify narrow was called correctly: narrow(dim=0, start=0, length=2)
|
||||
mock_narrow.assert_called_once_with(0, 0, 2)
|
||||
|
||||
# Verify data was copied
|
||||
self.assertTrue(torch.all(param == 1))
|
||||
|
||||
def test_weight_loader_tp_sharding_middle_rank(self):
|
||||
"""Test weight_loader with tensor parallelism sharding for middle rank"""
|
||||
# Configure mocks for rank 2 of 4
|
||||
self.mock_tp_rank.return_value = 2
|
||||
self.mock_tp_size.return_value = 4
|
||||
|
||||
param = torch.zeros(2, 5)
|
||||
loaded_weight = torch.ones(8, 5)
|
||||
|
||||
with patch.object(loaded_weight, "narrow", wraps=loaded_weight.narrow) as mock_narrow:
|
||||
self.weight_loader(param, loaded_weight)
|
||||
|
||||
# Verify narrow was called correctly: start = shard_size * rank = 2 * 2 = 4
|
||||
mock_narrow.assert_called_once_with(0, 4, 2)
|
||||
|
||||
self.assertTrue(torch.all(param == 1))
|
||||
|
||||
def test_weight_loader_tp_sharding_last_rank(self):
|
||||
"""Test weight_loader with tensor parallelism sharding for last rank"""
|
||||
# Configure mocks for rank 3 of 4
|
||||
self.mock_tp_rank.return_value = 3
|
||||
self.mock_tp_size.return_value = 4
|
||||
|
||||
param = torch.zeros(2, 5)
|
||||
loaded_weight = torch.ones(8, 5)
|
||||
|
||||
with patch.object(loaded_weight, "narrow", wraps=loaded_weight.narrow) as mock_narrow:
|
||||
self.weight_loader(param, loaded_weight)
|
||||
|
||||
# Verify narrow was called correctly: start = 2 * 3 = 6
|
||||
mock_narrow.assert_called_once_with(0, 6, 2)
|
||||
|
||||
self.assertTrue(torch.all(param == 1))
|
||||
|
||||
def test_weight_loader_shape_mismatch(self):
|
||||
"""Test weight_loader raises assertion error on shape mismatch"""
|
||||
self.mock_tp_rank.return_value = 0
|
||||
self.mock_tp_size.return_value = 2
|
||||
|
||||
param = torch.zeros(2, 3)
|
||||
loaded_weight = torch.ones(4, 4) # Different shape after sharding
|
||||
|
||||
# Mock narrow to return tensor with wrong shape
|
||||
with patch.object(loaded_weight, "narrow", return_value=torch.ones(2, 4)):
|
||||
with self.assertRaises(AssertionError) as context:
|
||||
self.weight_loader(param, loaded_weight)
|
||||
|
||||
# Verify error message contains expected information
|
||||
self.assertIn("Attempted to load weight", str(context.exception))
|
||||
self.assertIn("into parameter", str(context.exception))
|
||||
|
||||
def test_weight_loader_with_different_dtypes(self):
|
||||
"""Test weight_loader handles different dtypes correctly"""
|
||||
self.mock_tp_rank.return_value = 0
|
||||
self.mock_tp_size.return_value = 1 # No sharding
|
||||
|
||||
param = torch.zeros(2, 3, dtype=torch.float32)
|
||||
loaded_weight = torch.ones(2, 3, dtype=torch.float16)
|
||||
|
||||
self.weight_loader(param, loaded_weight)
|
||||
|
||||
# Verify data was copied and converted
|
||||
self.assertTrue(torch.all(param == 1))
|
||||
self.assertEqual(param.dtype, torch.float32)
|
||||
|
||||
|
||||
class TestAscendFAQuantAttentionMethodInit(unittest.TestCase):
|
||||
"""Test cases for AscendFAQuantAttentionMethod initialization"""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up test environment"""
|
||||
# Mock vllm_config
|
||||
self.config_patch = patch("vllm_ascend.quantization.methods.kv_c8.get_current_vllm_config")
|
||||
self.mock_get_config = self.config_patch.start()
|
||||
|
||||
# Create mock config with attributes
|
||||
self.mock_config = Mock()
|
||||
self.mock_hf_config = Mock()
|
||||
self.mock_hf_config.kv_lora_rank = 128
|
||||
self.mock_hf_config.qk_rope_head_dim = 64
|
||||
self.mock_config.model_config.hf_config = self.mock_hf_config
|
||||
self.mock_get_config.return_value = self.mock_config
|
||||
|
||||
# Import the class after patching
|
||||
from vllm_ascend.quantization.methods.kv_c8 import AscendFAQuantAttentionMethod
|
||||
|
||||
self.method_class = AscendFAQuantAttentionMethod
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up after each test"""
|
||||
self.config_patch.stop()
|
||||
|
||||
def test_init_with_full_config(self):
|
||||
"""Test initialization when config has all attributes"""
|
||||
method = self.method_class()
|
||||
self.assertEqual(method.kv_lora_rank, 128)
|
||||
self.assertEqual(method.qk_rope_head_dim, 64)
|
||||
|
||||
def test_init_without_both_attributes(self):
|
||||
"""Test initialization when config lacks both attributes"""
|
||||
delattr(self.mock_hf_config, "kv_lora_rank")
|
||||
delattr(self.mock_hf_config, "qk_rope_head_dim")
|
||||
|
||||
method = self.method_class()
|
||||
self.assertEqual(method.kv_lora_rank, 0)
|
||||
self.assertEqual(method.qk_rope_head_dim, 0)
|
||||
|
||||
|
||||
class TestAscendFAQuantAttentionMethodCreateWeights(unittest.TestCase):
|
||||
"""Test cases for create_weights method"""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up test environment"""
|
||||
# Mock vllm_config
|
||||
self.config_patch = patch("vllm_ascend.quantization.methods.kv_c8.get_current_vllm_config")
|
||||
self.mock_get_config = self.config_patch.start()
|
||||
|
||||
self.mock_config = Mock()
|
||||
self.mock_hf_config = Mock()
|
||||
self.mock_hf_config.kv_lora_rank = 128
|
||||
self.mock_hf_config.qk_rope_head_dim = 64
|
||||
self.mock_config.model_config.hf_config = self.mock_hf_config
|
||||
self.mock_get_config.return_value = self.mock_config
|
||||
|
||||
# Import the class
|
||||
from vllm_ascend.quantization.methods.kv_c8 import AscendFAQuantAttentionMethod
|
||||
|
||||
self.method_class = AscendFAQuantAttentionMethod
|
||||
|
||||
# Mock torch functions
|
||||
self.default_dtype_patch = patch("torch.get_default_dtype", return_value=torch.float32)
|
||||
self.mock_default_dtype = self.default_dtype_patch.start()
|
||||
|
||||
# Create a real nn.Module for testing
|
||||
self.layer = nn.Module()
|
||||
self.layer.num_heads = 32
|
||||
self.layer.num_kv_heads = 1
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up after each test"""
|
||||
self.config_patch.stop()
|
||||
self.default_dtype_patch.stop()
|
||||
|
||||
def test_create_weights_adds_submodules(self):
|
||||
"""Test that create_weights adds fa_q, fa_k, fa_v submodules"""
|
||||
method = self.method_class()
|
||||
|
||||
with patch("torch.empty") as mock_empty:
|
||||
mock_empty.return_value = torch.zeros(1, 1)
|
||||
|
||||
method.create_weights(self.layer)
|
||||
|
||||
# Verify submodules were added
|
||||
self.assertTrue(hasattr(self.layer, "fa_q"))
|
||||
self.assertTrue(hasattr(self.layer, "fa_k"))
|
||||
self.assertTrue(hasattr(self.layer, "fa_v"))
|
||||
|
||||
# Verify they are instances of nn.Module
|
||||
self.assertIsInstance(self.layer.fa_q, nn.Module)
|
||||
self.assertIsInstance(self.layer.fa_k, nn.Module)
|
||||
self.assertIsInstance(self.layer.fa_v, nn.Module)
|
||||
|
||||
def test_create_weights_creates_correct_tensors(self):
|
||||
"""Test that create_weights creates tensors with correct shapes and dtypes"""
|
||||
method = self.method_class()
|
||||
|
||||
# Track torch.empty calls
|
||||
empty_calls = []
|
||||
|
||||
def mock_empty(size, dtype=None):
|
||||
empty_calls.append((size, dtype))
|
||||
return torch.zeros(size, dtype=dtype if dtype else torch.float32)
|
||||
|
||||
with patch("torch.empty", side_effect=mock_empty):
|
||||
method.create_weights(self.layer)
|
||||
|
||||
# Verify tensor creations
|
||||
expected_calls = [
|
||||
((32, 1), torch.float32), # fa_q.scale
|
||||
((1, 1), torch.float32), # fa_k.scale
|
||||
((1, 1), torch.float32), # fa_v.scale
|
||||
((32, 1), torch.int8), # fa_q.offset
|
||||
((1, 1), torch.int8), # fa_k.offset
|
||||
((1, 1), torch.int8), # fa_v.offset
|
||||
]
|
||||
|
||||
# Compare without considering order
|
||||
self.assertEqual(len(empty_calls), len(expected_calls))
|
||||
for call in expected_calls:
|
||||
self.assertIn(call, empty_calls)
|
||||
|
||||
def test_create_weights_registers_parameters(self):
|
||||
"""Test that create_weights registers parameters with correct attributes"""
|
||||
method = self.method_class()
|
||||
|
||||
# Create real tensors for testing
|
||||
def create_tensor(*args, **kwargs):
|
||||
size = args[0] if args else kwargs.get("size", (1,))
|
||||
dtype = kwargs.get("dtype", torch.float32)
|
||||
return torch.zeros(*size, dtype=dtype)
|
||||
|
||||
with patch("torch.empty", side_effect=create_tensor):
|
||||
method.create_weights(self.layer)
|
||||
|
||||
# Import weight_loader for comparison
|
||||
from vllm_ascend.quantization.methods.kv_c8 import _fa_quant_weight_loader as weight_loader
|
||||
|
||||
# Verify each parameter exists and has weight_loader
|
||||
self.assertTrue(hasattr(self.layer.fa_q, "scale"))
|
||||
self.assertTrue(hasattr(self.layer.fa_q.scale, "weight_loader"))
|
||||
self.assertEqual(self.layer.fa_q.scale.weight_loader, weight_loader)
|
||||
self.assertFalse(self.layer.fa_q.scale.requires_grad)
|
||||
|
||||
self.assertTrue(hasattr(self.layer.fa_k, "scale"))
|
||||
self.assertTrue(hasattr(self.layer.fa_k.scale, "weight_loader"))
|
||||
|
||||
self.assertTrue(hasattr(self.layer.fa_v, "scale"))
|
||||
self.assertTrue(hasattr(self.layer.fa_v.scale, "weight_loader"))
|
||||
|
||||
self.assertTrue(hasattr(self.layer.fa_q, "offset"))
|
||||
self.assertTrue(hasattr(self.layer.fa_q.offset, "weight_loader"))
|
||||
self.assertEqual(self.layer.fa_q.offset.dtype, torch.int8)
|
||||
|
||||
|
||||
class TestAscendFAQuantAttentionMethodProcessWeights(unittest.TestCase):
|
||||
"""Test cases for process_weights_after_loading method"""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up test environment"""
|
||||
# Mock vllm_config
|
||||
self.config_patch = patch("vllm_ascend.quantization.methods.kv_c8.get_current_vllm_config")
|
||||
self.mock_get_config = self.config_patch.start()
|
||||
|
||||
self.mock_config = Mock()
|
||||
self.mock_hf_config = Mock()
|
||||
self.mock_hf_config.kv_lora_rank = 64
|
||||
self.mock_hf_config.qk_rope_head_dim = 32
|
||||
self.mock_config.model_config.hf_config = self.mock_hf_config
|
||||
self.mock_get_config.return_value = self.mock_config
|
||||
|
||||
# Import the class
|
||||
from vllm_ascend.quantization.methods.kv_c8 import AscendFAQuantAttentionMethod
|
||||
|
||||
self.method_class = AscendFAQuantAttentionMethod
|
||||
|
||||
# Create method instance with real layer
|
||||
self.method = self.method_class()
|
||||
|
||||
# Create a real nn.Module for testing
|
||||
self.layer = nn.Module()
|
||||
|
||||
# Create real tensors for fa_k
|
||||
self.fa_k_scale = torch.tensor([[2.0, 3.0, 4.0]], dtype=torch.float16) # Shape (1,3)
|
||||
self.fa_k_offset = torch.tensor([[1, 2, 3]], dtype=torch.int8) # Shape (1,3)
|
||||
|
||||
# Create fa_k module with parameters
|
||||
self.layer.fa_k = nn.Module()
|
||||
self.layer.fa_k.scale = nn.Parameter(self.fa_k_scale, requires_grad=False)
|
||||
self.layer.fa_k.offset = nn.Parameter(self.fa_k_offset, requires_grad=False)
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up after each test"""
|
||||
self.config_patch.stop()
|
||||
|
||||
def test_process_weights_with_single_value_scale(self):
|
||||
"""Test process_weights with single value scale"""
|
||||
# Create new layer with single value scale
|
||||
layer = nn.Module()
|
||||
layer.fa_k = nn.Module()
|
||||
layer.fa_k.scale = nn.Parameter(torch.tensor([[2.0]], dtype=torch.float16), requires_grad=False)
|
||||
layer.fa_k.offset = nn.Parameter(torch.tensor([[1]], dtype=torch.int8), requires_grad=False)
|
||||
|
||||
self.method.kv_lora_rank = 4
|
||||
self.method.process_weights_after_loading(layer)
|
||||
|
||||
self.assertEqual(layer.quant_kscale.shape, (1, 4))
|
||||
self.assertEqual(layer.quant_kscale.dtype, torch.float32)
|
||||
|
||||
|
||||
class TestIntegration(unittest.TestCase):
|
||||
"""Integration tests for the complete kv_c8 functionality"""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up test environment"""
|
||||
# Mock vllm_config
|
||||
self.config_patch = patch("vllm_ascend.quantization.methods.kv_c8.get_current_vllm_config")
|
||||
self.mock_get_config = self.config_patch.start()
|
||||
|
||||
self.mock_config = Mock()
|
||||
self.mock_hf_config = Mock()
|
||||
self.mock_hf_config.kv_lora_rank = 64
|
||||
self.mock_hf_config.qk_rope_head_dim = 32
|
||||
self.mock_config.model_config.hf_config = self.mock_hf_config
|
||||
self.mock_get_config.return_value = self.mock_config
|
||||
|
||||
# Mock distributed functions
|
||||
self.tp_rank_patch = patch("vllm_ascend.quantization.methods.kv_c8.get_tensor_model_parallel_rank")
|
||||
self.tp_size_patch = patch("vllm_ascend.quantization.methods.kv_c8.get_tensor_model_parallel_world_size")
|
||||
self.mock_tp_rank = self.tp_rank_patch.start()
|
||||
self.mock_tp_size = self.tp_size_patch.start()
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up after each test"""
|
||||
self.config_patch.stop()
|
||||
self.tp_rank_patch.stop()
|
||||
self.tp_size_patch.stop()
|
||||
|
||||
def test_complete_workflow(self):
|
||||
"""Test complete workflow from weight creation to processing"""
|
||||
from vllm_ascend.quantization.methods.kv_c8 import AscendFAQuantAttentionMethod
|
||||
|
||||
# Create method instance
|
||||
method = AscendFAQuantAttentionMethod()
|
||||
|
||||
# Create real layer
|
||||
layer = nn.Module()
|
||||
layer.num_heads = 32
|
||||
layer.num_kv_heads = 1
|
||||
|
||||
# Step 1: Create weights
|
||||
method.create_weights(layer)
|
||||
|
||||
# Verify weights were created with correct structure
|
||||
self.assertTrue(hasattr(layer, "fa_q"))
|
||||
self.assertTrue(hasattr(layer, "fa_k"))
|
||||
self.assertTrue(hasattr(layer, "fa_v"))
|
||||
|
||||
self.assertTrue(hasattr(layer.fa_q, "scale"))
|
||||
self.assertTrue(hasattr(layer.fa_q, "offset"))
|
||||
self.assertTrue(hasattr(layer.fa_k, "scale"))
|
||||
self.assertTrue(hasattr(layer.fa_k, "offset"))
|
||||
self.assertTrue(hasattr(layer.fa_v, "scale"))
|
||||
self.assertTrue(hasattr(layer.fa_v, "offset"))
|
||||
|
||||
# Step 2: Simulate weight loading
|
||||
self.mock_tp_rank.return_value = 0
|
||||
self.mock_tp_size.return_value = 1
|
||||
|
||||
# Create dummy weights
|
||||
q_scale = torch.randn(32, 1)
|
||||
k_scale = torch.randn(1, 1)
|
||||
v_scale = torch.randn(1, 1)
|
||||
q_offset = torch.randint(-128, 127, (32, 1), dtype=torch.int8)
|
||||
k_offset = torch.randint(-128, 127, (1, 1), dtype=torch.int8)
|
||||
v_offset = torch.randint(-128, 127, (1, 1), dtype=torch.int8)
|
||||
|
||||
# Load weights using weight_loader
|
||||
from vllm_ascend.quantization.methods.kv_c8 import _fa_quant_weight_loader as weight_loader
|
||||
|
||||
with torch.no_grad():
|
||||
weight_loader(layer.fa_q.scale, q_scale)
|
||||
weight_loader(layer.fa_k.scale, k_scale)
|
||||
weight_loader(layer.fa_v.scale, v_scale)
|
||||
weight_loader(layer.fa_q.offset, q_offset)
|
||||
weight_loader(layer.fa_k.offset, k_offset)
|
||||
weight_loader(layer.fa_v.offset, v_offset)
|
||||
|
||||
# Verify weights were loaded correctly
|
||||
self.assertTrue(torch.all(layer.fa_q.scale == q_scale))
|
||||
self.assertTrue(torch.all(layer.fa_k.scale == k_scale))
|
||||
self.assertTrue(torch.all(layer.fa_v.scale == v_scale))
|
||||
|
||||
# Step 3: Process after loading
|
||||
method.process_weights_after_loading(layer)
|
||||
|
||||
# Verify processed parameters
|
||||
self.assertTrue(hasattr(layer, "fak_descale"))
|
||||
self.assertTrue(hasattr(layer, "fak_offset"))
|
||||
self.assertTrue(hasattr(layer, "quant_kscale"))
|
||||
|
||||
|
||||
class TestC8KVScaleWeightLoader(TestBase):
|
||||
"""Tests for _c8_kv_scale_weight_loader in kv_c8.py."""
|
||||
|
||||
def setUp(self):
|
||||
from vllm_ascend.quantization.methods.kv_c8 import _c8_kv_scale_weight_loader
|
||||
|
||||
self.loader = _c8_kv_scale_weight_loader
|
||||
|
||||
def test_shape_match_copies_value(self):
|
||||
param = nn.Parameter(torch.ones(4, dtype=torch.float32), requires_grad=False)
|
||||
loaded = torch.tensor([1.0, 2.0, 3.0, 4.0])
|
||||
self.loader(param, loaded)
|
||||
self.assertTrue(torch.allclose(param.data, loaded.float()))
|
||||
|
||||
def test_shape_mismatch_resizes_param(self):
|
||||
param = nn.Parameter(torch.ones(1, dtype=torch.float32), requires_grad=False)
|
||||
loaded = torch.arange(8, dtype=torch.float32)
|
||||
self.loader(param, loaded)
|
||||
self.assertEqual(param.data.shape, (8,))
|
||||
self.assertTrue(torch.allclose(param.data, loaded))
|
||||
|
||||
def test_squeeze_before_compare(self):
|
||||
param = nn.Parameter(torch.ones(4, dtype=torch.float32), requires_grad=False)
|
||||
loaded = torch.arange(4, dtype=torch.float32).unsqueeze(0) # shape [1, 4]
|
||||
self.loader(param, loaded)
|
||||
self.assertEqual(param.data.shape, (4,))
|
||||
|
||||
def test_dtype_preserved_as_param_dtype(self):
|
||||
param = nn.Parameter(torch.ones(4, dtype=torch.float32), requires_grad=False)
|
||||
loaded = torch.arange(4, dtype=torch.float16)
|
||||
self.loader(param, loaded)
|
||||
self.assertEqual(param.data.dtype, torch.float32)
|
||||
|
||||
|
||||
class TestAscendC8KVCacheAttentionMethod(TestBase):
|
||||
"""Tests for AscendC8KVCacheAttentionMethod in kv_c8.py."""
|
||||
|
||||
def _make_method(self, is_kv_producer=False):
|
||||
from vllm_ascend.quantization.methods.kv_c8 import AscendC8KVCacheAttentionMethod
|
||||
|
||||
mock_config = MagicMock(spec=VllmConfig)
|
||||
if is_kv_producer:
|
||||
kv_config = MagicMock(spec=KVTransferConfig)
|
||||
kv_config.is_kv_producer = True
|
||||
mock_config.kv_transfer_config = kv_config
|
||||
else:
|
||||
mock_config.kv_transfer_config = None
|
||||
|
||||
with patch("vllm_ascend.quantization.methods.kv_c8.get_current_vllm_config", return_value=mock_config):
|
||||
return AscendC8KVCacheAttentionMethod(quant_description={}, prefix="model.layers.0.self_attn.attn")
|
||||
|
||||
def _make_layer_with_impl(self):
|
||||
layer = nn.Module()
|
||||
layer.impl = MagicMock()
|
||||
return layer
|
||||
|
||||
def test_create_weights_sets_kv_cache_torch_dtype(self):
|
||||
method = self._make_method(is_kv_producer=False)
|
||||
layer = self._make_layer_with_impl()
|
||||
method.create_weights(layer)
|
||||
self.assertEqual(layer.kv_cache_torch_dtype, torch.int8)
|
||||
|
||||
def test_create_weights_does_not_set_int8_when_kv_producer(self):
|
||||
method = self._make_method(is_kv_producer=True)
|
||||
layer = self._make_layer_with_impl()
|
||||
method.create_weights(layer)
|
||||
self.assertFalse(hasattr(layer, "kv_cache_torch_dtype"))
|
||||
|
||||
def test_create_weights_registers_scale_offset_params(self):
|
||||
method = self._make_method()
|
||||
layer = self._make_layer_with_impl()
|
||||
method.create_weights(layer)
|
||||
self.assertIsInstance(layer.k_cache_scale, nn.Parameter)
|
||||
self.assertIsInstance(layer.k_cache_offset, nn.Parameter)
|
||||
self.assertIsInstance(layer.v_cache_scale, nn.Parameter)
|
||||
self.assertIsInstance(layer.v_cache_offset, nn.Parameter)
|
||||
self.assertFalse(layer.k_cache_scale.requires_grad)
|
||||
self.assertFalse(layer.v_cache_offset.requires_grad)
|
||||
|
||||
def test_create_weights_initial_values(self):
|
||||
method = self._make_method()
|
||||
layer = self._make_layer_with_impl()
|
||||
method.create_weights(layer)
|
||||
self.assertEqual(layer.k_cache_scale.data.item(), 1.0)
|
||||
self.assertEqual(layer.v_cache_scale.data.item(), 1.0)
|
||||
self.assertEqual(layer.k_cache_offset.data.item(), 0.0)
|
||||
self.assertEqual(layer.v_cache_offset.data.item(), 0.0)
|
||||
|
||||
def test_create_weights_assigns_weight_loader(self):
|
||||
from vllm_ascend.quantization.methods.kv_c8 import _c8_kv_scale_weight_loader
|
||||
|
||||
method = self._make_method()
|
||||
layer = self._make_layer_with_impl()
|
||||
method.create_weights(layer)
|
||||
self.assertIs(layer.k_cache_scale.weight_loader, _c8_kv_scale_weight_loader)
|
||||
self.assertIs(layer.v_cache_scale.weight_loader, _c8_kv_scale_weight_loader)
|
||||
self.assertIs(layer.k_cache_offset.weight_loader, _c8_kv_scale_weight_loader)
|
||||
self.assertIs(layer.v_cache_offset.weight_loader, _c8_kv_scale_weight_loader)
|
||||
|
||||
def test_process_weights_after_loading_flattens(self):
|
||||
method = self._make_method()
|
||||
layer = nn.Module()
|
||||
layer.k_cache_scale = nn.Parameter(torch.ones(2, 4), requires_grad=False)
|
||||
layer.k_cache_offset = nn.Parameter(torch.zeros(2, 4), requires_grad=False)
|
||||
layer.v_cache_scale = nn.Parameter(torch.ones(2, 4), requires_grad=False)
|
||||
layer.v_cache_offset = nn.Parameter(torch.zeros(2, 4), requires_grad=False)
|
||||
method.process_weights_after_loading(layer)
|
||||
self.assertEqual(layer.k_cache_scale.data.dim(), 1)
|
||||
self.assertEqual(layer.k_cache_scale.data.shape[0], 8)
|
||||
self.assertEqual(layer.v_cache_offset.data.dim(), 1)
|
||||
|
||||
def test_apply_raises_runtime_error(self):
|
||||
method = self._make_method()
|
||||
layer = MagicMock()
|
||||
with self.assertRaises(RuntimeError):
|
||||
method.apply(layer, MagicMock(), MagicMock(), MagicMock(), None, None, None, None, None)
|
||||
|
||||
|
||||
class TestAscendC8AttentionBackendImplScales(TestBase):
|
||||
"""Tests for AscendC8AttentionBackendImpl scale helpers (scalar scales/offsets)."""
|
||||
|
||||
def setUp(self):
|
||||
self.original_dtype = torch.get_default_dtype()
|
||||
|
||||
def _make_impl(self, num_kv_heads=4, head_size=8):
|
||||
from vllm_ascend.attention.attention_v1 import AscendC8AttentionBackendImpl
|
||||
|
||||
impl = object.__new__(AscendC8AttentionBackendImpl)
|
||||
impl.num_heads = num_kv_heads
|
||||
impl.num_kv_heads = num_kv_heads
|
||||
impl.head_size = head_size
|
||||
impl.scale = 1.0
|
||||
impl.key_cache = None
|
||||
impl.value_cache = None
|
||||
return impl
|
||||
|
||||
def _make_layer(self, num_kv_heads=4, head_size=8):
|
||||
layer = nn.Module()
|
||||
shape = num_kv_heads * head_size
|
||||
layer.k_cache_scale = nn.Parameter(torch.ones(shape, dtype=self.original_dtype), requires_grad=False)
|
||||
layer.k_cache_offset = nn.Parameter(torch.zeros(shape, dtype=self.original_dtype), requires_grad=False)
|
||||
layer.v_cache_scale = nn.Parameter(torch.ones(shape, dtype=self.original_dtype), requires_grad=False)
|
||||
layer.v_cache_offset = nn.Parameter(torch.zeros(shape, dtype=self.original_dtype), requires_grad=False)
|
||||
return layer
|
||||
|
||||
@patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_rank", return_value=0)
|
||||
@patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_world_size", return_value=1)
|
||||
def test_prepare_c8_scales_runs_once(self, mock_tp_size, mock_tp_rank):
|
||||
impl = self._make_impl()
|
||||
layer = self._make_layer()
|
||||
impl._prepare_c8_scales(layer, torch.device("cpu"))
|
||||
self.assertTrue(hasattr(layer, "_c8_scales_prepared"))
|
||||
self.assertTrue(layer._c8_scales_prepared)
|
||||
|
||||
@patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_rank", return_value=0)
|
||||
@patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_world_size", return_value=1)
|
||||
def test_prepare_c8_scales_idempotent(self, mock_tp_size, mock_tp_rank):
|
||||
impl = self._make_impl()
|
||||
layer = self._make_layer()
|
||||
impl._prepare_c8_scales(layer, torch.device("cpu"))
|
||||
k_scale_after_first = layer._c8_k_scale.clone()
|
||||
layer.k_cache_scale.data = torch.ones(32, dtype=self.original_dtype) * 99
|
||||
impl._prepare_c8_scales(layer, torch.device("cpu"))
|
||||
self.assertTrue(torch.allclose(layer._c8_k_scale, k_scale_after_first))
|
||||
|
||||
@patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_rank", return_value=0)
|
||||
@patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_world_size", return_value=1)
|
||||
def test_prepare_c8_scales_creates_bnsd_shape(self, mock_tp_size, mock_tp_rank):
|
||||
num_kv_heads, head_size = 4, 8
|
||||
impl = self._make_impl(num_kv_heads, head_size)
|
||||
layer = self._make_layer(num_kv_heads, head_size)
|
||||
impl._prepare_c8_scales(layer, torch.device("cpu"))
|
||||
self.assertEqual(layer._c8_k_aq_scale_nz_bnsd.shape, (num_kv_heads, 1, head_size))
|
||||
self.assertEqual(layer._c8_v_aq_scale_nz_bnsd.shape, (num_kv_heads, 1, head_size))
|
||||
self.assertEqual(layer._c8_k_aq_scale_nz_bnsd.dtype, self.original_dtype)
|
||||
|
||||
@patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_rank", return_value=0)
|
||||
@patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_world_size", return_value=1)
|
||||
def test_quantize_kv_to_int8_output_dtype(self, mock_tp_size, mock_tp_rank):
|
||||
num_kv_heads, head_size = 4, 8
|
||||
impl = self._make_impl(num_kv_heads, head_size)
|
||||
layer = self._make_layer(num_kv_heads, head_size)
|
||||
impl._prepare_c8_scales(layer, torch.device("cpu"))
|
||||
num_tokens = 6
|
||||
key = torch.zeros(num_tokens, num_kv_heads, head_size, dtype=self.original_dtype)
|
||||
value = torch.zeros(num_tokens, num_kv_heads, head_size, dtype=self.original_dtype)
|
||||
key_q, value_q = impl._quantize_kv_to_int8(key, value, layer, num_tokens)
|
||||
self.assertEqual(key_q.dtype, torch.int8)
|
||||
self.assertEqual(value_q.dtype, torch.int8)
|
||||
self.assertEqual(key_q.shape, key.shape)
|
||||
|
||||
@patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_rank", return_value=0)
|
||||
@patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_world_size", return_value=1)
|
||||
def test_quantize_kv_to_int8_formula(self, mock_tp_size, mock_tp_rank):
|
||||
"""With scale=2.0, offset=0: q = round(x / 2)."""
|
||||
num_kv_heads, head_size = 1, 4
|
||||
impl = self._make_impl(num_kv_heads, head_size)
|
||||
layer = nn.Module()
|
||||
scale_val = torch.full((num_kv_heads * head_size,), 2.0, dtype=self.original_dtype)
|
||||
layer.k_cache_scale = nn.Parameter(scale_val.clone(), requires_grad=False)
|
||||
layer.k_cache_offset = nn.Parameter(
|
||||
torch.zeros(num_kv_heads * head_size, dtype=self.original_dtype), requires_grad=False
|
||||
)
|
||||
layer.v_cache_scale = nn.Parameter(scale_val.clone(), requires_grad=False)
|
||||
layer.v_cache_offset = nn.Parameter(
|
||||
torch.zeros(num_kv_heads * head_size, dtype=self.original_dtype), requires_grad=False
|
||||
)
|
||||
impl._prepare_c8_scales(layer, torch.device("cpu"))
|
||||
key = torch.full((1, num_kv_heads, head_size), 4.0, dtype=self.original_dtype)
|
||||
value = torch.full((1, num_kv_heads, head_size), 4.0, dtype=self.original_dtype)
|
||||
key_q, _ = impl._quantize_kv_to_int8(key, value, layer, 1)
|
||||
self.assertTrue(torch.all(key_q[0] == 2))
|
||||
|
||||
@patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_rank", return_value=0)
|
||||
@patch("vllm_ascend.attention.attention_v1.get_tensor_model_parallel_world_size", return_value=1)
|
||||
def test_dequant_paged_kv_to_dense_round_trip(self, mock_tp_size, mock_tp_rank):
|
||||
"""With scale=1, offset=0: dequant(int8) == float(int8)."""
|
||||
NZ_FMT_LAST_DIM = 32
|
||||
num_kv_heads, head_size = 2, 32 # head_size must be divisible by NZ_FMT_LAST_DIM
|
||||
block_size = 32
|
||||
num_blocks = 2
|
||||
nz_dim = head_size // NZ_FMT_LAST_DIM
|
||||
impl = self._make_impl(num_kv_heads, head_size)
|
||||
impl.key_cache = torch.empty(num_blocks, block_size, num_kv_heads, head_size)
|
||||
layer = self._make_layer(num_kv_heads, head_size)
|
||||
impl._prepare_c8_scales(layer, torch.device("cpu"))
|
||||
|
||||
# NZ 5D format: (num_blocks, num_kv_heads, nz_dim, block_size, NZ_FMT_LAST_DIM)
|
||||
key_int8 = torch.randint(
|
||||
-10, 10, (num_blocks, num_kv_heads, nz_dim, block_size, NZ_FMT_LAST_DIM), dtype=torch.int8
|
||||
)
|
||||
value_int8 = torch.randint(
|
||||
-10, 10, (num_blocks, num_kv_heads, nz_dim, block_size, NZ_FMT_LAST_DIM), dtype=torch.int8
|
||||
)
|
||||
seq_lens = [32, 32]
|
||||
block_table = torch.tensor([[0], [1]], dtype=torch.long)
|
||||
|
||||
dense_k, dense_v = impl._dequant_paged_kv_to_dense(
|
||||
key_int8, value_int8, block_table, seq_lens, torch.float32, layer
|
||||
)
|
||||
# Convert NZ 5D to ND 3D for expected comparison
|
||||
expected_k = key_int8.permute(0, 3, 1, 2, 4).contiguous().view(-1, num_kv_heads, head_size).float()
|
||||
self.assertEqual(dense_k.shape, (64, num_kv_heads, head_size))
|
||||
self.assertTrue(torch.allclose(dense_k, expected_k))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=2)
|
||||
23
tests/ut/quantization/methods/test_moe_logical_experts.py
Normal file
23
tests/ut/quantization/methods/test_moe_logical_experts.py
Normal file
@@ -0,0 +1,23 @@
|
||||
from types import SimpleNamespace
|
||||
|
||||
from vllm_ascend.quantization.methods.base import get_moe_num_logical_experts
|
||||
|
||||
|
||||
def test_get_moe_num_logical_experts_uses_vllm_config_field():
|
||||
layer = SimpleNamespace(moe_config=SimpleNamespace(num_logical_experts=128))
|
||||
|
||||
assert get_moe_num_logical_experts(layer, num_experts=130, global_redundant_expert_num=2) == 128
|
||||
|
||||
|
||||
def test_get_moe_num_logical_experts_falls_back_for_older_configs():
|
||||
layer = SimpleNamespace(moe_config=SimpleNamespace())
|
||||
|
||||
assert (
|
||||
get_moe_num_logical_experts(
|
||||
layer,
|
||||
num_experts=133,
|
||||
global_redundant_expert_num=2,
|
||||
num_shared_experts=3,
|
||||
)
|
||||
== 128
|
||||
)
|
||||
70
tests/ut/quantization/methods/test_registry.py
Normal file
70
tests/ut/quantization/methods/test_registry.py
Normal file
@@ -0,0 +1,70 @@
|
||||
from tests.ut.base import TestBase
|
||||
from vllm_ascend.quantization.methods.base import (
|
||||
AscendLinearScheme,
|
||||
AscendMoEScheme,
|
||||
)
|
||||
from vllm_ascend.quantization.methods.registry import (
|
||||
_SCHEME_REGISTRY,
|
||||
get_scheme_class,
|
||||
register_scheme,
|
||||
)
|
||||
|
||||
|
||||
class TestRegisterScheme(TestBase):
|
||||
def test_register_scheme(self):
|
||||
@register_scheme("TEST_QUANT_TYPE", "linear")
|
||||
class TestLinearScheme(AscendLinearScheme):
|
||||
def get_weight(self, input_size, output_size, params_dtype):
|
||||
return {}
|
||||
|
||||
def apply(self, layer, x, bias=None, tp_rank=0):
|
||||
return x
|
||||
|
||||
scheme_class = get_scheme_class("TEST_QUANT_TYPE", "linear")
|
||||
self.assertIs(scheme_class, TestLinearScheme)
|
||||
|
||||
def test_register_scheme_duplicate_raises(self):
|
||||
with self.assertRaises(ValueError):
|
||||
|
||||
@register_scheme("W8A8_DYNAMIC", "linear")
|
||||
class Duplicate:
|
||||
pass
|
||||
|
||||
|
||||
class TestGetSchemeClass(TestBase):
|
||||
def test_get_existing_scheme_class_existing_linear(self):
|
||||
cls = get_scheme_class("W8A8_DYNAMIC", "linear")
|
||||
self.assertIsNotNone(cls)
|
||||
self.assertTrue(issubclass(cls, AscendLinearScheme))
|
||||
|
||||
cls = get_scheme_class("W8A8_DYNAMIC", "moe")
|
||||
self.assertIsNotNone(cls)
|
||||
self.assertTrue(issubclass(cls, AscendMoEScheme))
|
||||
|
||||
cls = get_scheme_class("FAKQuant", "attention")
|
||||
self.assertIsNotNone(cls)
|
||||
|
||||
def test_get_nonexistent_scheme_class(self):
|
||||
cls = get_scheme_class("NONEXISTENT", "linear")
|
||||
self.assertIsNone(cls)
|
||||
cls = get_scheme_class("W8A8_DYNAMIC", "nonexistent")
|
||||
self.assertIsNone(cls)
|
||||
|
||||
def test_all_linear_schemes_subclass_ascend_linear_scheme(self):
|
||||
for (quant_type, layer_type), scheme_cls in _SCHEME_REGISTRY.items():
|
||||
if layer_type == "linear":
|
||||
self.assertTrue(
|
||||
issubclass(scheme_cls, AscendLinearScheme),
|
||||
f"{scheme_cls.__name__} for {quant_type}/{layer_type} should be subclass of AscendLinearScheme",
|
||||
)
|
||||
|
||||
def test_all_moe_schemes_subclass_ascend_moe_scheme(self):
|
||||
for (quant_type, layer_type), scheme_cls in _SCHEME_REGISTRY.items():
|
||||
if layer_type == "moe":
|
||||
self.assertTrue(
|
||||
issubclass(scheme_cls, AscendMoEScheme),
|
||||
f"{scheme_cls.__name__} for {quant_type}/{layer_type} should be subclass of AscendMoEScheme",
|
||||
)
|
||||
|
||||
def test_registry_not_empty(self):
|
||||
self.assertGreater(len(_SCHEME_REGISTRY), 0)
|
||||
63
tests/ut/quantization/methods/test_w4a16_mxfp4.py
Normal file
63
tests/ut/quantization/methods/test_w4a16_mxfp4.py
Normal file
@@ -0,0 +1,63 @@
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
from tests.ut.quantization.conftest_quantization import create_mock_ascend_config, create_mock_vllm_config
|
||||
from vllm_ascend.quantization.methods.w4a16_mxfp4 import AscendW4A16MXFP4FusedMoEMethod
|
||||
|
||||
|
||||
class TestAscendW4A16MXFP4MoEMethod(TestBase):
|
||||
num_experts = 8
|
||||
hidden_size = 128
|
||||
intermediate_size = 256
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a16_mxfp4.ensure_mxfp4_moe_available")
|
||||
@patch("vllm_ascend.quantization.methods.w4a16_mxfp4.get_current_vllm_config")
|
||||
@patch("vllm_ascend.quantization.methods.w4a16_mxfp4.get_ascend_config")
|
||||
@patch("vllm_ascend.quantization.methods.w4a16_mxfp4.get_ep_group")
|
||||
def setUp(self, mock_ep_group, mock_ascend, mock_vllm, mock_ensure):
|
||||
mock_vllm.return_value = create_mock_vllm_config()
|
||||
mock_ascend.return_value = create_mock_ascend_config()
|
||||
mock_ensure.return_value = None
|
||||
mock_ep_group.return_value = Mock()
|
||||
self.scheme = AscendW4A16MXFP4FusedMoEMethod()
|
||||
|
||||
@pytest.mark.skip("Execute after the issue is fixed")
|
||||
def test_get_weight_static_method(self):
|
||||
result = self.scheme.get_weight(self.num_experts, self.intermediate_size, self.hidden_size, torch.bfloat16)
|
||||
self.assertEqual(result["w13_weight"].dtype, torch.uint8)
|
||||
self.assertEqual(result["w2_weight"].dtype, torch.uint8)
|
||||
self.assertEqual(
|
||||
result["w13_weight"].shape, (self.num_experts, 2 * self.intermediate_size, self.hidden_size // 2)
|
||||
)
|
||||
self.assertEqual(result["w2_weight"].shape, (self.num_experts, self.hidden_size, self.intermediate_size // 2))
|
||||
|
||||
@pytest.mark.skip("Execute after the issue is fixed")
|
||||
def test_get_dynamic_quant_param_based_on_group_size(self):
|
||||
group_sizes = [16, 32, 64]
|
||||
for gs in group_sizes:
|
||||
self.scheme.group_size = gs
|
||||
result = self.scheme.get_dynamic_quant_param(
|
||||
self.num_experts, self.intermediate_size, self.hidden_size, torch.bfloat16
|
||||
)
|
||||
self.assertEqual(result["w13_weight_scale"].shape[2], self.hidden_size // gs)
|
||||
self.assertEqual(result["w13_weight_scale"].dtype, torch.uint8)
|
||||
self.assertEqual(result["w2_weight_scale"].dtype, torch.uint8)
|
||||
|
||||
@pytest.mark.skip("Execute after the issue is fixed")
|
||||
def test_process_weights_transposes_weights(self):
|
||||
layer = nn.Module()
|
||||
layer.w13_weight = nn.Parameter(torch.randint(0, 255, (8, 256, 64), dtype=torch.uint8), requires_grad=False)
|
||||
layer.w2_weight = nn.Parameter(torch.randint(0, 255, (8, 128, 128), dtype=torch.uint8), requires_grad=False)
|
||||
layer.w13_weight_scale = nn.Parameter(
|
||||
torch.randint(0, 255, (8, 256, 4), dtype=torch.uint8), requires_grad=False
|
||||
)
|
||||
layer.w2_weight_scale = nn.Parameter(torch.randint(0, 255, (8, 128, 8), dtype=torch.uint8), requires_grad=False)
|
||||
self.scheme.process_weights_after_loading(layer)
|
||||
self.assertEqual(layer.w13_weight.shape, (8, 128, 32))
|
||||
self.assertEqual(layer.w13_weight_scale.shape, (8, 4, 256))
|
||||
self.assertEqual(layer.w2_weight.shape, (8, 256, 16))
|
||||
self.assertEqual(layer.w2_weight_scale.shape, (8, 8, 128))
|
||||
140
tests/ut/quantization/methods/test_w4a4_mxfp4.py
Normal file
140
tests/ut/quantization/methods/test_w4a4_mxfp4.py
Normal file
@@ -0,0 +1,140 @@
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
from tests.ut.quantization.conftest_quantization import create_mock_ascend_config, create_mock_vllm_config
|
||||
from vllm_ascend.quantization.methods.w4a4_mxfp4 import (
|
||||
AscendW4A4MXFP4DynamicFusedMoEMethod,
|
||||
AscendW4A4MXFP4DynamicLinearMethod,
|
||||
)
|
||||
|
||||
|
||||
class TestAscendW4A4MXFP4LinearMethod(TestBase):
|
||||
@patch("vllm_ascend.quantization.methods.w4a4_mxfp4.ensure_mxfp4_linear_available")
|
||||
@patch("vllm_ascend.quantization.methods.w4a4_mxfp4.get_current_vllm_config")
|
||||
def setUp(self, mock_vllm, mock_ensure):
|
||||
mock_vllm.return_value = create_mock_vllm_config()
|
||||
mock_ensure.return_value = None
|
||||
self.scheme = AscendW4A4MXFP4DynamicLinearMethod()
|
||||
|
||||
def test_get_weight_various_input_sizes(self):
|
||||
for input_size in [64, 128, 256, 512]:
|
||||
result = self.scheme.get_weight(input_size, 128, torch.bfloat16)
|
||||
self.assertEqual(result["weight"].shape, (128, input_size // 2))
|
||||
self.assertEqual(result["weight"].dtype, torch.uint8)
|
||||
|
||||
def test_get_pergroup_param_based_on_group_size(self):
|
||||
group_sizes = [16, 32, 64]
|
||||
for gs in group_sizes:
|
||||
self.scheme.group_size = gs
|
||||
result = self.scheme.get_pergroup_param(256, 128, torch.bfloat16)
|
||||
self.assertEqual(result["weight_scale"].shape, (128, 256 // gs))
|
||||
self.assertEqual(result["weight_scale"].dtype, torch.uint8)
|
||||
|
||||
def test_process_weights_after_loading_transposes(self):
|
||||
layer = nn.Module()
|
||||
layer.weight = nn.Parameter(torch.randint(0, 255, (128, 128), dtype=torch.uint8), requires_grad=False)
|
||||
layer.weight_scale = nn.Parameter(torch.randint(0, 255, (128, 8), dtype=torch.uint8), requires_grad=False)
|
||||
self.scheme.process_weights_after_loading(layer)
|
||||
self.assertEqual(layer.weight.shape, (128, 128))
|
||||
self.assertEqual(layer.weight_scale.shape[0], 4)
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a4_mxfp4.torch_npu")
|
||||
def test_apply_3d_input(self, mock_npu):
|
||||
mock_npu.npu_dynamic_mx_quant.return_value = (
|
||||
torch.randint(0, 255, (32, 128), dtype=torch.uint8),
|
||||
torch.randint(0, 255, (32, 4), dtype=torch.uint8),
|
||||
)
|
||||
mock_npu.npu_quant_matmul.return_value = torch.randn(32, 1, 128)
|
||||
layer = MagicMock()
|
||||
layer.weight = MagicMock(data=torch.randint(0, 255, (128, 128), dtype=torch.uint8))
|
||||
layer.weight_scale = MagicMock(data=torch.randint(0, 255, (4, 128, 2), dtype=torch.uint8))
|
||||
x = torch.randn(32, 1, 256, dtype=torch.bfloat16)
|
||||
with patch.object(self.scheme, "group_size", 32):
|
||||
output = self.scheme.apply(layer, x)
|
||||
self.assertEqual(output.shape[0], 32)
|
||||
|
||||
|
||||
class TestAscendW4A4MXFP4MoEMethod(TestBase):
|
||||
num_experts = 8
|
||||
hidden_size = 128
|
||||
intermediate_size = 256
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a4_mxfp4.ensure_mxfp4_moe_available")
|
||||
@patch("vllm_ascend.quantization.methods.w4a4_mxfp4.get_current_vllm_config")
|
||||
@patch("vllm_ascend.quantization.methods.w4a4_mxfp4.get_ascend_config")
|
||||
def setUp(self, mock_ascend, mock_vllm, mock_ensure):
|
||||
mock_vllm.return_value = create_mock_vllm_config()
|
||||
mock_ascend.return_value = create_mock_ascend_config()
|
||||
mock_ensure.return_value = None
|
||||
self.scheme = AscendW4A4MXFP4DynamicFusedMoEMethod()
|
||||
|
||||
def test_get_weight_static_method(self):
|
||||
result = self.scheme.get_weight(self.num_experts, self.intermediate_size, self.hidden_size, torch.bfloat16)
|
||||
self.assertEqual(result["w13_weight"].dtype, torch.uint8)
|
||||
self.assertEqual(result["w2_weight"].dtype, torch.uint8)
|
||||
self.assertEqual(
|
||||
result["w13_weight"].shape, (self.num_experts, 2 * self.intermediate_size, self.hidden_size // 2)
|
||||
)
|
||||
self.assertEqual(result["w2_weight"].shape, (self.num_experts, self.hidden_size, self.intermediate_size // 2))
|
||||
|
||||
def test_get_dynamic_quant_param_based_on_group_size(self):
|
||||
group_sizes = [16, 32, 64]
|
||||
for gs in group_sizes:
|
||||
self.scheme.group_size = gs
|
||||
result = self.scheme.get_dynamic_quant_param(
|
||||
self.num_experts, self.intermediate_size, self.hidden_size, torch.bfloat16
|
||||
)
|
||||
self.assertEqual(result["w13_weight_scale"].shape[2], self.hidden_size // gs)
|
||||
self.assertEqual(result["w13_weight_scale"].dtype, torch.uint8)
|
||||
self.assertEqual(result["w2_weight_scale"].dtype, torch.uint8)
|
||||
|
||||
def test_process_weights_transposes_weights(self):
|
||||
layer = nn.Module()
|
||||
layer.w13_weight = nn.Parameter(torch.randint(0, 255, (8, 256, 64), dtype=torch.uint8), requires_grad=False)
|
||||
layer.w2_weight = nn.Parameter(torch.randint(0, 255, (8, 128, 128), dtype=torch.uint8), requires_grad=False)
|
||||
layer.w13_weight_scale = nn.Parameter(
|
||||
torch.randint(0, 255, (8, 256, 4), dtype=torch.uint8), requires_grad=False
|
||||
)
|
||||
layer.w2_weight_scale = nn.Parameter(torch.randint(0, 255, (8, 128, 8), dtype=torch.uint8), requires_grad=False)
|
||||
self.scheme.process_weights_after_loading(layer)
|
||||
self.assertEqual(layer.w13_weight.shape, (8, 64, 256))
|
||||
self.assertEqual(layer.w13_weight_scale.shape, (8, 2, 256, 2))
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a4_mxfp4.torch_npu")
|
||||
@patch("vllm_ascend.quantization.methods.w4a4_mxfp4._EXTRA_CTX")
|
||||
@patch("vllm_ascend.quantization.methods.w4a4_mxfp4.select_experts")
|
||||
def test_apply_full_params(self, mock_select, mock_ctx, mock_npu):
|
||||
tokens = 4
|
||||
layer = nn.Module()
|
||||
layer.w13_weight = nn.Parameter(torch.randint(0, 255, (8, 64, 256), dtype=torch.uint8), requires_grad=False)
|
||||
layer.w2_weight = nn.Parameter(torch.randint(0, 255, (8, 128, 128), dtype=torch.uint8), requires_grad=False)
|
||||
layer.w13_weight_scale = nn.Parameter(
|
||||
torch.randint(0, 255, (8, 64, 128, 2), dtype=torch.uint8), requires_grad=False
|
||||
)
|
||||
layer.w2_weight_scale = nn.Parameter(
|
||||
torch.randint(0, 255, (8, 128, 64, 2), dtype=torch.uint8), requires_grad=False
|
||||
)
|
||||
x = torch.randn(tokens, self.hidden_size, dtype=torch.bfloat16)
|
||||
router_logits = torch.randn(tokens, self.num_experts, dtype=torch.float32)
|
||||
topk_weights = torch.randn(tokens, 2)
|
||||
topk_ids = torch.randint(0, self.num_experts, (tokens, 2))
|
||||
mock_select.return_value = (topk_weights, topk_ids)
|
||||
mock_comm = Mock()
|
||||
mock_comm.fused_experts.return_value = torch.randn(tokens, self.hidden_size)
|
||||
mock_ctx.moe_comm_method = mock_comm
|
||||
mock_ctx.moe_comm_type = Mock()
|
||||
self.scheme.apply(
|
||||
layer,
|
||||
x,
|
||||
router_logits,
|
||||
top_k=2,
|
||||
renormalize=True,
|
||||
num_experts=self.num_experts,
|
||||
activation="silu",
|
||||
pertoken_scale=torch.randn(tokens),
|
||||
apply_router_weight_on_input=True,
|
||||
)
|
||||
mock_comm.fused_experts.assert_called_once()
|
||||
@@ -0,0 +1,202 @@
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
from vllm_ascend.quantization.methods.w4a4_mxfp4_flatquant import (
|
||||
MAX_SUPPORT_DIM,
|
||||
AscendW4A4MXFP4FlatQuantDynamicLinearMethod,
|
||||
get_decompose_dim,
|
||||
)
|
||||
|
||||
|
||||
class TestGetDecomposeDim(TestBase):
|
||||
"""Unit tests for the get_decompose_dim helper."""
|
||||
|
||||
def test_perfect_square_decomposition(self):
|
||||
self.assertEqual(get_decompose_dim(1024, 1), (32, 32))
|
||||
|
||||
def test_non_square_decomposition(self):
|
||||
left, right = get_decompose_dim(32, 1)
|
||||
self.assertEqual((left, right), (4, 8))
|
||||
self.assertEqual(left * right, 32)
|
||||
|
||||
def test_decomposition_product_equals_n(self):
|
||||
left, right = get_decompose_dim(256, 1)
|
||||
self.assertEqual(left * right, 256)
|
||||
|
||||
def test_raises_when_dim_sum_exceeds_max(self):
|
||||
n = (MAX_SUPPORT_DIM + 1) ** 2
|
||||
with self.assertRaisesRegex(ValueError, "should be less than"):
|
||||
get_decompose_dim(n, 1)
|
||||
|
||||
def test_fallback_when_left_times_m_exceeds_max(self):
|
||||
n = MAX_SUPPORT_DIM * MAX_SUPPORT_DIM
|
||||
left, right = get_decompose_dim(n, 2)
|
||||
self.assertEqual(left, MAX_SUPPORT_DIM)
|
||||
self.assertEqual(right, 2 * n // MAX_SUPPORT_DIM)
|
||||
|
||||
|
||||
class TestAscendW4A4MXFP4FlatQuantDynamicLinearMethod(TestBase):
|
||||
"""Unit tests for AscendW4A4MXFP4FlatQuantDynamicLinearMethod."""
|
||||
|
||||
input_size = 1024
|
||||
output_size = 64
|
||||
group_size = 32
|
||||
max_supported_tp = 4
|
||||
|
||||
def _build_method(self, tp_size=1, max_supported_tp=None, group_size=None):
|
||||
max_supported_tp = self.max_supported_tp if max_supported_tp is None else max_supported_tp
|
||||
group_size = self.group_size if group_size is None else group_size
|
||||
mock_vllm_config = Mock()
|
||||
mock_vllm_config.quant_config = Mock(
|
||||
quant_description={"group_size": group_size, "max_supported_tp": max_supported_tp}
|
||||
)
|
||||
with (
|
||||
patch("vllm_ascend.quantization.methods.w4a4_mxfp4_flatquant.ensure_mxfp4_flatquant_linear_available"),
|
||||
patch(
|
||||
"vllm_ascend.quantization.methods.w4a4_mxfp4_flatquant.get_current_vllm_config",
|
||||
return_value=mock_vllm_config,
|
||||
),
|
||||
patch(
|
||||
"vllm_ascend.quantization.methods.w4a4_mxfp4_flatquant.get_tensor_model_parallel_world_size",
|
||||
return_value=tp_size,
|
||||
),
|
||||
):
|
||||
return AscendW4A4MXFP4FlatQuantDynamicLinearMethod()
|
||||
|
||||
def setUp(self):
|
||||
self.method = self._build_method()
|
||||
|
||||
def test_init_default(self):
|
||||
self.assertEqual(self.method.group_size, self.group_size)
|
||||
self.assertEqual(self.method.max_supported_tp, self.max_supported_tp)
|
||||
self.assertEqual(self.method.tp_size, 1)
|
||||
|
||||
def test_init_raises_on_oversized_tp(self):
|
||||
with self.assertRaisesRegex(ValueError, "is not supported"):
|
||||
self._build_method(tp_size=8, max_supported_tp=4)
|
||||
|
||||
def test_get_weight(self):
|
||||
params = self.method.get_weight(self.input_size, self.output_size, torch.bfloat16)
|
||||
self.assertIn("weight", params)
|
||||
self.assertEqual(params["weight"].dtype, torch.uint8)
|
||||
self.assertEqual(params["weight"].shape, (self.output_size, self.input_size // 2))
|
||||
self.assertEqual(self.method.input_size, self.input_size)
|
||||
|
||||
def test_get_weight_raises_on_odd_input(self):
|
||||
with self.assertRaisesRegex(ValueError, "must be divisible by 2"):
|
||||
self.method.get_weight(127, self.output_size, torch.bfloat16)
|
||||
|
||||
def test_get_pertensor_param_non_row(self):
|
||||
self.method.get_weight(self.input_size, self.output_size, torch.bfloat16)
|
||||
params = self.method.get_pertensor_param(torch.bfloat16, layer_type="others")
|
||||
left_dim, right_dim = get_decompose_dim(self.input_size, 1)
|
||||
self.assertEqual(params["left_trans"].shape, (left_dim, left_dim))
|
||||
self.assertEqual(params["right_trans"].shape, (right_dim, right_dim))
|
||||
self.assertEqual(params["clip_ratio"].shape, (1,))
|
||||
self.assertEqual(params["left_trans"].dtype, torch.bfloat16)
|
||||
self.assertEqual(params["right_trans"].dtype, torch.bfloat16)
|
||||
self.assertEqual(params["clip_ratio"].dtype, torch.float32)
|
||||
|
||||
def test_get_pertensor_param_row(self):
|
||||
self.method.get_weight(self.input_size, self.output_size, torch.bfloat16)
|
||||
params = self.method.get_pertensor_param(torch.bfloat16, layer_type="row")
|
||||
origin_size = self.input_size * self.method.tp_size
|
||||
_, right_trans_dim = get_decompose_dim(
|
||||
origin_size // self.method.max_supported_tp, self.method.max_supported_tp
|
||||
)
|
||||
left_trans_dim = origin_size // right_trans_dim
|
||||
self.assertEqual(params["left_trans"].shape, (left_trans_dim, left_trans_dim))
|
||||
self.assertEqual(params["right_trans"].shape, (right_trans_dim, right_trans_dim))
|
||||
|
||||
def test_get_pergroup_param(self):
|
||||
params = self.method.get_pergroup_param(self.input_size, self.output_size, torch.bfloat16)
|
||||
self.assertIn("weight_scale", params)
|
||||
self.assertEqual(params["weight_scale"].dtype, torch.uint8)
|
||||
self.assertEqual(
|
||||
params["weight_scale"].shape,
|
||||
(self.output_size, self.input_size // self.group_size),
|
||||
)
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a4_mxfp4_flatquant.torch_npu")
|
||||
def test_apply(self, mock_torch_npu):
|
||||
layer = MagicMock()
|
||||
layer.left_trans = torch.randn(32, 32)
|
||||
layer.right_trans = torch.randn(32, 32)
|
||||
layer.aclnn_clip_ratio = 0.9
|
||||
layer.weight = MagicMock()
|
||||
layer.weight_scale = MagicMock()
|
||||
|
||||
batch = 8
|
||||
x = torch.randn(batch, self.input_size, dtype=torch.bfloat16)
|
||||
bias = torch.randn(self.output_size, dtype=torch.bfloat16)
|
||||
|
||||
mock_torch_npu.npu_kronecker_quant.return_value = (MagicMock(), MagicMock())
|
||||
expected_output = torch.randn(batch, self.output_size, dtype=torch.bfloat16)
|
||||
mock_torch_npu.npu_quant_matmul.return_value = expected_output
|
||||
|
||||
output = self.method.apply(layer, x, bias=bias)
|
||||
|
||||
mock_torch_npu.npu_kronecker_quant.assert_called_once()
|
||||
mock_torch_npu.npu_quant_matmul.assert_called_once()
|
||||
call_kwargs = mock_torch_npu.npu_quant_matmul.call_args.kwargs
|
||||
self.assertIs(call_kwargs["bias"], bias)
|
||||
self.assertEqual(call_kwargs["output_dtype"], torch.bfloat16)
|
||||
self.assertEqual(call_kwargs["group_sizes"], [1, 1, self.method.group_size])
|
||||
self.assertEqual(output.shape, (batch, self.output_size))
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w4a4_mxfp4_flatquant.torch_npu")
|
||||
def test_apply_preserves_input_shape(self, mock_torch_npu):
|
||||
layer = MagicMock()
|
||||
layer.left_trans = torch.randn(32, 32)
|
||||
layer.right_trans = torch.randn(32, 32)
|
||||
layer.aclnn_clip_ratio = 0.9
|
||||
x = torch.randn(2, 4, self.input_size, dtype=torch.bfloat16)
|
||||
mock_torch_npu.npu_kronecker_quant.return_value = (MagicMock(), MagicMock())
|
||||
mock_torch_npu.npu_quant_matmul.return_value = torch.randn(8, self.output_size, dtype=torch.bfloat16)
|
||||
output = self.method.apply(layer, x)
|
||||
self.assertEqual(output.shape, (2, 4, self.output_size))
|
||||
|
||||
def test_apply_dimension_mismatch_raises(self):
|
||||
layer = MagicMock()
|
||||
layer.left_trans = torch.randn(16, 16)
|
||||
layer.right_trans = torch.randn(16, 16)
|
||||
x = torch.randn(4, self.input_size)
|
||||
with self.assertRaisesRegex(ValueError, "dimension mismatch"):
|
||||
self.method.apply(layer, x)
|
||||
|
||||
def test_process_weights_after_loading(self):
|
||||
layer = MagicMock()
|
||||
layer.weight.data = torch.randint(0, 255, (self.output_size, self.input_size // 2), dtype=torch.uint8)
|
||||
|
||||
weight_scale_data = torch.randint(
|
||||
0, 255, (self.output_size, self.input_size // self.group_size), dtype=torch.uint8
|
||||
)
|
||||
layer.weight_scale.data = weight_scale_data
|
||||
layer.weight_scale.shape = weight_scale_data.shape
|
||||
layer.left_trans.data = torch.randn(32, 32, dtype=torch.bfloat16)
|
||||
layer.right_trans.data = torch.randn(32, 32, dtype=torch.bfloat16)
|
||||
layer.clip_ratio.data = torch.tensor([0.95])
|
||||
|
||||
self.method.process_weights_after_loading(layer)
|
||||
|
||||
# weight transposed: (output, input/2) -> (input/2, output)
|
||||
self.assertEqual(layer.weight.data.shape, (self.input_size // 2, self.output_size))
|
||||
# weight_scale view+transpose: (out, in/group) -> (in/group/2, out, 2)
|
||||
self.assertEqual(
|
||||
layer.weight_scale.data.shape,
|
||||
(self.input_size // self.group_size // 2, self.output_size, 2),
|
||||
)
|
||||
# left_trans is parameterized after a t().contiguous(); shape remains (32, 32)
|
||||
self.assertIsInstance(layer.left_trans, torch.nn.Parameter)
|
||||
self.assertEqual(layer.left_trans.shape, (32, 32))
|
||||
self.assertTrue(layer.left_trans.data.is_contiguous())
|
||||
# clip_ratio cast to float32, aclnn_clip_ratio set to its scalar value
|
||||
self.assertEqual(layer.clip_ratio.dtype, torch.float32)
|
||||
self.assertAlmostEqual(layer.aclnn_clip_ratio, 0.95, places=5)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(argv=["first-arg-is-ignored"], exit=False)
|
||||
169
tests/ut/quantization/methods/test_w8a8_mxfp8.py
Normal file
169
tests/ut/quantization/methods/test_w8a8_mxfp8.py
Normal file
@@ -0,0 +1,169 @@
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
from tests.ut.quantization.conftest_quantization import (
|
||||
create_mock_ascend_config,
|
||||
create_mock_vllm_config,
|
||||
create_mxfp_moe_layer,
|
||||
)
|
||||
from vllm_ascend.quantization.methods.w8a8_mxfp8 import (
|
||||
AscendW8A8MXFP8DynamicFusedMoEMethod,
|
||||
AscendW8A8MXFP8DynamicLinearMethod,
|
||||
)
|
||||
|
||||
|
||||
class TestAscendW8A8MXFP8LinearMethod(TestBase):
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_mxfp8.ensure_mxfp8_linear_available")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_mxfp8.get_current_vllm_config")
|
||||
def setUp(self, mock_vllm, mock_ensure):
|
||||
mock_vllm.return_value = create_mock_vllm_config()
|
||||
mock_ensure.return_value = None
|
||||
self.scheme = AscendW8A8MXFP8DynamicLinearMethod()
|
||||
|
||||
def test_get_weight_various_input_sizes(self):
|
||||
sizes = [(128, 64), (512, 256), (1024, 512)]
|
||||
for input_size, output_size in sizes:
|
||||
result = self.scheme.get_weight(input_size, output_size, torch.bfloat16)
|
||||
self.assertEqual(result["weight"].shape, (output_size, input_size))
|
||||
self.assertEqual(result["weight"].dtype, torch.float8_e4m3fn)
|
||||
|
||||
def test_get_pergroup_param_group_size_variations(self):
|
||||
group_sizes = [16, 32, 64, 128]
|
||||
for gs in group_sizes:
|
||||
self.scheme.group_size = gs
|
||||
result = self.scheme.get_pergroup_param(256, 128, torch.bfloat16)
|
||||
self.assertEqual(result["weight_scale"].shape, (128, 256 // gs))
|
||||
self.assertEqual(result["weight_scale"].dtype, torch.uint8)
|
||||
|
||||
def test_process_weights_stores_original_shapes(self):
|
||||
layer = nn.Module()
|
||||
layer.weight = nn.Parameter(torch.randn(128, 256).to(torch.float8_e4m3fn), requires_grad=False)
|
||||
layer.weight_scale = nn.Parameter(torch.randint(0, 255, (128, 8), dtype=torch.uint8), requires_grad=False)
|
||||
self.scheme.process_weights_after_loading(layer)
|
||||
self.assertTrue(hasattr(layer, "_mxfp8_original_shapes"))
|
||||
self.assertEqual(layer._mxfp8_original_shapes["weight"], (128, 256))
|
||||
self.assertTrue(layer._mxfp8_transformed)
|
||||
self.assertEqual(layer.weight_scale.shape, (4, 128, 2))
|
||||
self.assertTrue(layer.weight.data.is_contiguous())
|
||||
self.assertTrue(layer.weight_scale.data.is_contiguous())
|
||||
|
||||
def test_restore_after_process_returns_original_shape(self):
|
||||
layer = nn.Module()
|
||||
layer.weight = nn.Parameter(torch.randn(128, 256).to(torch.float8_e4m3fn), requires_grad=False)
|
||||
layer.weight_scale = nn.Parameter(torch.randint(0, 255, (128, 8), dtype=torch.uint8), requires_grad=False)
|
||||
original_weight_shape = layer.weight.shape
|
||||
original_scale_shape = layer.weight_scale.shape
|
||||
self.scheme.process_weights_after_loading(layer)
|
||||
self.scheme.restore_weights_for_rl_loading(layer)
|
||||
self.assertEqual(layer.weight.shape, original_weight_shape)
|
||||
self.assertEqual(layer.weight_scale.shape, original_scale_shape)
|
||||
self.assertFalse(layer._mxfp8_transformed)
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_mxfp8.torch_npu")
|
||||
def test_apply(self, mock_torch_npu):
|
||||
from vllm_ascend.device.mxfp_compat import FLOAT8_E8M0FNU_DTYPE
|
||||
|
||||
dynamic_scale = torch.randint(0, 255, (32, 8), dtype=torch.uint8)
|
||||
mock_torch_npu.npu_dynamic_mx_quant.return_value = (
|
||||
torch.randint(0, 255, (32, 256), dtype=torch.uint8),
|
||||
dynamic_scale,
|
||||
)
|
||||
mock_torch_npu.npu_quant_matmul.return_value = torch.randn(32, 128, dtype=torch.float16)
|
||||
layer = nn.Module()
|
||||
layer.weight = nn.Parameter(torch.randn(256, 128).to(torch.float8_e4m3fn), requires_grad=False)
|
||||
layer.weight_scale = nn.Parameter(torch.randint(0, 255, (4, 128, 2), dtype=torch.uint8), requires_grad=False)
|
||||
x = torch.randn(32, 1, 256, dtype=torch.float16)
|
||||
bias = torch.randn(128, dtype=torch.float16)
|
||||
output = self.scheme.apply(layer, x, bias)
|
||||
self.assertEqual(output.shape, (32, 1, 128))
|
||||
call_kwargs = mock_torch_npu.npu_quant_matmul.call_args.kwargs
|
||||
self.assertEqual(call_kwargs["bias"].dtype, torch.float32)
|
||||
self.assertEqual(call_kwargs["group_sizes"], [1, 1, self.scheme.group_size])
|
||||
self.assertEqual(call_kwargs["scale_dtype"], FLOAT8_E8M0FNU_DTYPE)
|
||||
self.assertEqual(call_kwargs["output_dtype"], torch.float16)
|
||||
|
||||
|
||||
class TestAscendW8A8MXFP8MoEMethod(TestBase):
|
||||
num_experts = 8
|
||||
hidden_size = 128
|
||||
intermediate_size = 256
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_mxfp8.ensure_mxfp8_moe_available")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_mxfp8.get_current_vllm_config")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_mxfp8.get_ascend_config")
|
||||
def setUp(self, mock_ascend, mock_vllm, mock_ensure):
|
||||
mock_vllm.return_value = create_mock_vllm_config()
|
||||
mock_ascend.return_value = create_mock_ascend_config()
|
||||
mock_ensure.return_value = None
|
||||
self.scheme = AscendW8A8MXFP8DynamicFusedMoEMethod()
|
||||
|
||||
def test_get_weight_various_expert_counts(self):
|
||||
for num_experts in [4, 8, 16]:
|
||||
result = self.scheme.get_weight(num_experts, self.intermediate_size, self.hidden_size, torch.bfloat16)
|
||||
self.assertEqual(result["w13_weight"].shape[0], num_experts)
|
||||
self.assertEqual(result["w2_weight"].dtype, torch.float8_e4m3fn)
|
||||
|
||||
def test_get_dynamic_quant_param_dtype_uint8(self):
|
||||
result = self.scheme.get_dynamic_quant_param(
|
||||
self.num_experts, self.intermediate_size, self.hidden_size, torch.bfloat16
|
||||
)
|
||||
self.assertEqual(result["w13_weight_scale"].shape, (8, 512, 4))
|
||||
self.assertEqual(result["w2_weight_scale"].dtype, torch.uint8)
|
||||
|
||||
def test_process_weights_stores_original_shapes(self):
|
||||
layer = create_mxfp_moe_layer(
|
||||
num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size
|
||||
)
|
||||
original_shape = layer.w13_weight.shape
|
||||
self.scheme.process_weights_after_loading(layer)
|
||||
self.assertTrue(hasattr(layer, "_mxfp8_original_shapes"))
|
||||
self.assertIn("w13_weight", layer._mxfp8_original_shapes)
|
||||
self.assertEqual(layer.w13_weight.shape, (original_shape[0], original_shape[2], original_shape[1]))
|
||||
self.assertFalse(layer.w13_weight.data.is_contiguous())
|
||||
self.assertFalse(layer.w2_weight.data.is_contiguous())
|
||||
self.assertFalse(layer.w13_weight_scale.data.is_contiguous())
|
||||
self.assertFalse(layer.w2_weight_scale.data.is_contiguous())
|
||||
|
||||
def test_restore_weights_for_rl_loading(self):
|
||||
layer = create_mxfp_moe_layer(
|
||||
num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size
|
||||
)
|
||||
original_w13_shape = layer.w13_weight.shape
|
||||
self.scheme.process_weights_after_loading(layer)
|
||||
self.assertNotEqual(layer.w13_weight.shape, original_w13_shape)
|
||||
self.scheme.restore_weights_for_rl_loading(layer)
|
||||
self.assertEqual(layer.w13_weight.shape, original_w13_shape)
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_mxfp8._EXTRA_CTX")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_mxfp8.select_experts")
|
||||
def test_apply_full_params(self, mock_select, mock_ctx):
|
||||
tokens = 4
|
||||
layer = create_mxfp_moe_layer(
|
||||
num_experts=self.num_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size
|
||||
)
|
||||
self.scheme.process_weights_after_loading(layer)
|
||||
layer.swiglu_limit = 1000000
|
||||
x = torch.randn(tokens, self.hidden_size, dtype=torch.bfloat16)
|
||||
router_logits = torch.randn(tokens, self.num_experts, dtype=torch.float32)
|
||||
topk_weights = torch.randn(tokens, 2)
|
||||
topk_ids = torch.randint(0, self.num_experts, (tokens, 2))
|
||||
mock_select.return_value = (topk_weights, topk_ids)
|
||||
mock_comm = Mock()
|
||||
mock_comm.fused_experts.return_value = torch.randn(tokens, self.hidden_size)
|
||||
mock_ctx.moe_comm_method = mock_comm
|
||||
mock_ctx.moe_comm_type = Mock()
|
||||
self.scheme.apply(
|
||||
layer,
|
||||
x,
|
||||
router_logits,
|
||||
top_k=2,
|
||||
renormalize=True,
|
||||
num_experts=self.num_experts,
|
||||
activation="silu",
|
||||
pertoken_scale=torch.randn(tokens),
|
||||
)
|
||||
mock_select.assert_called_once()
|
||||
mock_comm.fused_experts.assert_called_once()
|
||||
138
tests/ut/quantization/methods/test_w8a8_pdmix.py
Normal file
138
tests/ut/quantization/methods/test_w8a8_pdmix.py
Normal file
@@ -0,0 +1,138 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
from tests.ut.quantization.conftest_quantization import create_mock_vllm_config
|
||||
from vllm_ascend.quantization.methods import (
|
||||
AscendW8A8LinearMethod,
|
||||
AscendW8A8PDMixFusedMoeMethod,
|
||||
AscendW8A8PDMixLinearMethod,
|
||||
)
|
||||
|
||||
|
||||
class TestAscendW8A8PDMixLinearScheme(TestBase):
|
||||
def setUp(self):
|
||||
self.method = AscendW8A8LinearMethod()
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8LinearMethod")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8DynamicLinearMethod")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_pdmix.get_current_vllm_config")
|
||||
def test_get_weight_delegates_to_static(self, mock_vllm_config, mock_dynamic_cls, mock_static_cls):
|
||||
mock_vllm_config.return_value = create_mock_vllm_config(kv_transfer_config=None)
|
||||
mock_dynamic_instance = MagicMock()
|
||||
mock_dynamic_cls.return_value = mock_dynamic_instance
|
||||
mock_static_instance = MagicMock()
|
||||
mock_static_instance.get_weight.return_value = {"weight": torch.empty(128, 256, dtype=torch.int8)}
|
||||
mock_static_cls.return_value = mock_static_instance
|
||||
scheme = AscendW8A8PDMixLinearMethod()
|
||||
for input_size, output_size in [(64, 128), (256, 512), (1024, 2048)]:
|
||||
scheme.get_weight(input_size, output_size, torch.bfloat16)
|
||||
mock_static_instance.get_weight.assert_called_with(input_size, output_size, torch.bfloat16)
|
||||
mock_dynamic_instance.get_weight.assert_not_called()
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8LinearMethod")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8DynamicLinearMethod")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_pdmix.get_current_vllm_config")
|
||||
def test_get_pertensor_param_delegates_to_static(self, mock_vllm_config, mock_dynamic_cls, mock_static_cls):
|
||||
mock_vllm_config.return_value = create_mock_vllm_config(kv_transfer_config=None)
|
||||
mock_dynamic_instance = MagicMock()
|
||||
mock_dynamic_cls.return_value = mock_dynamic_instance
|
||||
mock_static_instance = MagicMock()
|
||||
mock_static_instance.get_pertensor_param.return_value = {}
|
||||
mock_static_cls.return_value = mock_static_instance
|
||||
scheme = AscendW8A8PDMixLinearMethod()
|
||||
scheme.get_pertensor_param(torch.bfloat16)
|
||||
mock_static_instance.get_pertensor_param.assert_called_once_with(torch.bfloat16)
|
||||
mock_dynamic_instance.get_pertensor_param.assert_not_called()
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8LinearMethod")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8DynamicLinearMethod")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_pdmix.get_current_vllm_config")
|
||||
def test_get_perchannel_param_delegates_to_static(self, mock_vllm_config, mock_dynamic_cls, mock_static_cls):
|
||||
mock_vllm_config.return_value = create_mock_vllm_config(kv_transfer_config=None)
|
||||
mock_dynamic_instance = MagicMock()
|
||||
mock_dynamic_cls.return_value = mock_dynamic_instance
|
||||
mock_static_instance = MagicMock()
|
||||
mock_static_instance.get_perchannel_param.return_value = {}
|
||||
mock_static_cls.return_value = mock_static_instance
|
||||
scheme = AscendW8A8PDMixLinearMethod()
|
||||
scheme.get_perchannel_param(128, torch.bfloat16)
|
||||
mock_static_instance.get_perchannel_param.assert_called_once_with(128, torch.bfloat16)
|
||||
mock_dynamic_instance.get_perchannel_param.assert_not_called()
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8LinearMethod")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8DynamicLinearMethod")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_pdmix.get_current_vllm_config")
|
||||
def test_apply_uses_static_for_kv_consumer(self, mock_vllm_config, mock_dynamic_cls, mock_static_cls):
|
||||
mock_vllm_config.return_value = create_mock_vllm_config(kv_transfer_config=None)
|
||||
mock_static_instance = MagicMock()
|
||||
mock_static_instance.apply.return_value = torch.randn(4, 128)
|
||||
mock_static_cls.return_value = mock_static_instance
|
||||
mock_dynamic_instance = MagicMock()
|
||||
mock_dynamic_cls.return_value = mock_dynamic_instance
|
||||
scheme = AscendW8A8PDMixLinearMethod()
|
||||
layer = MagicMock()
|
||||
layer.is_kv_consumer = True
|
||||
x = torch.randn(4, 256)
|
||||
scheme.apply(layer, x)
|
||||
mock_static_instance.apply.assert_called_once()
|
||||
mock_dynamic_instance.apply.assert_not_called()
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8LinearMethod")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8DynamicLinearMethod")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_pdmix.get_current_vllm_config")
|
||||
def test_apply_uses_dynamic_for_non_kv_consumer(self, mock_vllm_config, mock_dynamic_cls, mock_static_cls):
|
||||
mock_vllm_config.return_value = create_mock_vllm_config(kv_transfer_config=None)
|
||||
mock_dynamic_instance = MagicMock()
|
||||
mock_dynamic_instance.apply.return_value = torch.randn(4, 128)
|
||||
mock_dynamic_cls.return_value = mock_dynamic_instance
|
||||
mock_static_instance = MagicMock()
|
||||
mock_static_cls.return_value = mock_static_instance
|
||||
scheme = AscendW8A8PDMixLinearMethod()
|
||||
layer = MagicMock()
|
||||
layer.is_kv_consumer = False
|
||||
x = torch.randn(4, 256)
|
||||
scheme.apply(layer, x)
|
||||
mock_dynamic_instance.apply.assert_called_once()
|
||||
mock_static_instance.apply.assert_not_called()
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8LinearMethod")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_pdmix.AscendW8A8DynamicLinearMethod")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_pdmix.get_current_vllm_config")
|
||||
def test_process_weights_after_loading_sets_is_kv_consumer(
|
||||
self, mock_vllm_config, mock_dynamic_cls, mock_static_cls
|
||||
):
|
||||
mock_vllm_config.return_value = create_mock_vllm_config(kv_transfer_config=None)
|
||||
mock_static_instance = MagicMock()
|
||||
mock_static_cls.return_value = mock_static_instance
|
||||
mock_dynamic_instance = MagicMock()
|
||||
mock_dynamic_cls.return_value = mock_dynamic_instance
|
||||
scheme = AscendW8A8PDMixLinearMethod()
|
||||
layer = MagicMock()
|
||||
layer.weight_scale = MagicMock(data=torch.randn(128, 1, dtype=torch.bfloat16))
|
||||
scheme.process_weights_after_loading(layer)
|
||||
mock_static_instance.process_weights_after_loading.assert_called_once_with(layer)
|
||||
mock_dynamic_instance.process_weights_after_loading.assert_not_called()
|
||||
self.assertFalse(layer.is_kv_consumer)
|
||||
|
||||
|
||||
class TestAscendW8A8PDMixMoEScheme(TestBase):
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_mc2_group")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_current_vllm_config")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_ascend_config")
|
||||
def test_get_dynamic_quant_param(self, mock_ascend, mock_vllm, mock_mc2):
|
||||
mock_mc2.side_effect = AttributeError()
|
||||
mock_vllm.return_value = create_mock_vllm_config()
|
||||
mock_ascend.return_value = MagicMock(eplb_config=MagicMock(dynamic_eplb=False))
|
||||
scheme = AscendW8A8PDMixFusedMoeMethod()
|
||||
num_experts, intermediate_size_per_partition, hidden_sizes, params_dtype = 8, 256, 128, torch.bfloat16
|
||||
result = scheme.get_dynamic_quant_param(
|
||||
num_experts, intermediate_size_per_partition, hidden_sizes, params_dtype
|
||||
)
|
||||
# test adds extra params
|
||||
self.assertEqual(result["w2_deq_scale"].shape, (num_experts, hidden_sizes))
|
||||
self.assertEqual(result["w2_deq_scale"].dtype, torch.float32)
|
||||
self.assertEqual(result["w13_deq_scale"].shape, (num_experts, 2 * intermediate_size_per_partition))
|
||||
self.assertEqual(result["w2_input_offset"].dtype, torch.int8)
|
||||
self.assertEqual(result["w13_input_offset"].shape, (num_experts, 1))
|
||||
136
tests/ut/quantization/methods/test_w8a8fp8_dynamic.py
Normal file
136
tests/ut/quantization/methods/test_w8a8fp8_dynamic.py
Normal file
@@ -0,0 +1,136 @@
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
from tests.ut.quantization.conftest_quantization import (
|
||||
create_mock_ascend_config,
|
||||
create_mock_vllm_config,
|
||||
)
|
||||
from vllm_ascend.ascend_forward_context import MoECommType
|
||||
from vllm_ascend.quantization.methods.w8a8fp8_dynamic import (
|
||||
AscendW8A8FP8DynamicFusedMoEMethod,
|
||||
AscendW8A8FP8DynamicLinearMethod,
|
||||
)
|
||||
|
||||
|
||||
class TestAscendW8A8FP8DynamicLinearMethod(TestBase):
|
||||
def setUp(self):
|
||||
self.method = AscendW8A8FP8DynamicLinearMethod()
|
||||
|
||||
def test_act_quant_type(self):
|
||||
self.assertEqual(self.method.act_quant_type, torch.float8_e4m3fn)
|
||||
|
||||
def test_get_weight_various_sizes(self):
|
||||
sizes = [(64, 128), (256, 512), (1024, 2048)]
|
||||
for input_size, output_size in sizes:
|
||||
weight = self.method.get_weight(input_size, output_size, torch.bfloat16)
|
||||
self.assertEqual(weight["weight"].dtype, torch.float8_e4m3fn)
|
||||
self.assertEqual(weight["weight"].shape, (output_size, input_size))
|
||||
|
||||
def test_get_perchannel_param_dtype_variations(self):
|
||||
dtypes = [torch.bfloat16, torch.float16]
|
||||
for dtype in dtypes:
|
||||
params = self.method.get_perchannel_param(128, dtype)
|
||||
self.assertEqual(params["weight_scale"].dtype, torch.float32)
|
||||
self.assertEqual(params["weight_offset"].dtype, dtype)
|
||||
self.assertEqual(params["weight_scale"].shape, (128, 1))
|
||||
self.assertEqual(params["weight_offset"].shape, (128, 1))
|
||||
|
||||
|
||||
class TestAscendW8A8FP8FusedMoEMethod(TestBase):
|
||||
num_experts = 8
|
||||
hidden_size = 128
|
||||
intermediate_size = 128
|
||||
|
||||
@patch("torch.distributed.get_rank")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_mc2_group")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_ascend_config")
|
||||
def setUp(self, mock_ascend, mock_mc2, mock_rank):
|
||||
with patch("vllm_ascend.quantization.methods.w8a8_dynamic.get_current_vllm_config") as mock_vllm:
|
||||
mock_vllm.return_value = create_mock_vllm_config()
|
||||
mock_ascend.return_value = create_mock_ascend_config()
|
||||
mock_mc2.return_value = MagicMock(
|
||||
device_group=Mock(
|
||||
_get_backend=Mock(return_value=Mock(get_hccl_comm_name=Mock(return_value="test_comm")))
|
||||
)
|
||||
)
|
||||
mock_rank.return_value = 0
|
||||
self.quant_method = AscendW8A8FP8DynamicFusedMoEMethod()
|
||||
|
||||
def test_quant_type_is_w8a8fp8(self):
|
||||
from vllm_ascend.quantization.quant_type import QuantType
|
||||
|
||||
self.assertEqual(self.quant_method.quant_type, QuantType.W8A8FP8)
|
||||
|
||||
def test_get_weight_dtype_is_float8_e4m3fn(self):
|
||||
param_dict = self.quant_method.get_weight(
|
||||
self.num_experts, self.intermediate_size, self.hidden_size, torch.bfloat16
|
||||
)
|
||||
self.assertEqual(param_dict["w13_weight"].dtype, torch.float8_e4m3fn)
|
||||
self.assertEqual(param_dict["w2_weight"].dtype, torch.float8_e4m3fn)
|
||||
self.assertEqual(
|
||||
param_dict["w13_weight"].shape, (self.num_experts, 2 * self.intermediate_size, self.hidden_size)
|
||||
)
|
||||
self.assertEqual(param_dict["w2_weight"].shape, (self.num_experts, self.hidden_size, self.intermediate_size))
|
||||
|
||||
def test_get_weight_various_expert_counts(self):
|
||||
expert_counts = [4, 8, 16, 32]
|
||||
for num_experts in expert_counts:
|
||||
param_dict = self.quant_method.get_weight(
|
||||
num_experts, self.intermediate_size, self.hidden_size, torch.bfloat16
|
||||
)
|
||||
self.assertEqual(param_dict["w13_weight"].shape[0], num_experts)
|
||||
self.assertEqual(param_dict["w2_weight"].shape[0], num_experts)
|
||||
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_dynamic._EXTRA_CTX")
|
||||
@patch("vllm_ascend.quantization.methods.w8a8_dynamic.select_experts")
|
||||
def test_apply_uses_explicit_dispatch_and_mlp_args(self, mock_select_experts, mock_extra_ctx):
|
||||
tokens = 4
|
||||
hidden_size = self.hidden_size
|
||||
layer = torch.nn.Module()
|
||||
layer.w13_weight = torch.randn(
|
||||
self.num_experts, 2 * self.intermediate_size, hidden_size, dtype=torch.bfloat16
|
||||
).to(torch.float8_e4m3fn)
|
||||
layer.w2_weight = torch.randn(self.num_experts, hidden_size, self.intermediate_size, dtype=torch.bfloat16).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
layer.w13_weight_scale_fp32 = torch.ones(self.num_experts, 2 * self.intermediate_size, dtype=torch.float32)
|
||||
layer.w2_weight_scale = torch.ones(self.num_experts, hidden_size, dtype=torch.float32)
|
||||
layer.swiglu_limit = 1000000
|
||||
|
||||
x = torch.randn(tokens, hidden_size, dtype=torch.float32)
|
||||
router_logits = torch.randn(tokens, self.num_experts, dtype=torch.float32)
|
||||
topk_weights = torch.randn(tokens, 2, dtype=torch.float32)
|
||||
topk_ids = torch.randint(0, self.num_experts, (tokens, 2), dtype=torch.int64)
|
||||
mc2_mask = torch.tensor([1, 0, 1, 0], dtype=torch.bool)
|
||||
pertoken_scale = torch.randn(tokens, dtype=torch.float32)
|
||||
|
||||
mock_select_experts.return_value = (topk_weights, topk_ids)
|
||||
mock_comm = Mock()
|
||||
mock_comm.fused_experts.return_value = torch.randn(tokens, hidden_size, dtype=torch.float32)
|
||||
mock_extra_ctx.moe_comm_method = mock_comm
|
||||
mock_extra_ctx.moe_comm_type = MoECommType.ALLGATHER
|
||||
self.quant_method.multistream_overlap_gate = False
|
||||
self.quant_method.in_dtype = torch.float32
|
||||
|
||||
self.quant_method.apply(
|
||||
layer=layer,
|
||||
x=x,
|
||||
router_logits=router_logits,
|
||||
top_k=2,
|
||||
renormalize=True,
|
||||
num_experts=self.num_experts,
|
||||
activation="gelu",
|
||||
apply_router_weight_on_input=True,
|
||||
mc2_mask=mc2_mask,
|
||||
pertoken_scale=pertoken_scale,
|
||||
)
|
||||
|
||||
fused_experts_input = mock_comm.fused_experts.call_args.kwargs["fused_experts_input"]
|
||||
self.assertEqual(fused_experts_input.activation, "gelu")
|
||||
self.assertTrue(fused_experts_input.routing.apply_router_weight_on_input)
|
||||
self.assertIs(fused_experts_input.routing.mc2_mask, mc2_mask)
|
||||
self.assertIs(fused_experts_input.routing.pertoken_scale, pertoken_scale)
|
||||
self.assertIs(fused_experts_input.topk_weights, topk_weights)
|
||||
self.assertIs(fused_experts_input.topk_ids, topk_ids)
|
||||
130
tests/ut/quantization/test_compressed_tensors_config.py
Normal file
130
tests/ut/quantization/test_compressed_tensors_config.py
Normal file
@@ -0,0 +1,130 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from vllm.model_executor.layers.attention import Attention
|
||||
from vllm.model_executor.layers.fused_moe import FusedMoE
|
||||
from vllm.model_executor.layers.linear import RowParallelLinear, UnquantizedLinearMethod
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
from tests.ut.quantization.conftest_quantization import COMPRESSED_TENSORS_W8A8_CONFIG
|
||||
from vllm_ascend.ops.fused_moe.fused_moe import AscendUnquantizedFusedMoEMethod
|
||||
from vllm_ascend.quantization.compressed_tensors_config import AscendCompressedTensorsConfig
|
||||
from vllm_ascend.quantization.method_adapters import AscendFusedMoEMethod, AscendLinearMethod
|
||||
from vllm_ascend.quantization.methods import AscendW8A8DynamicFusedMoEMethod, AscendW8A8DynamicLinearMethod
|
||||
from vllm_ascend.utils import COMPRESSED_TENSORS_METHOD, vllm_version_is
|
||||
|
||||
|
||||
class TestAscendCompressedTensorsQuanType(TestBase):
|
||||
def setUp(self):
|
||||
self.config = AscendCompressedTensorsConfig(
|
||||
target_scheme_map={"Linear": {}},
|
||||
ignore=["lm_head"],
|
||||
quant_format="",
|
||||
config={},
|
||||
)
|
||||
|
||||
def _make_weight_quant(self, num_bits=8, strategy="channel", dynamic=False, symmetric=True, group_size=None):
|
||||
mock = MagicMock()
|
||||
mock.num_bits = num_bits
|
||||
mock.strategy = strategy
|
||||
mock.dynamic = dynamic
|
||||
mock.symmetric = symmetric
|
||||
mock.group_size = group_size
|
||||
return mock
|
||||
|
||||
def _make_input_quant(self, num_bits=8, strategy="tensor", dynamic=False, symmetric=True):
|
||||
mock = MagicMock()
|
||||
mock.num_bits = num_bits
|
||||
mock.strategy = strategy
|
||||
mock.dynamic = dynamic
|
||||
mock.symmetric = symmetric
|
||||
return mock
|
||||
|
||||
def test_detect_w8a8_static(self):
|
||||
weight = self._make_weight_quant(num_bits=8, strategy="channel", dynamic=False, symmetric=True)
|
||||
input_q = self._make_input_quant(num_bits=8, strategy="tensor", dynamic=False, symmetric=True)
|
||||
result = self.config._detect_quant_type(weight, input_q, "int-quantized")
|
||||
self.assertEqual(result, "W8A8")
|
||||
|
||||
def test_detect_w8a8_dynamic(self):
|
||||
weight = self._make_weight_quant(num_bits=8, strategy="channel", dynamic=False, symmetric=True)
|
||||
input_q = self._make_input_quant(num_bits=8, strategy="token", dynamic=True, symmetric=True)
|
||||
result = self.config._detect_quant_type(weight, input_q, "int-quantized")
|
||||
self.assertEqual(result, "W8A8_DYNAMIC")
|
||||
|
||||
def test_detect_w4a8_dynamic(self):
|
||||
weight = self._make_weight_quant(num_bits=4, strategy="channel", dynamic=False, symmetric=True)
|
||||
input_q = self._make_input_quant(num_bits=8, strategy="token", dynamic=True, symmetric=True)
|
||||
result = self.config._detect_quant_type(weight, input_q, "int-quantized")
|
||||
self.assertEqual(result, "W4A8_DYNAMIC")
|
||||
|
||||
def test_detect_w4a16(self):
|
||||
from compressed_tensors.quantization import QuantizationType
|
||||
|
||||
weight = MagicMock()
|
||||
weight.num_bits = 4
|
||||
weight.strategy = "group"
|
||||
weight.dynamic = False
|
||||
weight.type = QuantizationType.INT
|
||||
result = self.config._detect_quant_type(weight, None, None)
|
||||
self.assertEqual(result, "W4A16")
|
||||
|
||||
def test_detect_unsupported_raises(self):
|
||||
weight = self._make_weight_quant(num_bits=2, strategy="channel", dynamic=False, symmetric=True)
|
||||
input_q = self._make_input_quant(num_bits=2, strategy="tensor", dynamic=False, symmetric=True)
|
||||
with self.assertRaises(NotImplementedError):
|
||||
self.config._detect_quant_type(weight, input_q, "int_quantized")
|
||||
|
||||
|
||||
class TestAscendCompressedTensorsConfigGetQuantMethod(TestBase):
|
||||
def setUp(self):
|
||||
self.config = AscendCompressedTensorsConfig.from_config(COMPRESSED_TENSORS_W8A8_CONFIG)
|
||||
|
||||
@patch("vllm_ascend.quantization.method_adapters.AscendLinearMethod.__init__")
|
||||
def test_get_linear_quant_method(self, mock_method):
|
||||
mock_method.return_value = None
|
||||
layer = MagicMock(spec=RowParallelLinear)
|
||||
result = self.config.get_quant_method(layer, "model.layers.0.self_attn.q_proj")
|
||||
self.assertEqual(layer.ascend_quant_method, COMPRESSED_TENSORS_METHOD)
|
||||
self.assertTrue(isinstance(result, AscendLinearMethod))
|
||||
self.assertTrue(isinstance(layer.scheme, AscendW8A8DynamicLinearMethod))
|
||||
|
||||
def test_get_linear_unquantized_method(self):
|
||||
layer = MagicMock(spec=RowParallelLinear)
|
||||
result = self.config.get_quant_method(layer, "lm_head")
|
||||
self.assertEqual(layer.ascend_quant_method, COMPRESSED_TENSORS_METHOD)
|
||||
self.assertTrue(isinstance(result, UnquantizedLinearMethod))
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not vllm_version_is("0.23.0"),
|
||||
reason="Legacy FusedMoE quant method UT is only for vLLM 0.23.0.",
|
||||
)
|
||||
@patch("vllm_ascend.quantization.methods.AscendW8A8DynamicFusedMoEMethod.__init__")
|
||||
def test_get_moe_quant_method(self, mock_method):
|
||||
mock_method.return_value = None
|
||||
layer = MagicMock(spec=FusedMoE)
|
||||
layer.moe_config = {}
|
||||
result = self.config.get_quant_method(layer, "model.layers.0.mlp.experts")
|
||||
self.assertEqual(layer.ascend_quant_method, COMPRESSED_TENSORS_METHOD)
|
||||
self.assertTrue(isinstance(result, AscendFusedMoEMethod))
|
||||
self.assertTrue(isinstance(layer.scheme, AscendW8A8DynamicFusedMoEMethod))
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not vllm_version_is("0.23.0"),
|
||||
reason="Legacy FusedMoE quant method UT is only for vLLM 0.23.0.",
|
||||
)
|
||||
@patch("vllm_ascend.ops.fused_moe.fused_moe.AscendUnquantizedFusedMoEMethod.__init__")
|
||||
@patch("vllm_ascend.quantization.compressed_tensors_config.should_ignore_layer")
|
||||
def test_get_moe_unquantized_method(self, mock_ignore_layer, mock_method):
|
||||
mock_method.return_value = None
|
||||
mock_ignore_layer.return_value = True
|
||||
layer = MagicMock(spec=FusedMoE)
|
||||
layer.moe_config = {}
|
||||
result = self.config.get_quant_method(layer, "model.layers.0.mlp.experts")
|
||||
self.assertEqual(layer.ascend_quant_method, COMPRESSED_TENSORS_METHOD)
|
||||
self.assertTrue(isinstance(result, AscendUnquantizedFusedMoEMethod))
|
||||
|
||||
def test_no_quant_method(self):
|
||||
layer = MagicMock(spec=Attention)
|
||||
result = self.config.get_quant_method(layer, "attn")
|
||||
self.assertIsNone(result)
|
||||
192
tests/ut/quantization/test_method_adapters.py
Normal file
192
tests/ut/quantization/test_method_adapters.py
Normal file
@@ -0,0 +1,192 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch
|
||||
from vllm.model_executor.layers.fused_moe import FusedMoeWeightScaleSupported
|
||||
from vllm.model_executor.layers.linear import ColumnParallelLinear
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
from vllm_ascend.quantization.method_adapters import (
|
||||
AscendFusedMoEMethod,
|
||||
AscendKVCacheMethod,
|
||||
AscendLinearMethod,
|
||||
)
|
||||
from vllm_ascend.quantization.methods.base import AscendAttentionScheme, AscendLinearScheme, AscendMoEScheme
|
||||
|
||||
|
||||
class TestAscendLinearMethod(TestBase):
|
||||
@patch("vllm_ascend.quantization.method_adapters.enable_dsa_cp_with_layer_shard")
|
||||
def setUp(self, mock_enable_dsa_cp_with_layer_shard):
|
||||
self.mock_scheme = MagicMock(spec=AscendLinearScheme)
|
||||
self.mock_scheme.get_weight.return_value = {
|
||||
"weight": torch.empty(128, 256, dtype=torch.int8),
|
||||
"_packed_dim": 0,
|
||||
"_packed_factor": 0.1,
|
||||
}
|
||||
self.mock_scheme.get_pertensor_param.return_value = {
|
||||
"weight_scale_pertensor": torch.empty(1, 1, dtype=torch.int8),
|
||||
}
|
||||
self.mock_scheme.get_perchannel_param.return_value = {
|
||||
"weight_scale_perchannel": torch.empty(128, 1, dtype=torch.int8),
|
||||
}
|
||||
self.mock_scheme.get_pergroup_param.return_value = {
|
||||
"weight_scale_second": torch.empty(128, 2, dtype=torch.int8),
|
||||
"weight_offset_second": torch.empty(128, 2, dtype=torch.int8),
|
||||
"weight_scale_pergroup": torch.empty(128, 2, dtype=torch.int8),
|
||||
}
|
||||
self.method = AscendLinearMethod(self.mock_scheme)
|
||||
|
||||
@patch("vllm_ascend.quantization.method_adapters.PerTensorScaleParameter")
|
||||
def test_create_weights(self, mock_parameter):
|
||||
mock_parameter.return_value = torch.nn.Parameter(torch.empty(1, 1, dtype=torch.int8), requires_grad=False)
|
||||
layer = torch.nn.Module()
|
||||
weight_loader = MagicMock()
|
||||
self.method.create_weights(
|
||||
layer,
|
||||
input_size_per_partition=256,
|
||||
output_partition_sizes=[128],
|
||||
input_size=256,
|
||||
output_size=128,
|
||||
params_dtype=torch.bfloat16,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
# Check get_weight method
|
||||
self.mock_scheme.get_weight.assert_called_once_with(256, 128, torch.bfloat16)
|
||||
self.assertIn("weight", dict(layer.named_parameters()))
|
||||
self.assertNotIn("_packed_dim", dict(layer.named_parameters()))
|
||||
self.assertNotIn("_packed_factor", dict(layer.named_parameters()))
|
||||
self.assertEqual(layer.weight.input_dim, 1)
|
||||
self.assertEqual(layer.weight.output_dim, 0)
|
||||
self.assertEqual(layer.weight.packed_dim, 0)
|
||||
self.assertEqual(layer.weight.packed_factor, 0.1)
|
||||
|
||||
# Check per tensor param
|
||||
self.mock_scheme.get_pertensor_param.assert_called_once()
|
||||
self.assertTrue(layer.weight_scale_pertensor.ignore_warning)
|
||||
self.assertEqual(layer.weight_scale_pertensor.weight_loader, weight_loader)
|
||||
|
||||
# Check per channel param
|
||||
self.mock_scheme.get_perchannel_param.assert_called_once_with(128, torch.bfloat16)
|
||||
self.assertEqual(layer.weight_scale_perchannel.output_dim, 0)
|
||||
self.assertEqual(layer.weight_scale_perchannel.weight_loader, weight_loader)
|
||||
|
||||
# Check per group param
|
||||
self.mock_scheme.get_pergroup_param.assert_called_once()
|
||||
self.assertEqual(layer.weight_scale_pergroup.output_dim, 0)
|
||||
self.assertFalse(hasattr(layer.weight_scale_pergroup, "input_dim"))
|
||||
self.assertEqual(layer.weight_scale_second.input_dim, 1)
|
||||
self.assertEqual(layer.weight_offset_second.input_dim, 1)
|
||||
|
||||
def test_process_weights_after_loading_delegates(self):
|
||||
layer = torch.nn.Module()
|
||||
self.mock_scheme.process_weights_after_loading.return_value = None
|
||||
self.method.process_weights_after_loading(layer)
|
||||
self.mock_scheme.process_weights_after_loading.assert_called_once_with(layer)
|
||||
|
||||
def test_apply_delegates_to_scheme(self):
|
||||
layer = MagicMock(spec=ColumnParallelLinear)
|
||||
x = torch.randn(4, 256)
|
||||
self.mock_scheme.apply.return_value = torch.randn(4, 128)
|
||||
output = self.method.apply(layer, x)
|
||||
self.mock_scheme.apply.assert_called_once()
|
||||
self.assertEqual(output.shape, (4, 128))
|
||||
|
||||
|
||||
class TestAscendKVCacheMethod(TestBase):
|
||||
def setUp(self):
|
||||
self.mock_scheme = MagicMock(spec=AscendAttentionScheme)
|
||||
self.mock_scheme.create_weights.return_value = None
|
||||
self.mock_scheme.process_weights_after_loading.return_value = None
|
||||
self.method = AscendKVCacheMethod(self.mock_scheme)
|
||||
|
||||
def test_create_weights_delegates(self):
|
||||
layer = torch.nn.Module()
|
||||
self.method.create_weights(layer)
|
||||
self.mock_scheme.create_weights.assert_called_once_with(layer)
|
||||
|
||||
def test_process_weights_after_loading_delegates(self):
|
||||
layer = torch.nn.Module()
|
||||
self.method.process_weights_after_loading(layer)
|
||||
self.mock_scheme.process_weights_after_loading.assert_called_once_with(layer)
|
||||
|
||||
def test_apply_delegates(self):
|
||||
layer = torch.nn.Module()
|
||||
query = torch.randn(4, 8, 64)
|
||||
key = torch.randn(4, 8, 64)
|
||||
value = torch.randn(4, 8, 64)
|
||||
self.mock_scheme.apply.return_value = torch.randn(4, 8, 64)
|
||||
self.method.apply(
|
||||
layer,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
kv_cache=None,
|
||||
attn_metadata=None,
|
||||
attn_type=None,
|
||||
scale=1.0,
|
||||
output=None,
|
||||
)
|
||||
self.mock_scheme.apply.assert_called_once()
|
||||
|
||||
|
||||
class TestAscendFusedMoEMethod(TestBase):
|
||||
def setUp(self):
|
||||
self.mock_scheme = MagicMock(spec=AscendMoEScheme)
|
||||
self.mock_scheme.group_size = 0
|
||||
self.mock_moe_config = MagicMock()
|
||||
self.method = AscendFusedMoEMethod(self.mock_scheme, self.mock_moe_config)
|
||||
|
||||
def test_process_weights_after_loading_delegates(self):
|
||||
layer = torch.nn.Module()
|
||||
self.mock_scheme.process_weights_after_loading.return_value = None
|
||||
self.method.process_weights_after_loading(layer)
|
||||
self.mock_scheme.process_weights_after_loading.assert_called_once_with(layer)
|
||||
|
||||
def test_create_weights_registers_parameters(self):
|
||||
self.mock_scheme.get_weight.return_value = {
|
||||
"w13_weight": torch.empty(8, 256, 128, dtype=torch.int8),
|
||||
"w2_weight": torch.empty(8, 128, 256, dtype=torch.int8),
|
||||
}
|
||||
self.mock_scheme.get_dynamic_quant_param.return_value = {
|
||||
"w13_weight_scale_second": torch.empty(8, 256, 1, dtype=torch.bfloat16),
|
||||
"w2_weight_offset_second": torch.empty(8, 128, 1, dtype=torch.bfloat16),
|
||||
"w2_scale_bias": torch.empty(8, 128, 1, dtype=torch.bfloat16),
|
||||
"w13_weight_scale": torch.empty(8, 256, 1, dtype=torch.bfloat16),
|
||||
"w2_weight_offset": torch.empty(8, 128, 1, dtype=torch.bfloat16),
|
||||
}
|
||||
# per channel quantization
|
||||
layer = self.create_moe_weights()
|
||||
self.assertIn("w13_weight", dict(layer.named_parameters()))
|
||||
self.assertIn("w2_weight", dict(layer.named_parameters()))
|
||||
|
||||
self.assertEqual(layer.w13_weight_scale_second.quant_method, FusedMoeWeightScaleSupported.GROUP.value)
|
||||
self.assertEqual(layer.w2_weight_offset_second.quant_method, FusedMoeWeightScaleSupported.GROUP.value)
|
||||
self.assertEqual(layer.w2_scale_bias.quant_method, FusedMoeWeightScaleSupported.GROUP.value)
|
||||
self.assertEqual(layer.w13_weight_scale.quant_method, FusedMoeWeightScaleSupported.CHANNEL.value)
|
||||
self.assertEqual(layer.w2_weight_offset.quant_method, FusedMoeWeightScaleSupported.CHANNEL.value)
|
||||
|
||||
# per group quantization
|
||||
self.mock_scheme.group_size = 128
|
||||
layer = self.create_moe_weights()
|
||||
self.assertEqual(layer.w13_weight_scale.quant_method, FusedMoeWeightScaleSupported.GROUP.value)
|
||||
self.assertEqual(layer.w2_weight_offset.quant_method, FusedMoeWeightScaleSupported.GROUP.value)
|
||||
|
||||
def create_moe_weights(self):
|
||||
layer = torch.nn.Module()
|
||||
self.method.create_weights(
|
||||
layer,
|
||||
num_experts=8,
|
||||
hidden_size=128,
|
||||
intermediate_size_per_partition=256,
|
||||
params_dtype=torch.bfloat16,
|
||||
)
|
||||
return layer
|
||||
|
||||
def test_apply_method(self):
|
||||
layer = torch.nn.Module()
|
||||
x = torch.randn(8, 64)
|
||||
router_logits = torch.randn(8, 64)
|
||||
top_k = 3
|
||||
renormalize = True
|
||||
self.mock_scheme.apply.return_value = None
|
||||
self.method.apply(layer, x, router_logits, top_k, renormalize)
|
||||
self.mock_scheme.apply.assert_called_once()
|
||||
471
tests/ut/quantization/test_modelslim_config.py
Normal file
471
tests/ut/quantization/test_modelslim_config.py
Normal file
@@ -0,0 +1,471 @@
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.model_executor.layers.attention import Attention
|
||||
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
|
||||
from vllm.model_executor.layers.fused_moe import FusedMoE
|
||||
from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
|
||||
from vllm.model_executor.layers.linear import LinearBase
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
from vllm_ascend.ops.linear import AscendUnquantizedLinearMethod
|
||||
from vllm_ascend.quantization.modelslim_config import (
|
||||
MODELSLIM_CONFIG_FILENAME,
|
||||
AscendModelSlimConfig,
|
||||
)
|
||||
from vllm_ascend.utils import ASCEND_QUANTIZATION_METHOD, vllm_version_is
|
||||
|
||||
|
||||
class TestAscendModelSlimConfig(TestBase):
|
||||
def setUp(self):
|
||||
self.sample_config = {
|
||||
"weight": "INT8",
|
||||
"fa_quant_type": "C8",
|
||||
"layers.1.fa_k.scale": "C8",
|
||||
"layer1.weight": "INT8",
|
||||
"layer2.weight": "FLOAT",
|
||||
"fused_layer.weight": "FLOAT",
|
||||
"fused_layer.shard1.weight": "FLOAT",
|
||||
"fused_layer.shard2.weight": "FLOAT",
|
||||
"shard1.weight": "FLOAT",
|
||||
"shard2.weight": "FLOAT",
|
||||
}
|
||||
self.ascend_config = AscendModelSlimConfig(self.sample_config)
|
||||
self.ascend_config.packed_modules_mapping = None
|
||||
|
||||
def test_init(self):
|
||||
self.assertEqual(self.ascend_config.quant_description, self.sample_config)
|
||||
|
||||
def test_repr(self):
|
||||
repr_str = repr(self.ascend_config)
|
||||
self.assertTrue(repr_str.startswith("AscendModelSlimConfig:\n"))
|
||||
|
||||
def test_get_name(self):
|
||||
self.assertEqual(AscendModelSlimConfig.get_name(), ASCEND_QUANTIZATION_METHOD)
|
||||
|
||||
def test_get_supported_act_dtypes(self):
|
||||
supported_dtypes = AscendModelSlimConfig.get_supported_act_dtypes()
|
||||
self.assertEqual(len(supported_dtypes), 3)
|
||||
|
||||
def test_get_min_capability(self):
|
||||
with self.assertRaises(NotImplementedError):
|
||||
AscendModelSlimConfig.get_min_capability()
|
||||
|
||||
def test_get_config_filenames(self):
|
||||
filenames = AscendModelSlimConfig.get_config_filenames()
|
||||
self.assertEqual(filenames, [])
|
||||
|
||||
def test_from_config(self):
|
||||
config = AscendModelSlimConfig.from_config(self.sample_config)
|
||||
self.assertIsInstance(config, AscendModelSlimConfig)
|
||||
self.assertEqual(config.quant_description, self.sample_config)
|
||||
|
||||
@patch("torch.npu.is_available")
|
||||
def test_override_quantization_method(self, mock_is_available):
|
||||
# Test when NPU is available
|
||||
mock_is_available.return_value = True
|
||||
result = AscendModelSlimConfig.override_quantization_method(None, None)
|
||||
self.assertIsNone(result)
|
||||
hf_quant_cfg = {"quant_method": ""}
|
||||
result = AscendModelSlimConfig.override_quantization_method(hf_quant_cfg, None)
|
||||
self.assertEqual(result, "ascend")
|
||||
|
||||
# Test when NPU is not available
|
||||
mock_is_available.return_value = False
|
||||
result = AscendModelSlimConfig.override_quantization_method(None, None)
|
||||
self.assertIsNone(result)
|
||||
hf_quant_cfg = {"quant_method": ""}
|
||||
result = AscendModelSlimConfig.override_quantization_method(hf_quant_cfg, None)
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_get_quant_method_for_linear(self):
|
||||
mock_config = MagicMock()
|
||||
mock_config.model_config.hf_config.model_type = None
|
||||
linear_layer = MagicMock(spec=LinearBase)
|
||||
# Test skipped layer
|
||||
with (
|
||||
patch("vllm_ascend.quantization.modelslim_config.get_current_vllm_config", return_value=mock_config),
|
||||
patch.object(self.ascend_config, "is_layer_skipped_ascend", return_value=True),
|
||||
):
|
||||
method = self.ascend_config.get_quant_method(linear_layer, ".attn")
|
||||
self.assertIsInstance(method, AscendUnquantizedLinearMethod)
|
||||
|
||||
# Test quantized layer
|
||||
mock_scheme = MagicMock()
|
||||
with (
|
||||
patch.object(self.ascend_config, "is_layer_skipped_ascend", return_value=False),
|
||||
patch("vllm_ascend.quantization.modelslim_config.get_current_vllm_config", return_value=mock_config),
|
||||
patch("vllm_ascend.quantization.modelslim_config.create_scheme_for_layer", return_value=mock_scheme),
|
||||
patch(
|
||||
"vllm_ascend.quantization.method_adapters.AscendLinearMethod", return_value=MagicMock()
|
||||
) as mock_ascend_linear,
|
||||
):
|
||||
method = self.ascend_config.get_quant_method(linear_layer, ".attn")
|
||||
self.assertIs(method, mock_ascend_linear.return_value)
|
||||
mock_ascend_linear.assert_called_once_with(mock_scheme)
|
||||
|
||||
def test_get_quant_method_for_attention(self):
|
||||
attention_layer = MagicMock(spec=Attention)
|
||||
mock_config = MagicMock()
|
||||
mock_config.model_config.hf_config.model_type = None
|
||||
mock_scheme = MagicMock()
|
||||
with (
|
||||
patch("vllm_ascend.quantization.modelslim_config.get_current_vllm_config", return_value=mock_config),
|
||||
patch("vllm_ascend.quantization.modelslim_config.create_scheme_for_layer", return_value=mock_scheme),
|
||||
patch(
|
||||
"vllm_ascend.quantization.method_adapters.AscendKVCacheMethod", return_value=MagicMock()
|
||||
) as mock_ascend_kvcache,
|
||||
):
|
||||
# Test with fa_quant_type
|
||||
method = self.ascend_config.get_quant_method(attention_layer, ".attn")
|
||||
self.assertIs(method, None)
|
||||
method = self.ascend_config.get_quant_method(attention_layer, "layers.1.attn")
|
||||
self.assertIs(method, mock_ascend_kvcache.return_value)
|
||||
|
||||
def test_get_quant_method_for_c8_kv_cache_attention(self):
|
||||
c8_config = AscendModelSlimConfig(
|
||||
{
|
||||
"kv_cache_type": "C8",
|
||||
"model.layers.0.k_proj.kv_cache_scale": "C8",
|
||||
}
|
||||
)
|
||||
attention_layer = MagicMock(spec=AttentionLayerBase)
|
||||
mock_vllm_config = MagicMock()
|
||||
mock_vllm_config.model_config.hf_config.model_type = None
|
||||
|
||||
mock_vllm_config_for_kv_c8 = MagicMock()
|
||||
mock_vllm_config_for_kv_c8.kv_transfer_config = None
|
||||
|
||||
with (
|
||||
patch("vllm_ascend.quantization.modelslim_config.get_current_vllm_config", return_value=mock_vllm_config),
|
||||
patch(
|
||||
"vllm_ascend.quantization.methods.kv_c8.get_current_vllm_config",
|
||||
return_value=mock_vllm_config_for_kv_c8,
|
||||
),
|
||||
patch(
|
||||
"vllm_ascend.quantization.method_adapters.AscendKVCacheMethod", return_value=MagicMock()
|
||||
) as mock_kvcache,
|
||||
):
|
||||
method = c8_config.get_quant_method(attention_layer, "model.layers.0.self_attn.attn")
|
||||
self.assertIs(method, mock_kvcache.return_value)
|
||||
args, _ = mock_kvcache.call_args
|
||||
from vllm_ascend.quantization.methods.kv_c8 import AscendC8KVCacheAttentionMethod
|
||||
|
||||
self.assertIsInstance(args[0], AscendC8KVCacheAttentionMethod)
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not vllm_version_is("0.23.0"),
|
||||
reason="Legacy FusedMoE quant method UT is only for vLLM 0.23.0.",
|
||||
)
|
||||
def test_get_quant_method_for_fused_moe(self):
|
||||
fused_moe_layer = MagicMock(spec=FusedMoE)
|
||||
fused_moe_layer.moe = MagicMock(spec=FusedMoEConfig)
|
||||
fused_moe_layer.moe_config = MagicMock(spec=FusedMoEConfig)
|
||||
mock_config = MagicMock()
|
||||
mock_config.model_config.hf_config.model_type = None
|
||||
|
||||
# Test skipped layer
|
||||
with (
|
||||
patch.object(self.ascend_config, "is_layer_skipped_ascend", return_value=True),
|
||||
patch("vllm_ascend.quantization.modelslim_config.get_current_vllm_config", return_value=mock_config),
|
||||
patch(
|
||||
"vllm_ascend.ops.fused_moe.fused_moe.AscendUnquantizedFusedMoEMethod", return_value=MagicMock()
|
||||
) as mock_ascend_moe,
|
||||
):
|
||||
method = self.ascend_config.get_quant_method(fused_moe_layer, "moe_layer")
|
||||
self.assertIs(method, mock_ascend_moe.return_value)
|
||||
|
||||
# Test quantized layer
|
||||
mock_scheme = MagicMock()
|
||||
with (
|
||||
patch.object(self.ascend_config, "is_layer_skipped_ascend", return_value=False),
|
||||
patch("vllm_ascend.quantization.modelslim_config.get_current_vllm_config", return_value=mock_config),
|
||||
patch("vllm_ascend.quantization.modelslim_config.create_scheme_for_layer", return_value=mock_scheme),
|
||||
patch(
|
||||
"vllm_ascend.quantization.method_adapters.AscendFusedMoEMethod", return_value=MagicMock()
|
||||
) as mock_ascend_moe,
|
||||
):
|
||||
method = self.ascend_config.get_quant_method(fused_moe_layer, "moe_layer")
|
||||
self.assertIs(method, mock_ascend_moe.return_value)
|
||||
|
||||
def test_is_layer_skipped_ascend(self):
|
||||
# Test non-fused layer that should be quantized
|
||||
self.assertFalse(self.ascend_config.is_layer_skipped_ascend("layer1"))
|
||||
|
||||
# Test non-fused layer that should be skipped
|
||||
self.assertTrue(self.ascend_config.is_layer_skipped_ascend("layer2"))
|
||||
|
||||
# Test fused layer
|
||||
fused_mapping = {"fused_layer": ["shard1", "shard2"]}
|
||||
self.assertTrue(self.ascend_config.is_layer_skipped_ascend("fused_layer", fused_mapping))
|
||||
|
||||
# Test inconsistent fused layer shards
|
||||
bad_config = {"shard1.weight": "FLOAT", "shard2.weight": "INT8"}
|
||||
config = AscendModelSlimConfig(bad_config)
|
||||
with self.assertRaises(ValueError):
|
||||
config.is_layer_skipped_ascend("fused_layer", fused_mapping)
|
||||
|
||||
def test_init_with_default_config(self):
|
||||
config = AscendModelSlimConfig()
|
||||
self.assertEqual(config.quant_description, {})
|
||||
|
||||
def test_maybe_update_config_already_populated(self):
|
||||
# When quant_description is already populated, should be a no-op
|
||||
self.assertTrue(len(self.ascend_config.quant_description) > 0)
|
||||
self.ascend_config.maybe_update_config("/some/model/path")
|
||||
# quant_description should remain unchanged
|
||||
self.assertEqual(self.ascend_config.quant_description, self.sample_config)
|
||||
|
||||
def test_maybe_update_config_loads_from_file(self):
|
||||
config = AscendModelSlimConfig()
|
||||
self.assertEqual(config.quant_description, {})
|
||||
|
||||
quant_data = {"layer1.weight": "INT8", "layer2.weight": "FLOAT"}
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
config_path = os.path.join(tmpdir, MODELSLIM_CONFIG_FILENAME)
|
||||
with open(config_path, "w") as f:
|
||||
json.dump(quant_data, f)
|
||||
|
||||
config.maybe_update_config(tmpdir)
|
||||
|
||||
self.assertEqual(config.quant_description, quant_data)
|
||||
|
||||
def test_maybe_update_config_raises_when_file_missing(self):
|
||||
config = AscendModelSlimConfig()
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
with self.assertRaises(ValueError) as ctx:
|
||||
config.maybe_update_config(tmpdir)
|
||||
|
||||
error_msg = str(ctx.exception)
|
||||
self.assertIn("ModelSlim Quantization Config Not Found", error_msg)
|
||||
self.assertIn(MODELSLIM_CONFIG_FILENAME, error_msg)
|
||||
|
||||
def test_maybe_update_config_raises_with_json_files_listed(self):
|
||||
config = AscendModelSlimConfig()
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# Create a dummy json file that is NOT the config file
|
||||
dummy_path = os.path.join(tmpdir, "config.json")
|
||||
with open(dummy_path, "w") as f:
|
||||
json.dump({"dummy": True}, f)
|
||||
|
||||
with self.assertRaises(ValueError) as ctx:
|
||||
config.maybe_update_config(tmpdir)
|
||||
|
||||
error_msg = str(ctx.exception)
|
||||
self.assertIn("config.json", error_msg)
|
||||
|
||||
def test_maybe_update_config_non_directory_raises(self):
|
||||
config = AscendModelSlimConfig()
|
||||
|
||||
with self.assertRaises(ValueError) as ctx:
|
||||
config.maybe_update_config("not_a_real_directory_path")
|
||||
|
||||
error_msg = str(ctx.exception)
|
||||
self.assertIn("ModelSlim Quantization Config Not Found", error_msg)
|
||||
|
||||
def test_apply_extra_quant_adaptations_shared_head(self):
|
||||
config = AscendModelSlimConfig()
|
||||
config.quant_description = {
|
||||
"model.layers.0.shared_head.weight": "INT8",
|
||||
"transformer.shared_head.output.weight": "INT8",
|
||||
"transformer.shared_head.norm.weight": "INT8",
|
||||
}
|
||||
config._apply_extra_quant_adaptations()
|
||||
self.assertIn("model.layers.0.weight", config.quant_description)
|
||||
self.assertEqual(config.quant_description["model.layers.0.weight"], "INT8")
|
||||
self.assertIn("shared_head.head.weight", config.quant_description)
|
||||
self.assertIn("shared_head.norm.weight", config.quant_description)
|
||||
|
||||
def test_apply_extra_quant_adaptations_weight_packed(self):
|
||||
config = AscendModelSlimConfig()
|
||||
config.quant_description = {
|
||||
"model.layers.0.weight_packed": "INT8",
|
||||
}
|
||||
config._apply_extra_quant_adaptations()
|
||||
self.assertIn("model.layers.0.weight", config.quant_description)
|
||||
self.assertEqual(config.quant_description["model.layers.0.weight"], "INT8")
|
||||
|
||||
|
||||
class TestApplyVllmMapper(TestBase):
|
||||
def test_apply_mapper_with_populated_quant_description(self):
|
||||
config = AscendModelSlimConfig({"old_key.weight": "INT8"})
|
||||
mock_mapper = MagicMock()
|
||||
mock_mapper.apply_dict.return_value = {"new_key.weight": "INT8"}
|
||||
|
||||
config.apply_vllm_mapper(mock_mapper)
|
||||
|
||||
self.assertEqual(config.quant_description, {"new_key.weight": "INT8"})
|
||||
mock_mapper.apply_dict.assert_called_once_with({"old_key.weight": "INT8"})
|
||||
|
||||
|
||||
class TestQuantPrefixMapper(TestBase):
|
||||
def test_lm_head_maps_to_language_model_lm_head_when_quant_key_exists(self):
|
||||
config = AscendModelSlimConfig({"language_model.lm_head.weight": "FLOAT"})
|
||||
|
||||
prefix = config.quant_prefix_mapper("qwen3_5_moe", "lm_head")
|
||||
|
||||
self.assertEqual(prefix, "language_model.lm_head")
|
||||
|
||||
def test_lm_head_keeps_original_prefix_when_quant_key_exists(self):
|
||||
config = AscendModelSlimConfig(
|
||||
{
|
||||
"lm_head.weight": "FLOAT",
|
||||
"language_model.lm_head.weight": "FLOAT",
|
||||
}
|
||||
)
|
||||
|
||||
prefix = config.quant_prefix_mapper("qwen3_5_moe", "lm_head")
|
||||
|
||||
self.assertEqual(prefix, "lm_head")
|
||||
|
||||
def test_step3p5_mtp_maps_direct_and_step3p7_wrapped_quant_keys(self):
|
||||
cases = [
|
||||
(
|
||||
"model.layers.45.self_attn",
|
||||
"model.layers.45.self_attn.qkv_proj",
|
||||
),
|
||||
(
|
||||
"language_model.model.layers.45.self_attn",
|
||||
"language_model.model.layers.45.self_attn.qkv_proj",
|
||||
),
|
||||
]
|
||||
for quant_prefix, expected in cases:
|
||||
with self.subTest(quant_prefix=quant_prefix):
|
||||
config = AscendModelSlimConfig(
|
||||
{
|
||||
f"{quant_prefix}.q_proj.weight": "FLOAT",
|
||||
f"{quant_prefix}.k_proj.weight": "FLOAT",
|
||||
f"{quant_prefix}.v_proj.weight": "FLOAT",
|
||||
}
|
||||
)
|
||||
|
||||
prefix = config.quant_prefix_mapper(
|
||||
"step3p5_mtp",
|
||||
"model.layers.45.mtp_block.self_attn.qkv_proj",
|
||||
)
|
||||
|
||||
self.assertEqual(prefix, expected)
|
||||
|
||||
|
||||
class TestGetCacheScale(TestBase):
|
||||
def test_c8_kv_cache_type_k_proj_scale(self):
|
||||
config = AscendModelSlimConfig({"kv_cache_type": "C8"})
|
||||
result = config.get_cache_scale("model.layers.0.k_proj.kv_cache_scale")
|
||||
self.assertEqual(result, "model.layers.0.attn.k_cache_scale")
|
||||
result = config.get_cache_scale("model.layers.0.v_proj.kv_cache_offset")
|
||||
self.assertEqual(result, "model.layers.0.attn.v_cache_offset")
|
||||
|
||||
def test_no_match(self):
|
||||
config = AscendModelSlimConfig({"kv_cache_type": "FLOAT"})
|
||||
result = config.get_cache_scale("model.layers.0.k_proj.kv_cache_scale")
|
||||
self.assertIsNone(result)
|
||||
|
||||
config = AscendModelSlimConfig({"kv_cache_type": "C8"})
|
||||
result = config.get_cache_scale("model.layers.0.other_key")
|
||||
self.assertIsNone(result)
|
||||
|
||||
|
||||
class TestGetKvQuantDtype(TestBase):
|
||||
def test_enable_fa_quant(self):
|
||||
config = AscendModelSlimConfig(
|
||||
{
|
||||
"fa_quant_type": "C8",
|
||||
"layers.1.fa_k.scale": "C8",
|
||||
}
|
||||
)
|
||||
mock_model_config = MagicMock()
|
||||
mock_model_config.dtype = torch.float16
|
||||
# test mla
|
||||
mock_model_config.use_mla = True
|
||||
k_dtype, v_dtype = config.get_kv_quant_dtype("layers.1.attn", torch.float16, mock_model_config)
|
||||
self.assertEqual(k_dtype, torch.int8)
|
||||
self.assertEqual(v_dtype, torch.float16)
|
||||
|
||||
# test gqa
|
||||
mock_model_config.use_mla = False
|
||||
k_dtype, v_dtype = config.get_kv_quant_dtype("layers.1.attn", torch.float16, mock_model_config)
|
||||
self.assertEqual(k_dtype, torch.int8)
|
||||
self.assertEqual(v_dtype, torch.int8)
|
||||
|
||||
def test_enable_fa_quant_false(self):
|
||||
config = AscendModelSlimConfig({})
|
||||
mock_model_config = MagicMock()
|
||||
mock_model_config.dtype = torch.float16
|
||||
k_dtype, v_dtype = config.get_kv_quant_dtype("layers.1.attn", torch.float16, mock_model_config)
|
||||
self.assertEqual(k_dtype, torch.float16)
|
||||
|
||||
|
||||
class TestGetKvQuantSplitFactor(TestBase):
|
||||
@patch("vllm_ascend.quantization.modelslim_config.calc_split_factor")
|
||||
def test_enable_fa_quant_true(self, mock_calc_split_factor):
|
||||
mock_calc_split_factor.return_value = 2.0
|
||||
config = AscendModelSlimConfig(
|
||||
{
|
||||
"fa_quant_type": "C8",
|
||||
"layers.1.fa_k.scale": "C8",
|
||||
}
|
||||
)
|
||||
kv_head_dim_list = [64, 64]
|
||||
|
||||
result = config.get_kv_quant_split_factor("layers.1.attn", kv_head_dim_list)
|
||||
self.assertEqual(result, 2.0)
|
||||
mock_calc_split_factor.assert_called_once_with([64, 128])
|
||||
|
||||
@patch("vllm_ascend.quantization.modelslim_config.calc_split_factor")
|
||||
def test_enable_fa_quant_false(self, mock_calc_split_factor):
|
||||
mock_calc_split_factor.return_value = 1.0
|
||||
config = AscendModelSlimConfig({})
|
||||
kv_head_dim_list = [64, 64]
|
||||
|
||||
result = config.get_kv_quant_split_factor("layers.1.attn", kv_head_dim_list)
|
||||
self.assertEqual(result, 1.0)
|
||||
mock_calc_split_factor.assert_called_once_with([64, 64])
|
||||
|
||||
|
||||
class TestAddKvcacheQuantMetadata(TestBase):
|
||||
def test_with_fa_quant_type(self):
|
||||
config = AscendModelSlimConfig(
|
||||
{
|
||||
"fa_quant_type": "C8",
|
||||
"layers.1.fa_k.scale": "C8",
|
||||
"layers.2.fa_k.scale": "C8",
|
||||
}
|
||||
)
|
||||
config._add_kvcache_quant_metadata()
|
||||
|
||||
self.assertTrue(config.enable_fa_quant)
|
||||
self.assertIn(1, config.kvcache_quant_layers)
|
||||
self.assertNotIn(5, config.kvcache_quant_layers)
|
||||
self.assertFalse(config.enable_indexer_quant)
|
||||
self.assertEqual(config.indexer_quant_layers, [])
|
||||
|
||||
def test_with_indexer_quant_type(self):
|
||||
config = AscendModelSlimConfig(
|
||||
{
|
||||
"indexer_quant_type": "INT8",
|
||||
"layers.1.indexer.quant_type": "INT8",
|
||||
"layers.3.indexer.quant_type": "INT8",
|
||||
}
|
||||
)
|
||||
config._add_kvcache_quant_metadata()
|
||||
|
||||
self.assertFalse(config.enable_fa_quant)
|
||||
self.assertEqual(config.kvcache_quant_layers, [])
|
||||
self.assertTrue(config.enable_indexer_quant)
|
||||
self.assertIn(1, config.indexer_quant_layers)
|
||||
self.assertNotIn(5, config.indexer_quant_layers)
|
||||
|
||||
def test_with_neither_quant_type(self):
|
||||
config = AscendModelSlimConfig({})
|
||||
config._add_kvcache_quant_metadata()
|
||||
|
||||
self.assertFalse(config.enable_fa_quant)
|
||||
self.assertEqual(config.kvcache_quant_layers, [])
|
||||
self.assertFalse(config.enable_indexer_quant)
|
||||
self.assertEqual(config.indexer_quant_layers, [])
|
||||
45
tests/ut/quantization/test_quant_parser.py
Normal file
45
tests/ut/quantization/test_quant_parser.py
Normal file
@@ -0,0 +1,45 @@
|
||||
from unittest.mock import patch
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
from vllm_ascend.quantization.quant_parser import (
|
||||
get_rollback_quant_type,
|
||||
parse_mxfp_quant_params,
|
||||
parse_quant_moe_down_proj_params,
|
||||
)
|
||||
|
||||
|
||||
class TestGetRollbackQuantType(TestBase):
|
||||
def test_returns_down_proj_quant_type(self):
|
||||
config = {
|
||||
"model.layers.0.mlp.gate_proj": "W8A8_MXFP8",
|
||||
"model.layers.0.mlp.down_proj": "W4A4_MXFP4",
|
||||
}
|
||||
result = get_rollback_quant_type(config)
|
||||
self.assertEqual(result, "W4A4_MXFP4")
|
||||
|
||||
def test_returns_default_when_no_down_proj(self):
|
||||
config = {"model.layers.0.mlp.gate_proj": "W4A8_MXFP"}
|
||||
result = get_rollback_quant_type(config)
|
||||
self.assertEqual(result, "W8A8_MXFP8")
|
||||
|
||||
|
||||
class TestParseMxfpQuantParams(TestBase):
|
||||
def test_default_values(self):
|
||||
act, weight, scale, per_token, round_mode = parse_mxfp_quant_params()
|
||||
self.assertIsNotNone(act)
|
||||
self.assertIsNotNone(weight)
|
||||
self.assertIsNotNone(round_mode)
|
||||
|
||||
|
||||
class TestParseQuantMoeDownProjParams(TestBase):
|
||||
@patch("vllm_ascend.quantization.quant_parser.ensure_mxfp8_scale_dtype_available")
|
||||
def test_w8a8_mxfp8_uses_rint_round_mode(self, mock_ensure):
|
||||
mock_ensure.return_value = None
|
||||
act, weight, scale, per_token, round_mode = parse_quant_moe_down_proj_params("W8A8_MXFP8", "round")
|
||||
self.assertEqual(round_mode, "rint")
|
||||
|
||||
@patch("vllm_ascend.quantization.quant_parser.ensure_mxfp4_dtype_available")
|
||||
def test_w4a4_mxfp4_respects_parsed_round_mode(self, mock_ensure):
|
||||
mock_ensure.return_value = None
|
||||
act, weight, scale, per_token, round_mode = parse_quant_moe_down_proj_params("W4A4_MXFP4", "round")
|
||||
self.assertEqual(round_mode, "round")
|
||||
@@ -1,62 +1,193 @@
|
||||
import types
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from vllm.config import KVTransferConfig
|
||||
|
||||
from tests.ut.base import TestBase
|
||||
from vllm_ascend.quantization.utils import (ASCEND_QUANTIZATION_METHOD_MAP,
|
||||
get_quant_method)
|
||||
from tests.ut.quantization.conftest_quantization import FAKQUANT_CONFIG, W8A8_CONFIG
|
||||
from vllm_ascend.quantization import AscendCompressedTensorsConfig
|
||||
from vllm_ascend.quantization.modelslim_config import MODELSLIM_CONFIG_FILENAME, AscendModelSlimConfig
|
||||
from vllm_ascend.quantization.utils import (
|
||||
detect_quantization_method,
|
||||
enable_fa_quant,
|
||||
maybe_auto_detect_quantization,
|
||||
)
|
||||
from vllm_ascend.utils import ASCEND_QUANTIZATION_METHOD, COMPRESSED_TENSORS_METHOD
|
||||
|
||||
|
||||
class TestGetQuantMethod(TestBase):
|
||||
class TestDetectQuantizationMethod(TestBase):
|
||||
def test_returns_none_for_non_existent_path(self):
|
||||
result = detect_quantization_method("/non/existent/path")
|
||||
self.assertIsNone(result)
|
||||
|
||||
def setUp(self):
|
||||
self.original_quantization_method_map = ASCEND_QUANTIZATION_METHOD_MAP.copy(
|
||||
)
|
||||
for quant_type, layer_map in ASCEND_QUANTIZATION_METHOD_MAP.items():
|
||||
for layer_type in layer_map.keys():
|
||||
ASCEND_QUANTIZATION_METHOD_MAP[quant_type][
|
||||
layer_type] = types.new_class(f"{quant_type}_{layer_type}")
|
||||
def test_detects_modelslim(self):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
config_path = os.path.join(tmpdir, MODELSLIM_CONFIG_FILENAME)
|
||||
with open(config_path, "w") as f:
|
||||
json.dump({"layer.weight": "INT8"}, f)
|
||||
|
||||
def tearDown(self):
|
||||
# Restore original map
|
||||
ASCEND_QUANTIZATION_METHOD_MAP.clear()
|
||||
ASCEND_QUANTIZATION_METHOD_MAP.update(
|
||||
self.original_quantization_method_map)
|
||||
result = detect_quantization_method(tmpdir)
|
||||
self.assertEqual(result, ASCEND_QUANTIZATION_METHOD)
|
||||
|
||||
def test_linear_quant_methods(self):
|
||||
for quant_type, layer_map in ASCEND_QUANTIZATION_METHOD_MAP.items():
|
||||
if "linear" in layer_map.keys():
|
||||
prefix = "linear_layer"
|
||||
cls = layer_map["linear"]
|
||||
method = get_quant_method({"linear_layer.weight": quant_type},
|
||||
prefix, "linear")
|
||||
self.assertIsInstance(method, cls)
|
||||
def test_detects_compressed_tensors(self):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
config_path = os.path.join(tmpdir, "config.json")
|
||||
with open(config_path, "w") as f:
|
||||
json.dump({"quantization_config": {"quant_method": "compressed-tensors"}}, f)
|
||||
|
||||
def test_moe_quant_methods(self):
|
||||
for quant_type, layer_map in ASCEND_QUANTIZATION_METHOD_MAP.items():
|
||||
if "moe" in layer_map.keys():
|
||||
prefix = "layer"
|
||||
cls = layer_map["moe"]
|
||||
method = get_quant_method({"layer.weight": quant_type}, prefix,
|
||||
"moe")
|
||||
self.assertIsInstance(method, cls)
|
||||
result = detect_quantization_method(tmpdir)
|
||||
self.assertEqual(result, COMPRESSED_TENSORS_METHOD)
|
||||
|
||||
def test_with_fa_quant_type(self):
|
||||
quant_description = {"fa_quant_type": "C8"}
|
||||
method = get_quant_method(quant_description, ".attn", "attention")
|
||||
self.assertIsInstance(
|
||||
method, ASCEND_QUANTIZATION_METHOD_MAP["C8"]["attention"])
|
||||
def test_returns_none_for_no_quant(self):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
result = detect_quantization_method(tmpdir)
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_with_kv_quant_type(self):
|
||||
quant_description = {"kv_quant_type": "C8"}
|
||||
method = get_quant_method(quant_description, ".attn", "attention")
|
||||
self.assertIsInstance(
|
||||
method, ASCEND_QUANTIZATION_METHOD_MAP["C8"]["attention"])
|
||||
def test_returns_none_for_non_compressed_tensors_quant_method(self):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
config_path = os.path.join(tmpdir, "config.json")
|
||||
with open(config_path, "w") as f:
|
||||
json.dump({"quantization_config": {"quant_method": "gptq"}}, f)
|
||||
|
||||
def test_invalid_layer_type(self):
|
||||
quant_description = {"linear_layer.weight": "W8A8"}
|
||||
with self.assertRaises(NotImplementedError):
|
||||
get_quant_method(quant_description, "linear_layer", "unsupported")
|
||||
result = detect_quantization_method(tmpdir)
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_invalid_quant_type(self):
|
||||
quant_description = {"linear_layer.weight": "UNKNOWN"}
|
||||
with self.assertRaises(NotImplementedError):
|
||||
get_quant_method(quant_description, "linear_layer", "linear")
|
||||
def test_returns_none_for_config_without_quant_config(self):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
config_path = os.path.join(tmpdir, "config.json")
|
||||
with open(config_path, "w") as f:
|
||||
json.dump({"model_type": "llama"}, f)
|
||||
|
||||
result = detect_quantization_method(tmpdir)
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_returns_none_for_malformed_config_json(self):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
config_path = os.path.join(tmpdir, "config.json")
|
||||
with open(config_path, "w") as f:
|
||||
f.write("not valid json{{{")
|
||||
|
||||
result = detect_quantization_method(tmpdir)
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_modelslim_takes_priority_over_compressed_tensors(self):
|
||||
"""When both ModelSlim config and compressed-tensors config exist,
|
||||
ModelSlim should take priority."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
modelslim_path = os.path.join(tmpdir, MODELSLIM_CONFIG_FILENAME)
|
||||
with open(modelslim_path, "w") as f:
|
||||
json.dump({"layer.weight": "INT8"}, f)
|
||||
|
||||
config_path = os.path.join(tmpdir, "config.json")
|
||||
with open(config_path, "w") as f:
|
||||
json.dump({"quantization_config": {"quant_method": "compressed-tensors"}}, f)
|
||||
|
||||
result = detect_quantization_method(tmpdir)
|
||||
self.assertEqual(result, ASCEND_QUANTIZATION_METHOD)
|
||||
|
||||
|
||||
class TestMaybeAutoDetectQuantization(TestBase):
|
||||
def _make_vllm_config(self, model_path="/fake/model", quantization=None, revision=None):
|
||||
vllm_config = MagicMock()
|
||||
vllm_config.model_config.model = model_path
|
||||
vllm_config.model_config.quantization = quantization
|
||||
vllm_config.model_config.revision = revision
|
||||
return vllm_config
|
||||
|
||||
@patch("vllm_ascend.quantization.utils.detect_quantization_method", return_value=None)
|
||||
def test_no_detection_does_nothing(self, mock_detect):
|
||||
vllm_config = self._make_vllm_config()
|
||||
maybe_auto_detect_quantization(vllm_config)
|
||||
self.assertIsNone(vllm_config.model_config.quantization)
|
||||
|
||||
@patch("vllm_ascend.quantization.utils.detect_quantization_method", return_value=ASCEND_QUANTIZATION_METHOD)
|
||||
def test_user_specified_same_method_no_change(self, mock_detect):
|
||||
vllm_config = self._make_vllm_config(quantization=ASCEND_QUANTIZATION_METHOD)
|
||||
maybe_auto_detect_quantization(vllm_config)
|
||||
self.assertEqual(vllm_config.model_config.quantization, ASCEND_QUANTIZATION_METHOD)
|
||||
|
||||
@patch("vllm.config.VllmConfig._get_quantization_config", return_value=MagicMock())
|
||||
@patch("vllm_ascend.quantization.utils.detect_quantization_method", return_value=ASCEND_QUANTIZATION_METHOD)
|
||||
def test_auto_detect_sets_quantization_and_logs_info(self, mock_detect, mock_get_quant_config):
|
||||
"""When no --quantization is specified but ModelSlim config is found,
|
||||
the method should auto-set quantization and emit an INFO log."""
|
||||
vllm_config = self._make_vllm_config(model_path="/fake/quant_model", quantization=None)
|
||||
|
||||
with patch("vllm_ascend.quantization.utils.logger") as mock_logger:
|
||||
maybe_auto_detect_quantization(vllm_config)
|
||||
|
||||
self.assertEqual(vllm_config.model_config.quantization, ASCEND_QUANTIZATION_METHOD)
|
||||
mock_logger.info.assert_called_once()
|
||||
call_args = mock_logger.info.call_args[0]
|
||||
self.assertIn("Auto-detected quantization method", call_args[0])
|
||||
self.assertIn(ASCEND_QUANTIZATION_METHOD, call_args)
|
||||
self.assertIn("/fake/quant_model", call_args)
|
||||
|
||||
@patch("vllm_ascend.quantization.utils.detect_quantization_method", return_value=ASCEND_QUANTIZATION_METHOD)
|
||||
def test_user_mismatch_logs_warning(self, mock_detect):
|
||||
"""When user specifies a different method than auto-detected,
|
||||
a WARNING should be emitted and user's choice should be respected."""
|
||||
vllm_config = self._make_vllm_config(model_path="/fake/quant_model", quantization=COMPRESSED_TENSORS_METHOD)
|
||||
|
||||
with patch("vllm_ascend.quantization.utils.logger") as mock_logger:
|
||||
maybe_auto_detect_quantization(vllm_config)
|
||||
|
||||
self.assertEqual(vllm_config.model_config.quantization, COMPRESSED_TENSORS_METHOD)
|
||||
mock_logger.warning.assert_called_once()
|
||||
call_args = mock_logger.warning.call_args[0]
|
||||
self.assertIn("Auto-detected quantization method", call_args[0])
|
||||
self.assertIn(ASCEND_QUANTIZATION_METHOD, call_args)
|
||||
self.assertIn(COMPRESSED_TENSORS_METHOD, call_args)
|
||||
|
||||
@patch("vllm_ascend.quantization.utils.detect_quantization_method", return_value=None)
|
||||
def test_no_detection_emits_info_log(self, mock_detect):
|
||||
"""When no quantization is detected, an info log tells the user the model loads as float."""
|
||||
vllm_config = self._make_vllm_config(quantization=None)
|
||||
|
||||
with patch("vllm_ascend.quantization.utils.logger") as mock_logger:
|
||||
maybe_auto_detect_quantization(vllm_config)
|
||||
|
||||
mock_logger.info.assert_called_once()
|
||||
call_args = mock_logger.info.call_args[0]
|
||||
self.assertIn("No quantization signature detected", call_args[0])
|
||||
self.assertIn("/fake/model", call_args)
|
||||
mock_logger.warning.assert_not_called()
|
||||
self.assertIsNone(vllm_config.model_config.quantization)
|
||||
|
||||
@patch("vllm.config.VllmConfig._get_quantization_config", return_value=MagicMock())
|
||||
@patch("vllm_ascend.quantization.utils.detect_quantization_method", return_value=ASCEND_QUANTIZATION_METHOD)
|
||||
def test_passes_revision_to_detect(self, mock_detect, mock_get_quant):
|
||||
"""Verify that model revision is forwarded to detect_quantization_method."""
|
||||
vllm_config = self._make_vllm_config(model_path="org/model-name", revision="v1.0", quantization=None)
|
||||
maybe_auto_detect_quantization(vllm_config)
|
||||
mock_detect.assert_called_once_with("org/model-name", revision="v1.0")
|
||||
|
||||
|
||||
class TestEnableFaQuant(TestBase):
|
||||
def test_non_quantization_scenarios(self):
|
||||
# non quantization scene
|
||||
vllm_config = MagicMock()
|
||||
vllm_config.quant_config = None
|
||||
result = enable_fa_quant(vllm_config)
|
||||
self.assertFalse(result)
|
||||
|
||||
# CompressedTensors scene
|
||||
vllm_config.quant_config = AscendCompressedTensorsConfig({}, [], "", {})
|
||||
result = enable_fa_quant(vllm_config)
|
||||
self.assertFalse(result)
|
||||
|
||||
# non fa3 quant scene
|
||||
vllm_config.quant_config = AscendModelSlimConfig(W8A8_CONFIG)
|
||||
result = enable_fa_quant(vllm_config)
|
||||
self.assertFalse(result)
|
||||
|
||||
def test_fa3_quantization_scenario(self):
|
||||
vllm_config = MagicMock()
|
||||
vllm_config.quant_config = AscendModelSlimConfig(FAKQUANT_CONFIG)
|
||||
vllm_config.kv_transfer_config = KVTransferConfig(kv_connector="MultiConnector", kv_role="kv_consumer")
|
||||
result = enable_fa_quant(vllm_config)
|
||||
self.assertTrue(result)
|
||||
result = enable_fa_quant(vllm_config, layer_name="test_layer")
|
||||
self.assertFalse(result)
|
||||
|
||||
Reference in New Issue
Block a user