无锡网站建设mkdns,seo论坛的网址,云主机服务,去长沙旅游攻略及费用本专题介绍常见的深度学习分布式训练的并行策略和背后使用到的集合通信操作#xff0c;希望能帮助读者理解分布式训练的原理#xff0c;以及集合通信之于分布式训练的重要性和必要性。鉴于篇幅限制#xff0c;将拆分成三个部分展开讲述。 在上两回中#xff0c;我们介绍了D…本专题介绍常见的深度学习分布式训练的并行策略和背后使用到的集合通信操作希望能帮助读者理解分布式训练的原理以及集合通信之于分布式训练的重要性和必要性。鉴于篇幅限制将拆分成三个部分展开讲述。 在上两回中我们介绍了DPPPTPEP等多种并行策略及其通信模式详情请参见LINK和LINK。第三部分将会介绍一些更高阶的并行方式如序列并行SP上下文并行CP混合序列并行UlysessZeRO系列并行优化策略完全分片数据并行FSDP。并且在文章最后将汇总所有介绍过的并行方案与它们的通信模式帮助读者初步建立起分布式训练与集合通信的知识结构体系。
序列并行/上下文并行/混合序列并行
序列并行可以看作是先前介绍过的数据并行在Transformer大语言模型下的延伸和扩展它是更加细粒度的数据并行。大家知道Transformer语言类模型所处理的数据对象是序列Sequence一个batch中包含多个序列一个序列中包含多个token每个token由一个向量表示。序列并行Sequence Parallelism, SP即为将一个序列切开分片段到多个节点上并行处理的策略其主要目的是可以摆脱单卡存储限制训练超长上下文的大模型。
根据现有文献描述序列并行可以被部署在Transformer模型的两个的阶段中
一个是Attention阶段最初来自Colossal-AIMegatron-LM后来也提出了相似概念将之称为上下文并行Context Parallelism, CP下文我们将之统称为Attention阶段的序列并行除此之外还有DeepSpeed提出的混合序列并行Ulysess也在Attention阶段部分使用到了序列并行另一个是LayerNorm与Dropout阶段来自Megatron-LM。
虽然都是序列并行但由于处于不同的运算阶段它们的行为以及影响是不同的。前者的目的主要是减少数据存储压力打破模型输入序列长度sequence length的限制而后者则是为了与LayerNorm与Dropout前后相邻阶段的张量并行TP搭配使用减少存储压力。
下面详细介绍这两种序列并行及其通信模式。
1Attention阶段的序列并行
下图展示了Attention部分的基本计算范式以single-head为例。可以看到一个单头的attention模块主要涉及到5个矩阵乘法运算分别是(1)(2)(3)输入token矩阵I与权重矩阵相乘得到的Q/K/V矩阵(4)Q矩阵与K的转置矩阵KT相乘得到的A矩阵图中仅画出A矩阵经过softmax转换后得到的A′矩阵以及(5)V矩阵与A′矩阵相乘得到的O矩阵。
输入token矩阵I中的每一列代表一个token如图中红色箭头所示。所以序列并行也可以看作是矩阵I的列并行——将矩阵I中不同token的列向量拆分开并行计算以减少单节点上针对序列长度的存储压力。
下面我们来看将输入输出矩阵I、O按照列token的维度拆分开放在两个节点上并行计算即序列并行SP2对以上这5个矩阵乘法造成的影响。 对于矩阵乘法(1)(2)(3)输入矩阵I按列被分成和放置于两个节点上权重矩阵维持不变即每个节点上都需要复制一份计算得到的结果矩阵Q/K/V在不同节点内自然也会按照列被分开对于矩阵乘法(4)为了注意力矩阵A′可以按照列的维度分开在两个节点上矩阵Q即为本节点在矩阵乘法(1)中得到的Q而矩阵K需要收集其他所有节点在矩阵乘法(2)中得到的K对于矩阵乘法(5)为了输出矩阵O可以按照列的维度分开在两个节点上矩阵A′即为本节点在矩阵乘法(4)中得到的A′而矩阵V需要收集其他所有节点在矩阵乘法(3)中得到的V。 总结一下就是Attention阶段的序列并行在模型训练的前向传播中主要涉及到的集合通信操作有二一是计算注意力矩阵A′时对于所有序列并行节点上矩阵K的AllGather操作二是计算输出矩阵O时对于所有序列并行节点上矩阵V的AllGather操作如下图所示。值得提及对于Causal Attention不是所有的K和V的分块都需要被每个节点用到这使得通信不均衡效率低。解决方案一般是将token做互补分组这里不做细节展开感兴趣读者可以参见论文Striped Attention: Faster Ring Attention for Causal Transformers。 前向传播中K与V的AllGather对应反向传播中K与V梯度的ReduceScatter其原因可以用我们在系列推文二中的列并行级联原理来解释。
以K为例Attention阶段的SP中涉及到K的两次矩阵乘法(2)(4)其实可以看作是两级列并行的级联第一级是矩阵I的列并行第二级是矩阵Q的列并行。列并行级联之间的通信模式为前向AllGather反向ReduceScatter如下图。V同理。
前向传播AllGather图中的MM指代矩阵乘法 反向传播ReduceScatter 除此之外Attention阶段的序列并行涉及到的通信操作还有将权重复制到所有卡上的broadcast操作和最后不同SP卡上权重梯度同步的AllReduce操作与DP类似一个epoch进行一次由于每张卡中都有一套权重的copy。 以上就是Attention阶段序列并行的基本原理了以下再做四点补充
SP与TP的对比Ring Self-AttentionMegatron-LM的 CP与Colossal-AI的SP的差别DeepSpeed的Ulysess方案 第一Attention阶段的序列并行与TP的对比Attention阶段除了SP属于数据并行还有一种通用的并行方式是TP属于模型并行如下图。TP将Attention阶段的计算按照Multi-Head中Head的维度分开涉及到的通信主要是正向一次的各Head输出矩阵的AllReduce和反向一次的输入矩阵梯度的AllReduce而SP是按照输入数据的序列长度的维度分的涉及到的通信主要是正向一次的KV矩阵的AllGather与反向一次的KV矩阵梯度的ReduceScatter。值得提及SP在Multi-Head与Single-Head情况下的通信行为一致由于Head没有被分开不涉及跨Head的通信行为。 图片来源论文Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
第二Attention阶段的序列并行的Ring Self-AttentionRing Self-Attention是一种Attention阶段SP的常见实现方法。最早由Colossal-AI提出如下图。它的本质思想是将Ring AllGather的通信步骤拆分开用P2P的Send和Receive替换与计算任务做流水掩盖。具体来说即在分步完成K(V)矩阵AllGather操作的时候同步进行注意力矩阵A(输出矩阵O)的计算。当计算一个A/O分块的时间大于传输一个K/V分块的时间就不会引入额外的通信开销。
图片来源论文Sequence Parallelism: Long Sequence Training from System Perspective 第三Megatron-LM提出的 CP与Colossal-AI提出的SP的差别CP在SP之后被提出其主体思想与Colossal-AI的SP一致都是Attention阶段的序列并行。其主要优化点在于CP进一步利用了Flash Attention的方法对注意力矩阵进行了分块计算。结合上述Ring Attention的计算通信流水部署方案CP一次传输一组KV矩阵得到一个分块的输出矩阵O最后再整合降低存储与通信开销。 图片来源论文Context Parallelism for Scalable Million-Token Inference
具体说来之前介绍的SP在计算注意力矩阵A′时需要按列的维度进行softmax计算这也就要求只有当所有(列)的K被收集到以后(AllGather完成)单个节点上才能完成注意力矩阵从A到A′的softmax计算所以KV矩阵的传输是分成两个阶段进行的而CP用了Flash Attention则可以移除这个限制分块计算注意力矩阵A′进而得到输出矩阵分块O最后再整合其KV矩阵的传输是一个阶段完成的减低了通信开销并且大幅降低了访存。
除此之外CP还对大模型中常见的causal attention中序列不同位置的序列段负载不均衡的问题进行了优化。CP在将序列分成多份的同时会考虑到不同序列段计算量的不一样进行对称互补组队再按组分配到不同节点中并行处理这样可以提升硬件利用率对计算与通信做更好的流水掩盖。
第四除了SP与CPAttention阶段的序列并行还有DeepSpeed提出的Ulysess方案如下图。Ulysess将Attention阶段的计算分成了两个阶段第一阶段计算QKV矩阵与SP/CP相同都是按照序列长度的维度进行切分第二阶段计算AA′O矩阵与SP/CP不同Ulysess按照Multi-Head中Head的维度进行了切分每个节点计算出全部序列的部分Head的结果最后再整合结果回到序列切分的维度上去。
图片来源论文Deepspeed Ulysses: System Optimizations For Enabling Training Of Extreme Long Sequence Transformer Models
所以Ulysess需要在第二阶段开始前对不同节点上的计算得到的QKV三矩阵分属不同序列段所有Head依据第二阶段各节点负责的Head进行三次QKV各一次AlltoAll的通信操作下图示意了矩阵Q的AlltoAll操作以及在二阶段结束后对不同节点上计算得到的O矩阵分属不同Head所有序列段依据各节点负责的序列段再进行一次AlltoAll的通信操作。Ulysess这里的AlltoAll是负载均衡的所有节点收发相同的数据量不像之前介绍过的EP可能出现专家负载不均从而导致AlltoAll劣化为AlltoAllV。 值得注意的是使用Ulysess有一个前提就是Multi-Head的Head个数需要能够整除序列并行数。 图片来源论文Sequence Parallelism: Long Sequence Training from System Perspective
综上所有讨论Attention的序列并行涉及到的主要通信模式有两种一种是SP与CP的AllGather ReduceScatter采用Ring Attention的方法将通信分步骤打散成多个P2P的Send和Receive操作来实现与分块计算进行流水掩盖另一种是Ulysess的AlltoAll使用在Attention一阶段序列并行与二阶段Multi-Head并行可以看做TP级联的交接处以及二阶段的结束处。
2LayerNorm与Dropout阶段的序列并行 Megatron-LM提出的SP特指LayerNorm与Dropout阶段的序列并行如下图是一种与Transformer模型中其他阶段部署的TP结合使用的并行策略简称SPTP。其达到的主要效果是——在不增加通信量的同时一次AllReduce等价为一次ReduceScatterAllGather减少各节点对中间激活值的存储需求以及降低LayerNorm与Dropout的计算量。 图片来源论文Reducing Activation Recomputation in Large Transformer Models
Megatron-LM的SP的理论基础是LayerNorm与Dropout阶段的计算对于序列长度的维度是独立的即1LayerNorm是对每个token的特征向量进行归一化而不是对整个序列的所有token一起归一化2Dropout是以元素级别的方式进行的即对每个token向量的每个元素独立地应用Dropout不会跨token共享Dropout掩码而是为每个元素独立生成掩码。
于是顺理成章的LayerNorm与Dropout阶段不需要所有节点都包含完整序列的信息可以把序列切成多份放在不同的节点上做LayerNorm与Dropout的计算这样既降低了计算量也减轻了节点的存储压力。
对于通信SPTP与之前纯TP相比通信总量维持不变只是通信的模式发生了改变。具体来说就是对于一个Transformer层纯TP需要前反向各做两次AllReduce而SPTP需要前反向各做两次AllGatherReduceScatter。解释一下这背后原因我们知道Attention/FFN阶段的TP使用的都是列并行行并行的级联形式如下图。 图片来源论文Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
列行并行的数据流动示意图如下详细推导见本系列的第二篇推文。 这种级联形式的主要特点就是列并行的输入需要完整copy输入矩阵到每个节点行并行后各节点的输出只是partial sum要得到完整输出需要将各节点的输出进行element-wise的求和列行并行级联的中间不涉及通信操作。结合这些特点与LayerNorm/Dropout阶段SP的对数据的切分方式可以知道在Attention/FFN阶段之前需要插入AllGather来收集所有序列片段的数据而在Attention/FFN阶段之后只需要插入ReduceScatter来分片求和输出的partial sum以在不同节点上获得不同序列片段对应的完整输出。 上面介绍的是前向传播过程因为一层transformer layer中涉及两个列行并行级联模块Attention和FFN在每个模块之前都需要插入AllGather在每个模块之后都需要插入ReduceScatter所有SPTP一共是两次AllGatherReduceScatter。 对于反向传播我们知道AllGather和ReduceScatter在前反向传播上是对偶关系——前向AllGather对应反向ReduceScatter前向ReduceScatter对应反向AllGather所以SPTP在反向传播时同样也需要做两次AllGatherReduceScatter。 至此序列并行就介绍完了下表总结了典型三种序列并行的通信模式所有序列并行的通信模式在HCCL API中均已支持。 ZeRO系列并行/完全分片数据并行 ZeRO的全称是Zero Redundancy Optimizer零冗余优化由DeepSpeed提出用来解决大模型训练中内存开销大的问题其主旨思想是去除冗余存储以及用通信换存储。以普通数据并行DP为基线对应下图中的第一行ZeRO按存储需求由多到少衍生出了三种进阶版的并行策略简称ZeRO-1ZeRO-2ZeRO-3对应下图中的第二三四行。 完全分片数据并行Fully Sharded Data Parallel, FSDP是pytorch1.11中的特性其本质就是ZeRO-3。下面我们逐一来看。 在基线DP中图中的第一行每个节点都会维护一套完整的模型。这包含模型参数本身Parameters参数的梯度Gradients以及优化器的状态Optimizer States。优化器状态一般包含多个参数如momentum和variance并以高精度如FP32存储。所以优化器状态需要的内存空间是远大于参数和梯度的对应图中绿色块大于蓝色和橙色块。而ZeRO做的就是将这套完整的模型一点一点切分到所有并行的节点上去。为方便对比先把基线DP的训练流程列出来
各节点收到完整的模型Parameters以及各自批次的训练数据完成前向和反向传播得到各自的Gradients通过AllReduce整合各节点的Gradients得到聚合好的完整的Gradients各节点使用完整的Gradients和完整的Optimizer States更新完整的模型parameters。
ZeRO-1图中的第二行首先对存储需求量最大的优化器状态开刀。通过切分优化器状态到各节点ZeRO-1将内存需求降低了4倍单卡通信数据量提高1.5倍对比基线。这背后的训练流程如下
各节点收到完整的模型Parameters以及各自批次的训练数据完成前向和反向传播得到各自的Gradients(同基线)通过AllReduce整合各节点的Gradients得到聚合好的完整的Gradients(同基线)各节点使用完整的Gradients和部分的Optimizer States更新部分的模型Parameters通过AllGather收集各节点的部分更新了的Parameters得到完整更新的模型Parameters。
ZeRO-2图中的第三行在ZeRO-1的基础上接着对梯度开刀。通过切分梯度、优化器状态到各节点ZeRO-2将内存需求降低了8倍单卡通信数据量基本不变对比基线。这背后的训练流程如下
各节点收到完整的模型Parameters以及各自批次的训练数据完成前向和反向传播得到各自的Gradients(同基线)通过ReduceScatter整合各节点的Gradients得到聚合好的部分的Gradients随传随把不属于各自节点维护的Gradients丢弃各节点使用部分的Gradients和部分的Optimizer States更新部分模型Parameters通过AllGather收集各节点的部分更新了的Parameters得到完整更新的模型Parameters。
ZeRO-3 / FSDP图中的第四行在ZeRO-2的基础上继续对模型参数本身开刀。通过切分模型参数、梯度、和优化器状态到各节点ZeRO-3 / FSDP将内存需求降低了64倍单卡通信数据量提高1.5倍对比基线。这背后的训练流程如下
各节点收到部分的模型Parameters以及各自批次的训练数据做前向传播各节点通过AllGather收集完整的模型Parameters随算随把不属于自己维护的模型Parameters丢弃做反向传播各节点通过AllGather收集完整的模型Parameters随算随把不属于自己维护的模型Parameters丢弃反向传播完各节点得到各自的Gradients通过ReduceScatter整合各节点的Gradients得到聚合好的部分的Gradients随传随把不属于各自节点维护的Gradients丢弃各节点使用部分的Gradients和部分的Optimizer States更新部分模型Parameters无需再对其做AllReduce操作。 总结一下ZeRO和FSDP是DP的进阶版它们在DP的基础上加入了模型并行的思路。通过切分模型的优化器状态、梯度、模型参数本身大幅降低数据并行时模型的内存开销。ZeRO以通信代存储尽可能的让不同节点不存重复的数据不做重复的计算需要时通过通信从其他节点获得。ZeRO在LLM大模型训练中被广泛使用其涉及到的主要通信操作为ReduceScatter和AllGather所有ZeRO和FSDP并行的通信模式在HCCL API中均已支持。
总结
本系列推文详细介绍了大模型分布式训练的各种并行策略聚焦分析了这些并行策略中使用到的通信操作帮助读者理解分布式训练的原理与通信模式并且HCCL API对这些通信操作进行了全方位的支持如需了解详细信息请查阅HCCL API列表。 参考材料
https://www.zhihu.com/question/637961859
https://www.hiascend.com/doc_center/source/zh/Pytorch/60RC3/modthirdparty/asdevguide/mindspeed_0017.html
[并行训练]Context Parallelism的原理与代码浅析-腾讯云开发者社区-腾讯云
https://zhuanlan.zhihu.com/p/5502876106
https://zhuanlan.zhihu.com/p/618865052
Getting Started with Fully Sharded Data Parallel (FSDP2) — PyTorch Tutorials 2.7.0cu126 documentation
HCCL概述-HCCL集合通信库-CANN商用版8.1.RC1开发文档-昇腾社区
