入门下大模型训练中的MP、PP、TP等操作

近年来,我们在许多 NLP 基准任务上看到了更好的结果,这得益于更大规模的预训练语言模型。训练大型且深度的神经网络是一个挑战,因为它需要大量的 GPU 内存和漫长的训练时间。 然而,单个 GPU 工作器的内存有限,许多大型模型的大小已经超出了单个 GPU 的容量。有几种并行范式可以实现在多个 GPU 上进行模型训练,以及各种模型架构和内存节省设计,以帮助实现训练真正大型神经网络的可能性。

一、训练并行化

训练非常大型神经网络模型的主要瓶颈是对大量 GPU 内存的强烈需求,这远远超出了单个 GPU 机器所能承载的范围。除了模型权重(例如数十亿个浮点数)之外,存储中间计算输出(如梯度和优化器状态,例如 Adam 中的动量和变化)通常成本更高。此外,训练大型模型通常伴随着大型训练语料库,因此单个进程可能需要无限长的时间。 因此,需要并行化。并行化可以在不同的维度上发生,包括数据、模型架构和张量操作。

1、数据并行化

数据并行化(DP)最简单的方法是将相同的模型权重复制到多个工作器中,并将数据的一部分分配给每个工作器同时处理。

如果模型大小超过了单个 GPU 节点的内存,简单的 DP 就无法很好地工作。像 GeePS(Cui 等人,2016)这样的方法将暂时不用的参数卸载回 CPU,以便在模型太大而无法放入一台机器时与有限的 GPU 内存一起工作。数据交换传输应该在后端进行,不要干扰训练计算。

在每个小批量结束时,工作器需要同步梯度或权重,以避免过时。有两种主要的同步方法,各有优缺点。

  1. 批量同步并行(BSP):工作器在每个小批量结束时同步数据。它防止模型权重过时,具有良好的学习效率,但每台机器都必须暂停并等待其他机器发送梯度。

  2. 异步并行(ASP):每个 GPU 工作器异步处理数据,无需等待或停滞。然而,这很容易导致使用过时的权重,从而降低统计学习效率。尽管它增加了计算时间,但可能不会加快收敛到训练时间的速度。

介于两者之间的是每n 次迭代同步一次全局梯度。这个特性被称为“梯度累积”,自 Pytorch v1.5(Li 等人,2021)起在分布式数据并行(DDP)中实现。桶装梯度避免了立即的 AllReduce 操作,而是将多个梯度放入一个 AllReduce 中,以提高吞吐量。可以根据计算图进行计算和通信调度优化。

在这里插入图片描述

2、模型并行化

模型并行化(MP) 旨在解决模型权重无法放入单个节点的情况。计算和模型参数分布在多台机器上。与数据并行化不同,数据并行化中每个工作器托管整个模型的完整副本,MP 只在一个工作器上分配模型参数的一部分,因此内存使用和计算都减少了。

由于深度神经网络通常包含一系列垂直层,因此按层分割大型模型似乎很直接,将一小部分连续层分组为一个分区放在一个工作器上。然而,一个简单的实现,即让每个数据批次通过多个这样的工作器,存在顺序依赖,会导致大量等待时间的“气泡”和计算资源的严重未充分利用。

在这里插入图片描述

图 2. 一个简单的模型并行化设置,模型垂直分割为 4 个分区。由于顺序依赖,数据一次只能被一个工作器处理,导致大量闲置时间的“气泡”。(图片来源:Huang 等人,2019)

3、管道并行化

管道并行化(Pipeline Parallelism), 结合了模型并行化和数据并行化,以减少低效的时间“气泡”。主要思想是将一个小批量分割成多个微批次,并使每个阶段工作器同时处理一个微批次。注意,每个微批次需要两次传递,一次前向和一次反向。工作器之间的通信仅传输激活(前向)和梯度(反向)。这些传递的调度方式以及梯度的聚合方式在不同方法中有所不同。分区(工作器)的数量也称为管道深度。

