195 lines
9.6 KiB
Python
195 lines
9.6 KiB
Python
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import torch
|
|
from transformers import DeepseekV2Config
|
|
|
|
from vllm_ascend.eplb.adaptor.vllm_adaptor import EPLB_EXPERT_WEIGHT_NAMES, VllmEplbAdaptor
|
|
from vllm_ascend.quantization.quant_type import QuantType
|
|
|
|
|
|
class TestVllmAdaptor(unittest.TestCase):
|
|
def setUp(self):
|
|
VllmEplbAdaptor._registered_moe_layers = []
|
|
|
|
n_routed_experts = 256
|
|
self.mock_layer = MagicMock()
|
|
self.mock_layer.local_num_experts = n_routed_experts
|
|
self.mock_layer.ep_rank = 0
|
|
self.mock_layer.quant_type = QuantType.W8A8
|
|
self.mock_layer.w13_weight_list = [torch.randn(256, 128) for _ in range(n_routed_experts)]
|
|
self.mock_layer.w2_weight_list = [torch.randn(128, 256) for _ in range(n_routed_experts)]
|
|
self.mock_layer.w13_weight_scale_fp32_list = [torch.tensor([1.0]) for _ in range(n_routed_experts)]
|
|
self.mock_layer.w2_weight_scale_list = [torch.tensor([1.0]) for _ in range(n_routed_experts)]
|
|
self.mock_layer.w13_weight = torch.randn(n_routed_experts, 256, 128)
|
|
self.mock_layer.w2_weight = torch.randn(n_routed_experts, 128, 256)
|
|
self.mock_layer.moe_load = torch.randn(n_routed_experts)
|
|
self.mock_layer.global_expert_map = torch.arange(n_routed_experts * 4).reshape(n_routed_experts, 4)
|
|
self.mock_layer.get_log2phy_map.return_value = torch.arange(4)
|
|
self.mock_layer.clear_moe_load = MagicMock()
|
|
VllmEplbAdaptor.register_layer(self.mock_layer)
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.model.named_parameters.return_value = dict()
|
|
config = DeepseekV2Config(n_routed_experts=n_routed_experts)
|
|
mock_model.config = config
|
|
del mock_model.language_model
|
|
self.model = mock_model
|
|
num_dense_layers = getattr(config, "first_k_dense_replace", 0)
|
|
self.model.model.layers[num_dense_layers].mlp.experts.quant_type = QuantType.W8A8
|
|
|
|
self.mock_rank = patch("vllm_ascend.eplb.adaptor.vllm_adaptor.dist.get_rank", return_value=0).start()
|
|
self.mock_size = patch("vllm_ascend.eplb.adaptor.vllm_adaptor.dist.get_world_size", return_value=4).start()
|
|
|
|
@patch("torch.empty_like", return_value=torch.zeros(16, 32))
|
|
@patch("vllm_ascend.eplb.adaptor.vllm_adaptor.get_ascend_config")
|
|
def test_init_fp16(self, mock_get_config, mock_func):
|
|
mock_config = MagicMock()
|
|
mock_config.enable_fused_mc2 = 1
|
|
mock_get_config.return_value = mock_config
|
|
self.model.quant_config = None
|
|
adaptor = VllmEplbAdaptor(self.model)
|
|
self.assertEqual(adaptor.expert_weight_key_per_layer[0], (QuantType.NONE, True))
|
|
self.assertIs(adaptor.expert_param_per_layer[0][0][0], self.mock_layer.w13_weight_list[0])
|
|
self.assertIs(adaptor.expert_param_per_layer[0][0][1], self.mock_layer.w2_weight_list[0])
|
|
|
|
@patch("torch.empty_like", return_value=torch.zeros(16, 32))
|
|
@patch("vllm_ascend.eplb.adaptor.vllm_adaptor.get_ascend_config")
|
|
def test_init_w8a8(self, mock_get_config, mock_func):
|
|
mock_config = MagicMock()
|
|
mock_config.enable_fused_mc2 = 0
|
|
mock_get_config.return_value = mock_config
|
|
VllmEplbAdaptor(self.model)
|
|
|
|
@patch("torch.empty_like", return_value=torch.zeros(16, 32))
|
|
@patch("vllm_ascend.eplb.adaptor.vllm_adaptor.get_ascend_config")
|
|
def test_language_model_w8a8(self, mock_get_config, mock_func):
|
|
mock_config = MagicMock()
|
|
mock_config.enable_fused_mc2 = 0
|
|
mock_get_config.return_value = mock_config
|
|
model = MagicMock()
|
|
model.language_model = self.model
|
|
model.config.text_config = self.model.config
|
|
VllmEplbAdaptor(model)
|
|
|
|
def test_pp_eplb_adaptor_init_with_registered_layer(self):
|
|
"""PP+EPLB: adaptor picks up MoE layers registered via register_layer."""
|
|
VllmEplbAdaptor._registered_moe_layers = []
|
|
layer = MagicMock()
|
|
layer.local_num_experts = 4
|
|
layer.ep_rank = 0
|
|
layer.quant_type = QuantType.W8A8
|
|
layer.w13_weight_list = [torch.randn(256, 128) for _ in range(4)]
|
|
layer.w2_weight_list = [torch.randn(128, 256) for _ in range(4)]
|
|
layer.w13_weight_scale_fp32_list = [torch.tensor([1.0]) for _ in range(4)]
|
|
layer.w2_weight_scale_list = [torch.tensor([1.0]) for _ in range(4)]
|
|
layer.moe_load = torch.randn(4)
|
|
layer.global_expert_map = torch.arange(16).reshape(4, 4)
|
|
layer.get_log2phy_map.return_value = torch.arange(4)
|
|
VllmEplbAdaptor.register_layer(layer)
|
|
|
|
with patch("vllm_ascend.eplb.adaptor.vllm_adaptor.get_ascend_config") as mock_get_config:
|
|
mock_config = MagicMock()
|
|
mock_config.enable_fused_mc2 = 0
|
|
mock_get_config.return_value = mock_config
|
|
model = MagicMock()
|
|
model.quant_config = MagicMock()
|
|
model.config.first_k_dense_replace = 0
|
|
del model.language_model
|
|
adaptor = VllmEplbAdaptor(model)
|
|
|
|
self.assertEqual(adaptor.num_moe_layers, 1)
|
|
self.assertEqual(adaptor.num_local_experts, 4)
|
|
self.assertEqual(adaptor.ep_rank, 0)
|
|
|
|
@patch("vllm_ascend.eplb.adaptor.vllm_adaptor.get_ascend_config")
|
|
def test_init_mixed_quant_type_per_layer(self, mock_get_config):
|
|
mock_config = MagicMock()
|
|
mock_config.enable_fused_mc2 = 1
|
|
mock_get_config.return_value = mock_config
|
|
|
|
VllmEplbAdaptor._registered_moe_layers = []
|
|
num_local_experts = 2
|
|
w8a8_layer = MagicMock()
|
|
w8a8_layer.local_num_experts = num_local_experts
|
|
w8a8_layer.ep_rank = 0
|
|
w8a8_layer.quant_type = QuantType.W8A8
|
|
w8a8_layer.w13_weight_list = [torch.randn(2, 2) for _ in range(num_local_experts)]
|
|
w8a8_layer.w2_weight_list = [torch.randn(2, 2) for _ in range(num_local_experts)]
|
|
w8a8_layer.w13_weight_scale_fp32_list = [torch.randn(1) for _ in range(num_local_experts)]
|
|
w8a8_layer.w2_weight_scale_list = [torch.randn(1) for _ in range(num_local_experts)]
|
|
w8a8_layer.fused_w1_scale_list = [torch.randn(1) for _ in range(num_local_experts)]
|
|
w8a8_layer.fused_w2_scale_list = [torch.randn(1) for _ in range(num_local_experts)]
|
|
w8a8_layer.moe_load = torch.zeros(num_local_experts)
|
|
w8a8_layer.global_expert_map = torch.arange(num_local_experts * 4).reshape(num_local_experts, 4)
|
|
w8a8_layer.get_log2phy_map.return_value = torch.arange(4)
|
|
|
|
mxfp8_layer = MagicMock()
|
|
mxfp8_layer.local_num_experts = num_local_experts
|
|
mxfp8_layer.ep_rank = 0
|
|
mxfp8_layer.quant_type = QuantType.MXFP8
|
|
mxfp8_layer.w13_weight = torch.randn(num_local_experts, 2, 2)
|
|
mxfp8_layer.w2_weight = torch.randn(num_local_experts, 2, 2)
|
|
mxfp8_layer.w13_weight_scale = torch.randn(num_local_experts, 1)
|
|
mxfp8_layer.w2_weight_scale = torch.randn(num_local_experts, 1)
|
|
mxfp8_layer.moe_load = torch.zeros(num_local_experts)
|
|
mxfp8_layer.global_expert_map = torch.arange(num_local_experts * 4).reshape(num_local_experts, 4)
|
|
mxfp8_layer.get_log2phy_map.return_value = torch.arange(4)
|
|
|
|
VllmEplbAdaptor.register_layer(w8a8_layer)
|
|
VllmEplbAdaptor.register_layer(mxfp8_layer)
|
|
|
|
model = MagicMock()
|
|
model.quant_config = MagicMock()
|
|
model.config.first_k_dense_replace = 0
|
|
del model.language_model
|
|
adaptor = VllmEplbAdaptor(model)
|
|
|
|
w8a8_key = (QuantType.W8A8, True)
|
|
mxfp8_key = (QuantType.MXFP8, True)
|
|
self.assertEqual(adaptor.expert_weight_key_per_layer[0], w8a8_key)
|
|
self.assertEqual(adaptor.expert_weight_key_per_layer[1], mxfp8_key)
|
|
self.assertEqual(len(adaptor.buffer_tensor_list[w8a8_key][0]), len(EPLB_EXPERT_WEIGHT_NAMES[w8a8_key]))
|
|
self.assertEqual(len(adaptor.buffer_tensor_list[mxfp8_key][0]), len(EPLB_EXPERT_WEIGHT_NAMES[mxfp8_key]))
|
|
self.assertEqual(len(adaptor.expert_param_per_layer[0][0]), len(EPLB_EXPERT_WEIGHT_NAMES[w8a8_key]))
|
|
self.assertEqual(len(adaptor.expert_param_per_layer[1][0]), len(EPLB_EXPERT_WEIGHT_NAMES[mxfp8_key]))
|
|
|
|
@patch("vllm_ascend.eplb.adaptor.vllm_adaptor.get_ascend_config")
|
|
def test_reused_buffer_requires_same_expert_weight_shape(self, mock_get_config):
|
|
mock_config = MagicMock()
|
|
mock_config.enable_fused_mc2 = 0
|
|
mock_get_config.return_value = mock_config
|
|
|
|
VllmEplbAdaptor._registered_moe_layers = []
|
|
num_local_experts = 2
|
|
for weight_shape in [(2, 2), (3, 2)]:
|
|
layer = MagicMock()
|
|
layer.local_num_experts = num_local_experts
|
|
layer.ep_rank = 0
|
|
layer.quant_type = QuantType.W8A8
|
|
layer.w13_weight_list = [torch.randn(*weight_shape) for _ in range(num_local_experts)]
|
|
layer.w2_weight_list = [torch.randn(2, 2) for _ in range(num_local_experts)]
|
|
layer.w13_weight_scale_fp32_list = [torch.randn(1) for _ in range(num_local_experts)]
|
|
layer.w2_weight_scale_list = [torch.randn(1) for _ in range(num_local_experts)]
|
|
layer.moe_load = torch.zeros(num_local_experts)
|
|
layer.global_expert_map = torch.arange(num_local_experts * 4).reshape(num_local_experts, 4)
|
|
layer.get_log2phy_map.return_value = torch.arange(4)
|
|
VllmEplbAdaptor.register_layer(layer)
|
|
|
|
model = MagicMock()
|
|
model.quant_config = MagicMock()
|
|
model.config.first_k_dense_replace = 0
|
|
del model.language_model
|
|
|
|
with self.assertRaisesRegex(AssertionError, "EPLB expert weight shapes mismatch"):
|
|
VllmEplbAdaptor(model)
|
|
|
|
def tearDown(self):
|
|
self.mock_rank.stop()
|
|
self.mock_size.stop()
|
|
VllmEplbAdaptor._registered_moe_layers = []
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|