第 10 章 LLMMoE混合专家

第 10 章 MoE：用更少的算力装下更多知识

本章来源：本文整理自 LLMs-from-scratch 仓库的 ch04/07_moe bonus 材料，原书为 Sebastian Raschka《Build a Large Language Model (From Scratch)》。

这章要解决什么问题

前面几章我们一直在跟"算力和显存"较劲：KV Cache 是用显存换速度，GQA/MLA 是把 K/V 压小来省显存。这一章要解决的是另一个更根本的矛盾：

我们既希望模型"肚子里的知识多"（参数多），又希望它"每一步算得少"（计算量小、显存省）。

这两件事天生打架——参数越多，每一步要动用的矩阵就越大，算得就越慢。混合专家（Mixture of Experts，简称 MoE）就是来打这个圆场的：它让模型的总参数变得很大（能装更多知识），但每个 token 实际只用其中一小撮参数（计算量小）。

一句话先记住：

MoE 的核心是"稀疏激活"——参数都在，但不是每个 token 都用得着。

这章我们就把这个"稀疏激活"从"为什么"讲到"怎么写代码"。

用大白话讲原理

前馈层才是参数大户

先在脑子里回想第 4 章搭过的那个 transformer 块：它内部有两块子模块——注意力（attention）和前馈网络（FeedForward，缩写 FFN）。

很多人以为注意力是"贵重部件"，其实真正把参数量撑起来的是前馈层。原因很简单：注意力干的是"词和词之间交换信息"，它的矩阵尺寸跟嵌入维度 $d_{\text{emb}}$ 同量级；而前馈层要先把向量"撑大"到一个中间维度（通常是嵌入维度的 4 倍左右），做完非线性变换再压回来——它经手的矩阵是 $d_{\text{emb}} \times d_{\text{hidden}}$ 量级，比注意力大一个档次。

更关键的是：这个前馈层不是出现一次，而是每一层都有一套。而 transformer 块在 LLM 里要叠很多层——以 DeepSeek-V3 为例，它叠了 61 层。所以前馈层对总参数的"贡献"，是被层数放大过的。

记住这个结论：在 transformer 里，前馈层是当之无愧的"参数大户"。 而 MoE 的一切，正是绕着这个大户做文章。

用多个"专家"替换一个前馈层

既然前馈层是参数大户，那我们想装更多知识，最直接的做法就是——把这个前馈层换成好几个前馈层。

这就是 MoE 的核心动作：

把 transformer 块里"一个"前馈层，替换成"多个"前馈层，其中每一个前馈层就叫一个"专家（expert）"。

现在关键点来了：一个换多个，总参数量当然会大幅增加——这正是我们想要的，因为参数多意味着"容量（capacity）大"，训练时能装下更多知识。

但要是每来一个 token，所有专家都得跑一遍，那计算量也跟着暴涨，等于白忙。于是 MoE 的第二个关键动作出场了——

不是每个 token 都走所有专家，而是只走其中一小撮。

这就是"稀疏（sparse）"的含义：模块里参数一大堆（稀疏），但每次前向只激活其中一小部分。

对比一下两种前馈：

稠密前馈（dense FFN）：每个 token 都激活全部参数，一个不落。
稀疏 MoE（sparse MoE）：每个 token 只激活一部分专家，其余专家这一轮"不上班"。

打个比方：普通前馈层像是一家"全员上阵"的小公司，每人每个单子都得签字；MoE 则像是一家有几百号专家的大集团，每来一个项目，只按需叫上最对口的几位，其余的照常待命。公司规模（总参数）很大，但每个项目的人力成本（计算量）只跟被叫来的那几个人有关。

路由器：每个词只走几个专家

那问题来了：来了一个 token，到底该派哪个专家去接？

MoE 的答案是：训练一个路由器（router）来决定。路由器其实是个很小的线性层，它给每一个专家打一个分数，然后我们只挑分数最高的前几个——这几个就叫 $\text{top-}k$ 专家。用术语说：

