fix fp8 fakequant (#67715)

lixcli · web-flow · commit cc69be8ddcd0 · 2024-08-28T02:21:29.000-07:00
diff --git a/python/paddle/nn/quant/format.py b/python/paddle/nn/quant/format.py
@@ -23,6 +23,44 @@
 from ..layer.layers import Layer
 
 
+def fake_fp8_quant(input, scale, axis=-1, type='e4m3'):
+    # only support channelwise or tensorwise
+    if axis >= 0:
+        shape = [1] * len(input.shape)
+        shape[axis] = scale.numel()
+        scale = scale.reshape(shape)
+    inp = input.astype("float32")
+
+    if type == 'e4m3':
+        return paddle.cast(
+            (inp * 448 / scale).clip(-448, 448), "float8_e4m3fn"
+        ).astype(
+            input.dtype
+        )  # clip then cast
+    elif type == 'e5m2':
+        return paddle.cast(
+            (inp * 57344 / scale).clip(-57344, 57344), "float8_e5m2"
+        ).astype(
+            input.dtype
+        )  # clip then cast
+    else:
+        raise NotImplementedError("only support e4m3 or e5m2 now")
+
+
+def fake_fp8_dequant(input, scale, axis=-1, type='e4m3'):
+    # only support channelwise or tensorwise
+    if axis >= 0:
+        shape = [1] * len(input.shape)
+        shape[axis] = scale.numel()
+        scale = scale.reshape(shape)
+    if type == 'e4m3':
+        return (input.astype("float32") / 448 * scale).astype(input.dtype)
+    elif type == 'e5m2':
+        return (input.astype("float32") / 57344 * scale).astype(input.dtype)
+    else:
+        raise NotImplementedError("only support e4m3 or e5m2 now")
+
+
 class LinearQuanterDequanter(Layer):
     def __init__(self, quanter, dequanter):
         super().__init__()
@@ -129,6 +167,16 @@ def forward(self, input):
                         self._qmax,
                     )
                 return quant_weight.cast(input.dtype)
+
+            if self._qmax == 448:
+                return fake_fp8_quant(
+                    input, self._scales, self._quant_axis, type='e4m3'
+                )
+            elif self._qmax == 57344:
+                return fake_fp8_quant(
+                    input, self._scales, self._quant_axis, type='e5m2'
+                )
+
             return _C_ops.quantize_linear(
                 input.cast('float32'),
                 self._scales,
@@ -249,6 +297,15 @@ def forward(self, input):
                     )
                 return quant_dequant_weight.cast(input.dtype)
 
+            if self._qmax == 448:
+                return fake_fp8_dequant(
+                    input, self._scales, self._quant_axis, type='e4m3'
+                )
+            elif self._qmax == 57344:
+                return fake_fp8_dequant(
+                    input, self._scales, self._quant_axis, type='e5m2'
+                )
+
             return _C_ops.dequantize_linear(
                 input.cast('float32'),
                 self._scales,
