AccText2Img/w8a8.py at master · Qcompiler/AccText2Img · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157

import torch
import math
import torch.nn as nn
import sys
import gc
from torch.nn.parameter import Parameter
from torch.nn import init
from torch.nn import functional as F
import mixgemm
import mixlib
try:
    from EETQ import quant_weights, w8_a16_gemm
    memory_bound_eetq_linear = True
except:
    memory_bound_eetq_linear = False

# memory_bound_eetq_linear = False

def FindOutliers(Activation, sigma = None):

    if sigma is None:
        sigma = 20

    tmp = torch.unique(torch.where((  Activation.abs() > sigma ))[1])
    return tmp.to(torch.int32)


layer_id = 0
class MixLinear_GEMM(nn.Module):
    def __init__(self, in_features, out_features, bias = True,
                 device=None):
        super().__init__()
        global layer_id
        layer_id += 1
        self.layer_id = layer_id
        dtype = torch.float16

        factory_kwargs = {'device': device, 'dtype': dtype, "requires_grad": False}
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Parameter(torch.empty((out_features, in_features), **factory_kwargs), requires_grad = False)
        self.q_scale_col = torch.empty((1, out_features), **factory_kwargs)

        self.q_weight = torch.empty((1, 1), dtype = torch.int8)
        if bias:
            self.bias = Parameter(torch.empty(out_features, **factory_kwargs))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()
        self.weight_cache = None
        self.ind = None

        self.init = False
        self.quanted = False
        self.n_outliers = 0

        self.cnt = 0
        self.output = torch.zeros((1) , dtype = torch.float16)


        self.y1 = None
        self.reuse_output_because_of_zeros_input = False
        self.last_input = None
        self.cache_computed = False


        self.q_scale_col = None
        self.input_scales = None
        self.reuse_scaling_factor = False

        self.scale_max = None
        self.scale_min = None
        self.doing_estimation = True

    def reset_parameters(self) -> None:
        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
            init.uniform_(self.bias, -bound, bound)


    def forward(self, input):


        assert input.dtype == torch.float16

        if  self.init is not True:
            if self.layer_id == 1:
                print("I am init the weight do not disturb me and count me during time estimation")

            if len(input.shape) == 3:
                M = input.shape[0] * input.shape[1]
                self.input_scales = torch.zeros((M, 1), dtype = torch.float32, device = input.device)
            self.init = True

            # print("I should quant this layer")
            tmp = input.reshape(-1, input.shape[-1])


            self.weight.data = self.weight.data.cpu()
            tmp = self.weight.data


            # self.weight.data = self.weight.data.cpu()
            # del self.weight
            self.q_scale_col =   (torch.max(torch.abs(tmp), dim=1)[0].unsqueeze(1) / (127)).to(torch.float16).reshape((1,self.out_features))
            tmp  /= self.q_scale_col.T
            tmp = torch.clamp(tmp, -128, 127)

            self.q_weight = tmp.round().to(torch.int8).cuda()
            self.q_scale_col = self.q_scale_col.cuda().reshape((self.out_features))


            # 把 bias 打包到 scale 里面


            self.quanted = True
            self.weight.data = self.weight.data.cpu()

            del tmp


        input_ = input.reshape(-1, input.shape[-1])

        M = input_.shape[0]
        K = self.in_features
        N = self.out_features


        if not input_.is_contiguous():

            input_ = input_.contiguous()

        scaleRow = torch.zeros((M, 1) , dtype= torch.float16, device= input_.device)
        q_xcache = mixlib.FindRowScale(input_, scaleRow, M, K, 8)
        scaleRow = scaleRow.to(torch.float32)
        self.q_scale_col = self.q_scale_col.cuda().to(torch.float32)
        from vllm import _custom_ops as ops
        y1 =   ops.cutlass_scaled_mm(
                    q_xcache,
                    self.q_weight.T,
                    out_dtype=torch.float16,
                    scale_a=scaleRow,
                    scale_b=self.q_scale_col,
                    bias = self.bias
                )

        return y1.reshape(input.shape[:-1] + (self.out_features,))

    def extra_repr(self) -> str:
        return f'in_features={self.in_features}, out_features={self.out_features}, bias={self.bias is not None}'