← 返回首页
11_moe_viz.html

手搓 MoE:参数翻倍,算力不变

第 3 章的 GPT 里,FFN 是每个 token 都要过的同一个小 MLP。 MoE 把它拆开:多个专家 + 一个路由器, 每个 token 只走最对口的 2 个专家。这是 Mixtral / DeepSeek 这代模型"参数千亿、激活百亿"的核心机关。 本关用 1.35M 参数把它完整跑通 —— 大概算全球最小的 MoE 之一,但路由、分工、塌缩,和 671B 的 DeepSeek 是同一套机制。 配套代码 phase3-moe/09_moe.py。

STEP 1
全员上岗的 FFN
STEP 2
路由器:给 token 派单
STEP 3
分工是训出来的
STEP 4
偏心与塌缩
STEP 5
算账 & 落到代码
这一关 n_expert 4 每层专家数 top_k 2 每 token 走几个 expert_hidden 256 =2·n_embd·单个专家隐层 aux 0.01 负载均衡系数 底座 03 的 GPT 4 层·n_embd=128

① dense FFN:不管什么 token,全员上岗

第 3 章讲过:注意力负责 token 之间"沟通",FFN 负责每个 token 单独"思考"。 但这个思考很浪费 —— 随便点下面哪个 token,看它过 FFN 时点亮了多少参数:

↓ 进入同一个 FFN(隐层 4·n_embd = 512 个神经元,下图 1 点 = 4 个)
点一个 token 试试
一个平平无奇的空格,和一个信息量很大的字母,动用的是同一套全部 FFN 参数。 模型越大这笔账越吓人:GPT 这类架构里 FFN 占了参数的大约 2/3,而它对每个 token 都是全量计算。
为什么偏偏对 FFN 下手,不拆注意力?
注意力的工作是 token 之间的交互,拆开会破坏"谁都能看谁"的沟通; FFN 是逐位置独立的加工 —— 每个 token 各算各的,谁跟谁都不打招呼。 这种"天生并行、互不依赖"的结构才方便按 token 派给不同专家。 (对注意力做稀疏化的研究也有,但主流 MoE 都是拆 FFN。)
↳ 下一步:把这个 512 隐层的大 FFN 拆成 4 个 256 隐层的小 FFN,每个 token 只走 2 个 —— 谁来决定走哪 2 个?

② 路由器:一个线性层,给每个 token 派单

路由器只是一个 n_embd → n_expert 的线性层:token 向量进去,4 个分数出来,softmax 变成概率, 只留最高的 top_k=2 个,权重重新归一化。选一个 token,看它被派给谁 (下面的概率是 09_moe.py 训练后 layer 0 的真实路由统计,按字符类别汇总):

输出 = 选中专家的加权和:out = w·Expert_i(x) + w'·Expert_j(x),权重就是归一化后的路由概率。 没被选中的专家这个 token 完全不算 —— 这就是"参数在那儿,但不花算力"的稀疏激活。
top-k 是"挑选",挑选不可导,路由器怎么学?
梯度走的是权重那条路:被选中的专家,其输出乘的系数 w 来自 softmax,是可导的。 loss 对 w 的梯度会一路传回路由器 —— "这次派给 E1 效果好,下次给 E1 打分再高点"。 没被选中的专家确实拿不到梯度,这正是下一步"强者恒强"问题的根源。
派单的单位是 token 还是整句话?
是 token。同一句话里,空格可能走 E1、字母 e 走 E2 —— 每个位置独立决策。 代码里先把 (B,T,C) 摊平成 (B·T, C) 再路由,就是在强调"逐 token 派单"。
↳ 代码:09_moe.py 的 MoELayer.forward:router 打分 → topk → 归一化 → 加权求和
↳ 下一步:派单派了 5000 步之后,4 个专家真的形成分工了吗?

③ 分工是训出来的,不是设计出来的

没有任何人规定"E1 管空格、E2 管元音"。但训练 5000 步后,把 val 集的 token 按类别统计 "每类字符被派往哪个专家"(layer 0 真实数据),分工自己浮出来了:

每行加起来 = 100% · 颜色越深 = 该类字符越爱找这个专家 · 鼠标悬停看准确数值
每个专家最常接手的字符(同一次训练的真实统计)
这些分工是涌现的:换个随机种子,分工方案会不一样,但"有分工"这件事稳定出现 —— 路由器 + 专家在训练里互相适应:专家碰巧对某类 token 加工得好,路由器就更爱把这类 token 派给它,它就练得更专。
专家 = "数学专家 / 法律专家"吗?
不是。大模型里的专家分工大多发生在低层统计模式上(词形、标点、语言、代码 token 之类), 很少出现人类意义上的"学科专家"。Mixtral 论文自己也承认:看不出专家按主题分工。 本页字符级模型的"空白/元音"分工,就是这种统计分工的微缩版。
↳ 下一步:"接单多 → 练得好 → 接单更多"是个正反馈。它有一个危险的极端……

④ 偏心的尽头:专家塌缩,和治它的 aux loss

正反馈不加管束,路由器会把越来越多 token 塞给少数专家 —— 被冷落的专家拿不到梯度、逐渐报废,这叫专家塌缩。 要看清它,得用最脆的设定:top-1 路由(Switch Transformer 同款,aux loss 就是它提出的)。 下面是同一份代码、都用 top-1 跑的两次真实训练,唯一区别是负载均衡开关。点 ▶ 看 layer 0 的负载怎么随训练演化:

step 0
训练进度
–
val loss
–
最忙专家的负载
完美均衡 = 每个专家 25%(虚线为参考)· 数据:layer 0,每 500 步采一次
诚实交代:本关的 top-2 配置,关掉 aux 其实也没塌
实测(同 5000 步):top-2 + 4 专家就算 --aux 0,最终负载也只是 20/24/27/29% 的轻微偏, val loss 1.549 和开 aux 打平。原因:top-2 下每个 token 同时养 2 个专家,4 选 2 让谁都分得到梯度,天生抗塌。 专家越多、top_k 越小,饿死人的风险越大 —— 所以 top-1 的 Switch 最先撞上这个问题, 而 64 选 8 的细粒度 MoE 全都标配负载均衡。上面的 top-1 实验还有一个诚实数字:这个玩具规模下, 偏载对 val loss 也没有可测的伤害(1.597 vs 1.600)—— 塌缩真正的代价是买了的容量白买, 以及大规模多机训练时忙闲不均拖慢整批,这两笔账在大模型上才要命。
aux loss 的公式为什么长这样?
aux = n_expert · Σ f_e · P_e:f_e 是实际派给专家 e 的 token 占比(不可导), P_e 是路由器给 e 的平均概率(可导)。两者都均匀时乘积和最小(=1); 谁的 f 偏大,压它的 P 就最能降 loss —— 梯度顺着可导的 P 把"过热"的专家分数往下压。 这是 Switch Transformer 的经典设计,系数 aux_coef=0.01,轻轻管一下就够。
塌缩之后,MoE 退化成了什么?
退化成一个只有少数专家在干活的小 dense 模型:参数表上还是 4 个专家, 实际容量缩水一半甚至更多 —— 花了 MoE 的显存,买到的容量白买,最亏的局面。 所以工业级 MoE(Switch / Mixtral / DeepSeek)全都带某种负载均衡手段。
↳ 代码:09_moe.py 的 MoELayer.forward 末尾 aux loss · 复现本组对照:python 09_moe.py --top-k 1 --aux 0
↳ 下一步:分工有了、均衡有了,回头算总账 —— 这笔交易到底赚在哪?

⑤ 算账:用 dense 的算力,买 2 倍的容量

MoE 的本质是一笔交易:总参数(容量)↑,每 token 激活的计算量不变。 拨动两个滑块,看这笔账怎么随 n_expert / top_k 变化(模型 = 本关的 4 层 GPT):

4
2
总参数(要存)
每 token 激活
dense 对照(03)
0.81 M
MoE 省的是算力,不是显存 —— 总参数一个都不能少,全都得放在(显)存里待命; 省下的是每个 token 的 FLOPs。所以它适合"显存管够、算力金贵"的大规模训练与推理。 真实世界同款账本:Mixtral 8×7B 总 46.7B / 激活 12.9B;DeepSeek-V3 总 671B / 激活 37B。
达标本关实测(四次训练 · 各 5000 步 · tiny shakespeare)
Q&A常见疑问
MoE 的"专家"和多头注意力的"头",都是"多个小的",区别在哪?
头是全员上岗,专家是按需点名。多头注意力里,每个 token 都要过所有的头,只是各头看不同关系; MoE 里每个 token 只走 top-k 个专家,其余专家对它零计算。 一个是"并行分视角"(算力照付),一个是"稀疏选路"(算力打折),目的完全不同。
现在的 MoE 还是"4 选 2"这么糙吗?
本页教的是 Mixtral 同款的经典配方(少量大专家 + top-2),机制上够用,但前沿已经演进了一步: 细粒度专家 + 共享专家(DeepSeekMoE 提出,DeepSeek-V2/V3、Qwen 系 MoE 都在用)。 做法是把专家切得更碎(比如 64 个小专家选 top-8,而不是 8 个大专家选 2), 组合方式从 C(8,2)=28 种暴涨到 C(64,8)≈44 亿种,分工可以更精细; 再留 1~几个共享专家,所有 token 必走,兜底"人人都需要的通用加工", 让路由出去的部分专攻差异化。原理和本页完全一致,只是"切多碎、留不留必选项"这两个旋钮拧得更讲究。
既然容量翻倍这么香,为什么不是所有模型都用 MoE?
代价在工程:显存占用照总参数算、多机并行时 token 要跨设备找专家(通信开销)、 负载不均会拖慢整批训练、微调也比 dense 娇气。小模型(显存不紧张)用 dense 更省心; MoE 的甜区是"参数规模顶到显存上限、还想再大"的场景 —— 所以你在千亿级模型上才总看到它。
↳ 跑法:python 09_moe.py · 对照塌缩 python 09_moe.py --aux 0 · 训完自动打印负载 + 分工报告
🎉 手搓 MoE · 通关
你已经亲手把 dense GPT 的 FFN 换成"路由器 + 4 专家",看到了分工的涌现、塌缩的风险、和"参数翻倍算力不变"的账本。下次再看 Mixtral / DeepSeek 的参数表,你知道那两个数字是怎么来的了。