[llm]support lora merge (#7733)

* support lora merge

* add finetune

* support bf16 and lora merge

* update

* add lora merge

* version adatpt

* update

* fix

---------

Co-authored-by: root <[email protected]>

Author: lugimzzz

llm/docs/finetune.md

@@ -215,13 +215,16 @@ python merge_tp_and_pp_params.py \
 为了后续的**压缩**和**静态图推理**方便，我们提供LoRA参数合并脚本，可以将LoRA参数合并到主干模型并保存相应的权重。
 ```
 python merge_lora_params.py \
-    --model_name_or_path meta-llama/Llama-2-7b-chat \
-    --lora_path ./checkpoints/llama_lora_ckpts
+    --lora_path ./checkpoints/llama_lora_ckpts \
+    --merge_lora_model_path ./checkpoints/llama_lora_merge \
+    --device "gpu" \
+    --low_gpu_mem True
 ```
+
 <summary>&emsp; 脚本参数介绍</summary><div>
 
-- `model_name_or_path`: 必须，预训练模型名称或者本地的模型路径，用于热启模型和分词器，默认为None。
 - `lora_path`: LoRA参数和配置路径，对LoRA参数进行初始化，默认为None。
 - `merge_model_path`: 必须，合并参数后保存路径，默认为None。
 - `device`: 运行环境，默认为gpu。
+- `low_gpu_mem`:降低合参时候所需显存，默认为False。如果合参时显存不足，建议开启
 </div>

llm/finetune_generation.py

@@ -162,7 +162,6 @@ def main():
         else:
             # NOTE(gongenlei): new add autotuner_benchmark
             model = AutoModelForCausalLM.from_config(model_config, dtype=dtype)
-
     if training_args.do_train and model_args.neftune:
         # Inspired by https://github.com/neelsjain/NEFTune
         if hasattr(model, "get_input_embeddings"):
@@ -418,6 +417,7 @@ def neft_post_hook(module, input, output):
                 tensor_parallel_degree=training_args.tensor_parallel_degree,
                 dtype=dtype,
                 do_qat=quant_args.do_qat,
+                base_model_name_or_path=model_args.model_name_or_path,
             )
             model = LoRAModel(model, lora_config)
         else:

llm/llama/qlora_argument.json

@@ -0,0 +1,33 @@
+{
+    "model_name_or_path": "facebook/llama-7b",
+    "dataset_name_or_path": "./data",
+    "output_dir": "./checkpoints/llama_lora_ckpts",
+    "per_device_train_batch_size": 4,
+    "gradient_accumulation_steps": 4,
+    "per_device_eval_batch_size": 8,
+    "eval_accumulation_steps":16,
+    "num_train_epochs": 3,
+    "learning_rate": 3e-04,
+    "warmup_steps": 30,
+    "logging_steps": 1,
+    "evaluation_strategy": "epoch",
+    "save_strategy": "epoch",
+    "src_length": 1024,
+    "max_length": 2048,
+    "fp16": true,
+    "fp16_opt_level": "O2",
+    "do_train": true,
+    "do_eval": true,
+    "disable_tqdm": true,
+    "load_best_model_at_end": true,
+    "eval_with_do_generation": false,
+    "metric_for_best_model": "accuracy",
+    "recompute": true,
+    "save_total_limit": 1,
+    "tensor_parallel_degree": 1,
+    "pipeline_parallel_degree": 1,
+    "lora": true,
+    "zero_padding": false,
+    "use_flash_attention": false,
+    "weight_quantize_algo": "nf4"
+  }

llm/llama/wint8_lora_argument.json

