mindspore.unified_safetensors

View Source On Gitee
mindspore.unified_safetensors(src_dir, src_strategy_file, dst_dir, merge_with_redundancy=True, file_suffix=None, max_process_num=64, choice_func=None)[source]

Merge multiple safetensor files into a unified safetensor file.

Parameters

  • src_dir (str) – Source weight saving directory.

  • src_strategy_file (str) – Source weight segmentation strategy file.

  • dst_dir (str) – Target save directory.

  • merge_with_redundancy (bool, optional) – Whether the merged source weight files are de-duplicated and saved safetensors files. Default: True, indicating that the merged source weight files are complete.

  • file_suffix (str, optional) – Specify the filename suffix for merging safetensors files. Default: None, meaning all safetensors files in the source weight directory will be merged.

  • max_process_num (int) – Maximum number of processes. Default: 64.

  • choice_func (callable) – A callable function used to filter parameters or modify parameter names. The return value of the function must be of type str (string) or bool (boolean). Default: None.

Raises

ValueError – If the safetensors file of rank is missing.

Supported Platforms:

Ascend GPU CPU

Examples

>>> import mindspore as ms
>>> src_dir = "/usr/safetensors/llama31B/4p_safetensors/"
>>> src_strategy_file = "/usr/safetensors/llama31B/strategy_4p.ckpt"
>>> dst_dir = "/usr/safetensors/llama31B/merge_llama31B_4p/"
>>> ms.unified_safetensors(src_dir, src_strategy_file, dst_dir)