LLM並行訓練5-MoE並行

SunStriKE發表於2024-07-20
並行

前置知識

MOE(MixerOfExpert)

image-20240718194134029

moe的主要原理是替換attention層後的MLP層, 透過將不同型別的token按照門控單元計算出的機率分配給最大機率處理的專家網路處理, 對比單一MLP更適合處理複雜多樣化的資料集. 主要思想和整合學習感覺很像, 而且擴充套件性(遇到新的目標任務可以新增專家網路)和可解釋性(每個專家分開調整)都比較強. MOE前向步驟(以最簡單的top2 Expert為例):

  1. 門控網路, 輸入是attention的輸出, dim為(batch_size, tokens, emb_size), 輸出dim為(batch_size, tokens, experts_num)
topkgate_linear = nn.Linear(n_embed, num_experts) # 從emb_size->專家個數的對映, 根據這個線性層計算每個token進入各個專家網路的機率
logits = topkgate_linear(mh_output)
top_k_logits, top_k_indices = logits.topk(top_k, dim=-1)  # 從4個專家網路裡取top2,
zeros = torch.full_like(logits, float('-inf'))
sparse_logits = zeros.scatter(-1, top_k_indices, top_k_logits) #把除了top2剩餘的位置置-inf
gating_output= F.softmax(sparse_logits, dim=-1)   #softmax計算進入各個專家網路的機率
  1. 稀疏化Experts

這部分每個expert的網路結構都可以根據場景設計的不一樣, 因為在fp/bp計算的時候, 每個token都是隻進入了topk的網路進行計算, 剩餘的網路沒有計算. 大部分的引數都沒參與更新, 所以也被稱為稀疏化的dense. 因為這個特性也給moe網路的並行化改造提供了基礎.

expert負載不均勻問題

因為expert的機率純粹是訓練出來的引數決定的, 沒法用LALB類似的負載均衡策略強制使每個expert接收到的token是均勻的, 極有可能出現某幾個expert接收到了很多, 其他的基本沒啥token的問題..主要有以下這麼幾個解決辦法:

TokenBuffer: 給每個expert設定固定容量, 容量設定公式如下, 當這個expert收滿token後就不再接受token

\[𝑐𝑎𝑝𝑎𝑐𝑖𝑡𝑦=𝑚𝑎𝑥(\frac{𝑆}{𝐸}∗𝐾∗𝑐𝑎𝑝𝑎𝑐𝑖𝑡𝑦\_𝑓𝑎𝑐𝑡𝑜𝑟,𝑚𝑖𝑛\_𝑐𝑎𝑝𝑎𝑐𝑖𝑡𝑦) \]

\(E\)：expert_num \(S\) ：token數 \(K\): topK數

在deepspeed的實現裡處理token溢位的方法:

  • 這個token在top2的兩個expert只溢位了一個, 那麼把另一個沒溢位的expert的softmax權重設成1放到那個expert裡(這裡有點像hash線性探測那種方法haha)
  • 在兩個expert裡都溢位了, 那麼把這個token直接跳過expert透過殘差的方式直連到上層

Random Routing:

隨機路由的方法主要針對的是2nd的expert, 1st的直接發出去. deepspeed的2nd隨機選擇策略:

  1. 從隨機分佈中取樣expert_num個隨機數作為噪聲
  2. 把噪聲加到softmax的結果上, 另外把1st的mask掉(因為1st是必發的, 只需要再選一個最高的就夠了)
  3. 在剩下的裡面找一個最高的作為2nd expert.
  4. 因為隨機後的不一定兩個expert機率加和為1, 所以需要進行重新normalize \(P_0' = \frac{P_0}{P_0 + P_1}\) \(P_1' = \frac{P_1}{P_0 + P_1}\) 這裡算的機率用於把經過2個對應expert後的token結果進行加權平均.

輔助損失函式:

\[l_{\text {aux }}=\frac{1}{E} \sum_{e=1}^E \frac{c_e}{S} * m_e \]

  • \(C_e\) ：某expert的buffer中已經存下的token數量（該expert作為1st時接收到的token數）
  • \(m_e\) ：某expert的buffer中已經存下的token在該專家上的avg(weight)（token考慮範圍也是那些將該專家作為1st專家的token）, 加這個引數主要是為了讓這部分可導, 能夠進行bp

把這個loss加到主loss後面, 目標也是最小化這個輔助loss. 因為我們最理想的情況是每個expert作為1/2的1st和1/2的2nd, 而如果某個expert溢位時他作為1st的機率遠高於其他expert, 而溢位後根據上面處理溢位的方法會把2nd轉成1st, 使得1st變多. 所以我們minimize 輔助loss的時候其實就是讓網路學習時儘量避免溢位.

MOE並行(gshard)

施工中

參考

gshard論文: https://arxiv.org/pdf/2006.16668

Deepspeed-moe: https://github.com/microsoft/Megatron-DeepSpeed/blob/main/megatron/model/transformer.py

相關文章