物理流式:让”流式”名副其实

命中率结论已经拿到,但引擎其实还全常驻——专家仍占着 ~4GB 内存。研究不需要物理执行,但要真省内存,得走完最后一步。这是全项目最难的一块。

1. 选择性加载:非专家常驻,专家留盘

专家 3D 张量约占 OLMoE 九成参数。真省内存必须不把它们载入。streaming_model.ccStreamingModel::Load:

  • gguf_initno_alloc=true 只拿元数据 + 形状;
  • 非专家张量(注意力/norm/router/embd/lm_head)在常驻 ctx 建同形张量、从文件定位读其数据;
  • 专家张量(名含 _exps.weight)跳过,交给 ExpertReader 按需流式。

常驻内存从 ~4GB 降到 ~250MB。判断”是不是专家张量”就一句 strstr(name, "_exps.weight")——router 是 ffn_gate_inp 不含,正好排除。

2. 两段式前向:难在哪

难点在于——专家选择依赖运行时激活。你必须先算完这一层的路由、知道选了哪 8 个专家,才能去磁盘装载、再计算。所以一层没法一趟图算完,得拆成两段(streaming_forward.cc):

逐 token(T=1):
  段A 图:  注意力(带 KV cache)+ 残差 ─► ffn_norm ─► 路由 ─► 选中的全局专家 id
           └─ 读回 host:h_mid、weights、selected
  host:    caches[l].Ensure(selected)  ← miss 真从盘 pread 装进槽
           全局 id ─► slot id(重映射)
  段B 图:  在槽张量上 mul_mat_id(cur, slot_id)算专家 FFN ─► + 残差 ─► h_next

h 在两段之间走 host 中转——因为逐 token 处理,活值只有 [n_embd, 1],极小,这个中转很便宜。段 B 里 mul_mat_id 的专家张量是 expert_slot_cache有界槽张量 [n_embd, n_ff, C](C = 每层容量),selected 被重映射成 0..C 的槽号。约束:C ≥ 每步选中的不同专家数(逐 token 下 = n_expert_used = 8)。

3. 有界槽缓存

ExpertSlotCache::Ensure(ids) 是物理流式的心脏:命中即用;miss 则占槽(空槽优先,否则淘汰 tick 最小的)+ 从 ExpertReader 读入该专家的 gate/up/down 三类权重;返回 专家id → 槽号 供重映射。这里 P20 用 LRU;它同时也是第 9 章那些策略在物理世界里的一个实现。

4. 正确性锚点:streaming == resident

流式路径的输出必须与常驻逐 token 一致——因为槽里装的是和常驻专家完全相同的量化字节,同样的 mul_mat_id,结果理应一致。单测 StreamingMatchesResident 断言 StreamingGenerate == GreedyGenerateCached(norm_topk 两种、含淘汰路径)。

真机结果:

路径 输出 峰值内存
流式(每层缓存 16 专家) “…is Paris.” 1.2 GB
常驻 同上,逐 token 一致 3.0 GB

同一模型、逐 token 完全一致的输出,内存降到 ~40%。 1.2GB ≈ 250MB 非专家 + 16 层×16 槽×~3MB 专家。至此,nuthatch 是真正的流式 MoE 引擎,不是模拟器。

5. 提速 2×

流式每 token 有 ~34 趟小图,起初慢。瓶颈是两件蠢事:每趟 ggml_init 都新 malloc 一块 64MB compute buffer、还以 4 线程 compute——对 T=1 的极小图纯属浪费。改成复用一块 thread_local compute buffer + 极小图单线程(只有末端 lm_head 大乘保留多线程),省掉每 token 578 次线程池 spin,提速 (2.73s → 1.36s,12 token 含加载)。输出不变、parity 仍成立。

异步预取(把 miss 的磁盘读与 compute 重叠)还没做,是后续。

下一章:工程判断与适用边界


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