mindspore.parallel.unified_safetensors

View Source On Gitee
mindspore.parallel.unified_safetensors(src_dir, src_strategy_file, dst_dir, merge_with_redundancy=True, file_suffix=None, max_process_num=64, choice_func=None, split_dst_file=())[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, optional) – Maximum number of processes. Default: 64.

  • choice_func (callable, optional) – 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.

  • split_dst_file (tuple, optional) – execution, represented as a tuple containing two elements. The first element indicates the number of the current subtask, and the second element indicates the total number of tasks. This parameter supports splitting and executing tasks multiple times on a single machine, and also supports executing different subtasks on multiple machines respectively. Default: ().

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.parallel.unified_safetensors(src_dir, src_strategy_file, dst_dir)