num_experts：这个 MoE 模块里一共有几个专家。
top_k（代码里常叫 num_experts_per_tok）：每个 token 只挑几个专家上班。

于是每个 token 的流程就是：

路由器给所有专家打分；
选出分数最高（最"对口"）的 $k$ 个专家；
只让这 $k$ 个专家处理这个 token，再把它们的输出按分数（权重）加权求和，拼成这个 token 的最终输出。

再套回那个类比：一个 token 就是一个项目，路由器就是前台，它看一眼项目类型，只叫最对口的那几位专家出来；剩下的专家这一轮完全不参与计算。

共享专家与稀疏性带来的性价比

最后一块拼图是共享专家（shared expert）。

你可能立刻会问：既然每个 token 都只挑一小撮专家，那些天天都用到的基础规律（比如语法、常见搭配）岂不是每次都散落在不同专家身上、要重复学好几遍？

DeepSeek 的做法是：专门留一个"共享专家"，它对每一个 token 永远在岗、永远激活。 这样常见、重复的规律就由这个共享专家统一兜底，其余专家就能腾出精力去学更专业、更细分的东西。

这个点不是 DeepSeek 首创：2022 年的 DeepSpeed-MoE 论文就发现，加了共享专家后整体建模效果更好——原因正是"通用规律不用每个专家各学一遍，专家们就有了更多空间去学专精模式"。2024 年的 DeepSeekMoE 论文也沿用了这个设计。

现在把账合起来看，MoE 的性价比就出来了：总容量（总参数）靠"多堆专家"做大，而每个 token 的计算量只跟被激活的那几个专家有关。 于是我们得到一个"总参数很大、每 token 计算量却很小"的模型。

DeepSeek-V3 就是最好的例子：

每个 MoE 模块有 256 个专家，模型总参数 671B（6710 亿）；但每个 token 推理时只激活 9 个专家（1 个共享专家 + 路由器选的 8 个），实际只动用约 37B（370 亿）参数，而不是全部 671B。

用大白话说：账面上 6710 亿参数量级的"知识库"，每个词只看其中 370 亿——知识在，但不用全部搬出来，省得又快又划算。

10.1 从 SwiGLU 说起：MoE 里的"专家"长什么样

在写 MoE 之前，得先确认"一个专家"具体是个什么模块。答案是：一个专家就是一个前馈网络（FFN），形状和我们第 4 章写过的 FeedForward 一模一样——两层线性层夹一个激活函数。

不过上游的 MoE 示例脚本用的是 SwiGLU 前馈，而不是 GPT-2 传统的 GELU。（顺带解释两个词：GELU 第 4 章见过，是"平滑开关"式激活；SwiGLU 是 Swish 门控线性单元，Swish 激活函数的典型代表就是 SiLU。SwiGLU 用一条"门控"分支来控制信息放行多少，效果通常更好。）

一个 SwiGLU 前馈，可以写成：

$\text{SwiGLU}(x) = W_3\big(\text{SiLU}(W_1 x) \odot (W_2 x)\big)$

其中 $\odot$ 是逐元素相乘。对照着记忆：我们第 4 章写过的传统 GELU 前馈只有两个矩阵（升维、降维），所以在内存估算里算 2 份参数；而 SwiGLU 是"内容 × 门控"这种三条分支的组合，对应下面三个矩阵：

fc1：把 $d_{\text{emb}}$ 升到 $d_{\text{hidden}}$ （门控里的"内容"分支 $W_1$ ）；
fc2：把 $d_{\text{emb}}$ 升到 $d_{\text{hidden}}$ （门控里的"放行"分支 $W_2$ ）；
fc3：把 $d_{\text{hidden}}$ 压回 $d_{\text{emb}}$ （输出投影 $W_3$ ）。

所以在内存估算脚本里，SwiGLU 每个专家按 3 份矩阵计入（这也是后面算参数量时 swiglu 对应 3、gelu 对应 2 的原因）。

