本篇的第二部分将基于 DP、PP 和 TP,介绍 LLM 早期经典的 3D 并行方法。 在训练参数量巨大的模型时,任意一种并行化方法都无法有效地扩展到足够多的 GPU 上。利用 DP、PP 和 TP 三种并行化方法各自的特点,可以将它们组合在一起,形成 3D 并行(3D Parallelism),从而实现对超大规模模型的训练。

1. 张量并行:Tensor Parallelism

1.1 张量并行的基础:矩阵分块乘法

在深度神经网络的计算中,矩阵乘法是最常见的计算操作之一。假设我们有一个输入矩阵 ,一个权重矩阵 ,其中 表示输入序列的长度, 表示隐藏层的维度, 表示输出的维度。 是矩阵乘法的输出。

下面简要回顾一下矩阵乘法的两个视角:列视角标量视角,它们非常自然地为矩阵分块计算提供了思路。默认情况下,本文使用的“向量”均指列向量。

深刻理解矩阵乘法要优先使用“向量思维”,而不是“标量思维”。对矩阵乘法还不太熟悉的朋友可以参考 Gilbert Strang 的 线性代数课程 或专著 《Introduction to Linear Algebra》。

图1.1 矩阵与向量乘法的列视角

图 1.1 所示为一个大小为 (4, 3) 的矩阵 和一个大小为 (3, 1) 的向量 之间的乘积。在列视角下,该乘积可看作是矩阵 中各列向量的线性组合,线性组合的系数由向量 的每个元素确定。例如,向量 的第一个元素 作为 中第一列向量的系数,向量 的第二个元素 作为 中第二列向量的系数,以此类推。

列视角最大的优势在于,它表明两个矩阵之间的乘法实际上可以看作是矩阵与多个向量之间的乘法,如图 1.2 所示。

图1.2 矩阵与矩阵乘法的列视角

从图 1.2 可以看到,矩阵 与 矩阵 的乘积可以分解为矩阵 与 矩阵 的每一列向量的乘积,然后将这些乘积结果直接按列的顺序拼接起来就得到最终结果。矩阵 与某个向量的乘积可参考图 1.1 介绍的线性组合视角。

这意味着当计算矩阵 与 矩阵 的乘积时,我们可以独立地使用矩阵 与 矩阵 的每一列向量进行乘积计算,只在最后做一次合并。从并行计算的角度来看,这表明可以对矩阵 的列向量进行切分,放到不同的 GPU 上进行计算 与它的乘积,然后将每个 GPU 上的计算结果进行拼接,得到最终结果。为了便于后续描述,将这种切分方法称为“列切分”。对它总结一下:

  • 列切分:当计算 时,将矩阵 按列切分成多个子矩阵 ,然后将每个子矩阵放到不同的 GPU 上进行计算 ,最后将所有 拼接起来得到最终结果 。

注意 不要求只有一列。例如假设矩阵 有 8 列,且我们有 4 个 GPU,那么可以将矩阵 按列切分成 4 个子矩阵,每个子矩阵包含 2 列。

第二种视角是标量视角,这也是我们线性代数教材中常见的视角。标量视角在每个数的计算上定义了矩阵乘法的规则。简单来说,结果矩阵 的每个元素 由矩阵 的第 行与矩阵 的第 列的对应元素逐一相乘后求和得到,如图 1.3 所示。

图1.3 矩阵与矩阵乘法的标量视角

矩阵乘法的一个神奇性质是,标量视角下定义的矩阵乘法规则适用于分块矩阵的乘法。例如,将矩阵 按列切分成两个子矩阵 ;将矩阵 按行切分成 两个子矩阵 。那么矩阵 与 矩阵 的乘积可以表示为:

可以仔细对比一下式(1.1)和图 1.3 中的计算过程。式(1.1)成立的一个前提是各分块子矩阵之间的乘法是有定义的。也就是说,矩阵 的列数要等于矩阵 的行数,矩阵 的列数要等于矩阵 的行数。将这种切分方法称为“外积和切分”,对它总结一下:

  • 外积和切分:当计算 时,将矩阵 按列切分成多个子矩阵 ,将矩阵 按行切分成多个子矩阵 ,然后将下标不同的子矩阵放到不同的 GPU 上进行计算 ,最后将所有 求和得到最终结果 。

