[LoRA] add quick_lora

llm/argument.py

@@ -126,6 +126,12 @@ class ModelArgument:
     lora: bool = field(default=False, metadata={"help": "Whether to use LoRA technique"})
     lora_path: str = field(default=None, metadata={"help": "Initialize lora state dict."})
     lora_rank: int = field(default=8, metadata={"help": "Lora attention dimension"})
+    use_quick_lora: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to use quick lora, The use of Quick LoRa will only take effect when lora_dropout is set to 0."
+        },
+    )
     rslora: bool = field(default=False, metadata={"help": "Whether to use RsLoRA"})
     lora_plus_scale: float = field(default=1.0, metadata={"help": "Lora B scale in LoRA+ technique"})
 

llm/finetune_generation.py

@@ -112,6 +112,7 @@ def main():
         weight_double_quant=model_args.weight_double_quant,
         weight_double_quant_block_size=model_args.weight_double_quant_block_size,
     )
+
     if training_args.pipeline_parallel_degree > 1:
         if data_args.eval_with_do_generation and training_args.do_eval:
             raise ValueError("Plese set eval_with_do_generation to false in pipeline parallel mode.")
@@ -426,10 +427,12 @@ def neft_post_hook(module, input, output):
                 dtype=dtype,
                 do_qat=quant_args.do_qat,
                 base_model_name_or_path=model_args.model_name_or_path,
+                use_quick_lora=model_args.use_quick_lora,
             )
             model = LoRAModel(model, lora_config)
         else:
             model = LoRAModel.from_pretrained(model=model, lora_path=model_args.lora_path)
+
         model.print_trainable_parameters()
 
     def compute_metrics_do_generation(eval_preds):

paddlenlp/peft/lora/lora_config.py

@@ -18,6 +18,7 @@
 from typing import List, Optional, Union
 
 from ...utils.env import LORA_CONFIG_NAME
+from ...utils.log import logger
 
 
 @dataclass
@@ -77,6 +78,20 @@
     base_model_name_or_path: Optional[str] = field(
         default=None, metadata={"help": "The name of the base model to use."}
     )
+    use_quick_lora: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to use quick lora, The use of Quick LoRa will only take effect when lora_dropout is set to 0."
+        },
+    )
+
+    def __post_init__(self):
+        if self.use_quick_lora and self.lora_dropout > 0:
+            logger.warning(
+                "Quick LoRa is enabled, but lora_dropout is set to a non-zero value. "
+                "We will automatically set `use_quick_lora` to `False` to avoid potential inconsistencies."
+            )
+            self.use_quick_lora = False
 
     @property
     def __dict__(self):

paddlenlp/peft/lora/lora_layers.py

@@ -25,6 +25,8 @@
     RowParallelLinear,
 )
 
+from .lora_quick_layers import quick_lora
+
 if "npu" in paddle.device.get_all_custom_device_type():
     from .mc2_lora_npu import MC2LoRaColumnParallelLinear, MC2LoRaRowParallelLinear
 else:
@@ -42,6 +44,7 @@
         lora_alpha: int = 1,
         lora_dropout: float = 0.0,
         merge_weights: bool = True,
+        use_quick_lora: bool = False,
         rslora: bool = False,
         lora_plus_scale: float = 1.0,
         **kwargs
@@ -84,6 +87,11 @@
 
         # Freezing the pre-trained weight matrix
         self.weight.stop_gradient = True
+        self._use_quick_lora = use_quick_lora and lora_dropout == 0.0
+
+    @property
+    def use_quick_lora(self):
+        return self._use_quick_lora and self.training and not self.merged
 
     def train(self):
         super().train()
@@ -102,9 +110,13 @@
             self.merged = True
 
     def forward(self, input: paddle.Tensor, *args, **kwargs):
-        result = F.linear(x=input, weight=self.weight, bias=self.bias, name=self.name)
-        if not self.merged:
-            result += (self.lora_dropout(input) @ self.lora_A @ self.lora_B) * self.scaling
+        if self.use_quick_lora:
+            # Use the quick lora implementation
+            result = quick_lora(input, self.lora_A, self.lora_B, self.weight, self.bias, self.scaling)
+        else:
+            result = F.linear(x=input, weight=self.weight, bias=self.bias, name=self.name)
+            if not self.merged:
+                result += (self.lora_dropout(input) @ self.lora_A @ self.lora_B) * self.scaling
         return result
 
     def extra_repr(self):
