MoE 路由与专家流式加载

这是 colibrì 的心脏——让 744B 模型跑在 25 GB 内存上的地方。每层 MoE 从 256 个专家里选 8 个,绝大多数专家留在磁盘上按需读取。本文档拆解路由(moe)、专家从磁盘的流式加载(expert_load)、以及围绕它的多级缓存(LRU / 固定热存储 / tier 换入换出 / 会学习的缓存)。


1. moe:DeepSeek-V3 式稀疏路由

moe(1147–1265 行)x[S,hidden] 做稀疏 MoE。它的 5 个阶段:

1.1 sigmoid 路由(noaux_tc)

1156–1157 行:每个位置先 matmul(logit, xs, router) 算出每个专家的 logit,然后 sig[e]=sigmoidf(logit[e])选择用的分数加上专家偏置:choice[e] = sig[e] + router_bias[e]

这里有一个 DeepSeek “noaux_tc” 的精妙之处:偏置只影响选择,不影响权重被选中专家应用的门控权重是无偏的 sig[best](1163 行),而不是带偏置的 choice。偏置只用来引导选哪些专家(负载均衡),不污染实际计算。

1.2 top-k 选择与自适应 top-p

同时更新持久的 eusage 和近期的 eheat 热度计数(1173–1176 行)——这两个计数驱动”会学习的缓存”和 repin。

1.3 batch-union、缓存解析、分块计算

1.4 共享专家

1256–1262 行:每个 token 都过的共享专家(sh_gate/sh_up/SiLU/sh_down),加到输出上。它是常驻的(不流式)。

1.5 前几层稠密

first_dense 层不是 MoE 而是普通 MLP:dense_mlp(1267 行)。分派在 layer_forward(1360 行)l->sparse ? moe : dense_mlp

1.6 自适应旋钮

旋钮 出处 作用
g_topk 548 覆盖每 token 专家数(比 config 少 → 省磁盘)
g_topp 549 累积权重截断(自适应少读专家)
g_absorb 984 MLA 吸收(见 06)

2. expert_load:从磁盘流式读一个专家

expert_load(865–943 行)是磁盘流式的关键路径。一个专家有 gate/up/down 三个矩阵,目标是用尽量少的 pread 把它们读进一个连续的 slab

expert_load 流程:
  1. 构造张量名 model.layers.L.mlp.experts.E.{gate,up,down}_proj.weight
  2. 若无预量化容器 → 逐个 qt_from_disk 运行时量化(fallback)
  3. 容器路径:st_find 定位 3 个权重 + 3 个缩放张量
  4. 若 slot 的 slab 太小 → 重新分配(4K 对齐 posix_memalign)
  5. 按文件偏移排序,检测三个张量是否连续
  6. 连续 → 一次 pread 读全部(~19 MB)
       ├─ O_DIRECT 路径(g_direct):4K 对齐,走 st_direct_fd 双 fd
       └─ 缓冲路径:普通 pread
     不连续 → 3 次 pread
  7. 缩放 .qs 单独 pread 进 fslab
  8. DROP(g_drop)→ posix_fadvise(DONTNEED) 立即驱逐页
  9. QT 视图 g/u/d 指向 slab 内偏移,零拷贝

相关旋钮:g_drop(536)g_prefetch(539)g_direct(542)。O_DIRECT 的价值在 README benchmark 里很明显:VHDX 缓冲读 ~0.8 GB/s,O_DIRECT 能到 2.3+ GB/s。


3. 多级专家缓存

Model 结构(121–125 行)里的几组字段构成了缓存层级:

pin  (ESlot **pin, npin)      ← 固定热存储:从不驱逐,autopin 学到的热专家
ecache (ESlot **ecache, ecap) ← 每层 LRU:最近用过的专家,满了驱逐最冷的
ws  (ESlot ws[64])            ← working set:本次 forward 临时读进来的

再往下是 OS 页缓存(读过的磁盘页还在内存里,免费的 L2)和磁盘本身。

3.1 会学习的缓存(.coli_usage)

引擎记录每次使用真正路由到哪些专家,写进模型旁边的 .coli_usage,启动时自动把最热的固定进 RAM。

这就是 README 那句”colibrì 用得越多越快”的实现——命中 pin 的专家完全不碰磁盘。


4. tier.h:热度驱动的 repin 换入换出

固定的热专家不是一成不变的——会话进行中,tier.h 提供的热度换入换出可以把冷的固定专家换成热的(--repin N 开启)。

在 glm.c 里:

设计注释(1791–1800 行)总结了整套策略:启动从 .coli_usage autopin,回合间 repin,25%+4 迟滞,每次最多 4 换。持久的 .coli_usage 是长期信号,不被 decay。


5. pilot_prefetch:router-lookahead 磁盘预取

一个有趣的观察(README 实测):GLM-5.2 下一层的路由是可提前预测的——用 L+1 层的路由器作用于 L 层的 post-attention 状态,能召回 71.6% 的真 top-8。PILOT=1 用这一点,从一个专门的 I/O 线程预读下一层的专家。

旋钮:g_pilot/g_pilot_k(566–567 行)g_spec(550 行)g_looka(556 行)(LOOKA 只测量、零副作用,用来评估各种预测器的召回率)。

README 诚实地说:在磁盘已经 ~80% 饱和的开发机上 PILOT 测得中性;在计算/磁盘平衡的机器上才应该 overlap 出真实收益。


6. 小结

机制 出处 作用
sigmoid 路由(noaux_tc) moe:1156 偏置引导选择,权重用无偏 sigmoid
自适应 top-p moe:1166 --topp 少读专家,省 30–40% 磁盘
batch-union moe:1195 每个唯一专家只读一次
合并 pread 读专家 expert_load:865 一次 ~19 MB,可选 O_DIRECT
DROP 驱逐 expert_load:930 读完 fadvise(DONTNEED),RSS 可控
会学习的缓存 usage_load:2135 越用越快,热专家 autopin
tier repin repin_pass:1818, tier.h:8 回合间换热专家,25%+4 迟滞
router-lookahead 预取 pilot_prefetch:1320 预测下层专家,后台预读

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