初始化与权重加载

本文档跟踪从进程启动到”就绪可推理”的完整路径:读配置、索引磁盘上的 safetensors、把稠密权重加载常驻内存,而把路由专家留在磁盘上,最后从可用内存自动推算安全的缓存预算。核心是 model_init(706 行)

一个贯穿始终的事实:colibrì 不用 mmap/madvise。权重常驻与专家流式全部建立在 pread(定位读)+ posix_fadvise(WILLNEED/DONTNEED)之上,可选 O_DIRECT 双 fd。这是 st.h 头注释 明确的设计。


1. 配置解析

1.1 Cfg 结构

Cfg(53–61 行)保存 GLM-5.2 的全部超参数:

  • 核心维度:hidden, n_layers, n_heads, n_experts, topk, moe_inter, dense_inter
  • MLA/MoE 维度:first_dense, q_lora, kv_lora, qk_nope, qk_rope, qk_head, v_head, n_shared, vocab
  • 分组:n_group, topk_group, norm_topk
  • 停止 token:stop_ids[8], n_stop注释说 GLM-5.2 有 3 个 EOS
  • DSA:index_topk, index_nh, index_hd每层 idx_type[128](1=full/compute,0=shared/reuse)
  • 浮点:eps, theta, attn_scale, routed_scale

1.2 load_cfg

load_cfg(620 行)<snap>/config.json 读取(cfg_root(612 行) 打开并 slurp,gi(619 行) 是取整数的助手)。关键映射:

config.json 键 Cfg 字段 备注
num_attention_heads n_heads GLM-5.2 是 64,无 GQA
n_routed_experts n_experts 256
num_experts_per_tok topk 8
first_k_dense_replace first_dense 前几层是稠密 MLP
q_lora_rank / kv_lora_rank q_lora / kv_lora MLA 的 LoRA 秩
qk_nope_head_dim / qk_rope_head_dim qk_nope / qk_rope 无 RoPE / 有 RoPE 部分
routed_scaling_factor routed_scale 路由权重缩放

派生量:qk_head = qk_nope + qk_ropeattn_scale = 1/sqrt(qk_head)(652–653 行)。还有一个 CKR 宏(657 行)对每个维度做范围钳制,防止恶意 config。引擎硬断言 n_group==1(654 行)——这是 GLM-5.2 的假设,否则直接退出。


2. Model / Layer / 缓存结构

2.1 Layer

Layer(84–94 行)持有一层的所有常驻张量:

  • 归一化:in_ln, post_ln
  • MLA(量化 QT):q_a, q_b, kv_a, kv_b, o 加 f32 的 q_a_ln, kv_a_ln
  • Dense MLP:gate_proj, up_proj, down_proj
  • MoE:f32 的 router, router_bias;共享专家 sh_gate, sh_up, sh_down
  • sparse 标志:这一层是 MoE 还是 dense

注意:路由专家不在 Layer 里——它们留在磁盘上,通过独立的缓存结构按需加载。

2.2 ESlot:专家缓存槽

ESlot(100–101 行)是一个专家在内存里的样子:

typedef struct { int eid; QT g,u,d; uint8_t *slab; float *fslab;
                 int64_t slab_cap, fslab_cap; uint64_t used; } ESlot;

注释(96–99 行)解释:g/u/d(gate/up/down)三个 QT 都是同一块 slab视图——一个专家的三个矩阵通过一次合并 pread 读进 slabfslab 存缩放系数,used 是 LRU 时钟戳,*_cap 让 slot 能在不同大小的层之间复用。详见 07-MoE路由与专家流式加载

2.3 Model

Model(110–143 行)是顶层容器,除了 CfgLayer *L、embed/lm_head,还有几组关键子系统的状态:


3. model_init:大加载器

model_init(706–845 行)的流程:

  1. memset + load_cfg + st_init(707–708 行)——索引所有 safetensors 分片。
  2. 加载 embed / lm_head / final_norm(713–715 行)(embed/lm_head 默认 int8)。
  3. 逐层加载稠密权重724–754 行循环):
  4. MTP 头自动检测与加载(756–800 行):要求 3 个分片里存在完整张量集,MTP=0 可禁用。MTP 头是一个完整的 transformer 层,位于索引 n_layers(GLM-5.2 是 78)。
  5. DSA 自动检测与加载(803–828 行)index_topk>0 且每个 full 层都有 indexer 权重时启用,DSA=0 可禁用。
  6. resident_bytes 只累加稠密/embed/共享/MTP/DSA 张量(831–844 行)——专家被排除,从代码层面证明了它们是流式的。

4. 张量从磁盘到 QT

预量化容器的磁盘布局:每个权重 name 是 U8 打包字节,name.qs 是 F32 per-row 缩放——正是转换器输出的格式


5. 存储层 st.h:pread + fadvise

所有磁盘访问都在 st.h 里,全部基于 pread


6. RAM 安全预算:从可用内存推算缓存上限

colibrì 的一个亮点是不会触发 OOM-killer——它从 MemAvailable 诚实地推算峰值内存,据此设定专家缓存的大小。

这个预算的重要性在 README 的 benchmark 里体现得淋漓尽致:24 GB 内存的机器会被自动限制到每层 2 个专家槽,即使磁盘更快,decode 也保持冷态——”在小内存机器上,是 RAM cap 而非磁盘成为约束”。


7. 小结

阶段 出处 关键点
读配置 load_cfg:620 GLM-5.2 超参,硬断言 n_group==1
索引磁盘 st_init:101 12 万张量建哈希表,避免线性扫描
稠密常驻 model_init:724 注意力/router/共享专家/embed 进内存
专家留盘 model_init:742 只加载 router,专家不加载
MTP/DSA 检测 model_init:756, 803 按权重存在与否自动启用
RAM 预算 cap_for_ram:2329 从 MemAvailable 推算,绝不 OOM

This site uses Just the Docs, a documentation theme for Jekyll.