@@ -0,0 +1,33 @@
+{
+    "model_name_or_path": "facebook/llama-7b",
+    "dataset_name_or_path": "./data",
+    "output_dir": "./checkpoints/llama_lora_ckpts",
+    "per_device_train_batch_size": 4,
+    "gradient_accumulation_steps": 4,
+    "per_device_eval_batch_size": 8,
+    "eval_accumulation_steps":16,
+    "num_train_epochs": 3,
+    "learning_rate": 3e-04,
+    "warmup_steps": 30,
+    "logging_steps": 1,
+    "evaluation_strategy": "epoch",
+    "save_strategy": "epoch",
+    "src_length": 1024,
+    "max_length": 2048,
+    "fp16": true,
+    "fp16_opt_level": "O2",
+    "do_train": true,
+    "do_eval": true,
+    "disable_tqdm": true,
+    "load_best_model_at_end": true,
+    "eval_with_do_generation": false,
+    "metric_for_best_model": "accuracy",
+    "recompute": true,
+    "save_total_limit": 1,
+    "tensor_parallel_degree": 1,
+    "pipeline_parallel_degree": 1,
+    "lora": true,
+    "zero_padding": false,
+    "use_flash_attention": false,
+    "weight_quantize_algo": "weight_only_int8"
+  }

llm/merge_lora_params.py

@@ -12,45 +12,135 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import argparse
+import copy
+import os
 
 import paddle
 
 from paddlenlp.peft import LoRAConfig, LoRAModel
-from paddlenlp.transformers import AutoModelForCausalLM
+
+try:
+    from paddle.nn.quant import weight_dequantize, weight_quantize
+except:
+    weight_dequantize = None
+    weight_quantize = None
+try:
+    from paddlenlp.quantization.qlora import qlora_weight_quantize_dequantize
+except:
+    qlora_weight_quantize_dequantize = None
+
+from paddlenlp.quantization.quantization_config import QuantizationConfig
+from paddlenlp.transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from paddlenlp.transformers.utils import device_guard
+from paddlenlp.utils.env import CONFIG_NAME
 
 
 def parse_arguments():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--model_name_or_path", default=None, required=True, help="The directory of pretrained model.")
+    parser.add_argument("--model_name_or_path", default=None, help="The directory of pretrained model.")
     parser.add_argument(
         "--lora_path", default=None, required=True, help="The directory of LoRA parameters. Default to None"
     )
-    parser.add_argument("--merge_model_path", default=None, help="The directory of merged parameters. Default to None")
+    parser.add_argument(
+        "--merge_lora_model_path",
+        default=None,
+        required=True,
+        help="The directory of merged parameters. Default to None",
+    )
     parser.add_argument("--device", type=str, default="gpu", help="Device")
+    parser.add_argument(
+        "--low_gpu_mem", type=bool, default=False, help="Whether to use low gpu memory. Default to False"
+    )
     return parser.parse_args()
 
 
+def weight_process(name, quant_config, lora_config, state_dict):
+    weight = state_dict.pop(name + ".weight").cuda()
+    if quant_config.weight_quantize_algo is None:
+        pass
+    elif quant_config.weight_quantize_algo in ["nf4", "fp4"]:
+        weight = qlora_weight_quantize_dequantize(
+            weight,
+            quant_algo=quant_config.weight_quantize_algo,
+            double_quant=quant_config.weight_double_quant,
+            block_size=quant_config.weight_blocksize,
+            double_quant_block_size=quant_config.weight_double_quant_block_size,
+        )
+    elif quant_config.weight_quantize_algo in ["weight_only_int8"]:
+        out, scale = weight_quantize(weight, algo=quant_config.weight_quantize_algo)
+        weight = weight_dequantize(out, scale)
+    else:
+        raise ValueError(f"quant_config.weight_quantize_algo {quant_config.weight_quantize_algo} is not supported.")
+    lora_A = state_dict.pop(name + ".lora_A").cuda()
+    lora_B = state_dict.pop(name + ".lora_B").cuda()
+    scaling = lora_config.lora_alpha / lora_config.r
+    state_dict[name + ".weight"] = (weight + lora_A @ lora_B * scaling).cpu()
+
+
 def merge():
     args = parse_arguments()
     paddle.set_device(args.device)
