LLM:并行化训练

灵海之森 2024-10-09 17:31:01 阅读 54

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



声明

本文内容仅代表作者观点,或转载于其他网站,本站不以此文作为商业用途
如有涉及侵权,请联系本站进行删除
转载本站原创文章,请注明来源及作者。