weight_quantize api support gpu (#59090)

wwbitejotunn · web-flow · commit 3aff2c0abf03 · 2023-11-24T13:48:40.000+08:00
* weight quant gpu wint8

* fix scale type of weight ony quant

* fix cpu quant weight

* update

* float scale for weight-only8 gpu quant

* fix infermeta

* fix include

* add sm70

* fixup and andd sm70

* code_style fix
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -5123,9 +5123,10 @@ void WeightQuantizeInferMeta(const MetaTensor& x,
                              MetaTensor* out,
                              MetaTensor* scale) {
   PADDLE_ENFORCE_EQ(
-      ((arch == 80) || (arch == 70)),
+      ((arch == 80) || (arch == 86) || (arch == 70) || (arch == 75)),
       true,
-      phi::errors::InvalidArgument("Currently, arch only support 70, 80."));
+      phi::errors::InvalidArgument(
+          "Currently, arch only support 70, 75, 80, 86."));
 
   auto x_dims = x.dims();
   PADDLE_ENFORCE_EQ(
diff --git a/paddle/phi/kernels/cpu/weight_quantize_kernel.cc b/paddle/phi/kernels/cpu/weight_quantize_kernel.cc
@@ -30,9 +30,10 @@ void quant_compute(const DeviceContext& dev_ctx,
                    const std::string& algo,
                    const int32_t arch) {
   PADDLE_ENFORCE_EQ(
-      ((arch == 80) || (arch == 70)),
+      ((arch == 80) || (arch == 86) || (arch == 75) || (arch == 70)),
       true,
-      phi::errors::InvalidArgument("Currently, arch only support 70, 80."));
+      phi::errors::InvalidArgument(
+          "Currently, arch only support 70, 75, 80, 86."));
 
   const auto x_dims = x.dims();
   PADDLE_ENFORCE_EQ(
@@ -49,7 +50,8 @@ void quant_compute(const DeviceContext& dev_ctx,
   float* scale_data = scale->data<float>();
 
   DenseTensor x_int(out->type());
-  if (arch == 80) {
+
+  if ((arch == 80) || (arch == 75) || (arch == 86)) {
     x_int.Resize({static_cast<int64_t>(m), static_cast<int64_t>(n)});
   } else {
     // phi::Copy may change tensor meta info, here we transpose the quanted
@@ -81,8 +83,12 @@ void quant_compute(const DeviceContext& dev_ctx,
       // Note(Zhengzekang): In sm70, we only need RowMajor layout, just add bias
       // to make it unsigned.
       add_bias_and_interleave_inplace<bits>(x_int_data, num);
-      phi::Copy(dev_ctx, x_int, dev_ctx.GetPlace(), false, out);
-    } else if (arch == 80) {
+      // phi::Copy break the shape of int4 output, use naive copy;
+      // only left half of x_int data is valid in int4 mode
+      for (int i = 0; i < out->numel(); ++i) {
+        out_data[i] = x_int_data[i];
+      }
+    } else if ((arch == 80) || (arch == 75) || (arch == 86)) {
       permute_B_rows_for_mixed_gemm<bits>(
           int_processed_data, x_int_data, std::vector<size_t>{m, n});
       subbyte_transpose_impl<bits>(
diff --git a/paddle/phi/kernels/gpu/weight_quantize_kernel.cu b/paddle/phi/kernels/gpu/weight_quantize_kernel.cu
@@ -0,0 +1,84 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/common/datatype_traits.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/common_shape.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/impl/weight_quantize_kernel_gpu_impl.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void WeightQuantizeKernel(const Context& dev_ctx,
+                          const DenseTensor& x,
+                          const std::string& algo,
+                          const int32_t arch,
+                          DenseTensor* out,
+                          DenseTensor* scale) {
+  DenseTensor quanted_x;
+  dev_ctx.template Alloc<int8_t>(out);
+  dev_ctx.template Alloc<float>(scale);
+  size_t m = x.dims()[0];
+  size_t n = x.dims()[1];
+  quanted_x.Resize({static_cast<int64_t>(m), static_cast<int64_t>(n)});
+  dev_ctx.template Alloc<int8_t>(&quanted_x);
+  std::vector<int> weight_shape{static_cast<int>(x.dims()[0]),
+                                static_cast<int>(x.dims()[1])};
+  PADDLE_ENFORCE_EQ(
+      ((arch == 80) || (arch == 86) || (arch == 75) || (arch == 70)),
+      true,
+      phi::errors::InvalidArgument(
+          "Currently, arch only support 70, 75, 80, 86."));
+
+  if (algo == "llm.int8") {
+    std::vector<int> axis = {1, 0};
+    funcs::Transpose<Context, int8_t, 2> trans;
+    weight_quant_gpu<T, Context>(dev_ctx,
+                                 x.data<T>(),
+                                 quanted_x.data<int8_t>(),
+                                 scale->data<float>(),
+                                 weight_shape);
+    trans(dev_ctx, quanted_x, out, axis);
+  } else if (algo == "weight_only_int8") {
+    weight_quant_gpu<T, Context>(dev_ctx,
+                                 x.data<T>(),
+                                 quanted_x.data<int8_t>(),
+                                 scale->data<float>(),
+                                 weight_shape);
+    weight_permute_gpu<Context>(dev_ctx,
+                                quanted_x.data<int8_t>(),
+                                out->data<int8_t>(),
+                                weight_shape,
+                                arch);
+  } else if (algo == "weight_only_int4") {
+    phi::errors::Unimplemented(
+        "Weight quant gpu kernel currently don't support weight_only_int4 "
+        "algo, please use cpu version.");
+  } else {
+    phi::errors::Unimplemented(
+        "The algo must be in ['weight_only_int8', 'weight_only_int4', "
+        "'llm.int8'], but got[%s]",
+        algo);
+  }
+}
+}  // namespace phi
+
+PD_REGISTER_KERNEL(weight_quantize,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::WeightQuantizeKernel,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/impl/weight_quantize_kernel_gpu_impl.h b/paddle/phi/kernels/impl/weight_quantize_kernel_gpu_impl.h
@@ -0,0 +1,185 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "paddle/phi/backends/gpu/gpu_launch_config.h"
+#include "paddle/phi/common/float16.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/aligned_vector.h"
+#include "paddle/phi/kernels/funcs/common_shape.h"
+namespace phi {
+
+__global__ void weight_permute_kernel_wint8(const int8_t* input_data_dev,
+                                            int8_t* output_data_dev,
+                                            int numel,
+                                            int total_k,
+                                            int total_n) {
+  for (int linear_idx = blockIdx.x * blockDim.x + threadIdx.x;
+       linear_idx < numel;
+       linear_idx += blockDim.x * gridDim.x) {
+    int k_id = linear_idx / total_n;
+    int n_id = linear_idx % total_n;
+    constexpr int k_permute_const = 8;
+    int k_mod_16 = k_id % 16;
+    int temp_k_expr_1 = k_mod_16 - k_mod_16 / 8 * 8;
+    int temp_k_expr_2 = k_mod_16 / 8;
+    int permute_kk = temp_k_expr_1 + temp_k_expr_2 +
+                     (temp_k_expr_2 + 1) % 2 * k_mod_16 * 2 / 2 +
+                     temp_k_expr_1 * temp_k_expr_2 + k_id / 16 * 16;
+    int permute_index = permute_kk % 64 + permute_kk / 64 * 128 +
+                        64 * (n_id % 2) + total_k * 2 * (n_id / 2);
+    uint8_t shift_quant_weight = static_cast<uint8_t>(
+        static_cast<int32_t>(input_data_dev[linear_idx]) + 128);
+    output_data_dev[permute_index] =
+        *reinterpret_cast<int8_t*>(&shift_quant_weight);
+  }
+}
+
+/*
+For SM70 volta arch, weightonly int8 dequantize invoked in load global memory.
+So it only need interleave in K-dimension
+K_index: 0 1 2 3 -> 0 2 1 3
+*/
+__global__ void weight_interleave_add_bias_kernel_wint8(
+    const int8_t* input_data_dev,
+    int8_t* output_data_dev,
+    int numel,
+    int total_k,
+    int total_n) {
+  for (int linear_idx = blockIdx.x * blockDim.x + threadIdx.x;
+       linear_idx < numel;
+       linear_idx += blockDim.x * gridDim.x) {
+    int k_id = linear_idx / total_n;
+    int n_id = linear_idx % total_n;
+    constexpr int n_interleaved_factor = 4;
+    int n_interleave_group_id = n_id / n_interleaved_factor;
+    int n_interleave_id = n_id % n_interleaved_factor;
+    if (n_interleave_id == 1 || n_interleave_id == 2) {
+      /*
+      0001 xor 0011 -> 0010
+      0010 xor 0011 -> 0001
+      */
+      n_interleave_id ^= 3;
+    }
+    const int new_n_id =
+        n_interleave_group_id * n_interleaved_factor + n_interleave_id;
+    const int interleave_idx = k_id * total_n + new_n_id;
+
+    uint8_t shift_quant_weight = static_cast<uint8_t>(
+        static_cast<int32_t>(input_data_dev[linear_idx]) + 128);
+    output_data_dev[interleave_idx] =
+        *reinterpret_cast<int8_t*>(&shift_quant_weight);
+  }
+}
+
+template <typename GPUContext>
+void weight_permute_gpu(const GPUContext& dev_ctx,
+                        int8_t* input_data,
+                        int8_t* output_data,
+                        const std::vector<int>& shape,
+                        const int32_t arch) {
+  auto total_k = shape[0];
+  auto total_n = shape[1];
+  auto numel = total_k * total_n;
+  auto gpu_config = phi::backends::gpu::GetGpuLaunchConfig1D(dev_ctx, numel, 1);
+  int grid_size = gpu_config.GetGridSize();
+  int block_size = gpu_config.GetBlockSize();
+  if ((arch == 80) || (arch == 86) || (arch == 75)) {
+    weight_permute_kernel_wint8<<<grid_size, block_size>>>(
+        input_data, output_data, numel, total_k, total_n);
+  } else if (arch == 70) {
+    weight_interleave_add_bias_kernel_wint8<<<grid_size, block_size>>>(
+        input_data, output_data, numel, total_k, total_n);
+  }
+}
+template <typename T, int VectorSize = 8>
+__global__ void per_channel_quant_gpu(const T* weight_data,
+                                      int8_t* quanted_weight_data,
+                                      float* scale_data,
+                                      int total_k,
+                                      int total_vec_n) {
+  int n = blockIdx.x * blockDim.x + threadIdx.x;
+  if (n < total_vec_n) {
+    const int4* vec_weight_data_ptr =
+        reinterpret_cast<const int4*>(weight_data);
+    int2* vec_quanted_weight_data =
+        reinterpret_cast<int2*>(quanted_weight_data);
+    phi::AlignedVector<float, VectorSize> abs_max;
+#pragma unroll
+    for (int i = 0; i < VectorSize; ++i) {
+      abs_max[i] = static_cast<float>(0.0f);
+    }
+#pragma unroll
+    for (int k = 0; k < total_k; ++k) {
+      int linear_index = k * total_vec_n + n;
+      phi::AlignedVector<T, VectorSize> weight;
+      *reinterpret_cast<int4*>(&weight) = vec_weight_data_ptr[linear_index];
+#pragma unroll
+      for (int i = 0; i < VectorSize; ++i) {
+        abs_max[i] = fmaxf((abs_max[i]), fabsf((weight[i])));
+      }
+    }
+    phi::AlignedVector<float, VectorSize> scale;
+#pragma unroll
+    for (int i = 0; i < VectorSize; ++i) {
+      scale[i] = static_cast<float>(abs_max[i] / static_cast<float>(127.0f));
+    }
+    *reinterpret_cast<float4*>(scale_data + VectorSize * n) =
+        *reinterpret_cast<float4*>(&scale);
+    *reinterpret_cast<float4*>(scale_data + VectorSize * n + 4) =
+        *reinterpret_cast<float4*>(&(scale[4]));
+
+    for (int k = 0; k < total_k; ++k) {
+      phi::AlignedVector<int8_t, VectorSize> quanted_weight;
+      int linear_index = k * total_vec_n + n;
+      phi::AlignedVector<T, VectorSize> weight;
+      *reinterpret_cast<int4*>(&weight) =
+          *reinterpret_cast<const int4*>(vec_weight_data_ptr + linear_index);
+#pragma unroll
+      for (int i = 0; i < VectorSize; ++i) {
+        float scaled_weight =
+            (static_cast<float>(weight[i]) / static_cast<float>(abs_max[i])) *
+            static_cast<float>(127.0);
+        int8_t clipped_weight = static_cast<int8_t>(
+            lroundf(fmaxf(-127.0f, fminf(127.0f, scaled_weight))));
+        quanted_weight[i] = clipped_weight;
+      }
+      *reinterpret_cast<int2*>(vec_quanted_weight_data + linear_index) =
+          *reinterpret_cast<int2*>(&quanted_weight);
+    }
+  }
+}
+
+template <typename T, typename GPUContext>
+void weight_quant_gpu(const GPUContext& dev_ctx,
+                      const T* weight_data,
+                      int8_t* quanted_weight_data,
+                      float* scale_data,
+                      const std::vector<int>& shape) {
+  int total_k = shape[0];
+  int total_n = shape[1];
+  int numel = total_k * total_n;
+  constexpr int kWarpSize = 32;
+  constexpr int kBlockSize = 64;
+  constexpr int kWarpNum = kBlockSize / kWarpSize;
+  constexpr int kVectorSize = 128 / sizeof(T) / 8;
+  int vec_total_n = total_n / kVectorSize;
+  int kGridSize = max(vec_total_n / kBlockSize, static_cast<int>(1));
+  per_channel_quant_gpu<T, kVectorSize><<<kGridSize, kBlockSize>>>(
+      weight_data, quanted_weight_data, scale_data, total_k, vec_total_n);
+}
+
+}  // namespace phi
diff --git a/paddle/phi/kernels/impl/weight_quantize_kernel_impl.h b/paddle/phi/kernels/impl/weight_quantize_kernel_impl.h
@@ -37,7 +37,6 @@
 #include "paddle/phi/kernels/funcs/common_shape.h"
 
 namespace phi {
-
 template <typename T>
 inline T xabs(const T x) {
   return x < static_cast<T>(0.0) ? -x : x;
@@ -51,7 +50,7 @@ void per_channel_scale(
     for (size_t j = 0; j < m; ++j) {
       max = xabs(input[j * n + i]) > max ? xabs(input[j * n + i]) : max;
     }
-    scale[i] = static_cast<float>(max) / bound;
+    scale[i] = static_cast<float>(static_cast<float>(max) / bound);
   }
 }
 
@@ -67,7 +66,7 @@ void per_channel_quant(int8_t* output,
     const T* current_weight_row = input + ii * num_cols;
     for (size_t jj = 0; jj < bytes_per_out_col; ++jj) {
       if (quant_bit == 8) {
-        const float col_scale = scale[jj];
+        const float col_scale = static_cast<float>(scale[jj]);
         const float weight_elt = static_cast<float>(current_weight_row[jj]);
         const float scaled_weight = round(weight_elt / col_scale);
         const int8_t clipped_weight = static_cast<int8_t>(
@@ -79,7 +78,7 @@ void per_channel_quant(int8_t* output,
         for (int packed_idx = 0; packed_idx < 2; ++packed_idx) {
           const size_t input_idx = 2 * jj + packed_idx;
           if (input_idx < num_cols) {
-            const float col_scale = scale[input_idx];
+            const float col_scale = static_cast<float>(scale[input_idx]);
             const float weight_elt =
                 static_cast<float>(current_weight_row[input_idx]);
             const float scaled_weight = round(weight_elt / col_scale);
@@ -349,7 +348,6 @@ void interleave_column_major_tensor(int8_t* interleaved_quantized_tensor,
             interleave * base_vec_row +
             vec_rows_per_tile * (read_col % interleave) +
             vec_read_row % vec_rows_per_tile;
-
         const size_t read_offset =
             size_t(read_col) * num_vec_rows + vec_read_row;
         const size_t write_offset =
diff --git a/python/paddle/nn/quant/quantized_linear.py b/python/paddle/nn/quant/quantized_linear.py
@@ -68,8 +68,8 @@ def weight_quantize(x, algo="weight_only_int8", arch=None):
         arch = _get_arch_info()
 
     assert (
-        arch == 70 or arch == 80
-    ), "Currently weight_quantize only support SM70/80. "
+        arch == 70 or arch == 80 or arch == 86 or arch == 75
+    ), f"Currently weight_quantize only support SM70/75/80/86. but got {arch} "
 
     if in_dynamic_mode():
         return _C_ops.weight_quantize(x, algo, arch)
diff --git a/test/quantization/test_weight_only_linear.py b/test/quantization/test_weight_only_linear.py