+
     lora_config = LoRAConfig.from_pretrained(args.lora_path)
-    dtype = lora_config.dtype
-    lora_config.merge_weights = True
+    if lora_config.base_model_name_or_path is None:
+        if args.model_name_or_path is not None:
+            raise ValueError("We can not find a valid model_name_or_path.")
+        else:
+            lora_config.base_model_name_or_path = args.model_name_or_path
 
-    model = AutoModelForCausalLM.from_pretrained(
-        args.model_name_or_path,
-        dtype=dtype,
-    )
-    model = LoRAModel.from_pretrained(model=model, lora_path=args.lora_path, lora_config=lora_config)
-    model.eval()
-    if args.merge_model_path is None:
-        args.merge_model_path = args.lora_path
-
-    model_state_dict = model.model.state_dict()
-    for key in list(model_state_dict):
-        if "lora" in key:
-            del model_state_dict[key]
-    model.model.save_pretrained(args.merge_model_path, state_dict=model_state_dict)
+    if os.path.isfile(os.path.join(args.lora_path, CONFIG_NAME)):
+        config = AutoConfig.from_pretrained(args.lora_path)
+    elif args.model_name_or_path is not None:
+        config = AutoConfig.from_pretrained(args.model_name_or_path)
+    else:
+        raise ValueError(
+            f"We can not find config.json in lora_path: {args.lora_path} or find a valid model_name_or_path."
+        )
+    config.dtype = lora_config.dtype
+    if (
+        lora_config.dtype == "bfloat16" or config.quantization_config.weight_quantize_algo in ["nf4", "fp4"]
+    ) and args.device == "cpu":
+        raise ValueError("We can not apply bfloat16 or nf4/fp4 lora merge on cpu.")
+
+    if args.low_gpu_mem and args.device == "gpu":
+        quant_config = copy.deepcopy(config.quantization_config)
+        config.quantization_config = QuantizationConfig()
+        lora_config.merge_weights = False
+        with device_guard():
+            model = AutoModelForCausalLM.from_pretrained(
+                lora_config.base_model_name_or_path,
+                config=config,
+                low_cpu_mem_usage=True,
+            )
+            model = LoRAModel.from_pretrained(model=model, lora_path=args.lora_path, lora_config=lora_config)
+        model.eval()
+        model_state_dict = model.model.state_dict()
+        lora_name_list = []
+        for key in model_state_dict.keys():
+            if "lora_A" in key:
+                lora_name_list.append(key[:-7])
+        for name in lora_name_list:
+            weight_process(name, quant_config, lora_config, model_state_dict)
+    else:
+        model = AutoModelForCausalLM.from_pretrained(
+            lora_config.base_model_name_or_path,
+            config=config,
+            low_cpu_mem_usage=True,
+        )
+        lora_config.merge_weights = True
+        model = LoRAModel.from_pretrained(model=model, lora_path=args.lora_path, lora_config=lora_config)
+        model.eval()
+        model_state_dict = model.model.state_dict()
+        for key in list(model_state_dict):
+            if "lora" in key:
+                del model_state_dict[key]
+            if "quant" in key:
+                del model_state_dict[key]
+        model.model.config.quantization_config = QuantizationConfig()
+    model.model.save_pretrained(args.merge_lora_model_path, state_dict=model_state_dict)
+
+    tokenizer = AutoTokenizer.from_pretrained(lora_config.base_model_name_or_path)
+    tokenizer.save_pretrained(args.merge_lora_model_path)
 
 
 if __name__ == "__main__":

paddlenlp/peft/lora/lora_config.py

@@ -72,6 +72,9 @@ class LoRAConfig:
         },
     )
     do_qat: bool = field(default=False, metadata={"help": "Whether the lora model would do quant-aware training"})
