Safetensors 为什么支持按需读取权重切片

Safetensors 能按需读取权重切片，关键不在某个“智能加载器”，而在文件格式本身已经记录了每个张量的 dtype、shape 和 data_offsets。加载器先解析很小的文件头，就能知道目标张量位于后续字节缓冲区的哪一段；调用 get_slice 时先保留这份描述，等切片索引确定后再构造所需区域，而不是先把完整张量全部实例化。

官方地址：https://huggingface.co/docs/safetensors/index

最小用法是 safe_open 配合 get_slice。这项能力特别适合大 embedding、张量并行和多设备加载，但它不等于“任何切片都绝对零拷贝”，也不表示 GPU 可以直接把磁盘页当显存使用。

最小写法：先拿切片视图，再选择区域

from safetensors import safe_open

with safe_open("model.safetensors", framework="pt", device="cpu") as f:
    weight_slice = f.get_slice("embedding")  # 先取得延迟切片对象，不先构造完整权重
    vocab_size, hidden_dim = weight_slice.get_shape()

    # 只选择词表前 1024 行和隐藏维度前一半
    rows = min(1024, vocab_size)
    part = weight_slice[:rows, : hidden_dim // 2]

get_tensor("embedding") 表达的是“取得这个完整张量”，而 get_slice("embedding") 先返回可查询形状、可接受索引的切片对象。只有最后一行的索引明确后，绑定层才知道目标区域。对于完整张量很大、当前进程只负责一部分参数的场景，这能避免先创建完整张量再截取。

文件布局让权重位置可以直接计算

Safetensors 文件开头的 8 字节是一个小端无符号整数 N，表示 JSON 文件头的长度。接下来的 N 字节是 UTF-8 JSON，其中每个张量名称对应一组描述信息：数据类型、形状，以及相对于权重字节缓冲区的起止偏移。文件剩余部分就是连续字节缓冲区。

{
  "embedding": {
    "dtype": "F16",
    "shape": [32000, 4096],
    "data_offsets": [0, 262144000]
  }
}

JSON 不支持注释，上面的字段含义分别是：dtype 决定每个元素占用多少位，shape 决定逻辑维度，data_offsets 指向权重缓冲区中的半开区间 [BEGIN, END)。加载器组合这三项，就能验证字节数并建立“逻辑索引到文件位置”的映射。

这种布局与 pickle 的对象反序列化思路不同。Safetensors 不需要执行自定义 Python 对象来恢复权重，也不必扫描整份文件才能知道某个具名张量在哪里。官方格式还要求数据缓冲区被完整索引且不能留洞，从而让偏移关系保持明确。

get_slice 延迟到索引时才构造目标区域

safe_open 负责打开文件并读取元数据，get_slice 根据张量名称找到对应描述，get_shape 允许程序在实际构造张量前计算分片。最后的切片索引把逻辑范围映射到目标字节窗口，再由所选框架和设备创建结果张量。

例如把 embedding 的词表行平均分给两个进程，可以先读 shape，再各取一段连续行。连续行在 C 顺序布局中更容易对应连续字节区间，也通常比复杂的跨步选择更适合作为加载分片。

from safetensors import safe_open

rank = 1
world_size = 2

with safe_open("model.safetensors", framework="pt", device="cpu") as f:
    view = f.get_slice("embedding")
    vocab_size, hidden_dim = view.get_shape()

    # 按词表行计算当前进程负责的连续区间
    start = vocab_size * rank // world_size
    stop = vocab_size * (rank + 1) // world_size
    local_weight = view[start:stop, :hidden_dim]

这段代码表达的是分区策略，不要求为每种并行方案另存一套权重文件。改变 rank 或 world_size 就能得到不同范围；同一份 Safetensors 文件继续提供稳定的张量名称、形状和偏移。

零拷贝、惰性读取和少占内存不是一个概念

官方项目把 Safetensors 描述为支持 zero-copy 和 lazy loading，但也明确说明，没有任何模型格式能让数据凭空从磁盘出现在 GPU。三个概念需要分开：

• 惰性读取：可以只检查元数据，或只选择部分张量，而不先扫描、反序列化全部权重。
• CPU 零拷贝：当文件已在操作系统页缓存中，并通过内存映射建立张量视图时，可能不需要再复制一份完整 CPU 缓冲。
• GPU 加载：磁盘或 CPU 内存中的数据仍需进入 GPU 显存；按需切片的收益是只传目标区域，并可避免同时在 CPU 构造所有完整张量。

因此，“get_slice 完全不占内存”是错误说法。文件映射需要虚拟地址空间，访问页面会进入页缓存，结果张量也需要目标设备内存。更准确的结论是：它减少不必要的完整张量实例化和中间副本，并把实际访问限制在所选区域附近。

它为什么适合多 GPU 和大模型加载

在张量并行中，每个设备常常只需要权重矩阵的一部分；在流水线并行中，不同进程只负责若干层；在模型检查或转换工具中，程序可能只需要少量张量名称和 shape。Safetensors 的文件头让这些任务可以先制定分片计划，再加载实际所需数据。

对旧代码的影响并不复杂：如果现有逻辑始终调用 get_tensor 并在完整张量上做切片，功能仍然成立，只是内存峰值可能保留在完整张量级别。迁移时应把“切片决定”提前到 get_slice 之后、结果张量创建之前。

from safetensors import safe_open

with safe_open("model.safetensors", framework="pt", device=0) as f:
    view = f.get_slice("embedding")
    vocab_size, hidden_dim = view.get_shape()

    # 只把目标行区间构造成 GPU 张量，避免先加载完整 embedding
    gpu_part = view[: vocab_size // 4, :hidden_dim]

device=0 表示结果面向第一个 GPU，但不代表磁盘到显存没有数据传输。实际峰值和速度还会受到文件系统、页缓存、PCIe 或互连、框架绑定、切片形状以及同时加载数量的影响。

按需切片的几个边界

它不会自动做远程 Range 下载

Hugging Face 官方文档展示了用 HTTP Range 请求只获取文件头，从而查看模型的张量名称、类型、shape 和参数量，而无需下载全部权重。但本地 safe_open("model.safetensors") 的 get_slice 本身不是远程下载协议。远程仓库、缓存层或网络文件系统是否只取所需字节，要看外层存储方案。

布局是 C 顺序且不保存 stride

格式采用 C 或 row-major 顺序，保存的是打包后的连续张量，不保存任意 stride。连续行或连续块通常最符合文件布局；复杂跨步、转置后的逻辑选择可能需要多个片段、额外整理或由绑定层拒绝。不要把“支持切片”理解为所有高级索引都天然只读取一个连续范围。

小于 1 字节的数据类型有对齐限制

官方格式说明指出，小于 1 字节的 dtype 会让对齐和寻址变复杂；触发非对齐读取时，库可能直接报错。这类格式应按对应版本和 API 的实际支持范围设计切片，不要沿用 F16、BF16 或 F32 的字节计算假设。

切片不能解决所有内存峰值

如果后续代码立即把各切片拼回完整权重、同时缓存多个设备副本，或框架算子要求连续重排，峰值仍可能上升。按需读取只优化加载入口，模型初始化、量化、转置、重排和优化器状态还需要分别检查。

迁移时这样判断

• 只需要少数完整张量：继续用 get_tensor，代码最直接。
• 一个大张量只需要连续区域：优先 get_slice，先用 get_shape 计算边界。
• 不同进程负责不同层：先按名称选择张量，再决定是否对单个张量切片。
• 远程文件未落地：先确认缓存或文件系统是否支持范围访问，不要把本地 API 当下载器。
• 目标是 GPU：关注实际传输量和显存结果，不宣称磁盘直达显存的零拷贝。
• 使用特殊低位 dtype：检查对齐要求和当前绑定支持。

常见问题

get_slice 会立刻读取完整权重吗？

它先返回切片对象并允许查询 shape；实际索引确定目标区域后才构造结果，因此用途就是避免先取得完整张量。

data_offsets 是文件绝对偏移吗？

不是。官方格式规定它相对于文件头之后的张量字节缓冲区，BEGIN 是起点，END 是不包含在内的终点。

按需切片一定比完整加载快吗？

不一定。读取范围很小或内存受限时通常更有价值；如果最终仍需要整个张量，额外切片规划未必带来收益。应按实际并行策略和存储环境判断。

Safetensors 的安全性是否等于模型内容可信？

不是。格式避免像 pickle 那样通过权重文件执行任意自定义代码，但张量数值本身仍可能包含 NaN、Inf 或恶意构造的数据，模型来源和应用层校验仍然重要。

Safetensors 支持按需读取权重切片，归根结底是“元数据先行、字节位置明确、张量数据连续”。这让加载器可以先知道要什么，再只构造目标区域。正确使用时，它减少完整权重的中间实例化；理解 CPU、页缓存和 GPU 之间的边界，才能把这种能力转化为真实的加载收益。