大模型推理 - GPTQ 落地与优化

 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

# 处理int4权重，转换名字，保存成numpy格式, model_file为gptq int4量化后保存的权重文件
def _process_quant_file(model_config, model_file, dtype: torch.dtype, tp_size: int, save_dir: Path):
    state_dict = torch.load(model_file, map_location="cpu")
    for name in state_dict:
        param = state_dict[name]
        # Preprocess
        param_name = convert_parameter_name(name)
        if not param_name.endswith(("qzeros", "scales", "qweight")):
            # param = safe_transpose(param) 注意不要转置，因为gqtq已转置
            continue
        param = handle_exceptions(model_config, param_name, param)
        convert_and_save_parameter(param_name, param.detach().cpu().numpy(), tp_size, save_dir)

# qkv权重reorder
def reorder_qkv_weight_or_bias(model_config: PretrainedConfig,
                               name: str,
                               param: torch.nn.Parameter):
    if 'query_key_value' not in name:
        # Nothing to do for the non-eligible parameters.
        return param

    num_heads = model_config.n_head
    head_dim = model_config.hidden_size // model_config.n_head
    if "qzeros" in name:
        head_dim = head_dim // 8 # qzero 一行8个int4合并int32, 所以需要除8

    # (..., 3 x hidden) view as (..., num_heads, 3, head_dim)
    param = param.view(-1, num_heads, 3, head_dim)
    # permute to (..., 3, num_heads, head_dim)
    param = param.permute(0, 2, 1, 3)
    # final shape: weight=(hidden, 3, hidden) or bias=(3, hidden)
    if 'query_key_value.bias' in name:
        return param.reshape(3, num_heads * head_dim)
    return param.reshape(-1, 3, num_heads * head_dim)

 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

// src/fastertransformer/layers/attention_layers/DecoderSelfAttentionLayer.cc

template<typename T>
void DecoderSelfAttentionLayer<T>::forward(TensorMap*                output_tensors,
                                           TensorMap*                input_tensors,
                                           const AttentionWeight<T>* attention_weights)
{
    if (int8_mode_ == 1) {
        // int8 weight only gemm
    }
    else if (int8_mode_ == 10) { // gptq gemm
              gptq_runner_->gemm(attention_weights->query_weight.qweight, // int4的权重 buf 
                              {d_model_, 3 * local_hidden_units_}, // 权重的尺寸
                               attention_weights->query_weight.qscales,  // 量化系数 buf
                               attention_weights->query_weight.qzeros,   // 零点值 buf
                               attention_input, // 输入的matrix buf 
                               {batch_size, d_model_},  // 输入的尺寸
                               qkv_buf_, // 输出buf
                               stream_, cublas_wrapper_);
    }
    else { // fp16
        cublas_wrapper_->Gemm(CUBLAS_OP_N,
                                CUBLAS_OP_N,
                                3 * local_hidden_units_,  // n
                                batch_size,
                                d_model_,  // k
                                attention_weights->query_weight.kernel,
                                3 * local_hidden_units_,  // n
                                attention_input,
                                d_model_,  // k
                                qkv_buf_,
                                3 * local_hidden_units_ /* n */);
    }
}

  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

// int4 矩阵乘
// BATCH: 每BATCH行共享mat的还原,。m<16时，BATCH=m; m>16时BATCH=16, 然后blockIdx.z=16的倍数，余数单独调用一次kernel。
// TWIDTH: 每个block处理[BATCH, TWIDTH] * [TWIDTH, TWIDTH], 需根据显卡实际测试最优值。
// TBANK: 控制deq2的冗余，减少bank冲突，取值范围[1, 32]
// scalar_t: vec的数据类型，fp16

