LLM训练各种并行策略

最新推荐文章于 2025-08-04 10:30:53 发布

原创最新推荐文章于 2025-08-04 10:30:53 发布 · 951 阅读

28 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能 #数据库 #机器学习 #算法 #ocr

1.张量并行（TP）

张量并行，Tensor Parallelism，TP（也有叫Model Parallelism，MP的）：LLM中，有的tensor或者layer很大，单卡放不下（或者单卡不够放整个模型），那么就需要用TP把tensor分割成多块，每一块放在一张卡上，分别使用和计算。仅当某些算子需要完整的张量时，才会进行聚合。

1.1.分块矩阵乘法

TP的基本思想是对矩阵乘法进行拆分：

那么矩阵乘法有两种拆分方法：（1）对矩阵A按列拆分（上图上）（2）对矩阵A按行拆分（上图下）。

注意，当对矩阵A按行拆分的时候，也要对矩阵X进行列的拆分，保持维度的一致。

当对矩阵A按行拆分的时候，X和A都是concat的关系，backward的时候可以分别计算X拆分出来的小矩阵的梯度，然后再拼接起来就可以得到X的完整梯度。

而当对矩阵A按列进行拆分时，X同时参与了两块GPU上的前向计算，因此X的完整梯度等于两张卡上X的梯度相加。

更加详细的说明可以参考：图解大模型训练之：张量模型并行(TP)，Megatron-LM，https://zhuanlan.zhihu.com/p/622212228。

1.2.MLP层的TP

