聊聊序列并行Sequence parallelism

作者丨无恶不作来源丨https://zhuanlan.zhihu.com/p/653067104编辑丨GiantPandaCV随着大语言模型的发展，大模型的相关技术日益细化，在生成式 AI、蛋白质结构预测等多个领域开花结果。由此产生的技术痛点也越发明显，特别是在蛋白质预测、多模态基础模型、语音图像处理模型等需要前后文推理的模型，对长序列训练有明显的需求。与此同时，现有的分布式训练技术，主要分为：数据并行、张量并行、流水线并行等在数据维度正交切分的分布式训练技术。如下图0所示，分别对应对于 batch size，hidden_size 维度进行切分。对于长序列模型训练，由于 transformer 类模型 self-attention 部分序列长度和显存消耗呈现平方关系，上述三种均不能提供较为高效、稳定的训练支持。而对于长序列模型采用序列并行可以很好的解决显存溢出等问题，由此大模型训练技术在序列并行 ………………………………