[PIR-3][33-35] Adapt weight api (#62988)

zrr1999 · web-flow · commit 2a697580af38 · 2024-04-09T17:27:14.000+08:00
diff --git a/python/paddle/nn/quant/quantized_linear.py b/python/paddle/nn/quant/quantized_linear.py
@@ -75,8 +75,7 @@ def weight_quantize(x, algo="weight_only_int8", arch=None, group_size=-1):
     assert (
         group_size == -1 or group_size == 64 or group_size == 128
     ), f"Currently group_size only support -1/64/128. but got {group_size} "
-
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.weight_quantize(x, algo, arch, group_size)
     else:
         type = "weight_quantize"
@@ -129,7 +128,7 @@ def weight_dequantize(
         out_dtype, 'out_dtype', ['float16', 'bfloat16'], 'weight_dequantize'
     )
     out_dtype = convert_np_dtype_to_dtype_(out_dtype)
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.weight_dequantize(x, scale, algo, out_dtype, group_size)
     else:
         type = "weight_dequantize"
@@ -200,7 +199,7 @@ def weight_only_linear(
         group_size == -1 or group_size == 64 or group_size == 128
     ), f"Currently weight_quantize only support group size of -1, 64 or 128. but got {group_size} "
 
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         out = _C_ops.weight_only_linear(
             x, weight, bias, weight_scale, weight_dtype, arch, group_size
         )
diff --git a/test/quantization/test_weight_only_linear.py b/test/quantization/test_weight_only_linear.py
@@ -26,6 +26,7 @@
 from paddle import base
 from paddle.base import core
 from paddle.framework import set_default_dtype
+from paddle.pir_utils import IrGuard
 
 np.random.seed(123)
 paddle.seed(123)
@@ -155,59 +156,9 @@ def get_weight_only_linear_out(self):
         )
         return out.numpy()
 
-    def get_weight_only_linear_out_static(self):
-        paddle.enable_static()
-        main = base.Program()
-        start = base.Program()
-        with base.program_guard(main, start):
-            x = paddle.static.data("x", self.x.shape, dtype=self.x.dtype)
-
-            weight = paddle.static.data(
-                "weight", self.weight.shape, dtype=self.weight.dtype
-            )
-            bias = paddle.static.data(
-                "bias", self.bias.shape, dtype=self.bias.dtype
-            )
-            x_np = self.x.numpy()
-            weight_np = self.weight.numpy()
-            bias_np = self.bias.numpy()
-            if self.weight_scale is not None:
-                weight_scale = paddle.static.data(
-                    "weight_scale",
-                    self.weight_scale.shape,
-                    dtype=self.weight_scale.dtype,
-                )
-                weight_scale_np = self.weight_scale.numpy()
-            else:
-                weight_scale = None
-                weight_scale_np = None
-
-            out = Q.weight_only_linear(
-                x,
-                weight,
-                bias,
-                weight_scale,
-                self.weight_dtype,
-                group_size=self.group_size,
-            )
-            feed_dict = {
-                'x': x_np,
-                'weight': weight_np,
-                'bias': bias_np,
-                "weight_scale": weight_scale_np,
-            }
-            exe = base.Executor(paddle.CUDAPlace(0))
-            exe.run(start)
-            (out,) = exe.run(main, feed=feed_dict, fetch_list=[out])
-        paddle.disable_static()
-        return out
-
     def test_weight_only_linear(self):
         out_expect = self.get_linear_out()
-        if self.static:
-            out_real = self.get_weight_only_linear_out_static()
-        else:
-            out_real = self.get_weight_only_linear_out()
+        out_real = self.get_weight_only_linear_out()
 
         if self.dtype == "bfloat16":
             out_real = convert_uint16_to_float(out_real)
@@ -651,14 +602,102 @@ def config(self):
 
 
 @unittest.skipIf(
-    not core.is_compiled_with_cuda() or get_cuda_version() < 11020,
+    not core.is_compiled_with_cuda()
+    or get_cuda_version() < 11020
+    or paddle.device.cuda.get_device_capability()[0] < 8,
     "quantized_matmul requires CUDA >= 11.2 and CUDA_ARCH >= 8",
 )
 class WeightOnlyLinearTestCaseStatic(WeightOnlyLinearTestCase):
     def config(self):
         super().config()
         self.static = True
 
