PeftModel supports saving and loading safetensors (#2025)

Author: xing-yiren

mindnlp/peft/peft_model.py

@@ -23,6 +23,7 @@
 import mindspore
 from mindspore import Tensor
 from mindspore.train.serialization import _exec_save
+from mindnlp.core.serialization import safe_save_file
 
 from mindnlp.core import nn, ops
 from mindnlp.core.nn import functional as F
@@ -45,7 +46,7 @@
     LNTuningModel,
 )
 from .utils import (
-    # SAFETENSORS_WEIGHTS_NAME,
+    SAFETENSORS_WEIGHTS_NAME,
     TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING,
     WEIGHTS_NAME,
     PeftType,
@@ -124,7 +125,7 @@ def __init__(self, model, peft_config: PeftConfig, adapter_name="default"):
         # if hasattr(self.base_model, "config") and hasattr(self.base_model.config, "pretraining_tp"):
         #     self.base_model.config.pretraining_tp = 1
 
-    def save_pretrained(self, save_directory, **kwargs):
+    def save_pretrained(self, save_directory, safe_serialization=False, **kwargs):
         r"""
         This function saves the adapter model and the adapter configuration files to a directory, so that it can be
         reloaded using the [`LoraModel.from_pretrained`] class method, and also used by the [`LoraModel.push_to_hub`]
@@ -144,10 +145,17 @@ def save_pretrained(self, save_directory, **kwargs):
             output_dir = os.path.join(save_directory, adapter_name) if adapter_name != "default" else save_directory
             os.makedirs(output_dir, exist_ok=True)
 
-            _exec_save(
-                ckpt_file_name=os.path.join(output_dir, WEIGHTS_NAME),
-                data_list=output_state_dict,
-            )
+            if safe_serialization:
+                safe_output_state_dict = {k : Tensor(v[2]).reshape(v[0]) for k, v in output_state_dict.items()}
+                safe_save_file(
+                    safe_output_state_dict,
+                    os.path.join(output_dir, SAFETENSORS_WEIGHTS_NAME),
+                )
+            else:
+                _exec_save(
+                    ckpt_file_name=os.path.join(output_dir, WEIGHTS_NAME),
+                    data_list=output_state_dict,
+                )
 
             # save the config and change the inference mode to `True`
             if peft_config.base_model_name_or_path is None:

mindnlp/peft/utils/save_and_load.py

@@ -18,8 +18,10 @@
 
 import mindspore
 
+from mindnlp.core.serialization import safe_load_file
+
 from .peft_types import PeftType
-from .constants import WEIGHTS_NAME
+from .constants import WEIGHTS_NAME, SAFETENSORS_WEIGHTS_NAME
 
 def get_data_list(param_dict):
     """Get state dict of the Peft model for saving."""
@@ -198,11 +200,15 @@ def load_peft_weights(model_id: str,) -> dict:
     """
     path = model_id
 
-    filename = os.path.join(path, WEIGHTS_NAME)
-    if not os.path.exists(filename):
-        # TODO: add download logic later
-        raise ValueError(f"load peft model failed, peft model file: {filename} not exists.")
+    safe_filename = os.path.join(path, SAFETENSORS_WEIGHTS_NAME)
+    ckpt_filename = os.path.join(path, WEIGHTS_NAME)
 
-    adapters_weights = mindspore.load_checkpoint(filename)
+    if os.path.exists(safe_filename):
+        adapters_weights = safe_load_file(safe_filename)
+    elif os.path.exists(ckpt_filename):
+        adapters_weights = mindspore.load_checkpoint(ckpt_filename)
+    else:
+        # TODO: add download logic later
+        raise ValueError(f"load peft model failed, peft model file: neither {ckpt_filename} nor {safe_filename} was found.")
 
     return adapters_weights

mindnlp/transformers/models/qwen2/modeling_qwen2.py

@@ -801,7 +801,7 @@ def prepare_inputs_for_generation(
 
         if attention_mask is not None and position_ids is None:
             # create position_ids on the fly for batch generation
-            position_ids = attention_mask.int().cumsum(-1) - 1
+            position_ids = ops.cumsum(attention_mask.int(), -1) - 1
             position_ids = position_ids.masked_fill(attention_mask == 0, 1)
             if past_key_values:
                 position_ids = position_ids[:, -input_ids.shape[1] :]