63 lines
2.8 KiB
Python
63 lines
2.8 KiB
Python
|
"""Check installed Kunlun loader's configuration and CPU parameter contracts.
|
||
|
|
|
||
|
|
This uses synthetic layers and a mocked TP rank, not a full engine startup.
|
||
|
|
"""
|
||
|
|
import copy
|
||
|
|
from types import SimpleNamespace
|
||
|
|
from unittest.mock import patch
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from convert_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
|
||
|
|
|
||
|
|
cfg = KunlunCompressedTensorsConfig.from_config(copy.deepcopy(quantization_config()))
|
||
|
|
cfg.packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
|
||
|
|
linear = torch.nn.Linear(128, 128, bias=False)
|
||
|
|
scheme = cfg.get_scheme(linear, "model.layers.0.self_attn.qkv_proj")
|
||
|
|
# CPU-only parameter allocation; no distributed processes or accelerator work.
|
||
|
|
del linear._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(linear, [128], 128, torch.bfloat16, lambda *a, **k: None)
|
||
|
|
assert type(scheme.kernel).__name__ == "KunlunScaledMMLinearKernel"
|
||
|
|
assert linear.weight.dtype == torch.int8
|
||
|
|
assert tuple(linear.weight_scale.shape) == (128, 1)
|
||
|
|
print("LINEAR_PASS", type(scheme.kernel).__name__)
|
||
|
|
|
||
|
|
class RoutedExperts(torch.nn.Module):
|
||
|
|
pass
|
||
|
|
|
||
|
|
experts = RoutedExperts()
|
||
|
|
experts.ckpt_gate_proj_name = "w1"
|
||
|
|
experts.ckpt_down_proj_name = "w2"
|
||
|
|
experts.ckpt_up_proj_name = "w3"
|
||
|
|
experts.moe_config = SimpleNamespace(is_act_and_mul=True)
|
||
|
|
method = KunlunCompressedTensorsMoEMethod.get_moe_method(
|
||
|
|
cfg, experts, "model.layers.0.block_sparse_moe.experts")
|
||
|
|
assert type(method).__name__ == "KunlunCompressedTensorsW8A8Int8MoEMethod"
|
||
|
|
method.create_weights(experts, 2, 128, 128, torch.bfloat16)
|
||
|
|
assert tuple(experts.w13_weight.shape) == (2, 256, 128)
|
||
|
|
assert tuple(experts.w13_weight_scale.shape) == (2, 256, 1)
|
||
|
|
assert tuple(experts.w2_weight_scale.shape) == (2, 128, 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_PASS", type(method).__name__)
|
||
|
|
print("SCALE_PASS: checkpoint scale * 127 = runtime absmax")
|
||
|
|
assert cfg.get_scheme(torch.nn.Linear(128, 128), "lm_head") is None
|
||
|
|
assert cfg.get_scheme(torch.nn.Linear(128, 128), "model.layers.0.block_sparse_moe.gate") is None
|
||
|
|
print("IGNORE_PASS: lm_head and router gate")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|