模型並行-Megatron-LM

1.原理

  針對transformer模型的模型並行方法，利用矩陣乘法的特性，將每個層切開，放入不同GPU中。

對於注意力層，是按注意力頭進行切分; 對於MLP層，是按矩陣列進行切分。整個執行流程中，每個層的需要進行資料通訊，因為對於矩陣切分後的乘法要得到最終的結果需要每個分塊的結果進行矩陣加。

2.缺點

  使用多頭注意力的transformer的頭數、MLP隱藏層大小需要能夠被GPU數整除。每計算一個層就需要進行all reduce，導致不同層之間的計算是同步的，且有較大的通訊開銷。