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網路的並行化改造提供了基礎. 這塊裡使用了enisum(愛因斯坦求和)簡化表達矩陣運算

#為了方便後面計算, 把gating_output reshape成(batch_size * tokens, experts_num)
reshaped_gating_emb = mh_output.reshape(-1, emb_size)
#使用如下簡寫表示dim S: batch_size * tokens E: experts_num  C: expert_buffer  M: emb_size
#這一步把token填充到experts的buffer裡, 而且記錄每個token在buffer裡的填充位置, 以上圖的token T0為例, 對應的(E,C)矩陣為:
#[[P0, 0, 0, 0],
#[[P1, 0, 0, 0],
#[[0, 0, 0, 0],
#[[0, 0, 0, 0]] mask用於表示buffer裡的相關位置有沒有被zero_padding.
combine_w, dispatch_mask = Top2Gating(gating_output) # combine_w: (S, E, C)  dispatch_mask:(S, E, C)
# 按順序把emb填充到每個expert_buffer裡, mask裡為false的時候說明是padding, 跳過
dispatched_expert_input = einsum("SEC, SM -> ECM", dispatch_mask, reshaped_gating_emb)  #dispatched_expert_input: (E, C, M)
# 經過專家網路的FFN 前向計算
h = enisum("ECM, EMH -> ECH", dispatched_expert_input, Wi)  #experts的Wi層，尺寸為(E，M, H)，
h = relu(h)
expert_outputs = enisum("ECH, EHM -> ECM", h, Wo)  #experts的Wo層，尺寸為(E, H, M)

# 對token的top2機率進行加權求和
outputs = enisum("SEC, ECM -> SM", combine_w, expert_outputs)
outputs_reshape = outputs.reshape(input.shape) # 從(S, M)變成(seq_len, batch_size, emb_size)

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)

1. EP+DP

以2機16卡為例, 如果想使用4份資料並行(ep_dp_world_size), 每套專家可以被平均分為4份為例(ep_world_size), 這裡其實就是\(e_0+ e_1+e_2+e_3\), 那麼切分結構如下圖:

image-20240719204611622

在切分的時候, 主要需要遵循幾個原則:

  1. 在資料並行的時候, 我們儘量使同一套專家切分後儘量分佈在一臺機器內, 這樣在一個batch內對所有專家進行fp/bp時不需要跨機通訊, 經過專家網路fp後需要進行一次卡間all2all, 把每張卡的expert裡算的其他卡的token給傳回去進行加權求和
  2. 這張圖裡的\(e_0\)不代表只有一個專家, 而是1/4的專家, 比如專家總量為32時, 就代表著有8個專家在一張卡里
  3. 多套專家之間和傳統的DP處理方式一樣, 每套專家fp輸入不同batch的資料, 在bp時allReduce梯度

2. EP+DP+TP

對比上面的EP+DP, 其實就是多了\(e_0\)專家的W全部進行了縱切, 使得卡1,2可以進行張量並行. 在這個場景下網路計算的主要邏輯如下(以Node1為例):

  1. 在前面的attention層因為也採用了TP的方式, 而在同一個TP組內的兩張卡[g0, g1]...輸出在allReduce後是一致的.
  2. [g0, g2, g4, g6]和[g1, g3, g5, g7]內各自做1次all2all(這裡其實是處理相同的2份輸入token)，將token發給對應的expert進行fp計算。
  3. [g0, g1], [g2, g3], [g4, g5], [g6, g7]這幾個tp組各自透過AllReduce取得完整的輸出結果
  4. [g0, g2, g4, g6]和[g1, g3, g5, g7]進行ep_group all2all，把expert fp計算完畢的emb傳送回對應的卡用於加權求和，同時處理完成的emb可以直接輸入下一個attention層不再需要集合通訊.
image-20240719211544426

一般上EP+DP+TP就能滿足視訊記憶體限制, 在MoE訓練裡不會再引入流水線並行.

3. Deepspeed all2all最佳化

分層(Hierarchical) all2all

這個最佳化其實zeropp裡提過, 只是少了其中的量化反量化的環節. 通訊量從\(O(p) \Rightarrow O(G+p/G)\), p: 卡數 G: 機器數

image-20240720113426462

基於TP的all2all最佳化

基線的all2all在EP+DP+TP這節裡簡單講過, 因為g0和g1是在同一個TP單元裡, 當non-MoE的結果allReduce之後兩張卡上的結果是完全一樣的. 所以我們如果直接用all2all通訊時其實有一半的通訊是完全冗餘的.

image-20240720113859643

為了解決上面的冗餘資料通訊的問題, 最佳化後的all2all主要有以下幾步:

  1. 根據機器間的關係把資料切分重排序. 為啥這裡g0要把B和C換位置呢? 可以觀察下上面baseline是咋處理的, B傳給了g1, C傳給了g2, 而實際上g1自身就有一份一樣的B並不需要通訊. 只需要讓g0把C傳給g2就可以了. D由g1傳給g3.
  2. 對tp rank相同的卡組成一個新的all2all group, 比如[g0, g2], [g1,g3], 然後拿重排序的一半資料進行all2all通訊.
  3. 機器間的all2all完成後, 因為TP內部每個卡目前持有一半資料, allGather後再恢復順序, 就能完成最終的計算
  4. 通訊量從\(O(p) \Rightarrow O(p/L) + O(L)\) p:卡數 L:TP並行數

注意: deepspeed在實際實現時因為要把原來的4卡all2all改成2卡, 沒法和EP切分數保持一致導致改動成本較大. 所以把gating_output(E, C, M)的第二維切分成TP並行數. 然後還是正常進行4卡all2all, 就解決了冗餘通訊的同時還能減少改動成本.

image-20240720114231994

參考

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

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

einsum簡介: https://zhuanlan.zhihu.com/p/542625230

Deepspeed-moe論文: https://arxiv.org/abs/2201.05596

moe並行部落格: https://zhuanlan.zhihu.com/p/681154742

相關文章