add distributed

Author: lugimzzz

paddlenlp/quantization/qat_utils.py

@@ -67,8 +67,8 @@ def quantize(
         if act_scale is not None:
             if training:
                 scale = paddle.max(paddle.abs(target_x)) / qmax
-                if paddle.distributed.is_initialized():
-                    paddle.distributed.all_reduce(scale, op=paddle.distributed.ReduceOp.MAX)
+                if group is not None:
+                    paddle.distributed.all_reduce(scale, op=paddle.distributed.ReduceOp.MAX, group=group, sync_op=True)
                 if state < quantization_config.apply_online_actscale_step:
                     act_scale.set_value((state * act_scale + scale) / (state + 1))
                 else:
@@ -97,7 +97,8 @@ def quantize(
             scale = scale.squeeze(0) / hadamard_scale
         elif weight_quantize_algo in ["fp8linear"]:
             scale = paddle.max(paddle.abs(target_x)) / qmax
-            paddle.distributed.all_reduce(scale, op=paddle.distributed.ReduceOp.MAX)
+            if group is not None:
+                paddle.distributed.all_reduce(scale, op=paddle.distributed.ReduceOp.MAX, group=group, sync_op=True)
             quant_x = (target_x / scale).astype(quantization_config.fp8_format[tensor_type]).view("int8").T
             scale = scale / hadamard_scale
         else:
@@ -143,6 +144,7 @@ def int8_forward(
     state=0,
     training=False,
     act_scale=None,
+    group=None,
 ):
     quant_x, scale_x = quantize(
         x=x,
@@ -154,6 +156,7 @@ def int8_forward(
         act_scale=act_scale,
         state=state,
         training=training,
+        group=group,
     )
 
     out = paddle.matmul(quant_x, quant_w.T).astype(scale_w.dtype) * (scale_x * scale_w)
@@ -201,6 +204,7 @@ def fp8_forward(
     state=0,
     training=False,
     act_scale=None,
+    group=None,
 ):
     x_fp8, x_scale = quantize(
         x,
@@ -212,6 +216,7 @@ def fp8_forward(
         act_scale=act_scale,
         state=state,
         training=training,
+        group=group,
     )
     x_fp8 = x_fp8.view(quantization_config.fp8_format["activation"])
     w_fp8 = w_fp8.view(quantization_config.fp8_format["weight"])
@@ -368,6 +373,7 @@ def forward(
         training,
         act_scale,
         weight_quantize_algo,
+        group,
     ):
         quant_x, x_scale = None, None
         if weight_quantize_algo in ["fp8linear"]:
@@ -382,6 +388,7 @@ def forward(
                 state=state,
                 training=training,
                 act_scale=act_scale,
+                group=group,
             )
         else:
             output, quant_x, x_scale = int8_forward(
@@ -394,6 +401,7 @@ def forward(
                 state=state,
                 training=training,
                 act_scale=act_scale,
+                group=group,
             )
         ctx.quantization_config = quantization_config
         ctx.weight_quantize_algo = weight_quantize_algo

paddlenlp/quantization/quantization_linear.py

@@ -209,7 +209,7 @@ def quant_weight_linear(
 ):
     if weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]:
 
-        state, training, act_scale = act_state
+        state, training, act_scale, group = act_state
 
         return QATFunc.apply(
             x,
@@ -237,6 +237,14 @@ def quant_weight_linear(
         )
 
 
+def get_act_scale_group(is_row=False):
+    if paddle.distributed.is_initialized():
+        group = None
+    else:
+        group = None
+    return group
+
+
 class QuantizationLinear(nn.Layer):
     """Quantization Linear layer."""
 
@@ -290,6 +298,7 @@ def __init__(
                     shape=[], 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:
@@ -349,6 +358,7 @@ def __init__(
             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(
@@ -363,7 +373,7 @@ def forward(self, x):
             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)
+            act_state=(self.state, self.training, self.act_scale, self.group)
             if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]
             else None,
         )
@@ -455,6 +465,7 @@ def __init__(
                 )
                 self.act_scale.is_distributed = True if self.is_mp else 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:
@@ -469,6 +480,7 @@ def __init__(
             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:
@@ -495,7 +507,7 @@ def forward(self, x):
             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)
+            act_state=(self.state, self.training, self.act_scale, self.group)
             if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]
             else None,
         )
@@ -594,6 +606,7 @@ def __init__(
                 )
                 self.act_scale.is_distributed = True if self.is_mp else 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}")
 
@@ -607,6 +620,8 @@ def __init__(
                 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
@@ -628,7 +643,7 @@ def forward(self, x):
                 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)
+                act_state=(self.state, self.training, self.act_scale, self.group)
                 if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]
                 else None,
             )
@@ -656,7 +671,7 @@ def forward(self, x):
                 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)
+                act_state=(self.state, self.training, self.act_scale, self.group)
                 if self.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]
                 else None,
             )

paddlenlp/utils/optimizer.py

@@ -324,15 +324,21 @@ def _append_optimize_op(self, block, param_and_grad):
                 skip_update_param,
             )
             if skip_update_param:
-                if self.quantization_config.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]:
+                if param.weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]:
+                    if "parallel_quantization_linear" not in param.name:
+                        group = None
+                    elif param.weight_quantize_algo in ["a8w8linear", "a8w4linear"] and "row" in param.name:
+                        group = None
+                    else:
+                        group = self.mp_group
                     param[:], new_quant_scale = quantize(
                         x=master_weight.astype(quant_scale.dtype),
                         weight_quantize_algo=self.quantization_config.weight_quantize_algo,
                         tensor_type="weight",
                         quantization_config=self.quantization_config,
                         side="left",
                         apply_hadamard=self.quantization_config.apply_hadamard,
-                        group=None,
+                        group=group,
                     )
                     quant_scale.set_value(new_quant_scale)
                 else: