MoE 怎么实现？shape 怎么变？
capacity、load balancing loss、EP 的 all-to-all 是什么？

英文问法Implement an MoE layer and show the dimensions. · What is the capacity factor, and what happens to dropped tokens? · Derive the load-balancing loss. · How do dispatch and combine work under expert parallelism?

考过NVIDIA SWE（"Implement MoE and show the dimensions"，Glassdoor RVW105422785，年份未核实；Glassdoor 的另一条摘要里，同一场还问了 CUDA kernel optimization，追问 kernel fusion 和 launch overhead）· 美团（"手撕：用python实现moe"）· 阿里（"ep的dispatch和combine"）· 智谱（手撕 MoE 层的 dispatch）· 牛客（"推导MoE架构的负载均衡损失函数，如何避免专家坍缩"）

先弄清楚：每个概念是什么

MoE layer: 把 Transformer 里的 FFN 换成 \(E\) 个 expert FFN，再加一个 router，每个 token 只走其中 \(k\) 个。参数随 \(E\) 涨，每个 token 的计算只随 \(k\) 涨。例：Mixtral 8x7B，E = 8、k = 2。
Router 与 top-k gate: 一个 \(h \times E\) 的 linear 给每个 token 打 \(E\) 个分，softmax 后取最大的 \(k\) 个。输出是这 \(k\) 个 expert 输出按 gate 值的加权和。gate 值让 router 能拿到 gradient。
Capacity factor 与 token dropping: 每个 expert 一个 batch 最多收 \(C\) 个 token。超出的 token 被 drop，这一层不算，经 residual 直接传到下一层。没收满的位置 pad。
Load balancing loss: router 一开始偏爱某个 expert，它就学得更好、更被偏爱，最后少数 expert 包揽一切（expert collapse）。加一个 aux loss，分配越均匀越小。
Expert parallel（EP）: 不同 expert 放在不同卡上。dispatch：all-to-all 把每个 token 发到它的 expert 所在的卡；combine：反向 all-to-all 把结果发回原来的卡。

看图：token 怎么走，capacity 怎么卡

图 1：2 张卡，8 个 expert 每卡放 4 个。交叉的箭头是 dispatch、combine 里跨卡的那部分 token。

图 2：每根柱子是一个 expert 这个 batch 收到的 token（高度只是示意）。\(C = 1280\) 是平均负载 \(kT/E\) 的 cf = 1.25 倍。橙色虚线是超出 \(C\) 被 drop 的，灰色虚线是要 pad 的空位。

30 秒回答

An MoE layer replaces the FFN with E expert FFNs plus a router. The router applies softmax to x·W_r and picks the top-k experts per token; the output is the gate-weighted sum of their outputs. Each expert takes at most C = capacity factor × kT/E tokens; overflow tokens are dropped and pass through the residual. Left alone, routing collapses onto a few experts, so Switch adds α·E·Σ fᵢPᵢ, minimized by uniform routing. Under expert parallelism, dispatch is an all-to-all sending tokens to their experts' GPUs; combine is the reverse. Parameters grow with E, per-token compute only with k.

推导 / 关键公式

shape（T 个 token）

\(x\,[T,h] \to\) logits \([T,E] \to\) top-k 下标、gate \([T,k] \to\) dispatch \([E,C,h] \to\) expert FFN \(\to\) combine \(\to y\,[T,h]\)

\[ C = \text{cf}\cdot\frac{k\,T}{E} \]

load balancing loss（Switch）

\[ \mathcal{L}_{\text{aux}} = \alpha\, E \sum_{i=1}^{E} f_i\, P_i \]

\(f_i\)：top-1 分到 expert \(i\) 的 token 比例（argmax，不可导）；\(P_i\)：router 给 \(i\) 的平均概率（可导，gradient 走这里）。均匀时 \(f_i = P_i = 1/E\)，loss \(= \alpha\)；全挤到一个 expert，loss \(= \alpha E\)。Switch 用 \(\alpha = 10^{-2}\)。

参数 vs active 参数（Mixtral 8x7B）

