分布式预训练模型并行主要包括以下内容:
如果想了解更多关于 transformer 在 NLP/多模态/AIGC 的算法知识、分布式训练知识,以及如何在 TVM 上做 PTQ 量化和部署,可以关注作者 aaronxic。
当模型参数继续增大,除了用流水线并行之外,还可以用张量并行(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哟~