AI大模型并行训练详解：DP、PP、TP、EP全解析
快速阅读: AI计算主要依赖并行计算，包括数据并行、流水线并行、张量并行和专家并行，这些方法通过多GPU协作加速模型训练，有效提升计算效率，减少训练时间。
大家都知道,AI 计算(尤其是模型训练和推理),主要以并行计算为主。
AI 计算中涉及到的很多具体算法(例如矩阵相乘、卷积、循环层、梯度运算等),都需要基于成千上万的 GPU,以并行任务的方式去完成。这样才能有效缩短计算时间。
搭建并行计算框架,一般会用到以下几种常见的并行方式: Data Parallelism,数据并行 Pipeline Parallelism,流水线并行 Tensor Parallelism,张量并行 Expert Parallelism, 专家并行 接下来,我们逐一看看,这些并行计算方式的工作原理。
▉ DP(数据并行) 首先看看 DP,数据并行(Data Parallelism)。
AI 训练使用的并行,总的来说,分为数据并行和模型并行两类。刚才说的 PP(流水线并行)、TP(张量并行)和 EP(专家并行),都属于模型并行,待会再介绍。
这里,我们需要先大概了解一下神经网络的训练过程。简单来说,包括以下主要步骤: 1、前向传播:输入一批训练数据,计算得到预测结果。
2、计算损失:通过损失函数比较预测结果与真实标签的差距。
3、反向传播:将损失值反向传播,计算网络中每个参数的梯度。
4、梯度更新:优化器使用这些梯度来更新所有的权重和偏置(更新参数)。
以上过程循环往复,直到模型的性能达到令人满意的水平。训练就完成了。
我们回到数据并行。
数据并行是大模型训练中最为常见的一种并行方式(当然,也适用于推理过程)。
它的核心思想很简单,就是每个 GPU 都拥有完整的模型副本,然后,将训练数据划分成多个小批次(mini-batch),每个批次分配给不同的 GPU 进行处理。
数据并行的情况下,大模型训练的过程是这样的: 1、对数据进行均匀切割,发给不同的、并行工作的 GPU(Worker); 2、各 GPU 都拥有一样的模型以及模型参数,它们各自独立进行前向传播、反向传播,计算得到各自的梯度; 3、各 GPU 通过卡间通信,以 All-Reduce 的通信方式,将梯度推给一个类似管理者的 GPU(Server); 4、Server GPU 对所有梯度进行求和或者平均,得到全局梯度; 5、Server GPU 将全局梯度回传(broadcast 广播)到每个 Worker GPU,进行参数更新(更新本地模型权重)。更新后,所有 worker GPU 模型参数保持一致。
然后,再继续重复这样的过程,直至完成所有的训练。
再来一张图,帮助理解: 从下往上看 这里提到的 All-Reduce,也是一个 AI 领域的常见概念,字面意思是“全(All)-规约(Reduce)”,即:对所有节点的数据进行聚合(如求和、求最大值),并将最终结果分发到所有节点。(参考: 到底什么是 All-Reduce、All-to-All? ) 数据并行的优点,在于实现过程比较简单,能够显著加速大规模数据的训练过程,尤其适用于数据量远大于模型参数的场景。
数据并行的缺点,在于显存的限制。因为每个 GPU 上都有完整的模型副本,而当模型的规模和参数越大,所需要的显存就越大,很可能超过单个 GPU 的显存大小。
数据并行的通信开销也比较大。不同 GPU 之间需要频繁通信,以同步模型参数或梯度。而且,模型参数规模越大,GPU 数量越多,这个通信开销就越大。例如,对于千亿参数模型,单次梯度同步需传输约 2TB 数据(FP16 精度下)。
▉ ZeRO 这里要插播介绍一个概念 ——ZeRO(Zero Redundancy Optimizer,零冗余优化器)。
在数据并行策略中,每个 GPU 的内存都保存一个完整的模型副本,很占内存空间。那么,能否每个 GPU 只存放模型副本的一部分呢? 没错,这就是 ZeRo—— 通过对模型副本中的优化器状态、梯度和参数进行切分,来实现减少对内存的占用。
ZeRO 有 3 个阶段,分别是: ZeRO-1:对优化器状态进行划分。
ZeRO-2:对优化器状态和梯度进行划分 ZeRO-3:对优化器状态、梯度和参数进行划分。(最节省显存) 通过下面的图和表,可以看得更明白些: 根据实测数据显示,ZeRO-3 在 1024 块 GPU 上训练万亿参数模型时,显存占用从 7.5TB 降至 7.3GB / 卡。
值得一提的是,DP 还有一个 DDP(分布式数据并行)。传统 DP 一般用于单机多卡场景。而 DDP 能多机也能单机。这依赖于 Ring-AllReduce,它由百度最先提出,可以有效解决数据并行中通信负载不均(Server 存在瓶颈)的问题。
▉ PP(流水线并行) 再来看看模型并行。
刚才数据并行,是把数据分为好几个部分。模型并行,很显然,就是把模型分为好几个部分。不同的 GPU,运行不同的部分。(注意:业界对模型并行的定义有点混乱。也有的资料会将张量并行等同于模型并行。) 流水线并行,是将模型的不同层(单层,或连续的多层)分配到不同的 GPU 上,按顺序处理数据,实现流水线式的并行计算。
例如,对于一个包含 7 层的神经网络,将 1~2 层放在第一个 GPU 上,3~5 层放在第二个 GPU 上,6~7 层放在第三个 GPU 上。训练时,数据按照顺序,在不同的 GPU 上进行处理。
乍一看,流水并行有点像串行。每个 GPU 需要等待前一个 GPU 的计算结果,可能会导致大量的 GPU 资源浪费。
上面这个图中,黄色部分就是 Bubble (气泡)时间。气泡越多,代表 GPU 处于等待状态(空闲状态)越长,资源浪费越严重。
为了解决上述问题,可以将 mini-batch 的数据进一步切分成 micro-batch 数据。当 GPU 0 处理完一个 micro-batch 数据后,紧接着开始处理下一个 micro-batch 数据,以此来减少 GPU 的空闲时间。如下图(b)所示: 还有,在一个 micro-batch 完成前向计算后,提前调度,完成相应的反向计算,这样就能释放部分显存,用以接纳新的数据,提升整体训练性能。如上图(c)所示。
这些方法,都能够显著减少流水线并行的 Bubble 时间。
对于流水线并行,需要对任务调度和数据传输进行精确管理,否则可能导致流水线阻塞,以及产生更多的 Bubble 时间。
▉ TP(张量并行) 模型并行的另外一种,是张量并行。
如果说流水线并行是将一个模型按层「垂直」分割,那么,张量并行则是在一个层内「横向」分割某些操作。
具体来说,张量并行是将模型的张量(如权重矩阵)按维度切分到不同的 GPU 上运行的并行方式。
张量切分方式分为按行进行切分和按列进行切分,分别对应行并行(Row Parallelism)(权重矩阵按行分割)与列并行(Column Parallelism)(权重矩阵按列分割)。
每个节点处理切分后的子张量。最后,通过集合通信操作(如 All-Gather 或 All-Reduce)来合并结果。
张量并行的优点,是适合单个张量过大的情况,可以显著减少单个节点的内存占用。
张量并行的缺点,是当切分维度较多的时候,通信开销比较大。而且,张量并行的实现过程较为复杂,需要仔细设计切分方式和通信策略。
放一张数据并行、流水线并行、张量并行的简单对比: ▉ 专家并行 2025 年初 DeepSeek 爆红的时候,有一个词也跟着火了,那就是 MoE(Mixture of Experts,混合专家模型)。
MoE 模型的核心是“多个专家层 + 路由网络(门控网络)”。
专家层的每个专家负责处理特定类型的 token(如语法、语义相关)。路由网络根据输入 token 的特征,选择少数专家处理这个 token,其他专家不激活。
MoE 实现了任务分工、按需分配算力,因此大幅提升了模型效率。
专家并行(Expert Parallelism),是 MoE(混合专家模型)中的一种并行计算策略。它通过将专家(子模型)分配到不同的 GPU 上,实现计算负载的分布式处理,提高计算效率。
专家并行与之前所有的并行相比,最大的不同在于,输入数据需要通过一个动态的路由选择机制分发给相应专家,此处会涉及到一个所有节点上的数据重分配的动作。
然后,在所有专家处理完成后,又需要将分散在不同节点上的数据按原来的次序整合起来。
这样的跨片通信模式,称为 All-to-All。(再次参考: 到底什么是 All-Reduce、All-to-All? ) 专家并行可能存在负载不均衡的问题。某个专家所接收到的输入数据大于了其所能接收的范围,就可能导致 Tokens 不被处理或不能被按时处理,成为瓶颈。
所以,设计合理的门控机制和专家选择策略,是部署专家并行的关键。
▉ 混合并行 在实际应用中,尤其是训练万亿参数级别的超大模型时,几乎不会只使用单一的并行策略,而是采用多维度的混合并行(结合使用多种并行策略)。
例如: 数据并行 + 张量并行:数据并行处理批量样本,张量并行处理单样本的大矩阵计算。
流水线并行 + 专家并行:流水线并行划分模型层,专家并行划分层内专家模块。
更高级的,是 3D 并行,通过“数据并行 + 张量并行 + 流水线并行”,实现三重拆分,是超大模型训练的主流方案。
3D 并行 ▉ 最后的话 好啦,以上就是关于 DP、PP、TP、EP 等并行训练方式的介绍。大家都看懂了没? 并行计算方式其实非常复杂,刚才我们只是做了最简单的介绍。但在真实工作中,开发者无需了解具体的实现细节,因为业界提供了例如 DeepSpeed(微软开源,支持 3D 并行 +ZeRO 内存优化)、Megatron-LM(NVIDIA 开源,3D 并行的标杆)、FSDP 等开源软件,能够让开发者直接进行大语言模型训练。
小枣君之所以要专门介绍并行训练方式,其实更多是为了帮助大家深入地理解算力集群架构和网络的设计。
大家可以看到,不同的并行训练方式,有着不同的通信流量特点。
算力集群整体架构和网络设计,需要尽量去适配这些并行计算方式的流量特点,才能满足模型训推任务的要求,实现更高的工作效率。
比如说,数据并行,由于需要频繁同步梯度信息,对网络带宽要求较高,需要确保网络带宽能够满足大量梯度数据快速传输的需求,避免因带宽不足导致通信延迟,影响训练效率。
流水线并行,大模型的每一段,在不同的服务器上以流水线的方式逐步计算,涉及到多个服务器“串起来”,就建议部署在比较靠近的服务器上(尽量部署在叶脊网络的同一个 leaf 叶下)。
张量并行,通信数据量大,就建议部署在一台服务器的多个 GPU 上进行计算。
专家并行中,不同专家分配在不同 GPU 上,GPU 间需要交换中间计算结果等信息,其通信流量特点取决于专家的数量以及数据交互的频率等,也需要合理规划 GPU 间的连接方式和通信路径。
总之,在 GPU 算卡性能越来越难以提升的背景下,深入研究并行计算的设计,从架构和网络上挖掘潜力,是业界的必然选择。
随着 AI 浪潮的继续发展,以后是否还会出现其它的并行训练方式呢?让我们拭目以待吧! 本文来自微信公众号: 鲜枣课堂(ID:xzclasscom) ,作者:小枣君 广告声明:文内含有的对外跳转链接(包括不限于超链接、二维码、口令等形式),用于传递更多信息,节省甄选时间,结果仅供参考,IT之家所有文章均包含本声明。
(以上内容均由Ai生成)