直观来看,列切分和外积和切分的区别在于,列切分是将矩阵 切分成多个子矩阵,而外积和切分是将矩阵 和矩阵 都切分成多个子矩阵。

1.2 MLP 与 Self-Attention 的张量并行

标准 Transformer Block 主要包含一个 Self-Attention 模块和一个 MLP 模块。Megatron 的张量并行其实就是将上文介绍的列切分和外积和切分应用到 Self-Attention 模块和 MLP 模块中。我们先从相对简单的 MLP 模块开始介绍。

1.2.1 MLP 模块的张量并行

MLP 通常由两个全连接层构成。假设输入 矩阵为 ,第一个全连接层的权重矩阵为 ,第二个全连接层的权重矩阵为 ,其中 表示隐藏层的维度。MLP 的计算过程可以表示为:

其中 Act 表示激活函数,例如 ReLU 或 GeLU。

可以首先对式(1.2)中的矩阵乘法 使用列切分。以 2 张 GPU 的张量并行为例,将矩阵 按列切分成两个子矩阵 ,然后将它们分别放到 GPU 0 和 GPU 1 上进行计算:

因为激活函数通常是 Element-wise 的,所以可以分别对 和 进行激活函数计算,得到 。

此时可以使用一个 All-Gather 操作将 和 从两个 GPU 上收集到一起,得到 ,然后再用同样的方法对式(1.3)中的矩阵乘法 使用列切分,将矩阵 按列切分成两个子矩阵 ,然后将它们分别放到 GPU 0 和 GPU 1 上进行计算:

最后再使用一个 All-Gather 操作将 和 从两个 GPU 上收集到一起,得到最终 。上诉方法一共需要 2 次 All-Gather 操作。实际上可以注意到,当进行完式(1.4)和(1.5)的计算后, 和 已经分别在 GPU 0 和 GPU 1 上,并且它们从逻辑上形成了对 的列切分。这意味着只需要对矩阵 按照行进行切分,就能使用前文介绍的“外积和切分”法。具体来说,将矩阵 按行切分成两个子矩阵 ,然后将它们分别放到 GPU 0 和 GPU 1 上与对应的 和 进行计算:

最后再使用一个 All-Reduce 操作将 和 从两个 GPU 上收集到一起并求和,得到最终 。这样就只需要执行一次聚合操作(All-Reduce),而不是两次(All-Gather)。MLP 的张量并行示意如图 1.4 所示。

图1.4 MLP模块的张量并行

1.2.2 Self-Attention 模块的张量并行

Self-Attention 模块对每个 Head 进行如下计算:

其中 , 表示 Head 的数量, 分别表示第 个 Head 的 Query、Key、Value 的权重矩阵。这里为了简单起见,假设 QKV 的维度相等。利用每个 Head 的输出 ,Self-Attention 模块的最终输出为:

其中 是参数矩阵; 是所有 拼接的结果:

因为每个 Head 的输出 计算相互独立,所以可以将每个 Head 的 计算分配到不同的 GPU 上进行。同时注意到式(1.13)中的 实际上已经按照列划分了,因此可以按照 MLP 中同样的做法对矩阵 进行行切分,分成 H 个子矩阵 ,然后将它们分别放到与对应的 对应的 GPU 上进行计算:

最后再使用一个 All-Reduce 操作将所有结果收集到一起并求和,得到最终 。Self-Attention 的张量并行示意如图 1.5 所示。

图1.5 Self-Attention模块的张量并行。假设 Head 数为 2

2. 3D 并行:组合 TP、PP 与 DP

本小节将结合前两篇文章中介绍的 DP(Data Parallelism,数据并行)、PP(Pipeline Parallelism,流水线并行)和本篇的 TP(Tensor Parallelism,张量并行),介绍将它们组合在一起的 3D 并行方法。之所以称为 3D,是因为恰好有三种不同的并行方法,并且它们相对独立,每种方法就是一个 Dimension。

我们首先分析一下三种并行化方法各自的特点,然后再介绍它们的组合方案。

2.1 三种并行化方法的特点

DP 的扩展方法非常简单,只需要将数据划分成多份,然后将每份数据分配到不同的 GPU 上进行计算即可。但这也意味着 DP 的扩展和数据 Batch 是绑定的。要将 DP 扩展到更多的 GPU 上,通常也需要增大全局 Batch 的大小。但是过大的 Batch 会影响模型的收敛,因此虽然 DP 在理论上可以无限扩展,但实际上需要防止 DP 扩展到更大规模时因过大的 Batch Size 而导致的副作用。

