第 2 章 · 01 st.h safetensors 索引与范围读取

第 2 章 · 01 st.h safetensors 索引与范围读取

本节摘要:本节深潜 Colibrì 数据读取层的入口——c/st.h(852 行,header-only)。它解决两个问题:怎么解析 safetensors 格式(JSON 头描述每个 tensor 的偏移量,后面跟二进制体),以及怎么只读需要的字节而不是整个文件(用 pread + posix_fadvise(DONTNEED) 而不是 mmap,这样读过的页不留在进程 RSS 里,峰值内存只占稠密 + 缓存,而不是整个模型)。st.h 还支持 shard 分片(大模型拆多个文件)、镜像多 SSD、范围切片(st_read_slice_f32)——后者对专家流式加载至关重要:一个专家就是大 tensor 的一个子范围,只读它的 gate/down/up 三段。读完本节,你理解了 Colibrì 怎么"按需读字节"。

内容来源:原项目源码 c/st.h(852 行)、docs/FORMATS.md。

⚠️ 注意:st.h 用 pread 而不是 mmap 是有历史原因的——文件头注释明说这是 RSS bug 的修复:早期 mmap 版本会让读过的页常驻进程,峰值内存变成"整个模型"。pread + fadvise(DONTNEED) 让 OS 把读过的页立刻回收。这是 Colibrì 能在小内存机器上跑的关键修复之一。Windows 上 posix_fadvise 不可用,会退化成纯 buffered read。

学习目标

阅读完本节,你应当能够:

1. 描述 safetensors 格式(8 字节头长度 + JSON 头 + 二进制体)。
2. 解释 shards 结构体为什么维护 512 个 fd 和它们的 O_DIRECT "孪生" fd。
3. 说出 pread + posix_fadvise(DONTNEED) 相对 mmap 解决了什么 RSS bug。
4. 理解 shard 分片与镜像(mirror)的 fd 组织(主盘 + 多个副本)。
5. 解释 st_read_slice_f32 为什么对专家流式加载至关重要。
6. 读懂 dtype 编码(0=BF16/1=F16/2=F32/3=U8/I8/4=F8_E4M3/5=F8_E8M0/6=I64)。
7. 理解 st.h 为什么是 header-only、all static,被所有引擎 .c include。

一、safetensors 格式与 st.h 的定位

safetensors 是 Hugging Face 提出的安全权重格式,结构极简:

┌──────────────────────────────────────────────┐
│ 8 字节小端无符号整数:JSON 头的字节长度 N       │
├──────────────────────────────────────────────┤
│ N 字节 JSON 头:每个 tensor 的 {name: {dtype,  │
│   shape, data_offsets:[start,end), ...}}     │
│   + 可选 __metadata__                          │
├──────────────────────────────────────────────┤
│ 二进制体:所有 tensor 的原始字节,按头里的       │
│   data_offsets 寻址                             │
└──────────────────────────────────────────────┘
​

c/st.h:1-6 的头注释直白地标出它的角色:"Indicizzazione e lettura on-demand di tensori da piu' file safetensors"(多文件 safetensors 的按需索引与读取)。三个关键词:

