457 lines
24 KiB
Python
457 lines
24 KiB
Python
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,
|
|
)
|