@@ -123,6 +135,7 @@
         rslora: bool = False,
         lora_plus_scale: float = 1.0,
         merge_weights: bool = True,
+        use_quick_lora: bool = False,
         **kwargs
     ):
         RowParallelLinear.__init__(self, in_features, out_features, **kwargs)
@@ -171,6 +184,11 @@
 
         # Freezing the pre-trained weight matrix
         self.weight.stop_gradient = True
+        self._use_quick_lora = use_quick_lora and lora_dropout == 0.0
+
+    @property
+    def use_quick_lora(self):
+        return self._use_quick_lora and self.training and not self.merged
 
     def train(self):
         super().train()
@@ -194,33 +212,52 @@
         else:
             input_mp = x
 
-        # x @ W : [bz, in_f / ws] ===> [bz, out_f]
-        if "npu" in paddle.device.get_all_custom_device_type() and int(os.getenv("MC2", "0")):
-            output = MC2LoRaRowParallelLinear.apply(input_mp, self.weight, self.model_parallel_group)
-        else:
-            result_mp = F.linear(x=input_mp, weight=self.weight, name=self.name)
-
+        if self.use_quick_lora:
+            # Use the quick lora implementation
+            result_mp = quick_lora(
+                input_mp,
+                self.lora_A,
+                self.lora_B,
+                self.weight,
+                self.bias,
+                self.scaling,
+                is_row=True,
+                group=self.model_parallel_group,
+                world_size=self.world_size,
+            )
             output = mp_ops._mp_allreduce(
                 result_mp,
                 group=self.model_parallel_group,
                 use_calc_stream=True,
                 use_model_parallel=True,
             )
+        else:
+            # x @ W : [bz, in_f / ws] ===> [bz, out_f]
+            if "npu" in paddle.device.get_all_custom_device_type() and int(os.getenv("MC2", "0")):
+                output = MC2LoRaRowParallelLinear.apply(input_mp, self.weight, self.model_parallel_group)
+            else:
+                result_mp = F.linear(x=input_mp, weight=self.weight, name=self.name)
+                output = mp_ops._mp_allreduce(
+                    result_mp,
+                    group=self.model_parallel_group,
+                    use_calc_stream=True,
+                    use_model_parallel=True,
+                )
 
-        if not self.merged:
-            # x @ A: [bz, in_f/ ws] ===> [bz, r]
-            input_mp = self.lora_dropout(input_mp) @ self.lora_A
-            # all reduce to keep Lora B's gradient on different gpu consistent
-            input_dup = mp_ops._mp_allreduce(
-                input_mp,
-                group=self.model_parallel_group,
-                use_calc_stream=True,
-                use_model_parallel=True,
-            )
-            #  @ B: [bz, r] ===> [bz, out_f]
-            delta_mp = (input_dup @ self.lora_B) * self.scaling
-            output += delta_mp
-        output = output + self.bias if self.bias is not None else output
+            if not self.merged:
+                # x @ A: [bz, in_f/ ws] ===> [bz, r]
+                input_mp = self.lora_dropout(input_mp) @ self.lora_A
+                # all reduce to keep Lora B's gradient on different gpu consistent
+                input_dup = mp_ops._mp_allreduce(
+                    input_mp,
+                    group=self.model_parallel_group,
+                    use_calc_stream=True,
+                    use_model_parallel=True,
+                )
+                #  @ B: [bz, r] ===> [bz, out_f]
+                delta_mp = (input_dup @ self.lora_B) * self.scaling
+                output += delta_mp
+            output = output + self.bias if self.bias is not None else output
         return output
 
     def extra_repr(self):
@@ -240,6 +277,7 @@
         lora_plus_scale: float = 1.0,
         merge_weights: bool = True,
         lora_A_weight_attr: Optional[paddle.ParamAttr] = None,
+        use_quick_lora: bool = False,
         **kwargs
     ):
         ColumnParallelLinear.__init__(self, in_features, out_features, **kwargs)
@@ -286,6 +324,11 @@
 
         # Freezing the pre-trained weight matrix
         self.weight.stop_gradient = True
