昇思MindSpore2024年技术帖分享大会圆满结束！全年收获80+高质量技术帖， 2025年全新升级，推出“2025年昇思干货小卖部，你投我就收！”，活动继续每月征集技术帖。本期技术文章由社区开发者yide12输出并投稿。如果您对活动感兴趣，欢迎在昇思论坛投稿。

# 01

概述

深度学习模型规模快速增长，给模型存储和加载带来了挑战。SafeTensors格式解决了传统格式的安全隐患，支持高效的大规模模型存储和快速加载。

MindSpore完整集成了SafeTensors格式，并实现了独特的懒加载机制，让普通硬件也能处理大规模模型。

# 02

SafeTensors 格式概述

SafeTensors是新一代机器学习模型存储格式，解决了传统存储方式的局限性。

pickle格式存在安全问题，反序列化时可能执行恶意代码。SafeTensors采用标准化二进制格式，避免代码执行风险，通过精心设计的文件结构，天然支持零拷贝加载和内存映射。

SafeTensors 文件结构

SafeTensors文件分为三个部分：

1.文件头：记录文件的基本信息

2.目录页：列出所有参数的位置和大小

3.数据区：存放实际的模型参数

与传统格式的对比

特性	SafeTensors	Pickle格式
安全性	只存数据，无代码	可能包含恶意代码
加载速度	很快，按需加载	较慢，全部加载
内存使用	省内存	占用较多内存

# 03

安全性分析

SafeTensors格式就像给你的模型加了一把安全锁：

- 防止恶意代码：不会执行任何危险的代码

- 数据完整性：自动检查文件是否被篡改

- 格式验证：确保文件格式正确无误

# 04

MindSpore 中的 SafeTensors 集成

MindSpore已经完美支持SafeTensors格式，具有以下特点：

- 无缝切换：只需修改一个配置参数

- 完全兼容：支持所有常用的模型操作

- 自动优化：系统会自动选择最佳的加载策略

# 05

懒加载机制

5.1 什么是懒加载

懒加载（Lazy Loading）是一种延迟加载策略，就像看书一样，你不需要把整本书都记在脑子里，只需要知道每章在第几页，需要的时候再翻到那一页。

对于AI模型来说，懒加载意味着：

- 不一次性加载：不把整个模型都放进内存

- 按需获取：需要哪部分参数就加载哪部分

- 快速启动：几秒钟就能开始工作

5.2 为什么需要懒加载

传统加载方式的问题

想象一下，如果你要查一个词典里的单词，传统方式就像把整本词典都背下来：

1.内存爆炸：一个大型AI模型可能有几十GB甚至上百GB的参数

2.启动缓慢：加载完整模型可能需要几分钟甚至更久

3.资源浪费：很多参数可能根本用不到，但也要占用内存

4.硬件限制：普通电脑的内存根本装不下大型模型

懒加载的优势：

懒加载就像一个智能助手，只在你需要的时候提供帮助：

1.内存友好：只占用实际需要的内存空间

2.快速启动：几秒钟就能开始工作

3.硬件兼容：普通电脑也能运行大型模型

4.动态调整：根据实际使用情况调整内存占用

5.3 MindSpore的懒加载实现原理

MindSpore采用了多层次的懒加载架构，每一层都有其独特的作用：

内存映射技术（Memory Mapping）

内存映射就像在文件上贴了很多智能标签：

工作原理：

- 操作系统将文件内容映射到虚拟内存空间

- 程序访问数据时，系统自动从文件读取对应部分

- 不需要的数据不会占用物理内存

技术细节：

- 使用mmap系统调用建立文件到内存的映射

- 支持随机访问，可以快速跳转到任意位置

- 操作系统自动管理页面换入换出

实际效果：

- 一个100GB的模型文件，可能只占用几百MB的实际内存

- 访问速度接近内存访问速度

- 多个进程可以共享同一份映射

元数据预解析（Metadata Parsing）

这就像先看书的目录，了解每章的内容和位置：

解析内容：

