79 lines
3.7 KiB
Python
79 lines
3.7 KiB
Python
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))
|