+    base_model_name_or_path: Optional[str] = field(
+        default=None, metadata={"help": "The name of the base model to use."}
+    )
 
     @property
     def __dict__(self):

paddlenlp/peft/lora/lora_model.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import copy
 import math
 import os
 import re
@@ -226,14 +227,16 @@ def _convert_tensor_parallel(self, lora_state_dict):
         return lora_state_dict
 
     def save_pretrained(self, save_directory: str, merge_tensor_parallel: bool = False, **kwargs):
+        save_model_config = kwargs.get("save_model_config", True)
+
         if self.is_pipelinemodel:
             self.model._single_to_pp_mapping = None
-        if self.quantized and merge_tensor_parallel and self.model.config.tensor_parallel_degree > 1:
+        if self.quantized and merge_tensor_parallel and self.lora_config.tensor_parallel_degre > 1:
             merge_tensor_parallel = False
             logger.warning(
                 "Quantized strategy does not support merge_tensor_parallel. Set merge_tensor_parallel to False."
             )
-        if self.is_pipelinemodel and merge_tensor_parallel and self.model.config.tensor_parallel_degree > 1:
+        if self.is_pipelinemodel and merge_tensor_parallel and self.lora_config.tensor_parallel_degre > 1:
             merge_tensor_parallel = False
             logger.warning(
                 "Pipeline parallism does not support merge_tensor_parallel. Set merge_tensor_parallel to False."
@@ -247,18 +250,20 @@ def save_pretrained(self, save_directory: str, merge_tensor_parallel: bool = Fal
         ), f"Saving directory ({save_directory}) should be a directory, not a file"
         os.makedirs(save_directory, exist_ok=True)
 
-        if merge_tensor_parallel and self.model.config.tensor_parallel_degree > 1:
+        lora_config_to_save = LoRAConfig(**self.lora_config.to_dict())
+
+        if merge_tensor_parallel and lora_config_to_save.tensor_parallel_degree > 1:
             trainable_state_dict = self.get_trainable_state_dict()
             trainable_state_dict = self._merge_trainable_tensor_parallel(trainable_state_dict)
             if not is_main_process:
                 logger.info("Saving with merge_tensor_parallel, tensor_parallel_rank > 0 don't need save")
                 return
             if variant is not None and "tp" in variant:
                 variant = "_".join([x for x in variant.split("_") if "tp" not in x])
-            self.lora_config.tensor_parallel_degree = -1
+            lora_config_to_save.tensor_parallel_degree = -1
         else:
             trainable_state_dict = self.get_trainable_state_dict()
-            if self.model.config.tensor_parallel_degree > 1:
+            if lora_config_to_save.tensor_parallel_degree > 1:
                 if variant is None:
                     variant = weight_name_suffix()
 
@@ -269,8 +274,12 @@ def save_pretrained(self, save_directory: str, merge_tensor_parallel: bool = Fal
 
         # save lora config
         if is_main_process:
-            self.lora_config.save_pretrained(save_directory)
-        self.lora_config.tensor_parallel_degree = self.model.config.tensor_parallel_degree
+            lora_config_to_save.save_pretrained(save_directory)
+            if save_model_config:
+                model_config_to_save = copy.deepcopy(self.model.config)
+                if merge_tensor_parallel:
+                    model_config_to_save.tensor_parallel_degree = -1
+                model_config_to_save.save_pretrained(save_directory)
 
     def _find_and_replace_module(self, model, module_name, lora_config, enable_lora):
         parent_module = model
@@ -366,6 +375,7 @@ def _find_and_replace_module(self, model, module_name, lora_config, enable_lora)
                     r=lora_config.r,
                     lora_alpha=lora_config.lora_alpha,
                     lora_dropout=lora_config.lora_dropout,
+                    merge_weights=lora_config.merge_weights,
                 )
                 self.quantized = True
             elif ColumnParallelQuantizationLinear is not None and isinstance(module, ColumnParallelQuantizationLinear):