1.Megatron-LM簡介
由NVIDIA開發，提高大模型分布式并行訓練效率和線性度。加速訓練手段包括：數據并行DP,張量并行TP,流水線并行PP等。
本次主要分享其中的TP張量并行，結合代碼看看具體實現，下圖為Megatron-LM 張量并行思維導圖
2.TP并行原理及實現
張量我們有行和列兩種切分方式，按照我們分塊矩陣的計算的法則，將這這個變量A進行分別按行和列切分，如下圖所示。
其實通訊的代價還是很大的，所以在實際計算的時候，我們有多次連續的相乘，利用數學上的傳遞性和結合律，把數據通訊延后到最后一步，再進行一個計算同步和通訊，如下圖所示。
3.Megatron-LM TP代碼剖析
3.1 整體配置
transformer結構如下圖所示
整個模型配置具體實現過程如下圖所示。
實際上是在我們的sh腳本上面去做好整體的配置的，配置完之后，就會給到model_provider, 然后，通過觸發我們的model provider里面的GPT model, 通過這個方式，去注冊到整個網絡模型結構里面。
3.2 embedding并行
把一些語料輸進去，會通過一個embedding層處理，變成我們真正的hidden size或者影空間，再給到整個transformer layer，整個embedding層是通過注冊的方式，輸進去整個GPTModel，會調用到整個LanguageModelEmbedding, 這個LanguageModelEmbedding，真正的又會調到VocabParallelEmbedding, 就是先構建embedding層，就真正的去做一些并行切分的工作。
并行完之后真正輸進去給我們的網絡模型，是完整的shape [s,b,h]，當TP=2，其實切分成兩個NPU卡了，于是，在輸出的時候對output執行一個并行的策略allreduce，把剛才的兩個向量聚合起來，變成一個完整的向量，給transformer layer作為輸入。
3.3 LayerNorm并行
LayerNorm的方式或者位置有兩個，post和pre，也就是有個放在后面，有些放在前面，具體，是根據我們的網絡模型結構來去設置的，在我們的代碼里面，也會有一個具體的聲明，到底是放在前面還是后面，那一般來說，會在整個transformer 輸入之前放一個layerNorm, 是因為embedding之后數值有大有小，為了避免數值產生一個不均勻的情況下，會加一個layerNorm做一個歸一化，所以一開始的時候，加個layerNorm在前面。
整個layerNorm具體的實現，layerNorm在沒有開啟SP(sequnece并行)，沒有序列并行的時候，layerNorm是不需要做并行的，也是不需要切分的，所以它直接調一個算子TENORM,T專門針對transform做低比特或者做加速的一個庫，那這里面，就會調用TE里面的一個layerNorm, 或者TE里面的RMSNorm。
3.4 Attention并行
Z = Dropout(self-attention(XA),B)
Attention并行其實主要還是依賴于網絡模型的參數的。那網絡模型參數，就會通過這里面的設置，最核心的就是number-layers, hidden-size，num-attention-heads, 有了這幾個參數之后，在整個megatron LM里面，真正去執行的時候，是調用get_gpt_layer_local_spec這個注冊模型結構注冊進去的，在這里面，就會調用整個selfAttentionSubmodules來去實現我們的self attention層。最核心的就是么去做列的切分，實際我們的輸入是一個[s,b,h]，但是輸出是[s,b,3h]。因為我們會有QKV三個矩陣，然后再執行self attention。
整個函數里面，最核心的就是DotProductAttention，就會有各種各樣的QKV，而且輸進去的是一個[b,np,sq,sk]。首先，我們會利用DotProductAttention，將每個NPU里面的所的head進行合并，因為有非常多的head不可能每個單獨的去計算，所以，我們就會把整個[s,b,np,hn]映射成為[s,b*np,hn]，再輸進去我們整個結構里面去計算，那接著，很重要的就是第一個我們會去計算Q×K，然后，得到每一個attention的score，每一個attention score就是這里面，還沒有經過softmax，那這個時候我們的attention的score，就是[b*np,s,sk]，我們在真正在做softmax的計算的時候，又會把它做一次映射，映射成為[b,np,s,sk]這么一個shape。
接著，在整個DotProductAttention里面，再進行下一步的計算了，也就是每一個NPU上面的attention score去執行scale_mask _soft max, 我們就得到了整個attention的probs，這個時候的概率的一個shape是[b,np,s,sk]，再乘以一個V，因此把V映射成為[n,np,s,hn]，然后，跟剛才的attention_probs這一個輸出，進行一個具體的計算，得到我們整個attention probs*V再進行真正的輸出。為了讓計算更加方便，所以對shape變換，變成[s,b,hp]這么一個結構，進行一個輸出。
剛才講到的是最核心的core attention，也就是我們中間的這個模塊，真正的輸出變成[s,b,hp]了，接著，我們還是需要進行下一步的計算的，就是QKV之后，我們就進行一個列并行的切分了，那輸入的是[s,b,hp]，接著，我們輸出是[s,b,h]，現在，我們回到整個transformer結構的注冊里面，這個selfAttentionSubmodule里面。首先QKV的計算，是經過一個列切分的，接著去計算剛才講到的core attention 是怎么去計算的，那算完之后，再給到下一個線性層的時候，就需要執行一個行切分了.結合上圖來說，X按照列進行切分，給到兩個npu，權重A按照行切分為A1,A2，兩個NPU的輸出Y1Y2，其實是一個[s,b,h]，然后在真正輸出之前，需要執行一個allgather, 把它們兩個聚合起來，把[s,b,h]真正的輸出出來.
self attention之后，就給到dropout了，dropout之前還有一個add&norm相關的一些操作，所以整個注冊機制里面，我們就來到了在下一層了，get_bias_dropout_add相關的內容，因為整體的self attention的輸出[s,b,h]作為整個dropout的輸入，那dropout的輸出[s,b,h]作為整個attention結構的整體的輸出。實際代碼，使用的是一個大的融合算子，也就是get_bias_dropout_add。
3.5 Add&Layer Norm
再往上是有一個Add的操作，會把剛才Embedding輸入或者我們上一層的輸入，做一個殘差的連接，給到我們的Add，然后再給到下一個LayerNorm, 實際上去實現的時候，也是一個融合算子.
3.6 MLP并行
Z = Dropout(GeLU(XA),B)
MLP部分是通過get_gpt_layer_ammo_specific做一個注冊的，去構建我們的整個mlp的Submodules來去真正的去執行。真正執行的時候，我們還是回到了一個行切分這么一個計算的方式，整個MLP的輸入是[s,b,h]，輸出是[s,b,4h]。兩個LayerNorm，第一個LayerNorm就是變4倍，第二個LayerNorm就把它縮回來,整體的輸出保持一致.
第一個MLP的輸出是[s,b,4h]，因為要放在兩張卡，所以真正的輸出又還會再切一部分，就變成[s,b,2h]，給到Gelu進行一個計算，計算完之后，就來到了第二個線性層了。那接著，每一個線性層也是一樣的，我們的輸入按列切分，而里面的權重就按行切分，然后得到Y1Y2的輸出，Y1Y2的輸出就是[s,b,h]，經過一個allreduce的操作，把最后的輸入輸出統一成[s,b,h]。
