论文阅读笔记——Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsit

Switch Transformers 论文
Switch Transformer 在 Mixtral of Experts 的基础上，简化路由机制，仅选择一个专家，采用动态分配 token，结合数据并行+模型并行+专家并行，预训练速度相比 T5 提升 4-7 倍，支持万亿参数规模。

• 专家: 分布在不同的 device 上, 每个专家是一个 FFN 网络, 有独立的权重；
• 专家容量: 每个专家处理的 batch 大小, 公式为 $\frac{\text{tokens per batch}}{\text{num experts}}\cdot \text{capacity factor}$
• 容量系数: 计算专家容量时乘上系数, 可以为专家多分配一些 buffer 来改善 token 溢出的情况
  如果专家容量系数设置过小,，太多的 token 被路由到一个专家上,，会造成溢出的 token 走了残差分支直接传给下一层； 如果专家容量系数设置过大,，会造成内存和计算的浪费
  $$\text{expert capacity}=\left(\frac{\text{tokens per batch}}{\text{number of experts}}\right)\times \text{capacity factor}.$$
  对于 Switch Transformer，辅助损失会被加到总损失上，给定 i 从 1 到 N 的 N 个专家，一个 batch $\mathcal{B}$ 中有 T 个 token，损失计算如下：
  l o s s = α ⋅ N ⋅ ∑ i = 1 N f i ⋅ P i f i = 1 T ∑ x ∈ B 1 { a r g m a x p ( x ) = i } P i = 1 T ∑ x ∈ B p i ( x ) \begin{aligned} loss=\alpha · N · \sum_{i=1}^Nf_i·P_i \\f_i =\frac{1}{T}\sum_{x\in\mathcal{B}}1\{argmax\quad p(x)=i\} \\P_i=\frac{1}{T}\sum_{x\in\mathcal{B}}p_i(x) \end{aligned} loss=α⋅N⋅i=1∑N​fi​⋅Pi​fi​=T1​x∈B∑​1{argmaxp(x)=i}Pi​=T1​x∈B∑​pi​(x)​
  最终 $\alpha$ 被设置为 $\alpha =10^{-2}$ ，是因为这个值足够大，能够有效确保负载均衡；同时又足够小，不会对主要的交叉熵损失函数产生过大干扰。
  
• 数据并行：数据并行时，在前反向传播过程中没有通信，只在优化器阶段进行 core 的同步，weight 权重在所有 core 上复制一份，数据是所有 core 计算使分别处理一部分；
• 模型并行：模型并行时，weight 权重被 16 个 core 切分成 16 块，每个 core 上保留一部分；每个 core 都会处理全量的 $\mathcal{B}$ 个 token 数据；由于在 $d_{ff}$ 维度上进行了模型切分，计算 FFN 第二个矩阵乘法时（$ReLU（h）W_{out}$），每个 core 都会发送一个 $[B,d_{model}]$ 的数据，在前向和反向传播过程中都会调用 all reduce 操作；
• 数据并行/模型并行混合: 在上图中第三列, 每个 core 会分到 B / n 个token处理, 以及会分到 $d_{ff}/m$ 的权重的大小与中间 activation 产出. 在前向和反向通信时的 tensor 大小 $[B/n,d_{model}]$​这时 cores 会被划分为 4 组， 每一组 core 上有一份全量的模型权重，4 组模型权重做数据并行,，4组模型组内模型切分做模型并行;；数据切分在不同组模型处理数据不一样(用 4 种颜色区分), 在一个组内的 4 个 core 做模型并行, 数据会复制多份处理(例如一份蓝色数据复制成了多个小块)；
• 专家并行/数据并行混合：专家并行时每个专家的权重都不一样, 互相独立，所以模型权重颜色都不一样；下方数据并行，共同处理同一份数据，也就是一个大的蓝色块，跟第一列一样。这里 data 切分的是按数据并行度 n 来进行的，每个专家的输出 shape 大小是 $[n,B/n,E,C]$,，然后输出会跟输入 tensor 进行相乘 $einsum([n,B/n,dmodel],[n,B/n,E,C,dimension=[B/n])$ ​，得到结果大小为 $[n,E,C,d_{model}]$。 每个 core 都有各自的 expert, 对结果 $[E,C,d_{model}]$ 进行 all to all 操作, 从 n 维度切分变为 E维度切分。
• 专家并行/数据并行/模型并行: 从模型权重切分上, 与数据并行/模型并行混合一样, 区别在于 expert 专家权重各不相同，所以有 4 个不同色块，每个色块中 4 个 core 来进行模型切分；下面数据切分上, 与数据并行/模型并行混合一样, 每组处理不同的数据。 为了最优的模型设计，寻求 FLOPS 与模型参数的平衡, 增大专家数能增大参数量, 但 FLOPS 没有增长。为了增加 FLOPS, 增大了$d_{ff}$​ 的大小, 对应模型并行度 m 也会增大, 由于 N = n × m 是个定值, 所以 n 会对应减少。