83 lines
4.0 KiB
Python
83 lines
4.0 KiB
Python
"""CPU checks of installed Kunlun INT8 loader and BF16 MLA weight preparation.
|
|
|
|
Uses synthetic layers and mocked TP rank/size, not a full engine startup.
|
|
"""
|
|
import copy
|
|
import os
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
os.environ["VLLM_KUNLUN_ENABLE_INT8_BMM"] = "0"
|
|
|
|
import torch
|
|
from convert_deepseek_fp8_to_int8 import quantization_config
|
|
|
|
|
|
def main():
|
|
from vllm.platforms import current_platform
|
|
current_platform.pre_register_and_update()
|
|
from vllm_kunlun.quantization.compressed_tensors.compressed_tensors import KunlunCompressedTensorsConfig
|
|
from vllm_kunlun.quantization.compressed_tensors.compressed_tensors_moe import KunlunCompressedTensorsMoEMethod
|
|
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
|
|
cfg = KunlunCompressedTensorsConfig.from_config(copy.deepcopy(quantization_config()))
|
|
cfg.packed_modules_mapping = {
|
|
"gate_up_proj": ["gate_proj", "up_proj"],
|
|
"fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
|
|
}
|
|
for name, widths in [
|
|
("model.layers.0.self_attn.fused_qkv_a_proj", [128, 64]),
|
|
("model.layers.0.self_attn.q_b_proj", [128]),
|
|
("model.layers.0.self_attn.o_proj", [128]),
|
|
]:
|
|
layer = torch.nn.Linear(128, sum(widths), bias=False)
|
|
scheme = cfg.get_scheme(layer, name)
|
|
del layer._parameters['weight']
|
|
with patch('vllm.model_executor.parameter.get_tensor_model_parallel_rank', return_value=0), \
|
|
patch('vllm.model_executor.parameter.get_tensor_model_parallel_world_size', return_value=1):
|
|
scheme.create_weights(layer, widths, 128, torch.bfloat16, lambda *a, **k: None)
|
|
assert type(scheme.kernel).__name__ == 'KunlunScaledMMLinearKernel'
|
|
assert layer.weight.dtype == torch.int8
|
|
assert tuple(layer.weight_scale.shape) == (sum(widths), 1)
|
|
print('LINEAR_PASS', name, type(scheme.kernel).__name__)
|
|
|
|
class RoutedExperts(torch.nn.Module):
|
|
pass
|
|
|
|
experts = RoutedExperts()
|
|
experts.moe_config = SimpleNamespace(is_act_and_mul=True)
|
|
method = KunlunCompressedTensorsMoEMethod.get_moe_method(cfg, experts, 'model.layers.3.mlp.experts')
|
|
assert type(method).__name__ == 'KunlunCompressedTensorsW8A8Int8MoEMethod'
|
|
method.create_weights(experts, 2, 128, 128, torch.bfloat16)
|
|
assert tuple(experts.w13_weight_scale.shape) == (2, 256, 1)
|
|
with torch.no_grad():
|
|
experts.w13_weight_scale.fill_(0.25)
|
|
experts.w2_weight_scale.fill_(0.25)
|
|
method.process_weights_after_loading(experts)
|
|
assert experts.w13_weight_scale.flatten()[0].item() == 31.75
|
|
print('MOE_SCALE_PASS', type(method).__name__)
|
|
for name in ['lm_head', 'model.layers.3.mlp.gate', 'model.layers.0.self_attn.kv_b_proj',
|
|
'model.layers.0.mlp.gate_up_proj', 'model.layers.0.mlp.down_proj',
|
|
'model.layers.3.mlp.shared_experts.gate_up_proj',
|
|
'model.layers.3.mlp.shared_experts.down_proj']:
|
|
assert cfg.get_scheme(torch.nn.Linear(128, 128), name) is None
|
|
print('IGNORE_PASS: head, router, MLA kv_b_proj, dense MLPs, shared experts')
|
|
|
|
from vllm_kunlun.v1.attention.backends.mla.common import MLACommonImpl
|
|
projection = torch.nn.Linear(128, 256, bias=False, dtype=torch.bfloat16)
|
|
projection.quant_method = UnquantizedLinearMethod()
|
|
impl = SimpleNamespace(kv_b_proj=projection, kv_lora_rank=128,
|
|
num_heads=2, qk_nope_head_dim=64, v_head_dim=64)
|
|
with torch.inference_mode():
|
|
MLACommonImpl.process_weights_after_loading(impl, torch.bfloat16)
|
|
assert tuple(impl.W_UK_T.shape) == (2, 64, 128)
|
|
assert tuple(impl.W_UV.shape) == (2, 128, 64)
|
|
assert impl.W_UK_T.dtype == impl.W_UV.dtype == torch.bfloat16
|
|
expected = projection.weight.T.reshape(128, 2, 128)
|
|
assert torch.equal(impl.W_UK_T, expected[:, :, :64].permute(1, 2, 0))
|
|
assert torch.equal(impl.W_UV, expected[:, :, 64:].transpose(0, 1))
|
|
print('MLA_BF16_PREPARATION_PASS')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|