chore: import upstream snapshot with attribution

paddlenlp/quantization/quantization_linear.py

@@ -0,0 +1,699 @@
+# Copyright (c) 2025 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
+#
+#     http://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.
+
+import paddle
+import paddle.nn as nn
+from paddle.autograd import PyLayer
+from paddle.distributed import fleet
+from paddle.distributed.fleet.base import topology as tp
+from paddle.distributed.fleet.layers.mpu import mp_ops
+from paddle.distributed.fleet.utils.sequence_parallel_utils import (
+    AllGatherOp,
+    ReduceScatterOp,
+)
+from paddle.nn.quant import llm_int8_linear, weight_dequantize, weight_only_linear
+
+from paddlenlp.utils import infohub
+
+from .qat_utils import QATFunc
+
+try:
+    from .qlora import qlora_weight_dequantize, qlora_weight_linear
+except:
+    qlora_weight_linear = None
+    qlora_weight_dequantize = None
+
+QuantMapping = {
+    # (quant_dtype, quant_weight_bit)
+    "weight_only_int8": ("int8", 8),
+    "weight_only_int4": ("int4", 4),
+    "llm.int8": ("int8", 8),
+    "fp4": ("fp4", 4),
+    "nf4": ("nf4", 4),
+    "a8w8linear": ("int8", 8),
+    "a8w4linear": ("int8", 8),
+    "fp8linear": ("fp8", 8),
+}
+
+
+def quant_weight_forward(
+    x,
+    quant_weight,
+    bias,
+    quant_scale,
+    quant_state,
+    quant_dtype,
+    quantization_config,
+    weight_quantize_algo,
+    dtype,
+):
+    if weight_quantize_algo in ["weight_only_int8", "weight_only_int4"]:
+        output = weight_only_linear(
+            x=x,
+            weight=quant_weight,
+            bias=bias,
+            weight_scale=quant_scale,
+            weight_dtype=quant_dtype,
+            group_size=quantization_config.group_size,
+        )
+    elif weight_quantize_algo in ["llm.int8"]:
+        output = llm_int8_linear(x, quant_weight, bias, quant_scale, quantization_config.llm_int8_threshold)
+    elif weight_quantize_algo in ["fp4", "nf4"]:
+        output = qlora_weight_linear(
+            x=x,
+            quant_weight=quant_weight,
+            dtype=dtype,
+            state=quant_state if quantization_config.qlora_weight_double_quant else quant_scale,
+            quant_algo=weight_quantize_algo,
+            double_quant=quantization_config.qlora_weight_double_quant,
+            block_size=quantization_config.qlora_weight_blocksize,
+            double_quant_block_size=quantization_config.qlora_weight_double_quant_block_size,
+            bias=bias,
+        )
+
+    return output
+
+
+def dequant_weight(
+    quant_weight,
+    quantization_config,
+    weight_quantize_algo,
+    dtype,
+    quant_scale,
+    quant_state,
+    input_shape,
+):
+    if weight_quantize_algo in ["weight_only_int8", "weight_only_int4", "llm.int8"]:
+        quant_dequant_weight = weight_dequantize(
+            x=quant_weight,
+            scale=quant_scale,
+            algo=weight_quantize_algo,
+            out_dtype=dtype,
+            group_size=quantization_config.group_size,
+        )
+    elif weight_quantize_algo in ["fp4", "nf4"]:
+        quant_dequant_weight = (
+            qlora_weight_dequantize(
+                quant_weight=quant_weight,
+                quant_algo=weight_quantize_algo,
+                state=quant_state if quantization_config.qlora_weight_double_quant else quant_scale,
+                double_quant=quantization_config.qlora_weight_double_quant,
+                block_size=quantization_config.qlora_weight_blocksize,
+                double_quant_block_size=quantization_config.qlora_weight_double_quant_block_size,
+            )
+            .reshape([input_shape[-1], -1])
+            .cast(dtype)
+        )
+    return quant_dequant_weight
+
+
+class QuantizationLinearFunc(PyLayer):
+    @staticmethod
+    def forward(
+        ctx,
+        x,
+        quant_weight,
+        bias,
+        quant_scale,
+        quant_state,
+        quant_dtype,
+        quantization_config,
+        weight_quantize_algo,
+        dtype,
+    ):
+
+        output = quant_weight_forward(
+            x=x,
+            quant_weight=quant_weight,
+            bias=bias,
+            quant_scale=quant_scale,
+            quant_state=quant_state,
+            quant_dtype=quant_dtype,
+            quantization_config=quantization_config,
+            weight_quantize_algo=weight_quantize_algo,
+            dtype=dtype,
+        )
+        ctx.quant_dtype = quant_dtype
+        ctx.quantization_config = quantization_config
+        ctx.weight_quantize_algo = weight_quantize_algo
+        ctx.dtype = dtype
+        if ctx.weight_quantize_algo in ["fp4", "nf4"] and ctx.quantization_config.qlora_weight_double_quant:
+            qquant_scale, double_quant_scale, quant_scale_offset = quant_state
+            ctx.save_for_backward(x, quant_weight, bias, qquant_scale, double_quant_scale, quant_scale_offset)
+        else:
+            ctx.save_for_backward(x, quant_weight, bias, quant_scale)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        if ctx.weight_quantize_algo in ["fp4", "nf4"] and ctx.quantization_config.qlora_weight_double_quant:
+            x, quant_weight, bias, qquant_scale, double_quant_scale, quant_scale_offset = ctx.saved_tensor()
+            quant_state = (qquant_scale, double_quant_scale, quant_scale_offset)
+            quant_scale = None
+        else:
+            x, quant_weight, bias, quant_scale = ctx.saved_tensor()
+            quant_state = None
+
+        qdq_weight = dequant_weight(
+            quant_weight=quant_weight,
+            quantization_config=ctx.quantization_config,
+            weight_quantize_algo=ctx.weight_quantize_algo,
+            dtype=ctx.dtype,
+            quant_scale=quant_scale,
+            quant_state=quant_state,
+            input_shape=x.shape,
+        )
+
+        if not x.stop_gradient:
+            input_grad = paddle.matmul(grad_output, qdq_weight.T)
+        else:
+            input_grad = None
+
+        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:
+            bias_grad = grad_output.sum(axis=[0, 1])
+        else:
+            bias_grad = None
+
+        return input_grad, weight_grad, bias_grad
+
+
+def quant_weight_linear(
+    x,
+    quant_weight,
+    quant_dtype,
+    quantization_config,
+    weight_quantize_algo,
+    dtype,
+    quant_scale=None,
+    quant_state=None,
+    bias=None,
+    act_state=None,
+):
+    if weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]:
+
+        state, training, act_scale, group = act_state
+
+        return QATFunc.apply(
+            x,
+            quant_weight,
+            bias,
+            quant_scale,
+            quantization_config,
+            dtype,
+            state,
+            training,
+            act_scale,
+            weight_quantize_algo,
+            group,
+        )
+    else:
+        return QuantizationLinearFunc.apply(
+            x,
+            quant_weight,
+            bias,
+            quant_scale,
+            quant_state,
+            quant_dtype,
+            quantization_config,
+            weight_quantize_algo,
+            dtype,
+        )
+
+
+def get_act_scale_group(is_row=False):
+    if paddle.distributed.is_initialized():
+        if getattr(infohub, "scale_group") is None:
+            hcg = fleet.get_hybrid_communicate_group()
+            rank = hcg._dp_degree * hcg._sharding_degree
+            group_no_row = hcg.create_fuse_group(["data", "sharding"])[1] if rank > 1 else None
+            rank *= hcg._mp_degree
+            group_row = hcg.create_fuse_group(["data", "sharding", "model"])[1] if rank > 1 else None
+
+            setattr(infohub, "scale_group", [group_no_row, group_row])
+        group = infohub.scale_group[1] if is_row else infohub.scale_group[0]
+    else:
+        group = None
+    return group
+
+
+class QuantizationLinear(nn.Layer):
+    """Quantization Linear layer."""
+
+    def __init__(
+        self,
+        in_features,
+        out_features,
+        quantization_config,
+        weight_quantize_algo,
+        dtype,
+        bias_attr=None,
+        mp_moe=False,
+        is_distributed=False,
+    ):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.quantization_config = quantization_config
+        self.weight_quantize_algo = weight_quantize_algo
+        self._dtype = dtype
+        self.quant_dtype, self.quant_weight_bit = QuantMapping[self.weight_quantize_algo]
+        self.state = 0
+
+        # PaddlePaddle doesn't support 4bit data type, one 8bit data represents two 4bit data.
+        # paddle.nn.quant.weight_quantize will transpose in_features and out_features.
+        if self.weight_quantize_algo in [
+            "weight_only_int8",
+            "weight_only_int4",
+            "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],
+                dtype="int8",
+                is_bias=False,
+            )
+            if self.quantization_config.group_size == -1:
+                self.quant_scale = self.create_parameter(
+                    shape=[out_features] if self.weight_quantize_algo not in ["fp8linear"] else [1],
+                    dtype=self._dtype,
+                    is_bias=False,
+                )
+                self.quant_scale.stop_gradient = True
+            else:
+                # TODO(lugimzzz): support groupwise in next PR
+                raise NotImplementedError("Not yet support grouwise weightonly quantization.")
+            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)
+                )
+                self.act_scale.stop_gradient = True
+                self.group = get_act_scale_group()
+
+        elif self.weight_quantize_algo in ["fp4", "nf4"]:
+            if qlora_weight_linear is None:
+                raise ImportError(
+                    "Please run the following commands to install: qlora related package first\n"
+                    "1) git clone https://github.com/PaddlePaddle/PaddleSlim \n"
+                    "2) cd PaddleSlim && pip install -e .\n"
+                    "3) cd csrc &&  python ./setup_cuda.py install"
+                )
+            self.quant_weight = self.create_parameter(
+                shape=[out_features * in_features // 2, 1],
+                attr=paddle.nn.initializer.Constant(value=0),
+                dtype="uint8",
+                is_bias=False,
+            )
+            if self.quantization_config.qlora_weight_double_quant:
+                # quantized quant_scale
+                self.qquant_scale = self.create_parameter(
+                    shape=[in_features * out_features // self.quantization_config.qlora_weight_blocksize],
+                    dtype="uint8",
+                    is_bias=False,
+                )
+                # double quant_scale: quant_scale of quantized quant_scale
+                self.double_quant_scale = self.create_parameter(
+                    shape=[
+                        in_features
+                        * out_features
+                        // self.quantization_config.qlora_weight_blocksize
+                        // self.quantization_config.qlora_weight_double_quant_block_size
+                    ],
+                    dtype="float32",
+                    is_bias=False,
+                )
+                self.quant_scale_offset = self.create_parameter(
+                    shape=[],
+                    dtype="float32",
+                    is_bias=False,
+                )
+            else:
+                self.quant_scale = self.create_parameter(
+                    shape=[in_features * out_features // self.quantization_config.qlora_weight_blocksize],
+                    dtype="float32",
+                    is_bias=False,
+                )
+        else:
+            raise NotImplementedError(f"Not yet support weight_quantize_algo: {self.weight_quantize_algo}")
+        if bias_attr is False:
+            self.bias = None
+        else:
+            self.bias = self.create_parameter(
+                shape=[out_features],
+                attr=bias_attr,
+                dtype=self._dtype,
+                is_bias=True,
+            )
+        if mp_moe or is_distributed:
+            for p in self.parameters():
+                p.is_distributed = is_distributed
+                p.mp_moe = mp_moe
+        self.quant_weight.weight_quantize_algo = self.weight_quantize_algo
+
+    def forward(self, x):
+        output = quant_weight_linear(
+            x=x,
+            quant_weight=self.quant_weight,
+            quant_dtype=self.quant_dtype,
+            quantization_config=self.quantization_config,
+            weight_quantize_algo=self.weight_quantize_algo,
+            dtype=self._dtype,
+            quant_scale=self.quant_scale,
+            quant_state=(self.qquant_scale, self.double_quant_scale, self.quant_scale_offset)
+            if (self.weight_quantize_algo in ["fp4", "nf4"] and self.quantization_config.qlora_weight_double_quant)
+            else None,
+            bias=self.bias,
+            act_state=(self.state, self.training, self.act_scale, self.group)
+            if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]
+            else None,
+        )
+        if self.training:
+            self.state += 1
+        return output
+
+
+class ColumnParallelQuantizationLinear(nn.Layer):
+    """Quantization Linear layer with mp parallelized(column).
+    The code implementation refers to paddle.distributed.fleet.meta_parallel.ColumnParallelLinear.
+    https://github.com/PaddlePaddle/Paddle/blob/develop/python/paddle/distributed/fleet/layers/mpu/mp_layers.py#L310
+    Different from ColumnParallelLinear, this class keeps weight in INT8/INT4 with quant scale, and supports matrix
+    multiplication(weight_only_linear/llm_int8_linear) for input tensor(fp16/bf16) and quantized weight(INT8/INT4)
+    and bias addition if provided.
+    Notice: quantized weight shape is transposed of weight shape in ColumnParallelLinear.
+    """
+
+    def __init__(
+        self,
+        in_features,
+        output_size_per_partition,
+        quantization_config,
+        weight_quantize_algo,
+        dtype,
+        bias_attr=None,
+        gather_output=True,
+        mp_skip_c_identity=False,
+        mp_group=None,
+        sequence_parallel=False,
+    ):
+        super().__init__()
+        self.in_features = in_features
+        self.output_size_per_partition = output_size_per_partition
+        self.weight_quantize_algo = weight_quantize_algo
+        self.quantization_config = quantization_config
+        self._dtype = dtype
+        self.mp_skip_c_identity = mp_skip_c_identity
+        self.quant_dtype, self.quant_weight_bit = QuantMapping[self.weight_quantize_algo]
+        self.state = 0
+        self.model_parallel_group = (
+            tp._HYBRID_PARALLEL_GROUP.get_model_parallel_group() if mp_group is None else mp_group
+        )
+        self.world_size = (
+            tp._HYBRID_PARALLEL_GROUP.get_model_parallel_world_size() if mp_group is None else mp_group.nranks
+        )
+        self.is_mp = self.world_size > 1
+        self.gather_output = gather_output
+        self.sequence_parallel = sequence_parallel
+        if self.sequence_parallel and self.gather_output:
+            raise ValueError("Sequence parallel does not support gather_output")
+
+        # PaddlePaddle doesn't support Int4 data type, one Int8 data represents two Int4 data.
+        if self.weight_quantize_algo in [
+            "weight_only_int8",
+            "weight_only_int4",
+            "llm.int8",
+            "a8w8linear",
+            "a8w4linear",
+            "fp8linear",
+        ]:
+            self.quant_weight = self.create_parameter(
+                shape=[self.output_size_per_partition // 2, in_features]
+                if self.quant_dtype == "int4"
+                else [self.output_size_per_partition, in_features],
+                dtype="int8",
+                is_bias=False,
+            )
+            self.quant_weight.is_distributed = True if self.is_mp else False
+            if self.quant_weight.is_distributed:
+                self.quant_weight.split_axis = 0
+
+            if self.quantization_config.group_size == -1:
+                self.quant_scale = self.create_parameter(
+                    shape=[self.output_size_per_partition] if self.weight_quantize_algo not in ["fp8linear"] else [1],
+                    dtype=self._dtype,
+                    is_bias=False,
+                )
+                self.quant_scale.stop_gradient = True
+                if self.weight_quantize_algo in ["fp8linear", "a8w4linear", "a8w8linear"]:
+                    self.quant_scale.is_distributed = False
+                else:
+                    self.quant_scale.is_distributed = True if self.is_mp else False
+                if self.quant_scale.is_distributed:
+                    self.quant_scale.split_axis = 0
+            else:
+                # TODO(lugimzzz): support groupwise in next PR
+                raise NotImplementedError("Not yet support grouwise weightonly quantization.")
+            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)
+                )
+                self.act_scale.is_distributed = False
+                self.act_scale.stop_gradient = True
+                self.group = get_act_scale_group()
+        else:
+            raise NotImplementedError(f"Not yet support weight_quantize_algo: {self.weight_quantize_algo}")
+        if bias_attr is False:
+            self.bias = None
+        else:
+            self.bias = self.create_parameter(
+                shape=[self.output_size_per_partition],
+                attr=bias_attr,
+                dtype=self._dtype,
+                is_bias=True,
+            )
+            self.bias.is_distributed = True if self.is_mp else False
+            if self.bias.is_distributed:
+                self.bias.split_axis = 0
+        self.quant_weight.weight_quantize_algo = self.weight_quantize_algo
+
+    def forward(self, x):
+        if self.is_mp:
+            if self.sequence_parallel:
+                input_parallel = AllGatherOp.apply(x)
+            else:
+                input_parallel = mp_ops._c_identity(
+                    x,
+                    group=self.model_parallel_group,
+                    skip_c_identity_dynamic=self.mp_skip_c_identity,
+                )
+        else:
+            input_parallel = x
+
+        output_parallel = quant_weight_linear(
+            x=input_parallel,
+            quant_weight=self.quant_weight,
+            quant_dtype=self.quant_dtype,
+            quantization_config=self.quantization_config,
+            weight_quantize_algo=self.weight_quantize_algo,
+            dtype=self._dtype,
+            quant_scale=self.quant_scale,
+            quant_state=(self.qquant_scale, self.double_quant_scale, self.quant_scale_offset)
+            if (self.weight_quantize_algo in ["fp4", "nf4"] and self.quantization_config.qlora_weight_double_quant)
+            else None,
+            bias=self.bias,
+            act_state=(self.state, self.training, self.act_scale, self.group)
+            if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]
+            else None,
+        )
+        if self.training:
+            self.state += 1
+
+        if self.gather_output and self.is_mp:
+            output = mp_ops._c_concat(output_parallel, group=self.model_parallel_group)
+        else:
+            output = output_parallel
+        return output
+
+
+class RowParallelQuantizationLinear(nn.Layer):
+    """Quantization Linear layer with mp parallelized(row).
+    The code implementation refers to paddle.distributed.fleet.meta_parallel.RowParallelLinear.
+    https://github.com/PaddlePaddle/Paddle/blob/develop/python/paddle/distributed/fleet/layers/mpu/mp_layers.py#L517
+    Different from RowParallelLinear, this class keeps weight in INT8/INT4 with quant scale, and supports matrix
+    multiplication(weight_only_linear/llm_int8_linear) for input tensor(fp16/bf16) and quantized weight(INT8/INT4)
+    and bias addition if provided.
+    Notice: quantized weight shape is transposed of weight shape in RowParallelLinear.
+    """
+
+    def __init__(
+        self,
+        input_size_per_partition,
+        out_features,
+        quantization_config,
+        weight_quantize_algo,
+        dtype,
+        bias_attr=None,
+        input_is_parallel=False,
+        mp_skip_c_identity=False,
+        mp_group=None,
+        sequence_parallel=False,
+    ):
+        super().__init__()
+        self.input_size_per_partition = input_size_per_partition
+        self.out_features = out_features
+        self.quantization_config = quantization_config
+        self.weight_quantize_algo = weight_quantize_algo
+        self._dtype = dtype
+        self.mp_skip_c_identity = mp_skip_c_identity
+        self.quant_dtype, self.quant_weight_bit = QuantMapping[self.weight_quantize_algo]
+        self.state = 0
+
+        self.model_parallel_group = (
+            tp._HYBRID_PARALLEL_GROUP.get_model_parallel_group() if mp_group is None else mp_group
+        )
+        self.world_size = (
+            tp._HYBRID_PARALLEL_GROUP.get_model_parallel_world_size() if mp_group is None else mp_group.nranks
+        )
+        self.is_mp = self.world_size > 1
+        self.input_is_parallel = input_is_parallel
+        self.sequence_parallel = sequence_parallel
+        if not self.input_is_parallel and self.sequence_parallel:
+            raise ValueError("Sequence parallel only support input_is_parallel.")
+
+        # PaddlePaddle doesn't support Int4 data type, one Int8 data represents two Int4 data.
+        # paddle.nn.quant.weight_quantize will transpose in_features and out_features.
+        if self.weight_quantize_algo in [
+            "weight_only_int8",
+            "weight_only_int4",
+            "llm.int8",
+            "a8w8linear",
+            "a8w4linear",
+            "fp8linear",
+        ]:
+            self.quant_weight = self.create_parameter(
+                shape=[out_features // 2, self.input_size_per_partition]
+                if self.quant_dtype == "int4"
+                else [out_features, self.input_size_per_partition],
+                dtype="int8",
+                is_bias=False,
+            )
+            self.quant_weight.is_distributed = True if self.is_mp else False
+            if self.quant_weight.is_distributed:
+                self.quant_weight.split_axis = 1
+
+            if self.quantization_config.group_size == -1:
+                self.quant_scale = self.create_parameter(
+                    shape=[out_features] if self.weight_quantize_algo not in ["fp8linear"] else [1],
+                    dtype=self._dtype,
+                    is_bias=False,
+                )
+                self.quant_scale.stop_gradient = True
+                if self.weight_quantize_algo in ["fp8linear", "a8w4linear", "a8w8linear"]:
+                    self.quant_scale.is_distributed = False
+                else:
+                    self.quant_scale.is_distributed = True if self.is_mp else False
+                if self.quant_scale.is_distributed:
+                    self.quant_scale.split_axis = 0
+            else:
+                # TODO(lugimzzz): support groupwise in next PR
+                raise NotImplementedError("Not yet support grouwise weightonly quantization.")
+            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)
+                )
+                self.act_scale.is_distributed = False
+                self.act_scale.stop_gradient = True
+                self.group = get_act_scale_group(is_row=True)
+        else:
+            raise NotImplementedError(f"Not yet support weight_quantize_algo: {self.weight_quantize_algo}")
+
+        if bias_attr is False:
+            self.bias = None
+        else:
+            self.bias = self.create_parameter(
+                shape=[out_features],
+                attr=bias_attr,
+                dtype=self._dtype,
+                is_bias=True,
+            )
+
+        self.quant_weight.weight_quantize_algo = self.weight_quantize_algo
+
+    def forward(self, x):
+        if self.input_is_parallel or (not self.is_mp):
+            input_parallel = x
+        else:
+            # split last dim
+            input_parallel = mp_ops._c_split(x, group=self.model_parallel_group)
+
+        # with paddle.amp.auto_cast(enable=False):
+        if self.is_mp:
+            output_parallel = quant_weight_linear(
+                x=input_parallel,
+                quant_weight=self.quant_weight,
+                quant_dtype=self.quant_dtype,
+                quantization_config=self.quantization_config,
+                weight_quantize_algo=self.weight_quantize_algo,
+                dtype=self._dtype,
+                quant_scale=self.quant_scale,
+                quant_state=(self.qquant_scale, self.double_quant_scale, self.quant_scale_offset)
+                if (self.weight_quantize_algo in ["fp4", "nf4"] and self.quantization_config.qlora_weight_double_quant)
+                else None,
+                bias=None,
+                act_state=(self.state, self.training, self.act_scale, self.group)
+                if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]
+                else None,
+            )
+            if self.sequence_parallel:
+                output_ = ReduceScatterOp.apply(output_parallel)
+            else:
+                output_ = mp_ops._mp_allreduce(
+                    output_parallel,
+                    group=self.model_parallel_group,
+                    use_calc_stream=True,
+                    use_model_parallel=True,
+                    skip_c_identity_dynamic=self.mp_skip_c_identity,
+                )
+            output = output_ + self.bias if self.bias is not None else output_
+        else:
+            output = quant_weight_linear(
+                x=input_parallel,
+                quant_weight=self.quant_weight,
+                quant_dtype=self.quant_dtype,
+                quantization_config=self.quantization_config,
+                weight_quantize_algo=self.weight_quantize_algo,
+                dtype=self._dtype,
+                quant_scale=self.quant_scale,
+                quant_state=(self.qquant_scale, self.double_quant_scale, self.quant_scale_offset)
+                if (self.weight_quantize_algo in ["fp4", "nf4"] and self.quantization_config.qlora_weight_double_quant)
+                else None,
+                bias=self.bias,
+                act_state=(self.state, self.training, self.act_scale, self.group)
+                if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]
+                else None,
+            )
+        if self.training:
+            self.state += 1
+
+        return output