10.2 一个 MoE 模块：专家、路由器怎么摆

现在看 MoE 模块的整体结构。上游的 MoEFeedForward 其实就是把"路由器 + 多个专家"装进一个模块里。注意：这里没有为每个专家单独定义一个类，而是用三个 nn.ModuleList 分别存所有专家的三段权重——第 i 个专家的 SwiGLU 就由 fc1[i]、fc2[i]、fc3[i] 拼成：

class MoEFeedForward(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.num_experts_per_tok = cfg["num_experts_per_tok"]  # 每个 token 选几个专家（top_k）
        self.num_experts = cfg["num_experts"]                  # 专家总数
        self.emb_dim = cfg["emb_dim"]

        # 路由器：一个不带偏置的小线性层，给每个专家打分
        self.gate = nn.Linear(cfg["emb_dim"], cfg["num_experts"], bias=False)

        # 每个专家的三段权重（SwiGLU 的 fc1 / fc2 / fc3）
        self.fc1 = nn.ModuleList(
            [nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], bias=False)
             for _ in range(self.num_experts)]
        )
        self.fc2 = nn.ModuleList(
            [nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], bias=False)
             for _ in range(self.num_experts)]
        )
        self.fc3 = nn.ModuleList(
            [nn.Linear(cfg["hidden_dim"], cfg["emb_dim"], bias=False)
             for _ in range(self.num_experts)]
        )

几个要点：

路由器 gate：把 $d_{\text{emb}}$ 维的 token 映射成 num_experts 个分数，一条分数对应一个专家。它的参数量只有 $d_{\text{emb}} \times \text{num\_experts}$ ，跟专家本体比小到可以忽略。
三个 ModuleList：fc1[i]、fc2[i]、fc3[i] 合起来就是第 i 个专家的 SwiGLU 前馈。专家多了，这里就是"参数大户"的所在。
注意 MoEFeedForward 里每个专家用的是 hidden_dim，但 MoE 演示里为了"总参数公平对比"，会把这个 hidden_dim 调小（下面算账时会看到）。

10.3 路由器与 top_k：每个 token 只挑几个专家

重头戏是 forward。它的逻辑分三步：打分 → 选 top_k → 只让选中的专家干活并加权合并。先看开头这一段（打分和选择）：

    def forward(self, x):
        # x 的形状: (batch, seq_len, emb_dim)
        scores = self.gate(x)  # (b, seq_len, num_experts) —— 路由给每个专家打分
        topk_scores, topk_indices = torch.topk(
            scores, self.num_experts_per_tok, dim=-1
        )  # 每个 token 只取分数最高的 top_k 个专家
        topk_probs = torch.softmax(topk_scores, dim=-1)  # 把这几位的分数变成权重（和为 1）

逐行看：

scores = self.gate(x)：路由器给每一个 token、每一个专家打分，形状是 (batch, seq_len, num_experts)。
torch.topk(...)：沿着专家维度，取出每个 token 分数最高的 num_experts_per_tok（也就是 k）个专家，拿到它们的分数和编号。到这里，每个 token 该找谁就已经定下来了。
torch.softmax(topk_scores, ...)：把这 k 个分数归一化成一组权重（加起来等于 1），后面按这个权重加权合并各专家的输出。

这里有个容易踩的坑：softmax 是作用在被选中的 topk_scores 上，不是作用在全部专家分数上。换句话说，权重是在"已入选的这几位"之间再分配，这样每个 token 拿到的输出权重之和恒为 1，不会因为专家数变化而漂移。

10.4 稀疏地跑一次前向：只算被选中的专家

