斯坦福大学 CS336 Lecture 04 Detailed Explanation of the MoE Architecture

tianjiangming_

205人浏览 · 2026-08-05 21:59:11

tianjiangming_ · 2026-08-05 21:59:11 发布

1. MoE Architecture

1.1 Basics of MoE

左图是 Original Transformer 架构，而右图是采用了 MoE 架构的 Transformer 模型。唯一的不同就是相较于 Original Transformer，MoE 架构将每层稠密的 FFN 替换为一组并行且结构相同但参数独立的多组 FFN 。同时引入了一个可训练的门控网络作为 Router，为每个 Token 选择 Top_K 个专家模型（Token Choice，见 2.1），将每个专家模型的输出按门控加权求和。

1.2 Problems of MoE

问题：存储这些专家模型需要显存。基本上没有一块 GPU 能够放下所有数据，所以要将这些专家模型分散放在不同的 GPU 上。所以数据传输又成了新的问题。

在传统的 Data Parallelism 中，每张显卡都存储着所有的数据参数，并行计算得到结构后只需同步一次梯度。而在 MoE 架构中，由于 Expert Parallelism，数据需要跟着专家跑。例如：Token “猫”在 GPU 0 上，Router 计算得出“猫”应该去 GPU 3 上的“动物”专家模型。GPU 0 必须把 “猫” 的特征向量通过 NVLink （见Lecture 01）打包发送给 GPU 3 ，计算结束后还要将结果传回 GPU0 。

一个 Batch 里有几千个 Token，每个 Token 可能去不同的 GPU。这就会产生 All-to-All 通信（全网全互联）。当专家数量很多时，通信时间甚至超过了计算时间，导致 GPU 算得再快，也只能干等着数据从网线传过来。

1.3 Different Kinds of MoE

除了传统的修改 FFN 的 MoE 架构，还有将 MoE 作用在 Atten 层的架构，但训练结果难收敛，训练不稳定：

2. Design Space of MoE

2.1 Routing Function

2.1.1 Routing Types

MoE 架构的难点：将 Token 经由 Router 到 Expert 的这个过程是离散的非连续的，无法求梯度

其他方法还有类似于强化学习方法（路由目标的选取本质上属于离散决策问题，而强化学习恰好擅长处理这类离散决策任务）以及最优传输算法等。但这些方法带来的庞大计算量远超过带来的好处。

2.1.2 Top-K Routing in Detail

（哥们公式怎么是从下往上写的）

$s_{i,t} = Softmax_{i}(u_{t}^{l^{T}}e_{i}^{l})$

$g_{i,t} = \begin{cases} s_{i,t}, s_{i,t}\in TopK(\left \{ s_{j,t}|1\leqslant j \leqslant N\right \},K) \\ 0, otherwise& \end{cases}$

$h_{t}^{l} =\sum_{i=1}^{N}(g_{i,t}FFN_{i}(u_{t}^{l}))+u_{t}^{l}$

2.1.3 Recent Variations

a. Fine-grained Expert (细粒度专家)

不把整个 FFN 当作一个专家，而是将 FFN 的中间层切成更小的切片，每个切片是一个“细粒度专家”。同时，用更细的粒度配合更多的激活专家（如 Top-6），实现更精准的路径组合。

细粒度专家的问题：当专家过多且需要激活大量专家时，router 要将大量 token 分发到众多计算设备，通信开销可能会急剧增加。

b. Shared Expert (共享专家)

在所有路由专家之外，额外设置一个或多个“共享专家”，它们永远保持激活状态，对所有 Token 无条件服务。（有点像必修课和选修课）

2.2 Expert Sizes

(应该就是超参数设置相关，都是经验之谈，跟着大家的设置走就好了)

2.3 Traing Objectives

难点：稀疏性（Top-K）保证训练效率，但稀疏门控决策是不可微的

解决方法：强化学习（Reinforcement Learning）、随机近似（Stochastic Approximations）、启发式平衡损失（Heuristic Balancing Loss）。

2.3.1 Reinforcement Learning

最早最符合直觉的方法，但效果一般

附一个对强化学习的理解：通过“试错”以及“奖惩反馈”来学习在什么状态下该进行怎样的决策，使累积奖励最大化。重点：是否有连续决策，决策改变未来状态，反馈延迟。

2.3.2 Stochastic Approximations

（老哥解释没咋听懂，自己找了点其他资料）

随机性会带来一些效率损失。

