add qat_fp8linear (#10570)

paddlenlp/quantization/qat_utils.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from copy import deepcopy
+
 import paddle
 from paddle.autograd import PyLayer
 
@@ -19,6 +21,20 @@ from paddlenlp.utils import infohub
 
 from .hadamard_utils import random_hadamard_matrix
 
+try:
+    from transformer_engine import transformer_engine_paddle as tex
+    from transformer_engine.paddle.constants import FP8BwdTensors, FP8FwdTensors
+    from transformer_engine.paddle.cpp_extensions import fp8_gemm
+    from transformer_engine.paddle.layer.base import get_workspace
+
+    TE_DType = {
+        paddle.float8_e4m3fn: tex.DType.kFloat8E4M3,
+        paddle.float8_e5m2: tex.DType.kFloat8E5M2,
+    }
+    USE_FP8_GEMM = True
+except ImportError:
+    USE_FP8_GEMM = False
+
 
 def quantize_tensorwise(x, quantization_config=None, bit_length=8, state=0, training=False, act_scale=None):
     qmax = (1 << (bit_length - 1)) - 1
@@ -30,14 +46,18 @@ def quantize_tensorwise(x, quantization_config=None, bit_length=8, state=0, trai
 
     if act_scale is not None:
         if training:
-            scale = paddle.max(paddle.abs(target_x)) / qmax
-            act_scale.set_value((state * act_scale + scale) / (state + 1))
-            if state > quantization_config.skip_first_act_scale_step:
+            scale = paddle.max(paddle.abs(target_x)) / qmax + quantization_config.epsilon
+            if state < quantization_config.skip_first_act_scale_step:
+                act_scale.set_value((state * act_scale + scale) / (state + 1))
+            else:
+                act_scale.set_value(
+                    (1 - quantization_config.moving_rate) * act_scale + quantization_config.moving_rate * scale
+                )
                 scale = act_scale
         else:
             scale = act_scale
     else:
-        scale = paddle.max(paddle.abs(target_x)) / qmax
+        scale = paddle.max(paddle.abs(target_x)) / qmax + quantization_config.epsilon
 
     x_int8 = paddle.clip((target_x / scale).round(), qmin, qmax).astype("int8")
     return x_int8, scale
@@ -50,6 +70,61 @@ def dequantize_tensorwise(x_int8, scale, apply_hadamard=False):
     return x
 
 
+def fp8_quantize_tensorwise(x, tensor_type, quantization_config=None, state=0, training=False, act_scale=None):
+    assert tensor_type in ["weight", "activation", "grad_output"], "Only support weight, activation and grad_output"
+    fp8_format = quantization_config.fp8_format[tensor_type]
+    qmin, qmax = (-448, 448) if fp8_format == "float8_e4m3fn" else (-57344, 57344)
+    tensor_type_to_shape_index = {"weight": 0, "activation": -1, "grad_output": -2}
+
+    if quantization_config is not None and quantization_config.apply_hadamard:
+        if getattr(infohub, "hadamard") is None:
+            setattr(infohub, "hadamard", {})
+
+        hadamard_matrix_shape = x.shape[tensor_type_to_shape_index[tensor_type]]
+        if hadamard_matrix_shape in infohub.hadamard:
+            hadamard_matrix, block_size = infohub.hadamard[hadamard_matrix_shape]
+        else:
+            hadamard_matrix, block_size = random_hadamard_matrix(hadamard_matrix_shape, x.dtype, is_block=True)
+            infohub.hadamard[hadamard_matrix_shape] = (hadamard_matrix, block_size)
+        target_x = hadamard_matrix.T @ x if tensor_type in ["weight", "grad_output"] else x @ hadamard_matrix
+    else:
+        target_x = x
+        block_size = 1
+
+    if act_scale is not None:
+        if training:
+            scale = paddle.max(paddle.abs(target_x)) / qmax + quantization_config.epsilon
+            if state < quantization_config.skip_first_act_scale_step:
+                act_scale.set_value((state * act_scale + scale) / (state + 1))
+            else:
+                act_scale.set_value(
+                    (1 - quantization_config.moving_rate) * act_scale + quantization_config.moving_rate * scale
+                )
+                # scale = act_scale
+        else:
+            scale = act_scale
+    else:
+        scale = paddle.max(paddle.abs(target_x)) / qmax + quantization_config.epsilon
+
+    x_fp8 = target_x / scale
+    x_fp8 = x_fp8.astype(fp8_format).view("int8")
+    x_fp8 = x_fp8.T if tensor_type == "weight" else x_fp8
+    scale.stop_gradient = True
+    scale = scale / block_size if tensor_type in ["weight", "grad_output"] else scale
+    return x_fp8, scale
+
+
+def fp8_dequantize_tensorwise(x_fp8, scale, tensor_type, quantization_config=None):
+    x_fp8 = x_fp8.view(quantization_config.fp8_format[tensor_type])
+    x_fp8 = x_fp8.T if tensor_type == "weight" else x_fp8
+    x = x_fp8.astype(scale.dtype) * scale
+    if quantization_config.apply_hadamard:
+        hadamard_matrix_shape = x.shape[0] if tensor_type == "weight" else x.shape[-1]
+        hadamard_matrix, _ = infohub.hadamard[hadamard_matrix_shape]
+        x = hadamard_matrix @ x if tensor_type == "weight" else x @ hadamard_matrix.T
+    return x
+
+
 def quantize_channelwise(w, apply_hadamard=False, bit_length=8):
     qmax = (1 << (bit_length - 1)) - 1
     qmin = -1 * qmax - 1
@@ -77,7 +152,7 @@ def dequantize_channelwise(w_int8, scale, apply_hadamard=False):
     return w
 
 
-def a8w8_linear(
+def a8w8_forward(
     x, w_int8, w_scale=None, bias=None, dtype=None, quantization_config=None, state=0, training=False, act_scale=None
 ):
     x_int8, x_scale = quantize_tensorwise(
@@ -86,7 +161,180 @@ def a8w8_linear(
     out = paddle.matmul(x_int8, w_int8.T).astype(dtype) * (x_scale * w_scale.unsqueeze(0))
     if bias is not None:
         out += bias
-    return out
+    return out, x_int8, x_scale
+
+
+def a8w8_backward(ctx, x, grad_output, quant_weight, quant_scale, quant_x, x_scale):
+    if not ctx.x_stop_gradient:
+        if ctx.quantization_config.quant_input_grad:
+            grad_output_int8, grad_output_scale = quantize_tensorwise(grad_output * quant_scale)
+            input_grad = paddle.matmul(grad_output_int8, quant_weight).astype(ctx.dtype) * grad_output_scale
+            if ctx.quantization_config.apply_hadamard:
+                input_grad = input_grad @ infohub.hadamard[quant_weight.shape[-1]][0].T
+        else:
+            qdq_weight = dequantize_channelwise(
+                quant_weight, quant_scale, apply_hadamard=ctx.quantization_config.apply_hadamard
+            )
+            input_grad = paddle.matmul(grad_output, qdq_weight.T)
+    else:
+        input_grad = None
+
+    if not ctx.w_stop_gradient:
+        if len(x.shape) == 2:
+            weight_grad = paddle.matmul(x.transpose([1, 0]), grad_output)
+        else:
+            weight_grad = paddle.matmul(
+                x.reshape([-1, x.shape[-1]]).transpose([1, 0]), grad_output.reshape([-1, grad_output.shape[-1]])
+            )
+    else:
+        weight_grad = None
+
+    return input_grad, weight_grad
+
+
+def fp8_forward(
+    x, w_fp8, w_scale=None, bias=None, dtype=None, quantization_config=None, state=0, training=False, act_scale=None
+):
+    x_fp8, x_scale = fp8_quantize_tensorwise(
+        x,
+        tensor_type="activation",
+        quantization_config=quantization_config,
+        state=state,
+        training=training,
+        act_scale=act_scale,
+    )
+    x_fp8 = x_fp8.view(quantization_config.fp8_format["activation"])
+    w_fp8 = w_fp8.view(quantization_config.fp8_format["weight"])
+
+    if USE_FP8_GEMM:
+        x_shape = x_fp8.shape
+        x_fp8 = x_fp8.view((-1, x_fp8.shape[-1]))
+        fwd_scales = paddle.stack([x_scale.astype("float32"), w_scale.astype("float32")])
+        out, _ = fp8_gemm(
+            A=w_fp8,
+            A_scale_inv=fwd_scales,
+            A_fp8_tensor=FP8FwdTensors.GEMM1_WEIGHT,
+            A_dtype=TE_DType[w_fp8.dtype],
+            B=x_fp8,
+            B_scale_inv=fwd_scales,
+            B_fp8_tensor=FP8FwdTensors.GEMM1_INPUT,
+            B_dtype=TE_DType[x_fp8.dtype],
+            out_dtype=dtype,
+            workspace=get_workspace(),
+            bias=bias,
+            use_bias=True if bias is not None else False,
+            use_split_accumulator=True,
+        )
+        x_fp8 = x_fp8.view(x_shape)
+        out = out.view((*x_shape[:-1], -1))
+    else:
+        x = x_fp8.astype(dtype) * x_scale
+        w = w_fp8.astype(dtype) * w_scale
+        out = paddle.matmul(x, w.T).astype(dtype)
+        if bias is not None:
+            out += bias
+
+    return out, x_fp8, x_scale
+
+
+def fp8_backward(ctx, x, grad_output, quant_weight, quant_scale, quant_x, x_scale):
+    if not ctx.x_stop_gradient:
+        if ctx.quantization_config.quant_input_grad:
+            grad_output_fp8, grad_output_scale = fp8_quantize_tensorwise(
+                grad_output,
+                tensor_type="grad_output",
+                quantization_config=ctx.quantization_config,
+            )
+            grad_output_fp8 = grad_output_fp8.view(ctx.quantization_config.fp8_format["grad_output"])
+            quant_weight = quant_weight.view(ctx.quantization_config.fp8_format["weight"])
+            if USE_FP8_GEMM:
+                grad_output_shape = grad_output_fp8.shape
+                grad_output_fp8 = grad_output_fp8.view((-1, grad_output_fp8.shape[-1]))
+                fwd_scales = paddle.stack([x_scale.astype("float32"), quant_scale.astype("float32")])
+                bwd_scales = grad_output_scale[None].astype("float32")
+                input_grad, _ = fp8_gemm(
+                    A=quant_weight.T,
+                    A_scale_inv=fwd_scales,
+                    A_fp8_tensor=FP8FwdTensors.GEMM1_WEIGHT,
+                    A_dtype=TE_DType[quant_weight.dtype],
+                    B=grad_output_fp8,
+                    B_scale_inv=bwd_scales,
+                    B_fp8_tensor=FP8BwdTensors.GRAD_OUTPUT1,
+                    B_dtype=TE_DType[grad_output_fp8.dtype],
+                    out_dtype=ctx.dtype,
+                    workspace=get_workspace(),
+                    use_split_accumulator=True,
+                )
+                input_grad = input_grad.view((*grad_output_shape[:-1], -1))
+            else:
+                grad_output_ = grad_output_fp8.astype(ctx.dtype) * grad_output_scale
+                weight_ = quant_weight.astype(ctx.dtype) * quant_scale
+                input_grad = paddle.matmul(grad_output_, weight_).astype(ctx.dtype)
+            if ctx.quantization_config.apply_hadamard:
+                input_grad = infohub.hadamard[grad_output.shape[-2]][0] @ input_grad
+                input_grad = input_grad @ infohub.hadamard[quant_weight.shape[-1]][0].T
+        else:
+            qdq_weight = fp8_dequantize_tensorwise(
+                quant_weight, quant_scale, tensor_type="weight", quantization_config=ctx.quantization_config
+            )
+            input_grad = paddle.matmul(grad_output, qdq_weight.T)
+    else:
+        input_grad = None
+
+    if not ctx.w_stop_gradient:
+        if ctx.quantization_config.quant_weight_grad:
+            quantization_config_ = deepcopy(ctx.quantization_config)
+            quantization_config_.apply_hadamard = False
+            grad_output_fp8, grad_output_scale = fp8_quantize_tensorwise(
+                grad_output,
+                tensor_type="grad_output",
+                quantization_config=quantization_config_,
+            )
+            grad_output_fp8 = grad_output_fp8.view(ctx.quantization_config.fp8_format["grad_output"])
+            quant_x = quant_x.view(ctx.quantization_config.fp8_format["activation"])
+            if USE_FP8_GEMM:
+                quant_x = quant_x.view((-1, quant_x.shape[-1]))
+                grad_output_fp8 = grad_output_fp8.view((-1, grad_output_fp8.shape[-1]))
+                fwd_scales = paddle.stack([x_scale.astype("float32"), quant_scale.astype("float32")])
+                bwd_scales = grad_output_scale[None].astype("float32")
+                weight_grad, _ = fp8_gemm(
+                    A=grad_output_fp8.T,
+                    A_scale_inv=bwd_scales,
+                    A_fp8_tensor=FP8BwdTensors.GRAD_OUTPUT1,
+                    A_dtype=TE_DType[grad_output_fp8.dtype],
+                    B=quant_x.T,
+                    B_scale_inv=fwd_scales,
+                    B_fp8_tensor=FP8FwdTensors.GEMM1_INPUT,
+                    B_dtype=TE_DType[quant_x.dtype],
+                    out_dtype=ctx.dtype,
+                    workspace=get_workspace(),
+                    use_split_accumulator=True,
+                )
+            else:
+                grad_output_ = grad_output_fp8.astype(ctx.dtype) * grad_output_scale
+                x_ = quant_x.astype(ctx.dtype) * x_scale
+                if len(x_.shape) == 2:
+                    weight_grad = paddle.matmul(x_.transpose([1, 0]), grad_output_).astype(ctx.dtype)
+                else:
+                    weight_grad = paddle.matmul(
+                        x_.reshape([-1, x_.shape[-1]]).transpose([1, 0]),
+                        grad_output_.reshape([-1, grad_output_.shape[-1]]),
+                    ).astype(ctx.dtype)
+            if ctx.quantization_config.apply_hadamard:
+                hadamard_matrix, block_size = infohub.hadamard[quant_x.shape[-1]]
+                weight_grad = weight_grad / block_size
+                weight_grad = hadamard_matrix @ weight_grad
+        else:
+            if len(x.shape) == 2:
+                weight_grad = paddle.matmul(x.transpose([1, 0]), grad_output)
+            else:
+                weight_grad = paddle.matmul(
+                    x.reshape([-1, x.shape[-1]]).transpose([1, 0]), grad_output.reshape([-1, grad_output.shape[-1]])
+                )
+    else:
+        weight_grad = None
+
+    return input_grad, weight_grad
 
 
 class QATFunc(PyLayer):
@@ -103,53 +351,57 @@ class QATFunc(PyLayer):
         training,
         act_scale,
     ):
-        output = a8w8_linear(
-            x,
-            quant_weight,
-            w_scale=quant_scale,
-            bias=bias,
-            dtype=dtype,
-            quantization_config=quantization_config,
-            state=state,
-            training=training,
-            act_scale=act_scale,
-        )
+        quant_x, x_scale = None, None
+        if quantization_config.weight_quantize_algo in ["fp8linear"]:
+            output, quant_x, x_scale = fp8_forward(
+                x,
+                quant_weight,
+                w_scale=quant_scale,
+                bias=bias,
+                dtype=dtype,
+                quantization_config=quantization_config,
+                state=state,
+                training=training,
+                act_scale=act_scale,
+            )
+        else:
+            output, quant_x, x_scale = a8w8_forward(
+                x,
+                quant_weight,
+                w_scale=quant_scale,
+                bias=bias,
+                dtype=dtype,
+                quantization_config=quantization_config,
+                state=state,
+                training=training,
+                act_scale=act_scale,
+            )
         ctx.quantization_config = quantization_config
         ctx.dtype = dtype
-        ctx.save_for_backward(x, quant_weight, bias, quant_scale)
+        ctx.x_stop_gradient = x.stop_gradient
+        ctx.w_stop_gradient = quant_weight.stop_gradient
+        ctx.b_stop_gradient = bias.stop_gradient if bias is not None else True
+        ctx.save_for_backward(
+            x if not quantization_config.quant_weight_grad else None,
+            quant_weight,
+            bias,
+            quant_scale,
+            quant_x if quantization_config.quant_weight_grad else None,
+            x_scale if quantization_config.quant_weight_grad else None,
+        )
+
         return output
 
     @staticmethod
     def backward(ctx, grad_output):
-        x, quant_weight, bias, quant_scale = ctx.saved_tensor()
+        x, quant_weight, bias, quant_scale, quant_x, x_scale = ctx.saved_tensor()
 
-        if not x.stop_gradient:
-            if ctx.quantization_config.quant_input_grad:
-                x_int8, x_scale = quantize_tensorwise(grad_output * quant_scale)
-                input_grad = (
-                    paddle.matmul(x_int8, quant_weight).astype(ctx.dtype)
-                    @ infohub.hadamard[quant_weight.shape[-1]][0].T
-                    * x_scale
-                )
-            else:
-                qdq_weight = dequantize_channelwise(
-                    quant_weight, quant_scale, apply_hadamard=ctx.quantization_config.apply_hadamard
-                )
-                input_grad = paddle.matmul(grad_output, qdq_weight.T)
+        if ctx.quantization_config.weight_quantize_algo in ["fp8linear"]:
+            input_grad, weight_grad = fp8_backward(ctx, x, grad_output, quant_weight, quant_scale, quant_x, x_scale)
         else:
-            input_grad = None
+            input_grad, weight_grad = a8w8_backward(ctx, x, grad_output, quant_weight, quant_scale, quant_x, x_scale)
 
-        if not quant_weight.stop_gradient:
-            if len(x.shape) == 2:
-                weight_grad = paddle.matmul(x.transpose([1, 0]), grad_output)
-            else:
-                weight_grad = paddle.matmul(
-                    x.reshape([-1, x.shape[-1]]).transpose([1, 0]), grad_output.reshape([-1, grad_output.shape[-1]])
-                )
-        else:
-            weight_grad = None
-
-        if bias is not None and not bias.stop_gradient:
+        if not ctx.b_stop_gradient:
             bias_grad = grad_output.sum(axis=[0, 1])
         else:
             bias_grad = None

paddlenlp/quantization/quantization_config.py

@@ -17,6 +17,11 @@ import json
 from dataclasses import dataclass
 
 quant_inference_mapping = {"avg": "abs_max", "abs_max_channel_wise": "abs_max_channel_wise", "abs_max": "abs_max"}
+fp8_format_mapping = {
+    "hybrid": {"weight": "float8_e4m3fn", "activation": "float8_e4m3fn", "grad_output": "float8_e5m2"},
+    "e4m3": {"weight": "float8_e4m3fn", "activation": "float8_e4m3fn", "grad_output": "float8_e4m3fn"},
+    "e5m2": {"weight": "float8_e5m2", "activation": "float8_e5m2", "grad_output": "float8_e5m2"},
+}
 
 
 @dataclass
@@ -61,7 +66,11 @@ class QuantizationConfig:
         group_size=-1,
         apply_hadamard=True,
         quant_input_grad=False,
-        skip_first_act_scale_step=-1,
+        quant_weight_grad=False,
+        skip_first_act_scale_step=20,
+        moving_rate=0.01,
+        epsilon=1e-8,
+        fp8_format_type="hybrid",
         **kwargs,
     ):
         if weight_quantize_algo is not None:
@@ -77,6 +86,7 @@ class QuantizationConfig:
                         "fp4",
                         "a8w8linear",
                         "a8w4linear",
+                        "fp8linear",
                     ]
                     for algo in weight_quantize_algo
                 ):
@@ -92,6 +102,7 @@ class QuantizationConfig:
                 "fp4",
                 "a8w8linear",
                 "a8w4linear",
+                "fp8linear",
             ]:
                 raise ValueError(
                     f"weight_quantize_algo:{weight_quantize_algo} not in supported list ['weight_only_int8', 'weight_only_int4', 'llm.int8', 'a8w8', 'nf4', 'fp4']"
@@ -129,7 +140,11 @@ class QuantizationConfig:
         self.group_size = group_size
         self.apply_hadamard = apply_hadamard
         self.quant_input_grad = quant_input_grad
+        self.quant_weight_grad = quant_weight_grad
         self.skip_first_act_scale_step = skip_first_act_scale_step
+        self.moving_rate = moving_rate
+        self.epsilon = epsilon
+        self.fp8_format = fp8_format_mapping[fp8_format_type]
 
     def is_weight_quantize(self):
         if isinstance(self.weight_quantize_algo, dict):
@@ -143,6 +158,7 @@ class QuantizationConfig:
             "a8w8",
             "a8w8linear",
             "a8w4linear",
+            "fp8linear",
         ]:
             return True
         else:

paddlenlp/quantization/quantization_linear.py

@@ -40,6 +40,7 @@ QuantMapping = {
     "nf4": ("nf4", 4),
     "a8w8linear": ("int8", 8),
     "a8w4linear": ("int8", 8),
+    "fp8linear": ("fp8", 8),
 }
 
 
@@ -206,7 +207,7 @@ def quant_weight_linear(
     bias=None,
     act_state=None,
 ):
-    if weight_quantize_algo in ["a8w8linear", "a8w4linear"]:
+    if weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]:
 
         state, training, act_scale = act_state
 
@@ -258,6 +259,7 @@ class QuantizationLinear(nn.Layer):
             "llm.int8",
             "a8w8linear",
             "a8w4linear",
+            "fp8linear",
         ]:
             self.quant_weight = self.create_parameter(
                 shape=[out_features // 2, in_features] if self.quant_weight_bit == 4 else [out_features, in_features],
@@ -266,7 +268,7 @@ class QuantizationLinear(nn.Layer):
             )
             if self.quantization_config.group_size == -1:
                 self.quant_scale = self.create_parameter(
-                    shape=[out_features],
+                    shape=[out_features] if self.weight_quantize_algo not in ["fp8linear"] else [1],
                     dtype=self._dtype,
                     is_bias=False,
                 )
@@ -274,7 +276,7 @@ class QuantizationLinear(nn.Layer):
             else:
                 # TODO(lugimzzz): support groupwise in next PR
                 raise NotImplementedError("Not yet support grouwise weightonly quantization.")
-            if self.weight_quantize_algo in ["a8w8linear", "a8w4linear"]:
+            if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]:
                 self.act_scale = self.create_parameter(
                     shape=[], dtype=self._dtype, is_bias=False, default_initializer=nn.initializer.Constant(value=0.0)
                 )
