分布式预训练模型并行

当模型参数继续增大，除了用流水线并行之外，还可以用张量并行(Tensor Parallel)来缓解，就是以前常说的模型并行(Model Parallel)。这里拿最基础的矩阵乘举例。对于激活值矩阵A和权重矩阵B相乘，考虑对B做列切分或者行切分，会有两种方案方案1：按列切分。A不变，这样每一列独立计算，最后结果concat在一起方案2：按行切分。A对应列切分，这样两部分独立计算，最后结果sum在一起

这一切的起源是2017年发布的Attention Is All You Need([4])论文，之后基于大量语料的预训练模型百花齐放，比如：BERT(Bidirectional Encoder Representations from Transformers):Google在2018年提出，创新性的双向预训练并行获取上下文语义信息，以及掩码语言建模（MLM）让模型更好地推断语义信息。它开创了预训练语言表示范式，对自然语言处理产生了深远影响。参数规模：110M到340MGPT(Generative Pre-trained Transformer):OpenAI在2018年提出，开创了仅使用自回归语言建模作为预训练目标而无需额外监督信号。它展示了通过无监督大规模预训练获得的语言生成能力,对研究与应用都带来重大影响。参数规模：1750亿Large LAnguage Model Approach（LLAMA）:Meta在2021年提出，首个开源模型。为构建更大规模、更通用的语言模型提供了系统化的方法与工具。参数规模：十亿到千亿

最后小结一下，本文主要有4部分内容分布式通信原语包括了点对点通信和集合通信的方法。其中集合通信包括了一对多的Broadcast和Scatter，多对一的Reduce和Gather，多对多的AllReduce、AllGather、ReduceScatter和AllToAll3D并行包括了数据并行(DP)、流水线并行(PP)和模型并行(TP)DP的优势是计算和通信效率都很友好，但是权重的显存不友好，每张卡都有一份PP的问题是要求mini-batch里面batch size足够大才能掩盖住流水线带来的overhead。batch size如果过大，会增大激活显存的占用TP的优势是权重显存非常友好，没有冗余。但是计算和通信效率不友好，通信量要求很大，在超出了一个island的时候性能下降很快。ZeRO针对数据并行显存占用大的问题，借鉴了Parameter Server的思路，提出了ZeRO-1，ZeRO-2和ZeRO-3的优化。其中ZeRO-2让每张卡只维护一部分的梯度和优化器状态，显存占用减少到原来的$1/8$，通信带宽保持不变Alpa鉴了AI编译器的思路，对3D并行进行建模，用自动化搜索的方式得到了仅次于手工最优的并行策略PS：由于笔者小A并没有亲手撸过上述内容的所有细节，大部分是通过研究代码和精读优秀文章的方式bottom-up总结而来，本质上是个拾人牙慧的知识搬运工，所以终究是纸上谈兵。因此希望各方有实际经验的大佬猛锤，思维碰撞才生火花，真理越辩越明。如果想了解transformer在NLP/多模态/AIGC的算法知识，分布式训练的知识，以及如何在TVM上做PTQ量化和部署，可以关注我aaronxic哟~