直通估计器方法（Straight-Through Estimator，STE）：无噪声扰动和 τ 退火训练，退化的Gumbel。由于没有噪声，更容易出现专家死亡的现象。

2.3.3 Heuristic Balancing Loss

a. Expert_Level 负载平衡损失：

$loss_{Expert\_Level} = \alpha \cdot N \cdot \sum_{i = 1}^{N} f_{i} \cdot P_{i}$

其中 f_i 是专家被选中的频率，即一个 batch 中，分配给专家 i 的 Token 数量占总 Token 数量的比例；而 P_i 是专家被选中的概率，即通过路由网络输出给专家 i 的 Softmax 概率；α 是一个用来控制惩罚力度的超参数。

Aux Loss 的本质是：在“追求任务精度（主Loss）”和“保持专家多样性（Aux Loss）”之间找一个平衡点。

b. 除了每个专家分配到的 Token 平衡之外，每块设备上分配的 Token 数量也应该相同。例如现在有一个包含 64 个专家的 MoE 模型，部署在 8 块 GPU 上，每块 GPU 负责计算 8 个专家。Expert_level 的负载平衡损失保证 64 个专家收到的 Token 数相近；但可能激活的专家恰好集中在某一块 GPU 上，其余设备闲置。

$loss_{Device\_Level} = \alpha _{2} \sum_{i=1}^{D}d_{i} \cdot Q_{i}$

原理类似，d_i 指设备 i 上所有专家收到的 Token 总数占总 Token 数量的比例；Q_i 指设备 i 上所有专家的路由概率之和（即该设备被“期望”分配到的概率）。

c. DeepSeek v3 对 Expert_Level loss 进行了创新，在原本基础上新加了一个偏置项。