在 GPipe(Huang 等人,2019)中,多个微批次的梯度在结束时同步聚合并应用。同步梯度下降保证了学习的一致性和效率,无论工作器的数量如何。如图 3 所示,仍然存在气泡,但比图 2 中的要小得多。假设微批次均匀分割,且每个微批次的前向和反向传递各需一个时间单位,气泡的比例为:

GPipe 论文观察到,如果微批次的数量超过分区数量的 4 倍(应用激活重计算时),气泡开销几乎可以忽略不计。

在这里插入图片描述

图 3. GPipe 中管道并行化的示意图,包含 4 个微批次和 4 个分区。GPipe 在每个批次结束时同步聚合和更新跨设备的梯度。(图片来源:Huang 等人,2019)

GPipe 随着设备数量的增加几乎实现了线性吞吐量加速,尽管如果模型参数在工作器之间分布不均,则不总是能保证这一点。

PipeDream(Narayanan 等人,2019)安排每个工作器交替处理前向和反向传递(1F1B)。 PipeDream 将每个模型分区称为“阶段”,每个阶段工作器可以有多个副本以运行数据并行化。在此过程中,PipeDream 使用确定性轮询负载均衡策略在阶段的多个副本之间分配工作,以确保同一小批量的前向和反向传递发生在同一副本上。

在这里插入图片描述

图 4. PipeDream 中 1F1B 微批次调度的示意图。(图片来源:Harlap 等人,2018)

由于 PipeDream 在所有工作器之间没有每个批次结束时的全局梯度同步,1F1B 的原生实现很容易导致一个微批次的前向和反向传递使用不同版本的模型权重,从而降低学习效率。PipeDream 提出了一些设计来解决这个问题:

  1. 权重存储:每个工作器跟踪几个模型版本,并确保在给定一个数据批次时,前向和反向传递使用相同版本的权重。

  2. 垂直同步(可选):模型权重的版本与激活和梯度一起在阶段工作器之间流动。然后计算采用从上一个工作器传播的相应存储版本。这个过程保持了工作器之间的版本一致性。注意,这是异步的,与 GPipe 不同。

在训练开始时,PipeDream 首先分析模型中每一层的计算内存成本和时间,然后优化一个将层分割为阶段的解决方案,这是一个动态规划问题。

在这里插入图片描述

图 5. VGG16 在 ILSVRC12 上的结果。(顶部)准确率与时间。整数标记阶段工作器的数量。ASP = 异步并行 & BSP = 批量同步并行。(底部)不同并行化配置的训练时间加速。(图片来源:Harlap 等人,2018)

后来提出了两种 PipeDream 的变体,以减少存储模型版本所需的内存占用(Narayanan 等人,2021)。

PipeDream-flush 周期性地添加一个全局同步的管道刷新,就像 GPipe 一样。通过这种方式,它大大减少了内存占用(即只维护一个模型权重版本),以牺牲一些吞吐量为代价。

图 6. PipeDream-flush 中中管道调度的示意图。(图片来源:Narayanan 等人,2021)

PipeDream-2BW 仅维护两个模型权重版本,“2BW”是“双缓冲权重”的缩写。它每 n 个微批次生成一个新模型版本,且 n 应该大于管道深度 P。一个新更新的模型版本不能立即完全替换旧版本,因为一些剩余的反向传递仍然依赖于旧版本。总共只需要保存两个版本,因此内存成本大大降低。

在这里插入图片描述

图 7. PipeDream-2BW 中管道调度的示意图。(图片来源:Narayanan 等人,2021)

4、张量并行化

模型和管道并行化都是垂直分割模型。另一方面,我们可以水平分割一个张量操作的计算到多个设备上,称为张量并行化(TP)。

以 Transformer 为例,因其流行度而备受关注。Transformer 模型主要由 MLP 层和自注意力块组成。Megatron-LM(Shoeybi 等人,2020)采用了一种简单的方法来并行化 MLP 和自注意力的层内计算。