DP 的存储效率:传统 DP 需要在每个 GPU 上存储完整的模型参数,因此存储效率较低。在第一篇中介绍的 ZeRO 可以显著提升 DP 的存储效率,因此存储开销通常不是 DP 的瓶颈。

DP 的计算效率:DP 中每个 Worker 都需要执行一次完整的前向和反向计算,因此每个 Worker 的计算不会因 Worker 数量的增加而改变,始终恒定(假设 Per-Device Batch 大小固定)。但是随着模型的增大,整体通信开销也随之增大。在之前的 ZeRO 文中分析过,对于参数量为 的模型,最常见的 ZeRO-2 的通信开销为 Bytes。当模型非常大时(如 100B),模型每次 Step 都需要通信 200GB 的数据,尤其当需要跨节点通信时,通信开销会影响计算效率。此外,虽然理论上 ZeRO-2 的通信开销与 Worker 数量无关,但实际中随着 Worker 数量的增加,所有节点必须同步的问题非常容易受到慢节点、网络跳数等因素的影响,从而实际通信耗时增加。因此 DP 不适合大规模扩展,在小规模扩展上最具优势。

PP 将模型划分成多个阶段(Stage),将每个阶段分配到不同的 GPU 上进行计算。当收到一个 Batch 数据后,第一个 Stage 开始进行前向计算,计算完成后将结果(Activation)传递给下一个 Stage,然后下一个 Stage 开始计算,如此循环直到最后一个 Stage 计算完成。然后再从最后一个 Stage 开始进行反向计算,依次将梯度回传,直到第一个 Stage 计算完成。在实际应用中通常会使用更复杂的策略来尽量降低各个 GPU 的空闲时间。PP 的扩展上限受到模型中网络层数的限制。

PP 的存储效率:每个 Worker 只需要存储自己负责的 Stage 的参数,因此每个 Worker 的存储开销随着 Worker 数量的增加而线性减少。

PP 的计算效率:因为 PP 只需要在两个相邻 Stage 之间进行某一层 Activation/Gradient 的通信,所以 PP 的通信开销是三种并行化方法中最低。PP 的难点主要在于如何确保各个 Stage 的计算负载均衡,避免出现大量的空闲时间。

TP 将模型的张量划分成多个子张量,然后将每个子张量分配到不同的 GPU 上进行计算。

TP 的存储效率:每个 Worker 只需要存储自己负责的子张量,因此每个 Worker 的存储开销随着 Worker 数量的增加而线性减少。

TP 的计算效率:TP 的计算负载均衡性较好,因为每个 Worker 都需要参与每一层的计算,不会出现明显的空闲时间。但 TP 的通信开销较大,因为每一层的前向或反向计算都需要聚合操作(All-Gather 或 All-Reduce)。随着模型规模的增大,TP 的通信开销会显著增加,从而影响计算效率。因为巨大的通信开销,要绝对避免在不同节点上进行 TP。

2.2 3D Parallelism

根据上述三种并行化方法的特点,我们尝试分析一下它们合理的组合方式。为了简单起见,使用一个具体的示例来说明。假设模型的参数量是 512B。可以从三种并行化方法“受限程度”出发,优先决定受限程度高的方法的使用方式,然后再决定受限程度低的方法。

上述三种并行化方法中,TP 的受限程度最高,因为 TP 高昂的通信量导致其不能跨节点使用。其次是 PP,因为 PP 的理论扩展上限受到模型网络层数的限制,并且在实际中也要避免每个 Stage 过小而降低计算效率。最后是 DP,因为 DP 在理论上可以无限扩展,仅仅在规模非常大时会受到 Batch Size 和通信效率的影响。

因此可以先不考虑 DP,将模型使用 PP 划分为多个 Stage,然后在每个 Stage 内使用 TP 再次切分。

回到示例,假设首先将模型划分为 16 个 Stage,每个 Stage 的参数量为 32B,并且每个 Stage 放在一个 8 卡节点上。而在每个节点内,再使用 TP 将每个 Stage 切分成 8 份,每份参数量为 4B。我们将上述方案称为 16-way PP + 8-way TP。16-way PP + 8-way TP 一共需要 16 x 8 = 128 张 GPU,这 128 张 GPU 上的每个 GPU 保存的模型参数都不相同,共同构成一个完整的模型副本。为了便于描述,我们用 表示这 128 张 GPU。