+    def get_weight_only_linear_out_static(self):
+        paddle.enable_static()
+        main = paddle.static.Program()
+        start = paddle.static.Program()
+        with paddle.static.program_guard(main, start):
+            x = paddle.static.data("x", self.x.shape, dtype=self.x.dtype)
+
+            weight = paddle.static.data(
+                "weight", self.weight.shape, dtype=self.weight.dtype
+            )
+            bias = paddle.static.data(
+                "bias", self.bias.shape, dtype=self.bias.dtype
+            )
+            x_np = self.x.numpy()
+            weight_np = self.weight.numpy()
+            bias_np = self.bias.numpy()
+            if self.weight_scale is not None:
+                weight_scale = paddle.static.data(
+                    "weight_scale",
+                    self.weight_scale.shape,
+                    dtype=self.weight_scale.dtype,
+                )
+                weight_scale_np = self.weight_scale.numpy()
+            else:
+                weight_scale = None
+                weight_scale_np = None
+
+            out = Q.weight_only_linear(
+                x,
+                weight,
+                bias,
+                weight_scale,
+                self.weight_dtype,
+                group_size=self.group_size,
+            )
+            feed_dict = {
+                'x': x_np,
+                'weight': weight_np,
+                'bias': bias_np,
+                "weight_scale": weight_scale_np,
+            }
+            exe = base.Executor(paddle.CUDAPlace(0))
+            exe.run(start)
+            (out,) = exe.run(main, feed=feed_dict, fetch_list=[out])
+        paddle.disable_static()
+        return out
+
+    def test_weight_quantize_and_dequantize_pir(self, algo='weight_only_int8'):
+        with IrGuard():
+            weight = (
+                paddle.rand(shape=(4096, 12288), dtype='float16')
+                * 1
+                / math.sqrt(4096)
+            )
+
+            quant_weight, quant_scale = Q.weight_quantize(x=weight, algo=algo)
+            dequant_weight = Q.weight_dequantize(
+                quant_weight, quant_scale, algo=algo
+            )
+            exe = paddle.static.Executor(paddle.CUDAPlace(0))
+            res = exe.run(feed={}, fetch_list=[weight, dequant_weight])
+            np.testing.assert_allclose(res[0], res[1], rtol=1e-2, atol=1e-2)
+
+    def test_weight_quantize_and_dequantize_int4_pir(self):
+        self.test_weight_quantize_and_dequantize_pir(algo='weight_only_int4')
+
+    def test_weight_only_linear(self):
+        out_expect = self.get_linear_out()
+
+        out_real = self.get_weight_only_linear_out_static()
+        if self.dtype == "bfloat16":
+            out_real = convert_uint16_to_float(out_real)
+            out_expect = convert_uint16_to_float(out_expect)
+        np.testing.assert_allclose(
+            out_real, out_expect, rtol=self.rtol, atol=self.atol
+        )
+
+        with IrGuard():
+            out_real = self.get_weight_only_linear_out_static()
+        if self.dtype == "bfloat16":
+            out_real = convert_uint16_to_float(out_real)
+            out_expect = convert_uint16_to_float(out_expect)
+        np.testing.assert_allclose(
+            out_real, out_expect, rtol=self.rtol, atol=self.atol
+        )
+
 
 @unittest.skipIf(
     not core.is_compiled_with_cuda()
@@ -667,7 +706,9 @@ def config(self):
     "quantized_matmul requires CUDA >= 11.2 and CUDA_ARCH >= 8",
 )
 class WeightOnlyLinearBackwardAndWeightDequantizeTestCase(unittest.TestCase):
-    def test_weightonly_linear_backward(self):
+    def test_weightonly_linear_backward(
+        self, algo='weight_only_int8', weight_dtype='int8'
+    ):
         x = (
             paddle.rand(shape=(128, 4096), dtype='float16')
             * 1
@@ -683,16 +724,18 @@ def test_weightonly_linear_backward(self):
         )
 
         quant_weight, quant_scale = Q.weight_quantize(
-            x=weight.cuda(), algo='weight_only_int8'
+            x=weight.cuda(), algo=algo
+        )
+        dequant_weight = Q.weight_dequantize(
+            quant_weight.cuda(), quant_scale, algo=algo
         )
-        dequant_weight = Q.weight_dequantize(quant_weight.cuda(), quant_scale)
         np.testing.assert_allclose(weight, dequant_weight, rtol=1e-2, atol=1e-2)
 
         quant_out = Q.weight_only_linear(
             x=quant_x,
             weight=quant_weight,
             weight_scale=quant_scale,
-            weight_dtype="int8",
+            weight_dtype=weight_dtype,
         )
         out = paddle.matmul(x=x, y=weight)
         np.testing.assert_allclose(quant_out, out, rtol=1e-3, atol=1e-3)
@@ -701,6 +744,11 @@ def test_weightonly_linear_backward(self):
         out.backward()
         np.testing.assert_allclose(quant_x.grad, x.grad, rtol=1e-3, atol=1e-3)
 
+    def test_weightonly_linear_backward_int4(self):
+        self.test_weightonly_linear_backward(
+            algo='weight_only_int4', weight_dtype='int4'
+        )
+
 
 if __name__ == '__main__':
     unittest.main()