- 参数名称和层次结构

- 每个参数的数据类型和形状

- 参数在文件中的精确位置和大小

- 数据的字节序和对齐信息

存储结构：

文件头部信息：

├── 版本信息

├── 参数总数

├── 文件大小

└── 校验和

参数索引表：

├── 参数1: {名称, 类型, 形状, 偏移量, 大小}

├── 参数2: {名称, 类型, 形状, 偏移量, 大小}

└── …

优化策略：

- 索引表采用哈希表结构，支持O(1)查找

- 预计算常用参数的访问路径

- 缓存频繁访问的元数据

代理对象系统（Proxy Objects）

为每个参数创建一个"智能代理"：

代理对象的职责：

- 存储参数的元数据信息

- 管理实际数据的加载和释放

- 提供统一的访问接口

- 实现透明的数据获取

延迟实例化：

- 创建代理时不加载实际数据

- 首次访问时才触发数据加载

- 支持部分加载（只加载需要的切片）

内存管理：

- 自动检测数据使用情况

- 实现LRU缓存策略

- 支持手动释放不需要的数据

核心类：_fast_safe_open详解

这是MindSpore中实现懒加载的核心类，就像一个智能的文件管理器：

class _fast_safe_open:
    """
    Open a safetensors file and access its metadata and tensors efficiently.
    
    This function is designed to work similarly to `safetensors.safe_open`,
    providing a fast way to open and interact with safetensors files.
    """
    
    def __init__(self, filename, framework=None, device="cpu"):
        self.filename = filename
        self.framework = framework
        self.file = open(self.filename, "rb")
        self.file_mmap = mmap.mmap(self.file.fileno(), 0, access=mmap.ACCESS_COPY)
        try:
            self.base, self.tensors_decs, self.__metadata__ = read_metadata(self.file)
        except ValueError:
            raise ValueError(f"Fail to parse the input safetensors file: '{self.filename}'. "
                             f"Please check the correctness of the file.")
        self.tensors = OrderedDict()
        for key, info in self.tensors_decs.items():
            self.tensors[key] = PySafeSlice(info, self.file, self.base, self.file_mmap)
            self.tensors[key].key = key
    
    def __enter__(self):
        return self
   
    def __exit__(self, *args):
        self.file.close()
    
    def metadata(self):
        return self.__metadata__
    
    def keys(self):
        return list(self.tensors.keys())
   
    def get_tensor(self, name):
        return self.tensors[name].get()

初始化过程

1.文件验证：

• 检查文件格式和完整性

• 验证SafeTensors文件头

• 确认文件版本兼容性

2.元数据解析：

• 读取并解析文件头部信息

• 构建参数索引表

• 建立名称到位置的映射关系

3.内存映射建立：

• 创建文件到内存的映射

• 设置访问权限和保护模式

• 初始化页面管理器

# 06

性能优化

SafeTensors的懒加载机制可以显著提升性能：

6.1 内存优化

- 按需加载：只加载当前需要的参数

- 内存映射：直接从文件读取，不占用额外内存

6.2 加载速度优化

- 并行读取：支持多线程同时加载不同参数

- 格式优化：二进制格式比文本格式快很多

# 07

实际应用场景

大规模模型训练

训练大型AI模型时，SafeTensors就像一个智能仓库：

- 快速启动：几秒钟就能开始训练

- 节省内存：只加载正在使用的部分

- 分布式友好：多台机器可以各自加载需要的部分

# 08

总结

SafeTensors是一个让AI模型存储变得更安全、更快速的新技术。主要好处体现在下面几个方面：

- 更安全：不会有恶意代码，就像给模型加了安全锁

- 更快速：几秒钟就能开始使用，不用等很久

- 更省内存：只用需要的部分，其他暂时不管

- 更简单：使用起来很容易，不需要改很多代码

MindSpore已经完美支持SafeTensors，让每个人都能享受到这些好处。无论你是刚开始学AI，还是经验丰富的专家，都可以轻松使用这个技术。