MindSpore Transformers Safetensors权重详解

用户867573478982
2 阅读1分钟

​Safetensors是Hugging Face主导设计的安全、高效模型权重存储格式，核心解决传统CKPT、PT格式的安全隐患与加载效率问题，MindSpore Transformers已实现对Safetensors格式的完整支持，可无缝加载、保存、转换Safetensors权重，适配Qwen、Llama、DeepSeek等主流LLM，兼顾安全性、跨框架兼容性与昇腾硬件适配性。Safetensors权重的核心特性、格式原理，搭配MindSpore Transformers实操代码，覆盖权重加载、保存、格式转换及分布式场景处理，助力开发者高效管理大模型权重。

一、Safetensors权重核心特性与格式优势

与传统CKPT（MindSpore）、PT（PyTorch）权重格式相比，Safetensors权重具备三大核心优势，适配大模型开发全场景：一是安全可靠，采用标准化二进制格式，仅存储张量数据与元信息，避免pickle序列化带来的恶意代码注入风险，加载时自动校验文件完整性，防止篡改；二是高效加载，支持零拷贝加载与内存映射，无需完整读取文件即可访问张量，大幅提升大模型（如7B、13B）加载速度，降低内存占用；三是跨框架兼容，可直接在MindSpore、PyTorch等框架间通用，无需格式转换，同时支持完整权重与分布式分片权重两种存储形式。

Safetensors文件结构简洁，主要分为元数据区（记录张量形状、数据类型、偏移量）、偏移量表（映射张量与数据位置）、原始张量数据三部分，天然支持懒加载机制，可按需加载所需参数，适配资源受限场景。

二、环境搭建（适配Safetensors权重操作）

首先完成MindSpore Transformers环境搭建，确保支持Safetensors格式，代码适配昇腾NPU、GPU、CPU多硬件平台，可直接执行：

# 1. 创建虚拟环境并激活
conda create -n ms-safetensors python=3.8
conda activate ms-safetensors

# 2. 安装核心依赖（昇腾NPU为例，适配MindSpore 2.3.0）
pip install mindspore-npu==2.3.0 mindspore-transformers==1.8.0 datasets -i https://pypi.tuna.tsinghua.edu.cn/simple

# 3. 验证环境（确认Safetensors支持已启用）
python -c "import mindspore_transformers; print('Safetensors支持已启用')"

三、核心实操：Safetensors权重加载与保存

MindSpore Transformers提供简洁API，支持从本地或Hugging Face加载Safetensors权重，同时可将训练后的模型保存为Safetensors格式，代码适配昇腾NPU，兼顾基础与进阶场景：

from mindspore_transformers import AutoModelForCausalLM, AutoTokenizer
import mindspore as ms

# 1. 配置运行环境（昇腾NPU，动态图模式，支持懒加载）
ms.set_context(mode=ms.PYNATIVE_MODE, device_target="Ascend")

# 2. 加载Safetensors权重（两种场景：本地权重、Hugging Face在线加载）
# 场景1：加载本地Safetensors权重（以Qwen3-8B为例）
tokenizer = AutoTokenizer.from_pretrained("./qwen3-8b-safetensors")
model = AutoModelForCausalLM.from_pretrained(
    "./qwen3-8b-safetensors",
    load_ckpt_format="safetensors",  # 显式指定加载Safetensors格式
    device_target="Ascend"
)

# 场景2：从Hugging Face在线加载Safetensors权重（自动转换适配MindSpore）
# model = AutoModelForCausalLM.from_pretrained(
#     "Qwen/Qwen3-8B",
#     load_ckpt_format="safetensors",
#     device_target="Ascend"
# )

# 3. 验证权重加载（执行推理，确认模型可用）
prompt = "介绍MindSpore Transformers中Safetensors权重的优势"
inputs = tokenizer(prompt, return_tensors="ms")
outputs = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

# 4. 保存模型为Safetensors格式（训练后权重持久化）
model.save_pretrained(
    "./qwen3-8b-saved",
    save_ckpt_format="safetensors",  # 指定保存为Safetensors格式
    remove_redundancy=True  # 可选，去除冗余参数，减小文件体积
)
tokenizer.save_pretrained("./qwen3-8b-saved")

四、进阶实操：权重格式转换与分布式处理

实际开发中常需在CKPT与Safetensors格式间转换，同时应对分布式训练后的分片权重合并，MindSpore提供专用API实现高效处理，代码如下：

from mindspore import safetensors_to_ckpt, save_checkpoint
from mindspore_transformers import unified_safetensors

# 1. Safetensors与CKPT格式互转
# 场景1：Safetensors转CKPT（适配传统MindSpore训练流程）
safetensors_to_ckpt(
    safetensors_file="./qwen3-8b-safetensors/model.safetensors",
    save_path="./qwen3-8b.ckpt"
)

# 场景2：CKPT转Safetensors（提升加载效率与安全性）
save_checkpoint(
    model,
    "./qwen3-8b-converted.safetensors",
    save_ckpt_format="safetensors"
)

# 2. 分布式Safetensors权重合并（多卡训练后合并为完整权重）
unified_safetensors(
    src_dir="./distributed-safetensors",  # 分布式权重目录（含各rank分片）
    src_strategy_file="./train_strategy.ckpt",  # 分布式切分策略文件
    dst_dir="./unified-safetensors",  # 合并后完整权重目录
    max_process_num=4  # 并行合并，提升效率
)

# 3. 分布式权重加载（多卡推理场景，无需合并）
model = AutoModelForCausalLM.from_pretrained(
    "./distributed-safetensors",
    load_ckpt_format="safetensors",
    use_parallel=True,  # 启用并行加载
    auto_trans_ckpt=True  # 自动适配分布式切分策略
)

五、常见问题与优化技巧

1. 加载报错“格式不支持”：确保MindSpore Transformers版本≥1.8.0，显式指定load_ckpt_format="safetensors"；2. 加载速度慢：启用懒加载机制，避免一次性加载全部权重，同时利用内存映射特性提升效率；3. 权重体积过大：保存时启用remove_redundancy=True，去除优化器状态等冗余信息；4. 跨框架兼容问题：Hugging Face下载的Safetensors权重可直接加载，无需额外转换，MindSpore保存的Safetensors权重也可用于PyTorch框架。

avatar
文章
阅读
粉丝