@@ -353,7 +355,7 @@ class QuantizationLinear(nn.Layer):
             else None,
             bias=self.bias,
             act_state=(self.state, self.training, self.act_scale)
-            if self.weight_quantize_algo in ["a8w8linear", "a8w4linear"]
+            if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]
             else None,
         )
         if self.training:
@@ -412,6 +414,7 @@ class ColumnParallelQuantizationLinear(nn.Layer):
             "llm.int8",
             "a8w8linear",
             "a8w4linear",
+            "fp8linear",
         ]:
             self.quant_weight = self.create_parameter(
                 shape=[self.output_size_per_partition // 2, in_features]
@@ -426,7 +429,7 @@ class ColumnParallelQuantizationLinear(nn.Layer):
 
             if self.quantization_config.group_size == -1:
                 self.quant_scale = self.create_parameter(
-                    shape=[self.output_size_per_partition],
+                    shape=[self.output_size_per_partition] if self.weight_quantize_algo not in ["fp8linear"] else [1],
                     dtype=self._dtype,
                     is_bias=False,
                 )
@@ -437,7 +440,7 @@ class ColumnParallelQuantizationLinear(nn.Layer):
             else:
                 # TODO(lugimzzz): support groupwise in next PR
                 raise NotImplementedError("Not yet support grouwise weightonly quantization.")
-            if self.weight_quantize_algo in ["a8w8linear", "a8w4linear"]:
+            if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]:
                 self.act_scale = self.create_parameter(
                     shape=[], dtype=self._dtype, is_bias=False, default_initializer=nn.initializer.Constant(value=0.0)
                 )
@@ -484,7 +487,7 @@ class ColumnParallelQuantizationLinear(nn.Layer):
             else None,
             bias=self.bias,
             act_state=(self.state, self.training, self.act_scale)
-            if self.weight_quantize_algo in ["a8w8linear", "a8w4linear"]
+            if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]
             else None,
         )
         if self.training:
@@ -550,6 +553,7 @@ class RowParallelQuantizationLinear(nn.Layer):
             "llm.int8",
             "a8w8linear",
             "a8w4linear",
+            "fp8linear",
         ]:
             self.quant_weight = self.create_parameter(
                 shape=[out_features // 2, self.input_size_per_partition]
@@ -564,7 +568,7 @@ class RowParallelQuantizationLinear(nn.Layer):
 
             if self.quantization_config.group_size == -1:
                 self.quant_scale = self.create_parameter(
-                    shape=[out_features],
+                    shape=[out_features] if self.weight_quantize_algo not in ["fp8linear"] else [1],
                     dtype=self._dtype,
                     is_bias=False,
                 )
@@ -575,7 +579,7 @@ class RowParallelQuantizationLinear(nn.Layer):
             else:
                 # TODO(lugimzzz): support groupwise in next PR
                 raise NotImplementedError("Not yet support grouwise weightonly quantization.")
-            if self.weight_quantize_algo in ["a8w8linear", "a8w4linear"]:
+            if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]:
                 self.act_scale = self.create_parameter(
                     shape=[1], dtype=self._dtype, is_bias=False, default_initializer=nn.initializer.Constant(value=0.0)
                 )
@@ -616,7 +620,7 @@ class RowParallelQuantizationLinear(nn.Layer):
                 else None,
                 bias=None,
                 act_state=(self.state, self.training, self.act_scale)
-                if self.weight_quantize_algo in ["a8w8linear", "a8w4linear"]
+                if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]
                 else None,
             )
             if self.sequence_parallel:
@@ -644,7 +648,7 @@ class RowParallelQuantizationLinear(nn.Layer):
                 else None,
                 bias=self.bias,
                 act_state=(self.state, self.training, self.act_scale)
-                if self.weight_quantize_algo in ["a8w8linear", "a8w4linear"]
+                if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]
                 else None,
             )
         if self.training:

paddlenlp/quantization/quantization_utils.py

@@ -34,7 +34,7 @@ except:
     qlora_weight_quantize = None
 
 from ..utils.log import logger
-from .qat_utils import quantize_channelwise
+from .qat_utils import fp8_quantize_tensorwise, quantize_channelwise
 from .quantization_linear import (
     ColumnParallelQuantizationLinear,
     QuantizationLinear,
@@ -169,6 +169,13 @@ def convert_to_weight_quantize_state_dict(state_dict, name, quantization_config,
             act_scale = paddle.zeros([], dtype="bfloat16").cuda()
             act_scale.stop_gradient = True
             state_dict[act_scale_name] = act_scale
+        elif weight_quantize_algo in ["fp8linear"]:
+            quant_weight, quant_scale = fp8_quantize_tensorwise(
+                target_weight, tensor_type="weight", quantization_config=quantization_config
+            )
+            act_scale = paddle.zeros([], dtype="bfloat16").cuda()
+            act_scale.stop_gradient = True
+            state_dict[act_scale_name] = act_scale
         else:
             quant_weight, quant_scale = weight_quantize(
                 x=target_weight,
@@ -227,7 +234,14 @@ def convert_to_quantize_state_dict(state_dict, quantization_linear_list, quantiz
         if weight_quantize_algo is None:
             continue
         # Convert state dict
-        if weight_quantize_algo in ["weight_only_int8", "weight_only_int4", "llm.int8", "a8w8linear", "a8w4linear"]:
+        if weight_quantize_algo in [
+            "weight_only_int8",
+            "weight_only_int4",
+            "llm.int8",
+            "a8w8linear",
+            "a8w4linear",
+            "fp8linear",
+        ]:
             convert_to_weight_quantize_state_dict(state_dict, name, quantization_config, dtype, weight_quantize_algo)
         elif weight_quantize_algo in ["fp4", "nf4"]:
             convert_to_qlora_state_dict(state_dict, name, quantization_config, dtype, weight_quantize_algo)
@@ -260,7 +274,7 @@ def update_loaded_state_dict_keys(state_dict, quantization_linear_list, quantiza
             else:
                 state_dict.append(quant_scale_name)
                 weight_quantize_algo = parse_weight_quantize_algo(quantization_config, name)
-                if weight_quantize_algo in ["a8w8linear", "a8w4linear"]:
+                if weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]:
                     state_dict.append(act_scale_name)
 
         else:

paddlenlp/utils/optimizer.py

@@ -26,7 +26,12 @@ try:
 except:
     adamw_16bit_moment = None
 
-from ..quantization.qat_utils import dequantize_channelwise, quantize_channelwise
+from ..quantization.qat_utils import (
+    dequantize_channelwise,
+    fp8_dequantize_tensorwise,
+    fp8_quantize_tensorwise,
+    quantize_channelwise,
+)
 
 
 class AdamWMini(AdamW):
@@ -428,9 +433,14 @@ class AdamWQweight(AdamW):
             var_name = self._gen_master_weight_var_name(param)
             if param.name in self.quant_scale_mapping:
                 quant_scale = self.quant_scale_mapping[param.name]
-                var = dequantize_channelwise(
-                    param, quant_scale, apply_hadamard=self.quantization_config.apply_hadamard
-                ).astype("float32")
+                if self.quantization_config.weight_quantize_algo in ["fp8linear"]:
+                    var = fp8_dequantize_tensorwise(
+                        param, quant_scale, tensor_type="weight", quantization_config=self.quantization_config
+                    ).astype("float32")
+                else:
+                    var = dequantize_channelwise(
+                        param, quant_scale, apply_hadamard=self.quantization_config.apply_hadamard
+                    ).astype("float32")
             else:
                 var = paddle.cast(param, "float32")
             var.name = var_name
@@ -443,7 +453,7 @@ class AdamWQweight(AdamW):
         :param dtype: instance of core.VarDesc.VarType
         :return: True if dtype is one of fp16 or bf16, False otherwise
         """
-        if dtype == paddle.int8:
+        if dtype == paddle.int8 or dtype == paddle.float8_e4m3fn:
             return True
         assert isinstance(
             dtype, (core.VarDesc.VarType, core.DataType)
@@ -552,13 +562,21 @@ class AdamWQweight(AdamW):
             if quant_scale is None:
                 param[:] = p.astype(param.dtype)
             else:
-                if p.shape[1] == param.shape[0]:
-                    bit_length = 8
-                elif p.shape[1] / 2 == param.shape[0]:
-                    bit_length = 4
-                param[:], quant_scale[:] = quantize_channelwise(
-                    p.astype("bfloat16"), apply_hadamard=self.quantization_config.apply_hadamard, bit_length=bit_length
-                )
+                if self.quantization_config.weight_quantize_algo in ["fp8linear"]:
+                    param[:], new_quant_scale = fp8_quantize_tensorwise(
+                        p.astype("bfloat16"), tensor_type="weight", quantization_config=self.quantization_config
+                    )
+                    quant_scale.set_value(new_quant_scale)
+                else:
+                    if p.shape[1] == param.shape[0]:
+                        bit_length = 8
+                    elif p.shape[1] / 2 == param.shape[0]:
+                        bit_length = 4
+                    param[:], quant_scale[:] = quantize_channelwise(
+                        p.astype("bfloat16"),
+                        apply_hadamard=self.quantization_config.apply_hadamard,
+                        bit_length=bit_length,
+                    )
         else:
             param[:] = p
         moment1[:] = mom1