接下来是真正体现"稀疏"的地方——只对每个 token 选中过的专家做计算。先把张量摊平，方便按 token 处理：

        batch, seq_len, _ = x.shape
        x_flat = x.reshape(batch * seq_len, -1)  # 把 (batch, seq_len, ...) 摊成 (总token数, emb_dim)
        out_flat = torch.zeros(
            batch * seq_len, self.emb_dim, device=x.device, dtype=x.dtype
        )  # 输出缓冲，先全部填 0

        topk_indices_flat = topk_indices.reshape(-1, self.num_experts_per_tok)
        topk_probs_flat = topk_probs.reshape(-1, self.num_experts_per_tok)

        unique_experts = torch.unique(topk_indices_flat)  # 这一批里"被叫到"的专家编号（去重）

关键在于 unique_experts：一个 batch 里成千上万个 token，其实只会叫到少数几个专家。与其让每个专家对全部 token 算一遍、再挑出自己那份，不如反过来——遍历"被叫到的专家"，让它只处理点自己的那些 token：

        for expert_id_tensor in unique_experts:
            expert_id = int(expert_id_tensor.item())

            # 找出这个专家的"候选 token"：哪些 token 的 top_k 里包含它
            mask = topk_indices_flat == expert_id
            if not mask.any():
                continue
            token_mask = mask.any(dim=-1)
            selected_idx = token_mask.nonzero(as_tuple=False).squeeze(-1)
            if selected_idx.numel() == 0:
                continue

            # 只把属于它的 token 喂给这个专家（SwiGLU：SiLU(fc1(x)) * fc2(x)，再 fc3 投影回去）
            expert_input = x_flat.index_select(0, selected_idx)
            hidden = torch.nn.functional.silu(self.fc1[expert_id](expert_input)) * self.fc2[
                expert_id
            ](expert_input)
            expert_out = self.fc3[expert_id](hidden)

            # 取出这批 token 用到的"入选权重"（每个 token 在几位入选专家中各占一个槽位）
            mask_selected = mask[selected_idx]
            slot_indices = mask_selected.int().argmax(dim=-1, keepdim=True)
            selected_probs = torch.gather(
                topk_probs_flat.index_select(0, selected_idx), dim=-1, index=slot_indices
            ).squeeze(-1)

            # 按权重把专家输出累加到输出缓冲（index_add_ 支持同一个 token 被多个专家累加）
            out_flat.index_add_(0, selected_idx, expert_out * selected_probs.unsqueeze(-1))

        return out_flat.reshape(batch, seq_len, self.emb_dim)

这段就是"稀疏"的落地：

mask / selected_idx：挑出"确实要经过这个专家"的那批 token；
expert_input：只把这批 token 送进专家，其他 token 根本不参与这个专家的计算；
hidden = silu(fc1(x)) * fc2(x)：这就是 SwiGLU 的门控写法——fc1 走 SiLU 当"内容"，fc2 当"门"控制放行多少，两者逐元素相乘；最后 fc3 投影回 $d_{\text{emb}}$ 。一个专家的完整前馈，就这三下。

index_add_ 这一步很妙：因为一个 token 可能被多个专家处理，用"索引累加"就能自动把几位专家的加权输出叠加到同一个 token 上，不需要额外循环。最后再 reshape 回 (batch, seq_len, emb_dim)，形状和稠密前馈的输出完全一致——所以这个模块可以原地替换掉 transformer 块里的前馈层，其余结构一个字都不用改。

10.5 共享专家：一个永远在岗的"通用专家"

前面 10.2 的 MoEFeedForward 并没有包含共享专家——它是标准的地道 MoE。共享专家是 DeepSeek-V3 这类设计的增强项，思路是：

额外挂一个专家，它无视路由，对每个 token 永远激活；最后把这个共享专家的输出，加到路由器选出的那几个专家的加权输出上。

为什么它有用？前面讲过——通用、重复的规律交给共享专家统一处理，各路专家于是能专注在自己擅长的细分模式上，整体建模效果更好。

用最小的伪代码表示这个增量的思想（示意用，非上游脚本原样），大致是这样：