当上述方案已经接近极限时,为了继续扩大计算规模,只能尝试 DP 了。“极限”意味着没有办法继续增加 PP 的 Stage 数量。虽然 PP 的 Stage 数量理论上不能超过模型的网络层数,但实际应用中也不能让 PP 的每个 Stage 太小,否则会导致主要的计算开销全部集中在 Stage 之间的通信上,从而影响计算效率。

假设使用 2-way DP,并且有两组完全一样的 16-way PP + 8-way TP 的划分方案。将第一组的 GPU 记为 ,将第二组的 GPU 记为 。因为 DP 只能作用于参数相同的 GPU 之间,所以 和 构成一个 DP 组, 和 构成一个 DP 组,以此类推,一共构成 128 个 DP 组。在每个 DP 组内,可以使用 ZeRO 来进一步提升存储效率。因为每个 DP 组内的 GPU 数量通常较小,所以这种划分方法不会显著影响 ZeRO 的通信效率。

图 2.1 是一个 4-way DP + 4-way PP + 2-way TP 的示意图,大家可以结合上面的描述进行理解。早期经典的 3D 并行使用 TP + 1F1B PP + ZeRO DP 的组合方式。

图 2.1 3D 并行示意图

ZeRO DP 和 1F1B PP 的介绍请参考本系列前两篇文章。

3. 总结

本文介绍了张量并行和基于张量并行、流水线并行、数据并行的 3D 并行化方案。3D 并行是支撑早期大规模 LLM 训练的重要技术,其中稍微不那么直觉的部分是对数据并行的使用方法。 基于 3D 并行,后续逐渐演化出了 4D、5D 方案,不过这都是后话。

那么,如何系统的去学习大模型LLM?

作为一名深耕行业的资深大模型算法工程师,我经常会收到一些评论和私信,我是小白,学习大模型该从哪里入手呢?我自学没有方向怎么办?这个地方我不会啊。如果你也有类似的经历,一定要继续看下去!这些问题啊,也不是三言两语啊就能讲明白的。

所以我综合了大模型的所有知识点,给大家带来一套全网最全最细的大模型零基础教程。在做这套教程之前呢,我就曾放空大脑,以一个大模型小白的角度去重新解析它,采用基础知识和实战项目相结合的教学方式,历时3个月,终于完成了这样的课程,让你真正体会到什么是每一秒都在疯狂输出知识点。

由于篇幅有限,⚡️ 朋友们如果有需要全套 《2025全新制作的大模型全套资料》,扫码获取~
在这里插入图片描述

👉大模型学习指南+路线汇总👈

我们这套大模型资料呢,会从基础篇、进阶篇和项目实战篇等三大方面来讲解。
在这里插入图片描述
在这里插入图片描述

👉①.基础篇👈

基础篇里面包括了Python快速入门、AI开发环境搭建及提示词工程,带你学习大模型核心原理、prompt使用技巧、Transformer架构和预训练、SFT、RLHF等一些基础概念,用最易懂的方式带你入门大模型。
在这里插入图片描述

👉②.进阶篇👈

接下来是进阶篇,你将掌握RAG、Agent、Langchain、大模型微调和私有化部署,学习如何构建外挂知识库并和自己的企业相结合,学习如何使用langchain框架提高开发效率和代码质量、学习如何选择合适的基座模型并进行数据集的收集预处理以及具体的模型微调等等。
在这里插入图片描述

👉③.实战篇👈

实战篇会手把手带着大家练习企业级的落地项目(已脱敏),比如RAG医疗问答系统、Agent智能电商客服系统、数字人项目实战、教育行业智能助教等等,从而帮助大家更好的应对大模型时代的挑战。
在这里插入图片描述

👉④.福利篇👈

最后呢,会给大家一个小福利,课程视频中的所有素材,有搭建AI开发环境资料包,还有学习计划表,几十上百G素材、电子书和课件等等,只要你能想到的素材,我这里几乎都有。我已经全部上传到CSDN,朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】
在这里插入图片描述
相信我,这套大模型系统教程将会是全网最齐全 最易懂的小白专用课!!

Logo

DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。

更多推荐