h = 4096、SwiGLU 中间维 14336、32 层。每个 expert 3 × 4096 × 14336 ≈ 176M，每层 attention（8 个 KV head）≈ 42M，embedding + LM head 0.26B：

total = 32 × (42M + 8 × 176M) + 0.26B ≈ 46.7B
active = 32 × (42M + 2 × 176M) + 0.26B ≈ 12.9B

和论文的 47B、13B active 对得上，参数是 active 的 3.6 倍。

手撕代码（NumPy）

def moe(x, Wr, experts, k=2, cf=1.25):
    # x: [T, h]; Wr: [h, E]; experts[e]: [n, h] -> [n, h]
    T, E = x.shape[0], Wr.shape[1]
    p = softmax(x @ Wr)                         # router  [T, E]
    idx = np.argsort(-p, axis=1)[:, :k]         # top-k   [T, k]
    w = np.take_along_axis(p, idx, axis=1)      # gate    [T, k]
    w = w / w.sum(1, keepdims=True)
    C = int(cf * k * T / E)                     # capacity
    y = np.zeros_like(x)                        # [T, h]
    for e in range(E):
        t, j = np.nonzero(idx == e)             # tokens -> e
        t, j = t[:C], j[:C]                     # overflow dropped
        # dispatch [C, h] -> expert FFN -> combine
        y[t] += w[t, j, None] * experts[e](x[t])
    f = np.bincount(idx[:, 0], minlength=E) / T # token fraction
    aux = E * np.sum(f * p.mean(0))             # Switch aux loss
    return y, aux

必答点

router → top-k → dispatch → expert → combine，能写出每步 shape
\(C = \text{cf}\cdot kT/E\)：超出 drop（走 residual），不足 pad；cf 大少 drop，但多算多传
aux loss \(\alpha E\sum f_iP_i\)，均匀时最小；不加会 expert collapse
EP：dispatch、combine 两次 all-to-all，常是瓶颈，要和计算 overlap
参数随 \(E\)、计算随 \(k\)：Mixtral 47B 参数、13B active

追问

DeepSeek-V3 的 aux-loss-free 怎么做？

每个 expert 一个 bias \(b_i\)，只加在选 top-k 用的分数上，gate 值仍用原分数。每步看 load：过载的 \(b_i\) 减 \(\gamma\)，欠载的加 \(\gamma\)。不经过 loss，不干扰主目标（另留一个很小的 sequence-wise aux loss 防极端情况）。

router 为什么要 fp32？

Switch 说 bf16 会加剧 router softmax 的数值问题，没给更细的机制。全 bf16 训练直接发散，只把 router 内部 cast 到 fp32 就稳了，速度几乎不变。Megatron：--moe-router-dtype fp32。

延伸：kernel 太碎、launch overhead 大怎么办？

这题是自拟的延伸。按 expert 循环会发 \(E\) 个小 GEMM。先按 expert 排序 token（permute），用一个 grouped GEMM 算完所有 expert，再 unpermute（Megatron：--moe-grouped-gemm、--moe-permute-fusion）。

EP 的 all-to-all 怎么藏？

跨机 EP 时通信和计算差不多一样长（DeepSeek-V3 说约 1:1）。DualPipe 让两个 micro-batch 的计算和 all-to-all 交错（见 PP 页），DeepEP 之类的 dispatcher 去掉跨机重复 token、合并机内机间通信，见 EP communication。

来源Fedus et al. 2021（Switch）§2.1–2.2（pp.5–8，Eq. 1–6）、§2.4、Table 2（pp.8–10）· Lepikhin et al. 2020（GShard）§2.2（pp.5–6）、§3.3（p.11，AllToAll 分发）· Jiang et al. 2024（Mixtral）abstract、§2、Table 1（pp.1–3）、config.json · DeepSeek-V3 §2.1.2（pp.8–9）、§3.2.1（p.12）· Megatron Core: MoE

MoE 怎么实现？shape 怎么变？capacity、load balancing loss、EP 的 all-to-all 是什么？