上面展示的是矩阵乘法的TP。那么如果我们的计算不仅是Y=XA，而还有个激活函数呢，比如Y=ACT(XA）。把矩阵A按行切分的方式，需要在进入激活函数的计算前，同步各个GPU得到的Y，这就有不少的通讯量；而把A按列切分的方式则可以直接进行激活函数的计算。

那么再进一步，如果是MLP层，那么Y=ACT(XA)B，在上面的基础上又多了个B矩阵的计算，该怎么切分呢。理想的状况应该是尽量减少计算中的同步操作（从而减少通讯量），提升框架整体的计算效率。

基于前面的分析，我们可以对A按列切割，那么各个GPU得到的Y就是concat的关系，为了和各个小Y能够直接进行计算，那么B应该是按行切分：

1.3.Attention的TP

那么多头注意力如何做TP呢？先回顾一下多头注意力的计算，多头注意力本身就对Q、K、V在dim维度做了切分，然后concat起来。也就是说这多个头本身，天然就是可以并行，独立进行计算的。那么只需要把不同的注意力头放到不同的GPU上，我们就得到了多头注意力的TP了。

1.4.Embedding层的TP

最后还有embedding层。embedding层的做法是每块GPU维护一份embedding的子集，用id去gather向量的时候，各个GPU上分别获取，对于获取不到的id，则先用特殊向量比如零向量先表示，最后再allreduce各个GPU上的向量，替换掉零向量，就获得了完整的embedding输入了。

2.流水并行

流水并行，Pipeline Parallelism，PP：将网络按层切分，划分成多组，一张卡存一组。

TP是对模型宽度进行切分，而PP是对模型的高度进行切分。

# 假设模型有8层：L0~L7
# 两张卡：GPU0,GPU1
=====================   =====================
| L0 | L1 | L2 | L3 |   | L4 | L5 | L6 | L7 |
=====================   =====================
        GPU0                 GPU1

按这个思路，我们可以直接实现naive PP：假设模型有8层，把模型前4层放在一张卡，后4层放在另一张卡；前向的时候把中间激活数据从GPU0传给GPU1，反向的时候则把数据从GPU1传到GPU0。

naive PP的问题是，当GPU0在跑前向的时候，GPU1是没事干的，反过来也有一样的问题，这就导致GPU有大量的空闲时间在等数据。而且随着PP的GPU数量的提升，这个空闲率就越来越高。比如设置8卡的PP，那么GPU0在做前向计算的时候，GPU1到7都在休息。真所谓是一卡有难，七卡围观。这些GPU的空余等待时间叫bubble。

有N张卡的PP，卡的计算利用率就只有1/N。

那么怎么优化PP的GPU利用率呢。

一个自然的想法是，能不能在GPU0算下一个batch的前向数据时，让GPU1在算上一个batch数据的反向呢？是可以的，并且还可以把batch切分成更小的micro-batch，这样就能减少GPU的空闲等待时间。

这就是GPipe。GPipe单个batch进一步拆分为多个Micro-Batch，通过流水线调度不同Micro-Batch的前向和反向计算，减少设备空闲时间。

还有很多别的方案，比如Interleaved Pipeline、1-Forward-1-Backward等，可以看看大佬们的做法。

GPipe的Micro-Batch优化了bubble的问题，那还有显存问题呢。比如GPU1在接收来自GPU0的前向数据时，自己也还有反向传播的中间层数据，这么一来显存就很吃紧了。一个方法就是用activation checkpoint来减少显存的消耗。

实际上个人感觉流水并行是比较复杂的，也有很多不同的实现方法，可以看看框架大佬们的资料。

3.3D并行

3D = DP + TP + PP。

DP是对数据进行切分，TP是对模型宽度进行切分，而PP是对模型的高度进行切分。这三者是可以组合起来使用的。

层内使用TP，层间使用PP，多组TP+PP之间使用DP。一般来说DP可以跨机，而TP和PP的通讯更多，应尽量避免跨机。

看下来自Bloom论文的图：

每个白色方框表示一块GPU，每组机器有48块GPU，每组都复制了一份模型完整参数。左侧表示数据并行DP，有8组机器，每组输入一批数据；右侧图的竖向示意了PP过程，有12行，模型横跨了这12行GPU，例如模型有48层，则每4层放在一行中；右侧图横向示意了TP过程，一行4块GPU，表示这一行的模型参数被平摊到4块GPU上。

看下DeepSpeed博客的版本：

下图是个三维的3D并行示意图。每种颜色表示一个节点，每个节点有4块GPU。上面16张卡和下面16张卡分别是一组，每组输入一份数据，这是数据并行。上面一组16张卡，假设模型有32 layer，一组GPU中每个节点存放8layer，每个节点的输出作为下一个节点的输入，例如GPU0的输出是GPU8的输入，这就是流水线并行。每个节点执行模型并行，意思是每个layer被分成了4分，放到一个节点的4个卡上。

下图是对上图的拓展示意。模型有32 layer，每8个layer放到一个节点，黄色框是一个节点，包含4个GPU。每个节点执行模型并行/张量并行， MP-0、MP-1、MP-2、MP-3表示同一layer中的张量被切分成4份，分别放到4个GPU上。Rank 0 和Rank 1是数据并行。节点之间执行流水线并行，07layer放在第一个节点，以此类推，最后的2431layer放到最后一个节点。

4.序列并行

序列并行主要是解决LLM的输入数据长的问题。由于attention的计算复杂度是平方增长，中间激活值的量随着输入输出长度增长而暴增，naive attention实现的情况下，比如10k长度的序列所需的显存是1k长度的100倍。

前面TP和PP都是切模型，而序列并行就是切数据。

主流的实现有这三种，对比一下：

属性	Colossal-AI	Megatron-LM	DeepSpeed-Ulysses
核心目标	突破序列长度限制，支持超长序列（如114K Token）	减少LayerNorm和Dropout的显存占用，优化张量并行下的显存效率	高效支持超长序列（百万Token）和大模型训练，结合ZeRO-3参数分片
通信机制	环状通信（Ring Self-Attention），分块传递QKV，All-Gather聚合结果	All-Gather和Reduce-Scatter聚合序列分片的中间激活值	All-to-All转置QKV矩阵，将序列分片转换为注意力头分片
兼容性	兼容数据并行、流水线并行、张量并行	主要与张量并行结合使用	与ZeRO-3和数据并行结合，支持FlashAttention优化库
无损性验证	计算结果与单卡完全一致，实验验证Loss曲线和精度指标无差异	分布式与单卡输出的均方误差（MSE）为浮点误差量级（<1e-7）	生成文本的困惑度（Perplexity）与单卡一致，数学等价性通过矩阵分块转置严格保证

如何学习大模型 AI ？

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

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

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

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

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

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

在这里插入图片描述