template <int BATCH, int TWIDTH, int TBANK, typename scalar_t>
__global__ void VecQuant4MatMulKernel_perf(
    const  half2* __restrict__ vec,
    const    int* __restrict__ mat,
           half2* __restrict__ mul,
    const  scalar_t* __restrict__ scales,
    const    int* __restrict__ zeros,
    int vec_height,
    int height,
    int width,
    int zero_width
) 
{
  const int blockwidth2 = TWIDTH / 2;
  int h = TWIDTH / 8 * blockIdx.x;
  int w = TWIDTH * blockIdx.y + 2 * threadIdx.x;
  int w2 = w + 1;
  vec = vec + BATCH * blockIdx.z * vec_height;
  mul = mul + BATCH * blockIdx.z * width / 2;

  __shared__ half2 blockvec[BATCH][TWIDTH];
  #pragma unroll
  for (int b=0; b<BATCH; ++b) {
      blockvec[b][threadIdx.x] = vec[b * vec_height + blockIdx.x * blockwidth2 + threadIdx.x];
    }

  __shared__ half2 deq2[256][TBANK];
  int val = threadIdx.x / TBANK;
  int off = threadIdx.x % TBANK;
  for (; val < 256; val += blockwidth2 / TBANK) {
    deq2[val][off] = __halves2half2(
       __int2half_rn(val & 0xF), __int2half_rn(val >> 4)
    );
  }

  __syncthreads();

  int i = width * h + w;
  int g_h = h * 8;
  int k = 0;

  int z_w = w / 8; 
  int z_mod = (w % 8) * 4;

  // int z_w2 = w2 / 8; 
  int z_mod2 = (w2 % 8) * 4;

  // float res[BATCH] = {0};
  // float res2[BATCH] = {0};
  half2 resh[BATCH] = { {} };
  half2 resh2[BATCH] = { {} };

  unsigned int tmp;
  unsigned int tmp2;
  int g = g_h / GROUPSIZE;
  const int group_count = TWIDTH / GROUPSIZE;
  // #pragma unroll
  for (int gi=0; gi<group_count; ++gi) {
    float scale_f = scales[g * width + w];
    half2 scale = __float2half2_rn(scale_f);
    unsigned tmp_zero = as_unsigned(zeros[g * zero_width + z_w]);
    half2 zero = __float2half2_rn(-(scale_f * (((tmp_zero >> z_mod) & 0xF) + 1)));

    float scale_f2 = scales[g * width + w2];
    half2 scale2 = __float2half2_rn(scale_f2);
    half2 zero2 = __float2half2_rn(-(scale_f2 * (((tmp_zero >> z_mod2) & 0xF) + 1)));

    const int gk = GROUPSIZE / 2 * (gi+1);
    while (k < gk) { // blockvec一次取2个值

        tmp = as_unsigned(mat[i]);
        tmp2 = as_unsigned(mat[i + 1]);

        half2 m0 = __hfma2(deq2[(tmp >>  0) & 0xff][off], scale, zero);
        half2 m1 = __hfma2(deq2[(tmp >>  8) & 0xff][off], scale, zero);
        half2 m2 = __hfma2(deq2[(tmp >>  16) & 0xff][off], scale, zero);
        half2 m3 = __hfma2(deq2[(tmp >>  24) & 0xff][off], scale, zero);

        half2 m2_0 = __hfma2(deq2[(tmp2 >>  0) & 0xff][off], scale2, zero2);
        half2 m2_1 = __hfma2(deq2[(tmp2 >>  8) & 0xff][off], scale2, zero2);
        half2 m2_2 = __hfma2(deq2[(tmp2 >>  16) & 0xff][off], scale2, zero2);
        half2 m2_3 = __hfma2(deq2[(tmp2 >>  24) & 0xff][off], scale2, zero2);

          #pragma unroll
          for (int b=0; b<BATCH; ++b) {
              
              half2 v0 = blockvec[b][k + 0];
              half2 v1 = blockvec[b][k + 1];
              half2 v2 = blockvec[b][k + 2];
              half2 v3 = blockvec[b][k + 3];

              resh[b] = __hfma2(m0, v0, resh[b]);
              resh[b] = __hfma2(m1, v1, resh[b]);
              resh[b] = __hfma2(m2, v2, resh[b]);
              resh[b] = __hfma2(m3, v3, resh[b]);

              resh2[b] = __hfma2(m2_0, v0, resh2[b]);
              resh2[b] = __hfma2(m2_1, v1, resh2[b]);
              resh2[b] = __hfma2(m2_2, v2, resh2[b]);
              resh2[b] = __hfma2(m2_3, v3, resh2[b]);

          }
          i += width;
          k += 4;
        }

    g = g + 1;    
  }

  w = blockwidth2 * blockIdx.y + threadIdx.x;

  for (int b=0; b<BATCH; ++b) {
    float resa = __half2float(resh[b].x) + __half2float(resh[b].y);
    float resb = __half2float(resh2[b].x) + __half2float(resh2[b].y);
    
    half2 result = make_half2(resa, resb);
    atomicAdd(&mul[b * width / 2 + w], result);
    // atomicAdd(&mul[b * width + w2], res2[b]);
  }
}