LLM分布式预训练浅析

发布于 2024-7-4 10:45
浏览
0收藏

1)模型参数过大,如GLM 130B模型参数需要520GB(130B*4bytes)的显存,参数梯度为520GB,优化器状态需要1040GB,共计2096GB的显存,即使显存最大的GPU也放不下大模型的模型参数;

2)即使可以将模型参数放在一张GPU上,对海量数据进行预训练也需要耗费很长时间,单卡如A100训练130B参数量的模型大约需要232年。分布式预训练技术,就是在这样的背景下应运而生,它允许模型在多个计算节点上并行训练,大幅度缩短了训练时间,本文总结了目前常见的分布式训练的常见方法。

一、分布式训练

1.数据并行

数据并行指的是数据集切分为多份,每张GPU分配不同的数据进行训练,每个进程运行完整模型,在每个进程上独立运行前向和反向过程,最后通过AllReduce操作,将多张卡上的梯度进行同步。

LLM分布式预训练浅析-AI.x社区

▲ 图1 数据并行


2.张量并行

张量并行是指将一个独立的层划分到不同的GPU上,例如用两块GPU实现一个并行的GEMM矩阵乘法运算Y=X*W,可以将X和W分别沿着列方向和行方向切分成两份,通过矩阵乘法的分块规则,分别在每个GPU上做子矩阵乘法,然后将两个子矩阵乘法结果相加得到最终结果

LLM分布式预训练浅析-AI.x社区

针对transformer模型MLP的张量并行实现方式图2所示,第一个GEMM输入X不做切分,直接通过f操作拷贝到两个GPU上,权重A沿列方向进行切分,每个GPU上保存一半的运算结果,通过GELU运算得到Y1和Y2。第二个GEMM由于Y1和Y2已经分别存在于两个GPU上,相当于已经做了沿列方向的切分,所以对权重B做行方向的切分,分别进行矩阵乘法得到Z1和Z2,g操作是一个AllReduce操作,用来实现公式中的加法,AllReduce之后,每个GPU上都有一份完整的Z矩阵。

LLM分布式预训练浅析-AI.x社区

▲ 图2 MLP张量并行

针对transformer模型Atttention的张量并行如图3所示,QKV的三个矩阵的权重shape均为是[dim, head_hum*head_dim],考虑到多头注意力天然的并行性,可以将三个权重沿着列方向进行切分,每张卡上权重的shape为[dim, head_hum*head_dim/N],每张卡上不同部分head的softmax等计算式互相独立,分别得到Y1和Y2,再对输出的权重 B沿着行方向切分,得到计算结果Z1和Z2,再通过g做AllReduce操作得到最终的Z。可以看到,每张卡上做的仍然是一个完整的Attention运算,只是head number变为了原来的1/N。

LLM分布式预训练浅析-AI.x社区

▲ 图3 Attention 张量并行


3.流水线并行

流水线并行是指将不同的layer划分到不同的GPU上,为了提升节点的利用率,将一个batch,继续拆分为 M 个micro batch,当节点1计算完第一个micro batch, 节点2开始计算后,节点 1可以继续计算第二个micro batch,对于P个节点,M个micro batch的流水线并行,流水线气泡占比为(P-1)/M,为减少气泡占比,需保证M>>P。


LLM分布式预训练浅析-AI.x社区

▲ 图4 流水线并行


4.序列并行

序列并行是在 张量的基础上,将 Transformer 层中的 LayerNorm 以及 Dropout 的输入按输入长度(Sequence Length)维度进行了切分,使得各个设备上面只需要做一部分的 Dropout 和 LayerNorm,LayerNorm 和 Dropout 的计算被平摊到了各个设备上,减少了计算资源的浪费,进一步降低了显存开销。

LLM分布式预训练浅析-AI.x社区

▲ 图5 序列并行

4.激活重计算

激活重计算是指在前向计算时,只把检查点节点保存在显存中,检查点以外的中间值全部被丢弃,在方向传播计算梯度时,从最近的检查点开始,其余节点重新计算一次前向,激活重计算需要额外的计算,是一个以时间换空间的策略。


5.ZeRO零冗余优化器

数据并行在通信和计算效率上表现良好,但存在内存冗余问题,ZeRO通过分区参数和状态来减少冗余,使每个GPU只保留部分数据。图6展示了启用内存优化后各设备的内存消耗,其中Ψ代表模型大小,K是优化器状态的内存乘数,Nd是数据并行度。

LLM分布式预训练浅析-AI.x社区

▲ 图6 ZeRO零冗余优化

假设模型大小为Ψ=75亿,基于Adam优化器的混合精度训练,数据并行度为Nd=64,K=12,我们在数据并行的基础上,分三步对显存进行优化

Pos(优化器状态优化)

我们在每个GPU中保存全部的参数和梯度,但是只保存1/Nd的优化器状态,这会导致总的显存消耗变为2Ψ+2Ψ+KΨ/Nd , 在图示假设下为显存占用31.4GB。

Pos+g(优化器状态和梯度优化)

优化器状态优化的基础上增加对梯度内存的优化,即每个GPU中只保存1/Nd的梯度,这会导致总的显存消耗变为2Ψ+(2Ψ+KΨ)/Nd , 在图示假设下为显存占用为16.6GB。

Pos+g+p(优化器状态、梯度以及参数优化)

优化器状态和梯度优化的基础上增加对参数显存的优化,即每个GPU中只保存1/Nd的参数,这将导致总的内存消耗变为(2Ψ+2Ψ+KΨ)/Nd ,在图示假设下为显存占用为1.9GB。

二、讨论与展望

未来的LLM可能会采用更先进的并行策略,以进一步减少通信开销并提高计算效率。例如探索使用分层并行策略,将模型的不同层次分布在不同的节点上进行训练,以减少节点间的通信需求。此外,还可以结合使用数据并行和模型并行的方法,根据硬件环境和模型规模灵活调整并行策略。分布式预训练技术的发展为大规模智算平台提供了有效支撑,但仍然面临非常大挑战,通过精心设计解决方案并不断探索新的并行策略,我们可以进一步提高分布式预训练的效率和可扩展性,为构建更大规模的模型提供支持。

参考文献

[1] Korthikanti, Vijay, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, Bryan Catanzaro. 《Reducing Activation Recomputation in Large Transformer Models》.

[2] Narayanan, Deepak, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand. 《Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM》. Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, 1–15. St. Louis Missouri: ACM, 2021. https://doi.org/10.1145/3458817.3476209.

[3] Shoeybi, Mohammad, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, Bryan Catanzaro. 《Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism》.

[4] Rajbhandari, Samyam, Jeff Rasley, Olatunji Ruwase和Yuxiong He. 《ZeRO: Memory Optimizations Toward Training Trillion Parameter Models》. arXiv, 2020年5月13日. http://arxiv.org/abs/1910.02054.

本文转载自​AI遇见云​,作者: 何武 ​​


收藏
回复
举报
回复
相关推荐