mindspore.unified_safetensors

View Source On Gitee
mindspore.unified_safetensors(src_dir, src_strategy_file, dst_dir)[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.

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)