test_linear.py 1.84 KB
Newer Older
chenxl's avatar
chenxl committed
1
2
3
4
5
6
7
8
#!/usr/bin/env python
# coding=utf-8
'''
Description  :  
Author       : chenht2022
Date         : 2024-07-25 10:32:05
Version      : 1.0.0
LastEditors  : chenht2022 
9
LastEditTime : 2024-08-06 10:36:59
chenxl's avatar
chenxl committed
10
11
12
13
14
15
16
17
Copyright (c) 2024 by KVCache.AI, All Rights Reserved. 
'''
import os, sys
import time
sys.path.append(os.path.dirname(__file__) + '/../build')
import cpuinfer_ext
import torch

18
19
20
21
22
23
24
25
26
27
input_size = 16384
output_size = 5120
stride = 32
group_max_len = 1024
proj_type = 1 # ggml_type::GGML_TYPE_F16
hidden_type = 1 # ggml_type::GGML_TYPE_F16
qlen = 30
layer_num = 10
CPUInfer = cpuinfer_ext.CPUInfer(48)
validation_iter = 100
chenxl's avatar
chenxl committed
28

29
with torch.inference_mode(mode=True):
chenxl's avatar
chenxl committed
30
31
32
33
    linears = []
    projs = []
    for _ in range(layer_num):
        proj = torch.randn((output_size, input_size), dtype=torch.float16, device = "cuda").to("cpu").contiguous()
34
        config = cpuinfer_ext.linear.LinearConfig(input_size, output_size, stride, group_max_len, proj.data_ptr(), proj_type, hidden_type)
chenxl's avatar
chenxl committed
35
36
37
38
39
40
41
        linear = cpuinfer_ext.linear.Linear(config)
        projs.append(proj)
        linears.append(linear)

    # validation
    for i in range(validation_iter):
        linear = linears[i % layer_num]
42
43
        input = torch.randn((qlen, input_size), dtype=torch.float16).contiguous()
        output = torch.empty((qlen, output_size), dtype=torch.float16).contiguous()
chenxl's avatar
chenxl committed
44
45
        input = input / 100

46
47
48
49
50
51
52
        CPUInfer.submit(
            linear.forward(
                qlen,
                input.data_ptr(),
                output.data_ptr()
            )
        )
chenxl's avatar
chenxl committed
53
54
55
56
57
58
59
60
61
62
        CPUInfer.sync()
        # print('cpuinfer output', output)

        proj = projs[i%layer_num]
        t_output = torch.mm(input, proj.t())
        # print('torch output', t_output)

        diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output))
        print('diff = ', diff)
        assert(diff < 0.001)