+        self._use_quick_lora = use_quick_lora and lora_dropout == 0.0
+
+    @property
+    def use_quick_lora(self):
+        return self._use_quick_lora and self.training and not self.merged
 
     def train(self):
         super().train()
@@ -304,22 +347,37 @@
             self.merged = True
 
     def forward(self, input: paddle.Tensor):
-        if "npu" in paddle.device.get_all_custom_device_type() and int(os.getenv("MC2", "0")):
-            res_mp = MC2LoRaColumnParallelLinear.apply(input, self.weight, self.model_parallel_group)
-            result_mp = res_mp + self.bias
+        if self.use_quick_lora:
+            # Use the quick lora implementation
+            input_mp = mp_ops._c_identity(input, group=self.model_parallel_group) if self.is_mp else input
+            result_mp = quick_lora(
+                input_mp,
+                self.lora_A,
+                self.lora_B,
+                self.weight,
+                self.bias,
+                self.scaling,
+                is_column=True,
+                group=self.model_parallel_group,
+                world_size=self.world_size,
+            )
         else:
-            input_mp = mp_ops._c_identity(input, group=self.model_parallel_group)
-            result_mp = F.linear(x=input_mp, weight=self.weight, bias=self.bias, name=self.name)
-
-        if not self.merged:
-            input_a = self.lora_dropout(input) @ self.lora_A
             if "npu" in paddle.device.get_all_custom_device_type() and int(os.getenv("MC2", "0")):
-                tmp = MC2LoRaColumnParallelLinear.apply(input_a, self.lora_B, self.model_parallel_group)
-                delta_mp = tmp * self.scaling
+                res_mp = MC2LoRaColumnParallelLinear.apply(input, self.weight, self.model_parallel_group)
+                result_mp = res_mp + self.bias
             else:
-                input_a_mp = mp_ops._c_identity(input_a, group=self.model_parallel_group)
-                delta_mp = (input_a_mp @ self.lora_B) * self.scaling
-            result_mp += delta_mp
+                input_mp = mp_ops._c_identity(input, group=self.model_parallel_group)
+                result_mp = F.linear(x=input_mp, weight=self.weight, bias=self.bias, name=self.name)
+
+            if not self.merged:
+                input_a = self.lora_dropout(input) @ self.lora_A
+                if "npu" in paddle.device.get_all_custom_device_type() and int(os.getenv("MC2", "0")):
+                    tmp = MC2LoRaColumnParallelLinear.apply(input_a, self.lora_B, self.model_parallel_group)
+                    delta_mp = tmp * self.scaling
+                else:
+                    input_a_mp = mp_ops._c_identity(input_a, group=self.model_parallel_group)
+                    delta_mp = (input_a_mp @ self.lora_B) * self.scaling
+                result_mp += delta_mp
 
         if self.gather_output and self.is_mp:
             result = mp_ops._c_concat(result_mp, group=self.model_parallel_group)

paddlenlp/peft/lora/lora_model.py

@@ -385,6 +385,7 @@ def _find_and_replace_module(self, model, module_name, lora_config, enable_lora)
                     rslora=lora_config.rslora,
                     lora_plus_scale=lora_config.lora_plus_scale,
                     bias_attr=False if module.bias is None else None,
+                    use_quick_lora=lora_config.use_quick_lora,
                 )
             if isinstance(module, nn.Conv2D):
                 lora_module = LoRAConv2D(
@@ -422,6 +423,7 @@ def _find_and_replace_module(self, model, module_name, lora_config, enable_lora)
                             negative_slope=math.sqrt(5), nonlinearity="leaky_relu"
                         )
                     ),
+                    use_quick_lora=lora_config.use_quick_lora,
                 )
                 # Lora column parallel will spilt lora B matrix
                 self.add_lora_split_mapping(module_name + ".lora_B", is_column=True)
@@ -444,6 +446,7 @@ def _find_and_replace_module(self, model, module_name, lora_config, enable_lora)
                     rslora=lora_config.rslora,
                     lora_plus_scale=lora_config.lora_plus_scale,
                     merge_weights=lora_config.merge_weights,
+                    use_quick_lora=lora_config.use_quick_lora,
                 )
                 # Lora column parallel will spilt lora A matrix
                 self.add_lora_split_mapping(module_name + ".lora_A", is_column=False)