一文搞懂模型训练的并行方式

最新推荐文章于 2024-09-01 18:39:24 发布

小天才学习机打游戏

最新推荐文章于 2024-09-01 18:39:24 发布

阅读量661

点赞数 7

文章标签：爬虫人工智能架构 prompt chatgpt

本文链接：https://blog.csdn.net/m0_59164520/article/details/141726861

版权

随着模型的规模和复杂度不断增加，训练这些大规模模型需要消耗大量的计算资源和时间。为了提高训练效率，各种模型训练并行方式应运而生。

一、数据并行（Data Parallelism）

数据并行是最简单且常用的并行方式之一。**它的基本思想是将相同的模型权重复制到多个设备（如多个 GPU）上，然后将训练数据分成多个子集，每个设备处理其中一部分数据。**这样，多个设备可以同时进行前向传播、反向传播和参数更新，从而加快训练速度。

在数据并行中，通常采用分布式数据并行（Distributed Data Parallel，DDP）架构。在这种架构下，负载分散在每个 GPU 节点上，通信成本是恒定的，与 GPU 数量无关。例如，PyTorch 的 DDP 引入了 bucket 机制，在 backward 时进行梯度更新，当一个 bucket 内部的梯度计算完成后直接开始进行 all-reduce 操作，而不需要等到 backward 计算结束，由此提升了吞吐量和训练效率。

然而，数据并行也存在一些问题，例如主 GPU 可能会负载较大，导致负载不均衡的情况。

二、张量并行（Tensor Parallelism）

由于模型越来越大，单个设备保存模型参数、梯度和优化器越来越难。因为深度学习主要是矩阵计算，而矩阵计算可以分块计算，因此可以将模型参数拆成若干份，每份单独计算，以减少显存占用。

张量并行则是将张量（tensor）沿特定维度分成若干部分，在不同的设备上分别计算。以 Transformer 结构为例，对于 MLP 层中的第一个线性层，可以将权重矩阵沿其行拆分，或者将输入沿其列拆分，以实现并行计算。在反向传播中也需要相应的 all-reduce 操作来汇总梯度。

除了简单的 1D 并行，还有更复杂的 2D、2.5D 以及 3D 并行方式。例如，Colossal-AI 在 SUMMA（可扩展的通用矩阵乘法算法）的基础上引入了更多的张量并行形式。通过这种方式，可以更高效地利用多个设备的计算能力，提高训练速度。

三、流水线并行（Pipeline Parallelism）

流水线并行是一种将模型并行与数据流水线相结合的方法。**模型按层分割成若干块，分配给不同的设备。**在前向传递过程中，每个设备处理完一个微批次（micro-batch）后，将中间激活传递给下一个设备，然后立即开始处理下一个微批次，从而减少空闲时间。

然而，流水线并行仍然存在 bubble 时间，即各设备等待前一个设备处理完成的时间。为了减少 bubble 时间所占比例，可以使 micro-batch 拆分的数量远大于流水线并行数。此外，像 Pipedream 和 Megatron-LM 提出的非交错式和交错式 1f1b 等策略进一步提高了效率。

Megatron-LM 还提出了 PTD-P（Inter-node Pipeline Parallelism, Intra-node Tensor Parallelism, and Data Parallelism），结合了流水线并行、张量并行和数据并行三种策略来训练万亿级别参数量的模型。

四、多种并行方式的组合

在实际应用中，通常会组合多种并行方式以达到更好的效果。例如：

（一）DP + PP（Data Parallelism + Pipeline Parallelism）

将数据并行与流水线并行结合使用。如下图所示，dprank0 看不到 gpu2，dprank1 看不到 gpu3。对于数据并行而言，只有 gpu0 和 gpu1，并向它们供给数据。gpu0 使用流水线并行将部分负载转移到 gpu2，gpu1 同理将负载转移到 gpu3。这里每个维度至少需要 2 个 gpu，所以至少需要 4 个 gpu。

（二）3D 并行（DP + PP + TP）

为了更高效地训练，可以将流水线并行（PP）、张量并行（TP）和数据并行（DP）相结合，这被业界称为 3D 并行。在此模式下，每个维度至少需要 2 个 gpu，因此至少需要 8 个 gpu 才能实现完整的 3D 并行。

（三）Zero-DP + PP + TP

Zero 是 DeepSpeed 的主要功能之一，是数据并行的超级可伸缩增强版，它也可以与流水线并行（PP）、张量并行（TP）结合使用。当 Zero-DP 与 PP 和 TP 结合时，通常只启用 Zero 阶段 1（只对优化器状态进行分片）。

而启用 Zero 阶段 2 会对梯度进行分片，阶段 3 会对模型权重进行分片。虽然理论上可以将阶段 2 与流水线并行一起使用，但会对性能产生不良影响，因为每个 micro-batch 都需要额外的 reduce-scatter 通信来在分片之前聚合梯度，增加了潜在的显著通信开销。并且，由于 PP 已经减少了梯度大小，在此基础上进行梯度分片节省的内存有限。

此外，还可以采用其他组合方式，如 DP + TP 进行组合、PP + TP 进行组合，或者使用 Zero3 代替 DP + PP + TP。Zero3 本质上是 DP + MP 的组合，且无需对模型进行过多改造，使用相对更方便。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述