# 共享专家：永远对所有 token 生效
shared_out = shared_expert(x_flat)          # 一个固定的 SwiGLU 前馈，不经过路由
out_flat = out_flat + shared_out            # 与路由专家的加权输出相加

这样对应到 DeepSeek-V3，就是"1 个共享专家 + 路由器选的 8 个 = 每个 token 激活 9 个专家"。

10.6 算笔账：MoE 到底省了多少

光看代码还不够直观，上游提供了一个内存估算脚本 memory_estimator_moe.py，专门把"稠密前馈"和"MoE"的参数量摆在一起比。先看它怎么算：

稠密前馈参数量： $\text{(矩阵份数)} \times d_{\text{emb}} \times d_{\text{hidden}}$ 。GELU 是 2 份，SwiGLU 是 3 份。
路由器参数量： $d_{\text{emb}} \times \text{num\_experts}$ 。
MoE 每 token 激活参数量： $\text{路由器} + \text{top\_k} \times \text{单专家参数}$ 。

脚本里就是这么直译成代码的：calc_ffn_params 把"矩阵份数 × $d_{\text{emb}}$ × $d_{\text{hidden}}$ "算出来（SwiGLU 取 3 份、GELU 取 2 份），calc_router_params 则返回 $d_{\text{emb}} \times \text{num\_experts}$ 。

现在跑一个例子：嵌入维度 7168、中间维度 14336、SwiGLU、8 个专家、每个 token 激活 2 个（top_k=2），并让 MoE 总参数和稠密前馈对齐（--match_dense）。跑出来的结果里有两组关键数字：

稠密前馈约 3.08 亿参数（0.62 GB），前向时会全部激活；
换成总参数几乎一样（也是约 3.08 亿，0.62 GB）的 MoE 后，每个 token 只激活 约 0.77 亿参数（0.15 GB）——大约只有原来的四分之一。此时脚本为对齐总参数，自动把单专家中间维度缩到了 moe_hidden_size = 1792。

也就是说，在总参数量持平的前提下，MoE 把每 token 的实际计算量压到了 $\text{top\_k} / \text{num\_experts} = 2/8 = 1/4$ 左右。这不是巧合，而是 MoE 内存收益的来源：

$\text{激活比例} \approx \frac{\text{top\_k}}{\text{num\_experts}}$

（前提是 MoE 的总参数和稠密前馈对齐；match_dense 做的就是自动把每个专家的 hidden_dim 缩小，让"专家数 × 单专家参数 + 路由器"刚好等于稠密前馈的参数。上面的例子里 moe_hidden_size = 1792 就是这个缩小后的单专家中间维度。）

换个角度看这条公式：在总参数固定的前提下，专家数越多、每个专家越小，激活比例就越低，省得就越多。这解释了为什么现代 MoE 动辄上百个专家——专家越切越细，"大容量、低激活"的性价比就越突出。

真实模型里的速度/显存代价

内存省了，代价是什么？上游还给了两个能实际跑起来的脚本 gpt_with_kv_ffn.py（稠密前馈）和 gpt_with_kv_moe.py（MoE），在同一套 GPT 上对比。为了公平比较，MoE 版要把单专家尺寸调小：设 32 个专家，就把中间维度从 32768 除以 32 降到 1024，并且每个 token 只激活 2 个专家。跑下来的结果是：稠密前馈平均每次调用约 0.759 ms、激活内存约 0.19 MB（峰值 0.75 MB）；稀疏 MoE 平均每次约 1.555 ms、激活内存约 0.04 MB（峰值 0.11 MB）。

对比结果很说明问题：

显存：稠密前馈的激活内存约 0.19 MB（峰值 0.75 MB），MoE 只留约 0.04 MB（峰值 0.11 MB）——大约省了 4~5 倍；
耗时：MoE 约 1.555 ms，稠密约 0.759 ms——前馈计算时间大约翻了一倍。

一句话总结这笔交易：MoE 用大约 2 倍的（前馈）计算时间，换来大约 4~5 倍的激活内存节省。 这里有额外开销的一部分原因是路由本身耗时，而且上游这个实现也不一定最优。

