MOE(Mixture-of-Experts)

一句话：训练多个专家模型组合成一个更大的模型，在前向过程中动态选择TopK个专家。

这样做可以在训练时也避免各个领域知识互相打架导致模型难以收敛。

相比一个很大的模型什么都会什么都不太对，也就是学杂了，拆成小模型然后需要哪个用哪个来保证样样精通更容易提高模型的通用性和准确性。

组成

Gate 就是一个全连接打个分投个票，根据当前input得到他对于各个专家的相关性，然后选取topk个专家

所谓专家，也只是很简单的MLP升维再降维（当然可以换成任何网络结构）

还有共享专家，就是不走Gate打分，相当于普通的MLP。但是相比普通的MLP，拆成多个专家后计算量少很多。

专家们算完结果，最后拿到一起用Gate打的分做个加权。

得到输出，给到下一层的RN+MLA

MOE

EP(Expert Parallelism)

DeepSeek V3 有256个专家+32个共享专家，MOE中Topk = 8。

在DeepSeek 的MOE 中，使用EP，把全部256个专家分散到所有卡上，每个卡各自仍保留完整的32个共享专家。

根据Gate的结果选取8个专家，然后All-to-All把hidden_state传到选中的专家（可能就在当前卡，可能在同一个机器，也可能在其他机器上）。算完后再All-to-All收集回原来的卡上。

EP

通信开销

每层都有2次 All-to-All 通信，使用DeepEP实现。其中把hidden state根据topk 发送出去叫dispatch，Dtype是FP8，MLP后收回来叫combine，Dtype是BF16

通信量计算

(1Byte+2Bytes)×32×9×7168÷50GB/s=120.96us(1\text{Byte} + 2\text{Bytes}) \times 32 \times 9 \times 7168 \div 50\text{GB/s} = 120.96\text{us}

Here, dispatch uses FP8 (1 byte), while combine uses BF16 (2 bytes), and the hidden size of each token is approximately 7K. The factor 9 indicates that each token is ransferred to 8 routed experts and 1 shared expert.

profile发现通信开销是大头。

Micro batch Overlap (TBO, two-batch-overlap)

一层可以分为几个部分

  • MLA
  • Gate + select
  • share expert
  • dispatch
  • MLP
  • combine

其中通信有dispatch和combine，而且combine的耗时是dispatch的两倍，是前面小后面大。计算量都不大，MLA略多于MLP，是前面大后面小。

EP

考虑把大块通信和大块计算重叠，那么把输入平均切成两个小batch，错开顺序使得一个batch的combine和另一个batch的MLA重叠， 一个batch的dispatch和另一个batch的share expert重叠。

TBO overlap

MLA切成了两个部分，一部分是计算QK+RMS Norm，另一部分是Attention。细分计算步骤更灵活的重叠。