• 索引(Indicizzazione):启动时解析 JSON 头,建 tensor 名 → 偏移量的哈希表。
• 按需(on-demand):只读需要的字节,不一次性加载整个文件。
• 多文件(piu' file):大模型拆成多个 shard(分片),st.h 把它们当成一个虚拟大文件。

st.h 把这两件事做到极致:它绝不 mmap,只用 pread(带偏移量的读),读完后用 posix_fadvise(DONTNEED) 告诉内核"这页我用完了,可以回收"——这样进程 RSS 不会随读取累积。

二、shards 结构体:索引 + fd 池

st.h:42-74 定义核心结构 shards,它是 st.h 的全部状态:

42  typedef struct {
43      st_tensor *t;            /* tensor 索引数组 */
44      int        n, cap;       /* tensor 数量 */
45      int        fds[512];     /* 每个 shard 文件的 fd */
46      int        dfds[512];    /* O_DIRECT 孪生 fd(惰性打开),-2 = 还没试过 */
47      char      *paths[512];   /* shard 文件路径,与 fds 并行 */
48      int64_t    sizes[512];   /* 每个 shard 的字节大小 */
49      int        nfd;          /* 已打开的 shard 数 */
50  #define ST_MAX_MIR 4         /* 主盘之外最多 4 个读副本(多 SSD) */
51      int        mfds[ST_MAX_MIR][512];  /* MIRROR: 副本 r+1 的 fd,-1 = 缺 */
52      int        mdfds[ST_MAX_MIR][512]; /* 副本的 O_DIRECT 孪生 fd */
53      int        nmirror[ST_MAX_MIR];    /* 副本 r+1 接受了多少文件 */
54      int        nrep;         /* 已注册副本数(0 = 镜像未激活) */
55      int       *hidx;         /* 名字 -> 索引 的开放寻址哈希表 */
56      int        hcap;
        /* ... FORMAT METADATA STAMP 字段(fmt_name/fmt_val)... */
73  } shards;
​

几个关键点逐个解读:

• fds[512]:最多 512 个 shard 文件。GLM-5.2 744B 拆成几十个 shard,每个几 GB。
• dfds[512](O_DIRECT 孪生):st.h:160-161 的注释解释——同一文件除了普通 fd,还惰性打开一个 O_DIRECT fd 用于绕过页缓存(实测 ext4-in-VHDX 上 buffered read 卡在 0.8 GB/s,O_DIRECT 能到 2.3+ GB/s)。-2 表示还没试过,-1 表示不可用。
• mfds[ST_MAX_MIR][512]:多 SSD 镜像——最多 4 个读副本,每个副本独立一套 fd。这是第 5 章双 SSD 带宽翻倍的底层支撑。
• hidx 哈希表:st.h:55-57 的注释极其关键——GLM 有约 120k 个 tensor(256 expert × 78 layer × 3 矩阵 × 2),线性扫描会花几十秒/token(在第一次真实 run 上测出来的),所以必须有哈希表。st.h:76-80 是 FNV 哈希实现。

三、st_tensor:单个 tensor 的索引项

st.h:31-40 定义单个 tensor 的索引项:

31  typedef struct {
32      char   *name;       /* tensor 名,如 "layers.0.mlp.experts.42.gate.weight" */
33      int     fd;         /* 这个 tensor 在哪个 shard 文件里 */
34      int64_t off;        /* tensor 数据在文件里的绝对字节偏移 */
35      int64_t nbytes;     /* tensor 数据的字节长度 */
36      int     dtype;      /* 0=BF16 1=F16 2=F32 3=U8/I8 4=F8_E4M3 5=F8_E8M0 6=I64 */
37      int64_t numel;      /* 元素个数 */
38      int     rank;       /* 维度数 */
39      int64_t shape[ST_MAX_RANK];  /* 各维大小,最多 8 维 */
40  } st_tensor;
​

dtype 的编码值得记——st.h:82-99 的 st_dtype_code 函数把字符串映射成整数。注意 0/1/2/3 是早期就有的(BF16/F16/F32/U8-I8),4/5/6 是后加的(F8_E4M3/F8_E8M0/I64):

94      if (!strcmp(s, "F8_E4M3") || !strcmp(s, "F8_E4M3FN") ||
95          !strcmp(s, "float8_e4m3fn")) return 4;
96      if (!strcmp(s, "F8_E8M0") || !strcmp(s, "F8_E8M0FNU")) return 5;
97      if (!strcmp(s, "I64") || !strcmp(s, "U64")) return 6;
​

st.h:88-93 的注释解释为什么 4/5/6 是新加的:在它们之前,st_init 在 DeepSeek V4 checkpoint 的第一个 I64 tensor 处就 exit(1),根本读不到权重。新加的 dtype 只被索引和 st_read_raw 路径处理,浮点读取路径按名字拒绝而不是踩进去。st.h:104-111 的 st_dtype_esz 是唯一知道每种 dtype 几字节/元素的地方——之前 dtype 字节公式在三个地方重复(dtype==2 ? 4 : 2),加 I64 时会静默错成"2 字节",所以集中到一个函数。

四、pread 而不是 mmap:RSS bug 修复

st.h:1-6 头注释用意大利语写的,翻译过来:

多文件 safetensors 的按需索引与读取。等价于 engine.py 的 Shards,但:

• 用 pread(不用 mmap)+ posix_fadvise(DONTNEED) → 读过的页不留在进程里。这是 RSS bug 的修复:这样峰值内存保持"稠密 + 缓存",而不是整个模型(见 memory mmap-rss-bug)。
• 出口总转成 float32(支持 BF16/F16/F32)。

关键句是"le pagine NON restano residenti nel processo"(页不留在进程里)。这是为什么 st.h 不用 mmap:

• mmap 把文件页映射到进程地址空间,读过的页会留在进程 RSS 里(直到主动 madvise(MADV_DONTNEED) 或被内存压力驱逐)。在 744B 模型上,边推理边 mmap,几轮下来 RSS 就涨到几十 GB,把高速缓存挤光。
• pread 是显式 read,读完内核可以立即回收(配合 fadvise(DONTNEED)),进程 RSS 只反映"当前在用的页",不累积。

st.h:808-813 的 st_read_raw 是这哲学的典型实现——读 raw 量化字节:

808  static void st_read_raw(shards *S, const char *name, void *out, int drop) {
809      st_tensor *t = st_find(S, name);
810      if (!t) { fprintf(stderr, "missing tensor: %s\n", name); exit(1); }
811      st_pread_full(t->fd, out, t->nbytes, t->off, "pread raw");
812      if (drop) posix_fadvise(t->fd, t->off, t->nbytes, POSIX_FADV_DONTNEED);
813  }
​

注意三个细节:(1) 用 st_find 走哈希表找 tensor(不是线性扫描);(2) st_pread_full 是分块 pread 循环(st.h:261-281,分块是因为单次 pread 上限约 2^31 字节,bf16 大 tensor 会超);(3) drop=1 时调 fadvise(DONTNEED) 把刚读的页标记为可回收。这条 drop 参数贯穿所有 st.h 读取函数。

五、范围切片:专家流式加载的关键

对 MoE 路由专家来说,一个专家不是一个独立 tensor,而是大 tensor 的一个子范围。例如 GLM 把一层的 256 个专家打包成 experts.gate.weight 这种 [256, ...] 大 tensor,单个专家是第 expert 个切片。st.h:831-849 的 st_read_slice_f32 就是干这个的:

831  static void st_read_slice_f32(shards *S, const char *name,
832          int64_t elem_off, int64_t n_elems, float *out, int drop) {
833      st_tensor *t = st_find(S, name);
834      if (!t) { fprintf(stderr, "missing tensor: %s\n", name); exit(1); }
835      if (t->dtype >= 3) { fprintf(stderr, "slice %s: tensor is %s — not a float tensor\n",
836              name, st_dtype_name(t->dtype)); exit(1); }
837      int esz = st_dtype_esz(t->dtype);
838      if (elem_off < 0 || n_elems < 0 || elem_off > t->numel ||
839          n_elems > t->numel - elem_off) { /* 切片必须在 tensor 内 */
840          fprintf(stderr, "slice %s [%lld,+%lld) out of tensor bounds (numel %lld)\n",
841                  name, (long long)elem_off, (long long)n_elems, (long long)t->numel); exit(1); }
842      int64_t boff = t->off + elem_off * esz, nb = n_elems * esz;
843      void *raw = malloc(nb);
844      st_pread_full(t->fd, raw, nb, boff, "pread slice");
845      /* ... 按 dtype 转 f32(BF16/HF16/F32)... */
848      free(raw);
849      if (drop) posix_fadvise(t->fd, boff, nb, POSIX_FADV_DONTNEED);
850  }
​

st.h:828-830 的意大利语注释解释用途:"Serve per gli expert fusi di GLM (un tensore = blocco [E, ...]): si legge il solo expert richiesto via pread del sotto-range, niente lettura dell'intero blocco."(用于 GLM 的融合专家:一个 tensor 是 [E, ...] 块,只读需要的那个专家的子范围,不读整个块。)

注意两个安全细节:(1) elem_off > t->numel 和 n_elems > t->numel - elem_off 用"减法避免溢出"(注释里的 #1 指代这条边界 bug);(2) dtype >= 3(量化字节)直接拒绝,因为 slice 是 float 路径。一个专家的三个矩阵 gate/down/up 就是三次切片调用,各自只读自己那段字节——这就是"只读需要的字节"。

💡 深潜要点:st.h 的精髓是只读需要的字节,读完即弃。范围切片让"读一个专家"等价于"在大 tensor 上 pread 一段字节",这正是第 3 章 expert_store.h 流式加载的底层基石。如果 st.h 一次性读整个大 tensor,专家缓存就退化成"层缓存",稀疏性优势全没。

六、shard 分片与镜像:多 SSD 的雏形

大模型单个 safetensors 文件存不下(GLM-5.2 744B 约 370GB 专家 + 9.9GB 稠密),必须分片。st.h 的处理逻辑在 st_init_multi(st.h:449)里:扫描目录下所有 *.safetensors,每个文件单独解析 JSON 头,把所有 tensor 项合并到一个全局 S->t 数组,文件 fd 记到 S->fds[nfd]。tensor 的 data_offsets 是文件内偏移,所以跨 shard 的 tensor 各自指向自己的 fd。

镜像(多 SSD)的 fd 组织在第二节已经讲过(mfds/mdfds)。关键约束在 st.h:184-191 注释里:一个镜像文件只有在 size 和 safetensors 头都与主盘 byte-identical 时才被接受——这样 data_offsets 天然一致,每个 pread 在任意副本上都有效。缺失或分歧的文件静默留在主盘(部分镜像可行——小 SSD 只放部分 shard 也帮忙)。镜像永不写入(.coli_usage/.coli_kv 只在主盘)。这部分第 5 章详讲。

本节要点回顾

1. safetensors 格式:8 字节头长度 + JSON 头(tensor 偏移量)+ 二进制体。
2. shards 结构:512 个 shard fd + O_DIRECT 孪生 fd + 最多 4 套镜像 fd + FNV 名字哈希(120k tensor 必须哈希)。
3. pread 不 mmap:这是 RSS bug 的修复,fadvise(DONTNEED) 让读过的页不累积在进程。
4. dtype 编码:0=BF16/1=F16/2=F32/3=U8-I8/4=F8_E4M3/5=F8_E8M0/6=I64;st_dtype_esz 是唯一字节数源。
5. 范围切片:st_read_slice_f32 让"读一个专家"= "大 tensor 上 pread 一段字节",是专家流式加载的底层基石。
6. 镜像约束:size + 头 byte-identical,data_offsets 天然一致;部分镜像可行,镜像永不写。

下一节,我们从"读字节"上升到"解字节"——c/quant.h 怎么解码量化容器,以及 docs/FORMATS.md 怎么治理格式编号冲突。