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]。