$g_{i,t}^{'} = \begin{cases} & s_{i,t}, s_{i,t} + b_{i} \in TopK(\{s_{j,t}+b_{j}|1\leq j\leq N_{r}\},K_{r})\\ & 0, otherwise \end{cases}$

在传统的 Softmax 得分基础上，为每个专家各引入一个微调系数 b_i。在某次训练中，若某个专家分配得到的 Token 数量不够多，系统就会以某个学习率 γ 对 b_i 进行增量微调，提高其得分；反之进行减量调整。这个方法称为无辅助损失平衡机制，auxiliary loss free balancing

2.3.4 Training MoEs ---- the System Side

现有的几种并行方式：

a. Transformer 自带的 Token 并行和多头并行。

b. 数据并行：每张GPU（或每个进程）都拥有完整的、一模一样的模型副本（包括所有专家）。将整个Batch的数据切分成多份，每张GPU拿一部分数据进行前向和反向计算。缺点：显存瓶颈。每张卡都必须装下整个模型。

c.模型并行：

① 张量并行：将单个网络层（如FFN、Attention的权重矩阵）按列或按行切分到多张GPU上。计算时需要频繁通信来拼接结果。

② 流水线并行：按模型的层进行切分。例如：GPU 0 负责第1-10层，GPU 1 负责第11-20层。

d. 专家并行

① 将不同的专家（Experts）放到不同的GPU上。所有GPU依然拥有共享的专家以及非专家层（如Attention层、LayerNorm等）的完整副本

② 当单个设备部署了多个专家后，由于计算稀疏性，例如 Token 0 与 Expert 0 交互，Token 1 和 Token 2 分别与 Expert 1 和 Expert 2 交互——可看做三个小型系数矩阵乘法运算。

2.3.5 Stochasticity of MoE Models

即使把数学上的温度设为 0（消除算法随机性），也无法消除硬件拥堵导致的 Token 丢弃（工程随机性）

事故起因：路由失衡，多个 Token 涌向了同一专家。

处理方式：该专家所在的 GPU 显存溢出，于是直接舍弃超出的 Token。

数据结果：该 Token 没有进过运算，经过残差网络后就原封不动前往下一层。

最终结果：由于 Batch 里其他数据影响了设备能接收的 Token 数量，导致同样的输入得到不同的结果。

3. Issues with MoEs

3.1 z-loss

如果将路由分数看作是一个向量， $log\sum e^{x}$ 项大致相当于 “这个向量里最大的那个分数 + 一个小的冗余项”。所以这个 Loss 实际上在惩罚 “路由分数的绝对值过大”。（和 Lecture 3 里的 z_loss 原理上一样）

3.2 Upcycling

在一个已经训练好的稠密模型基础上，复制多份 FFN 作为初始专家，在此基础上继续训练。

优点就是训练计算量小很多，但缺点就是基础 FFN 网络相同，使得路由器很难区分各个专家，同时各个专家也很难学到独特的知识。

4. DeepSeek MoE v1-v2-v3

4.1 v1

总参数量 16B，实际激活参数量为 2.8B

专家组由两个共享专家以及 64 个细粒度专家构成，其中每次有 4 个专家激活，故每个 Token 会经过 6 个专家的运算。routing function 为Top-K 方法，损失函数包含了前文提到的 Expert_Level 和 Device_Level 的 z-loss 。

4.2 v2

总参数量 236B，实际激活参数量为 21B

模型架构无变化。相较于 v1 做了两层优化：

1. 正如之前所说的细粒度专家的问题，由于模型参数太大，为了减少通信开销，在直接采用 Top-K 方法进行 routing 前，首先对设备进行筛选，选出前 M 个最优设备；然后再在这 M 个设备上，进行 Top-K 。

2. 引入了通信均衡损失函数：确保每个设备发送和接收的Token数量大致相等。

设备均衡 vs 通信均衡：前者关心的是 GPU 设备的算力，后者关心的是 GPU 之间的网络通信带宽。注意计算力均衡不等于通信量均衡。

4.3 v3

4.3.1 MoE Part

总参数量 671B，实际激活参数量为 37B

模型架构无变化，相较于 v2 做了两层优化（但好像把通信均衡给删了）

1. Top-K 稍微有些许不同（从下往上看）

2. 见 2.3.3 c 无辅助损失平衡机制，auxiliary loss free balancing

4.3.2 Non-MoE Part

a. 多头潜在注意力（Multi-Head Latent Attention,MLA）:

对 KV_cache 进行的一种不同于 GQA 和 MQA 思路的优化：不直接存储原始的KV，而是将输入通过一个“下投影”矩阵，压缩成一个维度更低的向量 c_KV。当需要 KV 时，再将这个 c_KV 进行上投影。在实际推理时，缓存的是这个压缩后的c_KV，而不是庞大的KV，大大降低了显存占用。

问题：这样可能会引入新的计算量，因为原本不需要这个放缩过程。

解决：对于 c_K 上投影到 K 的矩阵，反正也是要和 Q 进行计算，Q 也是通过 W_q 和 x 计算得到的，那不如把这个上投影矩阵和 W_q 乘到一起，就可以节省计算。

$Q K^{T} = (x W_{q})(c_k W_{Up}^{k})^{T} =x W_{q} (W_{Up}^{k})^{T}(c_{k})^{T} = x W_{new\_q}(c_{k})^{T}$

至于 c_V 上投影到 V 确实引入了新的计算量。

新问题：RoPE 阻碍了该过程。RoPE 直接作用于 QK 上。带位置信息的计算无法合并：

$Q K^{T} = Rotate((x W_{q})) Rotate((c_k W_{Up}^{k})^{T} )$

$= x W_{q}R_{q} R_{k}(W_{Up}^{k})^{T}(c_{k})^{T}$

继续解决：将 K 分为两部分，一部分为内容部分（来自压缩的 c_KV）和位置部分（专门用来承载 RoPE）。前者可以用上述解决办法，后者只能老老实实先上投影，再继续计算。

b. MTP（Multi-Token Prediction）

让模型不仅学习预测下一个Token，还同时学习预测未来的多个Token。

亚马逊云科技技术品牌专区

更多推荐

基于STM32的全屋环境智能监测与预警终端设计

亚马逊云科技技术品牌专区

Keras：面向人类的深度学习库

Keras3支持多后端（TensorFlow、PyTorch、JAX），通过设置环境变量切换，兼容指定版本。pip安装后可用Sequential API快速构建CNN模型，如MNIST分类示例，实

亚马逊云科技技术品牌专区

AI行业日报｜2026-09-25：5个热点事件

本期行业动态聚焦于端侧算力重构、自主智能体工作流、开发工具经济模型、网络隔离架构设计与公共部门合规审查。以下结合一线技术场景进行结构化梳理，供 AI 研究员与软件工程团队参考。 PrismML 正将其轻量级语言模型适配至搭载高通处理器的智能眼镜终端 1 。该平台的核心战略指向开源权重模型的本地化运行，旨在最大化挖掘穿戴设备现有算力的利用效率 1 。对于嵌入式与边缘计算开发者而言，此类微型模型向低功