Safetensors 分片权重怎样按索引加载

我第一次处理 Safetensors 分片模型时，最容易犯的错是看文件名：先打开 model-00001-of-00008.safetensors，再猜下一个文件应该放哪些参数。真正可靠的入口其实是索引文件。分片文件只是承载数据，model.safetensors.index.json 才记录参数名与文件之间的对应关系。

官方资料入口：https://huggingface.co/docs/

先看索引文件而不是猜分片编号

一个常见的分片目录会同时放着配置文件、多个 .safetensors 文件和 model.safetensors.index.json。索引通常有两个关键部分：metadata 描述权重总大小，weight_map 把每个参数名映射到实际分片文件。参数名相同，文件名就不应该靠排序或编号推断。

import json
from pathlib import Path

# 读取索引，只关注参数到分片的映射关系。
root = Path("./checkpoint")
index_path = root / "model.safetensors.index.json"
index = json.loads(index_path.read_text(encoding="utf-8"))
weight_map = index["weight_map"]

# 用一个明确的参数名查询它所在的真实分片文件。
parameter_name = "model.layers.0.self_attn.q_proj.weight"
shard_name = weight_map.get(parameter_name)
if shard_name is None:
    raise KeyError(f"索引中没有参数：{parameter_name}")

print("参数所在分片：", root / shard_name)

这里的查询结果只是文件定位信息，不代表张量已经加载到内存。若索引能读到、但目标分片不存在，问题是目录不完整；若分片存在但参数名查不到，先检查模型版本或参数命名，而不是继续增加分片文件。

完整模型优先让 Transformers 自动恢复

如果目标是把整个 checkpoint 放进一个已经定义好的模型，手动遍历每个分片通常没有必要。Transformers 的 from_pretrained() 会根据目录中的索引找到分片，并把它们装载到对应模型结构中；需要直接给已有模型填充权重时，也可以使用 load_sharded_checkpoint()。

from transformers import AutoModel

# 目录中应包含 config、索引和索引引用的全部 safetensors 分片。
model = AutoModel.from_pretrained(
    "./checkpoint",
    device_map="cpu",
)

# 如果已经有按同一 config 创建的模型，可改用：
# from transformers.modeling_utils import load_sharded_checkpoint
# load_sharded_checkpoint(existing_model, "./checkpoint", prefer_safe=True)

这个路径适合推理、继续转换或需要完整 state_dict 的场景。它的好处是分片边界由官方加载逻辑处理，代价是最终仍然要为完整模型准备足够的内存。索引能减少“把所有文件一次性读成一个临时字典”的冲动，但不会让完整模型本身变小。

需要按需读取时按分片分组

做权重检查、抽取少数层，或只想读取某些参数时，我会先把目标参数按 weight_map 分组。这样同一个分片只打开一次，读完后关闭文件；没有必要把整个目录的所有张量都读进内存。

import json
from collections import defaultdict
from pathlib import Path
from safetensors import safe_open

root = Path("./checkpoint")

# 索引负责把参数名映射到真实分片，不能按文件名顺序猜测。
index = json.loads(
    (root / "model.safetensors.index.json").read_text(encoding="utf-8")
)
weight_map = index["weight_map"]
requested = [
    "model.layers.0.self_attn.q_proj.weight",
    "model.layers.0.self_attn.k_proj.weight",
]

# 先按文件聚合参数，避免同一个分片被重复打开。
groups = defaultdict(list)
for name in requested:
    shard = weight_map.get(name)
    if shard is None:
        raise KeyError(f"索引中没有参数：{name}")
    groups[shard].append(name)

loaded = {}
for shard, names in groups.items():
    shard_path = root / shard
    if not shard_path.is_file():
        raise FileNotFoundError(f"缺少索引引用的分片：{shard_path}")

    # safe_open 负责打开单个 safetensors 文件，按名称取出目标张量。
    with safe_open(str(shard_path), framework="pt", device="cpu") as handle:
        for name in names:
            loaded[name] = handle.get_tensor(name)

print("已读取参数：", sorted(loaded))

safe_open 的关键价值是把“文件选择”和“参数选择”分开：索引决定打开哪些文件，get_tensor() 决定从当前文件取哪个键。示例把张量放在 CPU，实际使用 GPU、MPS 或其他设备时，应结合当前 safetensors 与 PyTorch 版本的设备支持来调整，不要把设备迁移和索引解析混在一起排错。

报错时先区分文件、索引和模型结构

分片加载失败时，我会按三个层次看，而不是先删除索引重试。

尤其要注意，索引只负责告诉加载器“哪个参数在哪个文件”，它不会修复模型结构差异，也不会替你补齐缺失分片。把索引、配置和权重看成同一个发布单元，排查会快很多。

按任务选择加载粒度

如果我要运行完整模型，选择 from_pretrained()；如果只检查几个层或抽取少量张量，选择“解析索引、按分片分组、safe_open 读取”。前者维护成本低，后者更节省无关张量的读取，但需要自己处理参数名、文件存在性和输出字典。

我的经验是：索引文件存在时，不要把分片当成一组需要手工拼接的二进制块。它们是由参数映射组织起来的多个文件。对完整模型让 Transformers 负责，对局部张量才手动按 weight_map 取数，这个边界最不容易出错。

常见问题

可以按 00001、00002 的顺序读取分片吗？

不建议。分片编号主要是文件命名信息，参数归属以 weight_map 为准；即使顺序看起来连续，也不代表某个参数一定在预想的文件里。

只有一个 safetensors 文件，还需要索引吗？

单文件场景通常可以直接用 safetensors.torch.load_file() 或框架加载。索引的主要作用是描述多个分片之间的参数归属，不能为了形式给单文件强行拼出一份映射。

按需读取后能直接当成完整模型权重吗？

只有当你读取了模型需要的全部键，并按模型结构调用 load_state_dict() 时才可以这样做。只取几个张量更适合分析、转换或局部处理，不等于完整 checkpoint。