MOE

• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding

  • 模型設計：

    

  特點:

  (1) 專家容量 expert capacity

  • 強制規定了每一個expert所負責處理的token數量有一個最大值，這個最大值就叫專家容量；（2N/E）

  （2）分組分配 Local group dispatching

  • 給所有輸入token分成了G組，不同的組并行處理，每個組相應地也把組內專家容量變成2N/EG。這樣做相當于在前向推理時，把大的batch拆分成的 batch，每個小的batch就是一個group。這樣做的好處是通訊的時候只需要在每個group內進行就可以了，減少了通訊量。而進行反向計算的時候這些group可以合起來一起用，相當于進行了gradient accumulation。

  （3）輔助損失函數 Auxiliary loss

  • 損失函數如下：
    $$l_{aux}=\frac{1}{E}\sum _{e=1}^E\frac{c_e}{S}{m}_e$$
    $ S$是token數， $E$是專家數， $c_e$是分配給第$e$ 個專家的token數， $m_e$是第$e$ 個expert在$S$ 個token中獲得的平均權重;

  （4）隨機路由 Random routing

  • 如果top-1專家的權重很高，而第二個專家的權重如果較小，那很有可能只用第一個專家就足夠解決問題了。

• Switch Transformer: Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

  • 模型設計

  

  • 特點：

    (1) 專家容量 expert capacity
    $$expertcapacity=(\frac{tokensperbatch}{numbersofexperts})?capacityfactor.$$
    (2) 負載均衡損失
    $$loss=a?N?\sum _{i?1}^N{f}_i?P_i$$
    $f_i$表示被分配到第 $i$個expert的token數
    f i = 1 T ∑ x ∈ B 1 { a r g m a x ( p ( x ) ) = i } f_i = \frac{1}{T}\sum_{x \in B}1\{argmax(p(x)) = i \} fi?=T1?x∈B∑?1{argmax(p(x))=i}
    $P_i$表示整個batch每個token分配給第$i$ 個expert的概率的總和
    $$P_i=\frac{1}{T}\sum _{x\in B}{p}_i(x)$$
    負載損失函數設計的思路與上述GShard 完全一致，只不過系數不太相似。

  • 主要工作：

    （1）MoE to dense：把訓出來的效果較好的MoE模型蒸餾到dense模型，在壓縮MoE模型99%的參數的情況下，效果還是比直接訓練dense模型好。

    （2）訓練和微調技術：

    • 首次使用bf16成功訓練MoE模型：

      routing function內部使用單精度，其他部分使用半精度。

    • 更適合MoE結構的模型初始化

    • 增加的專家正則化，改善了稀疏模型的微調和多任務訓練

  ? （3）訓練框架：結合數據、模型和專家并行性，訓練了超過1T參數的MoE模型。

• ST-MoE: DESIGNING STABLE AND TRANSFERABLE SPARSE EXPERT MODELS

  • 思想： 穩定訓練

    1. 乘性操作對模型穩定性和效果的影響

       • GELU Gated Linear Units
       • RMSNorm
       • 增加dense層
       • 增加一個bias

       2. noise對模型穩定性和效果的影響

          給router的輸入logits乘以一個在[1e-2, 1e2]之間的均勻隨機變量來添加噪音。

    2. 限制激活值和梯度值對模型穩定性和效果的影響

       提出新的損失：router z-loss損失
       $$L_z(x)=\frac{1}{B}\sum _{i=1}^B(log\sum _{j=1}^N{e}_j^{(i)}{)}^2$$
       $B$ 是token的數量， $E$是專家數， $x\in R^{B?N}$是router的輸入。

    ? z-loss會對進入router的較大的logits值進行懲罰，以達到盡量減少進入指數函數的較大誤差的目的。