Transformer 中的 MLP 层包含一个 GEMM(通用矩阵乘法)后跟一个非线性 GeLU 转换。让我们按列分割权重矩阵 A得到

自注意力块根据上述分割并行运行 GEMM,使用Query、Key和Value),然后通过另一个 GEMM 结合它们以产生注意力头结果。


在这里插入图片描述

图 8. Megatron-LM 提出的 Transformer 关键组件的张量并行化示意图。(图片来源:Shoeybi 等人,2020)

Narayanan 等人(2021)结合了管道、张量和数据并行化,并采用了一种新的管道调度策略,将其方法称为 PTD-P。与其仅将一组连续的层(“模型块”)定位在设备上,每个工作器可以被分配多个较小连续子集的层块(例如,设备 1 有层 1、2、9、10;设备 2 有层 3、4、11、12;每个都有两个模型块)。一个批次中的微批次数量应能被工作器数量整除。如果有 M 个模型块每个工作器,管道气泡时间可以通过 GPipe 调度的倍数减少。

在这里插入图片描述

图 9.(顶部)PipeDream-flush 中默认的 1F1B 管道调度。(底部)交错的 1F1B 管道调度。第一模型块为深色,第二块为浅色。(图片来源:Narayanan 等人,2021)

参考:https://lilianweng.github.io/posts/2021-09-25-train-large/


二、如何系统学习掌握AI大模型?

AI大模型作为人工智能领域的重要技术突破,正成为推动各行各业创新和转型的关键力量。抓住AI大模型的风口,掌握AI大模型的知识和技能将变得越来越重要。

学习AI大模型是一个系统的过程,需要从基础开始,逐步深入到更高级的技术。

这里给大家精心整理了一份全面的AI大模型学习资源,包括:AI大模型全套学习路线图(从入门到实战)、精品AI大模型学习书籍手册、视频教程、实战学习、面试题等,资料免费分享

1. 成长路线图&学习规划

要学习一门新的技术,作为新手一定要先学习成长路线图方向不对,努力白费

这里,我们为新手和想要进一步提升的专业人士准备了一份详细的学习成长路线图和规划。可以说是最科学最系统的学习成长路线。
在这里插入图片描述

2. 大模型经典PDF书籍

书籍和学习文档资料是学习大模型过程中必不可少的,我们精选了一系列深入探讨大模型技术的书籍和学习文档,它们由领域内的顶尖专家撰写,内容全面、深入、详尽,为你学习大模型提供坚实的理论基础(书籍含电子版PDF)

在这里插入图片描述

3. 大模型视频教程

对于很多自学或者没有基础的同学来说,书籍这些纯文字类的学习教材会觉得比较晦涩难以理解,因此,我们提供了丰富的大模型视频教程,以动态、形象的方式展示技术概念,帮助你更快、更轻松地掌握核心知识

在这里插入图片描述

4. 2024行业报告

行业分析主要包括对不同行业的现状、趋势、问题、机会等进行系统地调研和评估,以了解哪些行业更适合引入大模型的技术和应用,以及在哪些方面可以发挥大模型的优势。

在这里插入图片描述

5. 大模型项目实战

学以致用 ,当你的理论知识积累到一定程度,就需要通过项目实战,在实际操作中检验和巩固你所学到的知识,同时为你找工作和职业发展打下坚实的基础。

在这里插入图片描述

6. 大模型面试题

面试不仅是技术的较量,更需要充分的准备。

在你已经掌握了大模型技术之后,就需要开始准备面试,我们将提供精心整理的大模型面试题库,涵盖当前面试中可能遇到的各种技术问题,让你在面试中游刃有余。

在这里插入图片描述

全套的AI大模型学习资源已经整理打包,有需要的小伙伴可以微信扫描下方CSDN官方认证二维码,免费领取【保证100%免费

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值