还有一个反直觉的现象：整个生成过程在两种情况下的 GPU 峰值显存都差不多（约 11.5 GB）。因为两版加载的权重数量一样、KV Cache 也一样大，而这俩才是显存的大头——MoE 省下的是中间激活，不是权重本身。另外，这里没有用 batch（为了示例简单），如果一次处理更多 token（batch 大于 1），稀疏带来的省显存效果会更明显。

关键概念

MoE（Mixture of Experts，混合专家）：把 transformer 块里的单个前馈层，替换成多个前馈层（专家），并由路由器只激活其中少数几个。
专家（expert）：MoE 里一个独立的子模块，通常就是一个前馈网络（本文里是 SwiGLU 前馈）。
路由器（router / gate）：一个很小的线性层，为每个 token 给各专家打分，决定该激活谁。
top_k（num_experts_per_tok）：每个 token 只激活的专家个数；num_experts 是专家总数。
稀疏（sparse）与稠密（dense）：稠密前馈每个 token 用全部参数，稀疏 MoE 每 token 只用一部分参数；稀疏带来"总参数大、单位计算量小"。
共享专家（shared expert）：对每个 token 永远激活的那个专家，负责兜底通用规律，让其余专家专注细分模式。
容量（capacity）：模型能装下的知识量，主要由总参数量决定；MoE 用"多堆专家"做大容量。
激活比例：约等于 $\text{top\_k}/\text{num\_experts}$ ，是 MoE 显存/计算收益的来源。
SwiGLU：一种门控前馈（ $\text{SiLU}(W_1x)\odot(W_2x)$ 再经 $W_3$ ），需要三个线性层，MoE 专家常采用。

动手练习

改配置看收益：用 memory_estimator_moe.py 固定 --emb_dim 和 --hidden_dim，把 --num_experts 从 4、8、16 一路调大（--top_k 保持不变），观察 MoE ACTIVE/Token 和 MoE TOTAL params 各自怎么变。验证一下那句话：专家数越多，激活比例越低、省得越多。
对齐总参数：给 memory_estimator_moe.py 加不加 --match_dense 各跑一遍，对比 Per-expert params 和 moe_hidden_dim，想清楚"要让 MoE 总参数和稠密前馈对齐，单专家的中间维度该缩到多少"。
亲手跑一次对比：按 10.6 的两条命令分别运行 gpt_with_kv_ffn.py 和 gpt_with_kv_moe.py（没有 GPU 也能跑，只是没有显存数据），亲眼看看"前馈耗时"和"前馈激活内存"这两个数字的变化。（提醒：模型没训练，生成的文本是乱码，这里只看性能数字。想看训练好的 MoE，可以翻作者在 ch05 里给的 Qwen3-MoE 示例。）
想一个设计题：如果让你给 DeepSeek-V3 那样的配置"记账"——256 个专家 + 1 个共享专家、每个 token 激活 9 个——你能用本章的激活比例公式，估算出"每 token 激活参数占总参数的比例"大约是多少吗？

第 10 章 MoE：用更少的算力装下更多知识

第 10 章 MoE：用更少的算力装下更多知识

这章要解决什么问题

用大白话讲原理

前馈层才是参数大户

用多个"专家"替换一个前馈层

路由器：每个词只走几个专家

共享专家与稀疏性带来的性价比

10.1 从 SwiGLU 说起：MoE 里的"专家"长什么样

10.2 一个 MoE 模块：专家、路由器怎么摆

10.3 路由器与 top_k：每个 token 只挑几个专家

10.4 稀疏地跑一次前向：只算被选中的专家

10.5 共享专家：一个永远在岗的"通用专家"

10.6 算笔账：MoE 到底省了多少

真实模型里的速度/显存代价

关键概念

动手练习

📑 从零构建大语言模型（LLMs from Scratch 通俗精讲版）