fix merge lora params (#7082)

Co-authored-by: lugimzzz <63761690+lugimzzz@users.noreply.github.com>

llm/causallm/data.py

@@ -97,7 +97,7 @@ def convert_example_common(example, tokenizer, data_args, is_test=True, intokens
         seq_length = len(input_ids)
         if intokens:
             features["position_ids"] = list(range(seq_length))
-            features["attention_mask"] = np.tril(np.ones((seq_length, seq_length), dtype="bool"))
+            features["attention_mask"] = np.tril(seq_length, seq_length, dtype="bool")
 
         return features
 

paddlenlp/peft/lora/lora_layers.py

@@ -170,7 +170,7 @@ class RowParallelLoRALinear(RowParallelLinear):
             input_mp = x
 
         # x @ W : [bz, in_f / ws] ===> [bz, out_f]
-        result_mp = F.linear(x=input_mp, weight=self.weight, bias=self.bias, name=self.name)
+        result_mp = F.linear(x=input_mp, weight=self.weight, name=self.name)
 
         output = mp_ops._mp_allreduce(
             result_mp,
@@ -192,7 +192,7 @@ class RowParallelLoRALinear(RowParallelLinear):
             #  @ 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):

paddlenlp/peft/lora/lora_model.py

@@ -183,7 +183,7 @@ class LoRAModel(nn.Layer):
             lora_state_dict[name] = action(tensor)
         return lora_state_dict
 
-    def save_pretrained(self, save_directory: str, merge_tensor_parallel: bool = False, **kwargs):
+    def save_pretrained(self, save_directory: str, merge_tensor_parallel: bool = True, **kwargs):
         variant = kwargs.get("variant", None)
         is_main_process = kwargs.get("is_main_process", paddle.distributed.get_rank() == 0)
 
@@ -214,7 +214,7 @@ class LoRAModel(nn.Layer):
         # 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
+        self.lora_config.tensor_parallel_degree = self.model.config.tensor_parallel_degree
 
     def _find_and_replace_module(self, model, module_name, lora_config, enable_lora):
         parent_module = model

paddlenlp/peft/prefix/prefix_model.py

@@ -303,7 +303,7 @@ class PrefixModelForCausalLM(paddle.nn.Layer):
 
         return prefix_model
 
-    def save_pretrained(self, save_directory: str, merge_tensor_parallel: bool = False, **kwargs):
+    def save_pretrained(self, save_directory: str, merge_tensor_parallel: bool = True, **kwargs):
         variant = kwargs.get("variant", None)
         is_main_process = kwargs.get("is_main_process", paddle.distributed.get_rank() == 0)
 
@@ -349,9 +349,13 @@ class PrefixModelForCausalLM(paddle.nn.Layer):
         # save prefix config & past key values
         if is_main_process:
             self.prefix_config.save_pretrained(save_directory)
-            self.prefix_config.tensor_parallel_degree = self.model.config.tensor_parallel_degree
             np.save(os.path.join(save_directory, PAST_KEY_VALUES_FILE_NAME), past_key_values)
 
+        if self.model.base_model_prefix == "chatglm2":
+            self.prefix_config.tensor_parallel_degree = -1
+        else:
+            self.prefix_config.tensor_parallel_degree = self.model.config.tensor_parallel_degree
+
     def set_state_dict(self, state_dict):
         self.prefix_encoder.set_state_dict(state_dict)
         logger.info("Load prefix weight successfully")

paddlenlp/trainer/trainer.py

@@ -49,12 +49,7 @@ from tqdm.auto import tqdm
 
 from ..data import DataCollator, DataCollatorWithPadding, default_data_collator
 from ..peft import LoRAModel, PrefixModelForCausalLM
-from ..transformers.model_utils import (
-    PretrainedModel,
-    _add_variant,
-    paddlenlp_load,
-    unwrap_model,
-)
+from ..transformers.model_utils import PretrainedModel, _add_variant, unwrap_model
 from ..transformers.tokenizer_utils import PretrainedTokenizer
 from ..utils.batch_sampler import DistributedBatchSampler as NlpDistributedBatchSampler
 from ..utils.env import LORA_WEIGHTS_NAME, PADDLE_WEIGHTS_NAME, PREFIX_WEIGHTS_NAME
@@ -414,23 +409,33 @@ class Trainer:
             else:
                 weight_name = PADDLE_WEIGHTS_NAME
 
-            if not os.path.isfile(
-                os.path.join(resume_from_checkpoint, _add_variant(weight_name, self.args.weight_name_suffix))
-            ):
-                raise ValueError(f"Can't find a valid checkpoint at {resume_from_checkpoint}")
+            if resume_from_checkpoint is not None and self.args.dataset_rank == 0:
+                if (isinstance(self.model, LoRAModel) and self.model.lora_config.tensor_parallel_degree > 1) or (
+                    isinstance(self.model, PrefixModelForCausalLM)
+                    and self.model.prefix_config.tensor_parallel_degree > 1
+                ):
+                    file_path = os.path.join(resume_from_checkpoint, weight_name)
+                    state_dict = paddle.load(file_path, return_numpy=True)
+                    state_dict = self.model._convert_tensor_parallel(state_dict)
+                else:
+                    file_path = os.path.join(
+                        resume_from_checkpoint, _add_variant(weight_name, self.args.weight_name_suffix)
+                    )
+                    if not os.path.isfile(file_path):
+                        raise ValueError(f"Can't find a valid checkpoint at {resume_from_checkpoint}, no {file_path}")
 
-            logger.info(f"Loading model from {resume_from_checkpoint} .")
+                    logger.info(f"Loading model from {resume_from_checkpoint} .")
 
-            # We load the model state dict on the CPU to avoid an OOM error.
-            state_dict = paddle.load(
-                os.path.join(resume_from_checkpoint, _add_variant(weight_name, self.args.weight_name_suffix)),
-                return_numpy=True,
-            )
-            # If the model is on the GPU, it still works!
-            self._set_state_dict_in_model(state_dict)
+                    # We load the model state dict on the CPU to avoid an OOM error.
+                    state_dict = paddle.load(file_path, return_numpy=True)
 
-            # release memory
-            del state_dict
+                # If the model is on the GPU, it still works!
+                self._set_state_dict_in_model(state_dict)
+
+                # release memory
+                del state_dict
+            elif resume_from_checkpoint is not None:
+                logger.info(f"not loading ckpt :{self.args.dataset_rank}")
 
     def train(
         self,
@@ -1800,17 +1805,21 @@ class Trainer:
 
         merge_tensor_parallel = merge_tensor_parallel and self.args.use_hybrid_parallel
 
-        if (
-            not isinstance(self.model, PretrainedModel)
-            and not isinstance(self.model, LoRAModel)
-            and not isinstance(self.model, PrefixModelForCausalLM)
-        ):
+        if isinstance(self.model, LoRAModel) or isinstance(self.model, PrefixModelForCausalLM):
+            self.model.save_pretrained(
+                output_dir,
+                variant=self.args.weight_name_suffix,
+                merge_tensor_parallel=True,
+                is_main_process=self.args.should_save,
+            )
+        elif not isinstance(self.model, PretrainedModel):
             if isinstance(unwrap_model(self.model), PretrainedModel):
                 unwrap_model(self.model).save_pretrained(
                     output_dir,
                     merge_tensor_parallel=merge_tensor_parallel,
                     variant=self.args.weight_name_suffix,
                     is_main_process=self.args.should_save,
+                    max_shard_size=self.args.max_shard_size,
                 )
             else:
                 logger.info("Trainer.model is not a `PretrainedModel`, only saving its state dict.")
@@ -1828,6 +1837,7 @@ class Trainer:
                 merge_tensor_parallel=merge_tensor_parallel,
                 variant=self.args.weight_name_suffix,
                 is_main_process=self.args.should_save,
+                max_shard_size=self.args.max_shard_size,
             )
 
         if self.args.should_save:
@@ -1848,7 +1858,7 @@ class Trainer:
             os.path.join(checkpoint, SCHEDULER_NAME)
         ):
             # Load in optimizer and scheduler states
-            self.optimizer.set_state_dict(paddlenlp_load(os.path.join(checkpoint, optimizer_name), map_location="cpu"))
+            self.optimizer.set_state_dict(paddle.load(os.path.join(checkpoint, optimizer_name)))
 
             self.lr_scheduler.set_state_dict(paddle.load(os.path.join(checkpoint, SCHEDULER_NAME)))
             if self.do_grad_scaling and os.path.isfile(os.path.join(checkpoint, SCALER_NAME)):

paddlenlp/trainer/training_args.py

@@ -313,6 +313,8 @@ class TrainingArguments:
             scripts](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/examples) for more details.
         flatten_param_grads (`bool`, *optional*):
             Whether use flatten_param_grads method in optimizer, only used on NPU devices. Default is `False`.
+        max_shard_size (`str`, *optional*):
+            The max shard size for saving model state dict param files. Default is `"10GB"`.
     """
 
     output_dir: str = field(
@@ -639,6 +641,10 @@ class TrainingArguments:
         default=True,
         metadata={"help": "Whether use lazy data processing."},
     )
+    max_shard_size: Optional[str] = field(
+        default="10GB",
+        metadata={"help": "The max shard size for saving model state dict param files."},
+    )
 
     def __post_init__(self):
         env_local_rank = int(os.environ.get("PADDLE_RANK_IN_NODE", -1))

paddlenlp/transformers/conversion_utils.py

@@ -37,7 +37,7 @@ from numpy import allclose, ndarray, transpose
 from paddle import Tensor
 from paddle.nn import Layer
 
-from paddlenlp.utils.distributed import distributed_gather
+from paddlenlp.utils.distributed import distributed_gather, distributed_allgather
 from paddlenlp.utils.env import CONFIG_NAME, PADDLE_WEIGHTS_NAME, PYTORCH_WEIGHTS_NAME
 from paddlenlp.utils.import_utils import (
     is_package_available,
@@ -1008,7 +1008,10 @@ class ConversionMixin:
         for key in state_dict.keys():
             tensor = state_dict[key]
             if key in name_action_mappings:
-                ret = distributed_gather(tensor, group=mp_group, offload=True)
+                if paddle.is_compiled_with_xpu():
+                    ret = distributed_allgather(tensor, group=mp_group, offload=True)
+                else:
+                    ret = distributed_gather(tensor, group=mp_group, offload=True)
                 action = name_action_mappings.pop(key)
                 tensor = action(ret) if is_dst else None
             else:

paddlenlp/transformers/generation_utils.py

@@ -1891,7 +1891,11 @@ def TopKProcess(probs, top_k, min_tokens_to_keep):
 
 
 def TopPProcess(probs, top_p, min_tokens_to_keep):
+    org_dtype = probs.dtype
+    if org_dtype == paddle.bfloat16:
+        probs = paddle.cast(probs, paddle.float32)
     sorted_indices = paddle.argsort(probs, descending=True)
+
     if isinstance(sorted_indices, tuple):
         sorted_probs, sorted_indices = sorted_indices
     else:
@@ -1917,4 +1921,7 @@ def TopPProcess(probs, top_p, min_tokens_to_keep):
     )
     condition = paddle.cast(condition, "bool").reshape(probs.shape)
     probs = paddle.where(condition, paddle.full_like(probs, 0.0), probs)
+
+    if org_dtype == paddle.bfloat16:
+        probs = paddle.cast(probs, paddle.bfloat16)
     return probs