LLM:并行化训练
llama3.1论文称使用了:张量并行(TP)、流水线并行(PP)、上下文并行(CP)和数据并行(DP)。
并行化可以先粗分为:
1.数据并行
2.模型并行:张量并行,流水线并行
3.上下文并行(llama3.1)
4.deepspeed (专门提出)
所以主要就是 数据并行(DP)、张量并行(TP)、流水线并行(PP)。还有deepspeed。
1.数据并行
数据并行是最常见的并行形式,因为它很简单。在数据并行训练中,数据集被分割成几个碎片,每个碎片被分配到一个设备上。这相当于沿批次(Batch)维度对训练过程进行并行化。每个设备将持有一个完整的模型副本,并在分配的数据集碎片上进行训练。在反向传播之后,模型的梯度将会聚合(All Reduce),以便在不同设备上的模型参数能够保持同步。
DP(Data Parallelism):早期数据并行模式,一般采用参数服务器(Parameters Server)这一编程框架。实际中多用于单机多卡。
DDP(Distributed Data Parallelism):分布式数据并行,采用Ring AllReduce的通讯方式,多用于多机多卡场景。
可见:https://www.bilibili.com/video/BV1mm42137X8/?share_source=copy_web&vd_source=29af710704ae24d166ca951b4c167d53
2.张量并行
张量并行训练是将一个张量沿特定维度分成 N 块,每个设备只持有整个张量的 1/N,同时不影响计算图的正确性。
以一般的矩阵乘法为例,假设我们有 C = AB。我们可以将B沿着列分割成 [B0 B1 B2 … Bn],每个设备持有一列。然后我们将 A 与每个设备上 B 中的每一列相乘,我们将得到 [AB0 AB1 AB2 … ABn] 。此刻,每个设备仍然持有一部分的结果,例如,设备(rank=0)持有 AB0。为了确保结果的正确性,我们需要收集全部的结果,并沿列维串联张量。通过这种方式,我们能够将张量分布在设备上,同时确保计算流程保持正确。
也可以对这两个矩阵都进行切分。
3.流水线并行
流水线并行的核心思想是,模型按层分割成若干块,每块都交给一个设备。
在前向传播过程中,每个设备将中间的激活传递给下一个阶段。
在后向传播过程中,每个设备将输入张量的梯度传回给前一个流水线阶段。
该方案有个很大的缺陷是GPU的利用率太低,当一个GPU进行计算的时候,其他层的GPU都是闲置的。为了解决这个缺陷,谷歌提出了GPipe 流水线并行(Pipeline model parallesim ),下图来自GPipe官方论文。
方案引用了一个chunks的超参,将每一个mini-batch数据划分到chunks个micro-batches上。
4.deepspeed
官方文档:https://www.deepspeed.ai/
模型训练时的大概流程
正向传播时(FWD),上一层fp16的激活值和fp16的参数参与了计算,得到这一层的fp16激活值
反向传播时(BWD),本层的激活值和参数参与计算,得到fp16的梯度
参数更新时,fp16的梯度以及fp32的参数副本,momentum和variance参与计算,最终算出更新后的fp32参数、momentum和variance,然后将fp32的参数转化为fp16进行本层的参数更新
训练过程中的显存占用分析:
训练大模型时通常会采用AdamW优化器,并用混合精度训练来加速训练,基于这个前提分析显存占用。
来自:https://zhuanlan.zhihu.com/p/624740065
来自:https://zhuanlan.zhihu.com/p/639872915
按我的理解:fp16的模型参数,fp16的梯度,fp32的优化器状态:模型参数、动量、方差。最终就是2+2+4*(1+1+1)=16。
推理过程中的显存占用分析
deepspeed详解
可以看到baseline的参数估计就是我理解的这种,条件是提到的那些,k是优化器状态,2是模型参数,2是梯度。
ZeRO技术优化重点便是优化Model States 显存占用。在整个训练过程中, States只是会在对应算子计算时候才会被用到,optimizer states只在最终做update时才用到。数据并行中,gradients只在最后做AllReduce和updates时才用到。参数W只在做forward和backward的那一刻才用到,ZeRO技术针对以上做了以下优化:
Stage 0: 禁用所有类型的分片,仅使用 DDP (Distributed Data Parallel) ,对应下图的Baseline情况;
Stage 1: 把优化器状态(optimizer states) 分片到每个数据并行的工作进程(每个GPU)下,对应下图Pos
Stage 2: 把优化器状态(optimizer states) + 梯度(gradients) 分片到每个数据并行的工作进程(每个GPU)下,对应下图Pos + Pg的情况;
Stage 3: 把优化器状态(optimizer states) + 梯度(gradients) + 模型参数(parameters) 分片到每个数据并行的工作进程(每个GPU)下,Pos + Pg +Pp的情况;
Optimizer Offload: 在Stage2的基础上,把梯度和优化器状态下沉到CPU内存或硬盘上
Param Offload: 在Stage3的基础上,把模型参数下沉到CPU内存或硬盘上
deepspeed的实现需要设置一个config文件,在训练时一般使用zero2。
参考:
https://zhuanlan.zhihu.com/p/598714869
https://zhuanlan.zhihu.com/p/688873027
https://zhuanlan.zhihu.com/p/685472786
https://zhuanlan.zhihu.com/p/624740065
【动画理解Pytorch 大模型分布式训练技术 DP,DDP,DeepSpeed ZeRO技术】
https://zhuanlan.zhihu.com/p/639872915
原文地址:https://blog.csdn.net/qq_43814415/article/details/140857209
免责声明:本站文章内容转载自网络资源,如本站内容侵犯了原著者的合法权益,可联系本站删除。更多内容请关注自学内容网(zxcms.com)!