当前位置:首页 > 科技资讯 > 正文

AI并行计算方式详解:从数据并行到专家并行

众所周知,人工智能计算,特别是在模型训练和推理阶段,主要依赖于并行计算技术。

在AI计算中,许多核心算法(如矩阵乘法、卷积操作、循环层处理、梯度计算等)需要借助成千上万的GPU,通过并行任务的方式执行,以显著缩短计算时间。

搭建并行计算框架时,通常会采用以下几种常见的并行策略:

  • Data Parallelism,数据并行
  • Pipeline Parallelism,流水线并行
  • Tensor Parallelism,张量并行
  • Expert Parallelism, 专家并行

接下来,我们将逐一剖析这些并行计算方式的工作原理。

DP(数据并行)

首先关注数据并行(Data Parallelism)。

AI训练所使用的并行策略,总体上可分为数据并行和模型并行两大类。前面提到的PP(流水线并行)、TP(张量并行)和EP(专家并行)均属于模型并行范畴,稍后会详细介绍。

AI并行计算方式详解:从数据并行到专家并行 数据并行  流水线并行 张量并行 专家并行 第1张

在此之前,我们有必要简要回顾神经网络的训练流程。简单来说,训练过程包含以下几个主要步骤:

AI并行计算方式详解:从数据并行到专家并行 数据并行  流水线并行 张量并行 专家并行 第2张

1、前向传播:输入一批训练数据,计算得出预测结果。

2、计算损失:通过损失函数比较预测结果与真实标签的差距。

3、反向传播:将损失值反向传播,计算网络中每个参数的梯度。

4、梯度更新:优化器使用这些梯度来更新所有权重和偏置(参数更新)。

上述过程循环往复,直至模型性能达到理想水平,训练即告完成。

现在回到数据并行的话题。

数据并行是大模型训练中最常用的并行方式之一(当然也适用于推理过程)。

其核心思想非常直观:每个GPU都持有完整的模型副本,同时将训练数据划分为多个小批次(mini-batch),每个小批次分配给不同的GPU进行处理。

在数据并行模式下,大模型训练的流程如下:

AI并行计算方式详解:从数据并行到专家并行 数据并行  流水线并行 张量并行 专家并行 第3张

1、对数据进行均匀切割,分发给并行工作的不同GPU(Worker);

2、各GPU拥有相同的模型及参数,独立进行前向传播和反向传播,计算各自的梯度;

3、各GPU通过卡间通信,以All-Reduce方式将梯度推送给一个类似管理者的GPU(Server);

4、Server GPU对所有梯度进行求和或平均,得到全局梯度;

5、Server GPU将全局梯度广播(broadcast)到每个Worker GPU,用于参数更新(更新本地模型权重)。更新后,所有worker GPU的模型参数保持一致。

之后不断重复此过程,直至训练全部完成。

再借助一张图来加深理解:

AI并行计算方式详解:从数据并行到专家并行 数据并行  流水线并行 张量并行 专家并行 第4张

从下往上看

这里提到的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分为三个阶段:

ZeRO-1:对优化器状态进行划分。

ZeRO-2:对优化器状态和梯度进行划分。

ZeRO-3:对优化器状态、梯度和参数进行划分(最节省显存)。

通过下图和表格可以更清晰地理解:

AI并行计算方式详解:从数据并行到专家并行 数据并行  流水线并行 张量并行 专家并行 第5张

AI并行计算方式详解:从数据并行到专家并行 数据并行  流水线并行 张量并行 专家并行 第6张

根据实测数据显示,ZeRO-3在1024块GPU上训练万亿参数模型时,显存占用从7.5TB降至7.3GB/卡。

值得一提的是,DP还有一个DDP(分布式数据并行)变体。传统DP一般用于单机多卡场景,而DDP支持多机多卡,这得益于百度最先提出的Ring-AllReduce技术,它能有效解决数据并行中通信负载不均(Server成为瓶颈)的问题。

AI并行计算方式详解:从数据并行到专家并行 数据并行  流水线并行 张量并行 专家并行 第7张

PP(流水线并行)

接下来探讨模型并行。

数据并行是将数据拆分,而模型并行显然是把模型拆分为多个部分,由不同的GPU分别运行不同的部分(注意:业界对模型并行的定义存在一定混淆,有些资料会将张量并行等同于模型并行)。

流水线并行是将模型的不同层(单层或连续多层)分配到不同的GPU上,按顺序处理数据,形成流水线式的并行计算。

AI并行计算方式详解:从数据并行到专家并行 数据并行  流水线并行 张量并行 专家并行 第8张

例如,对于一个包含7层的神经网络,可将1~2层放在第一个GPU,3~5层放在第二个GPU,6~7层放在第三个GPU。训练时,数据按顺序在不同GPU上依次处理。

乍看之下,流水线并行有些类似串行:每个GPU需要等待前一个GPU的计算结果,可能导致大量GPU资源闲置。

AI并行计算方式详解:从数据并行到专家并行 数据并行  流水线并行 张量并行 专家并行 第9张

上图中黄色部分即Bubble(气泡)时间。气泡越多,表示GPU处于等待(空闲)状态的时间越长,资源浪费越严重。

为解决这一问题,可以将mini-batch数据进一步切分成micro-batch。当GPU 0处理完一个micro-batch后,立即开始处理下一个micro-batch,从而减少GPU的空闲时间。如下图所示(b):

AI并行计算方式详解:从数据并行到专家并行 数据并行  流水线并行 张量并行 专家并行 第10张

此外,在一个micro-batch完成前向计算后,可提前调度其反向计算,从而释放部分显存,接纳新数据,提升整体训练性能。如上图(c)所示。

这些方法能显著减少流水线并行的Bubble时间。

对于流水线并行,需要精确管理任务调度和数据传输,否则可能导致流水线阻塞,产生更多Bubble时间。

TP(张量并行)

模型并行的另一种形式是张量并行。

如果说流水线并行是将模型按层“垂直”分割,那么张量并行则是在一个层内“横向”分割某些操作。

AI并行计算方式详解:从数据并行到专家并行 数据并行  流水线并行 张量并行 专家并行 第11张

具体而言,张量并行是将模型的张量(如权重矩阵)按维度切分到不同GPU上并行运行。

切分方式分为按行切分和按列切分,分别对应行并行(Row Parallelism)(权重矩阵按行分割)与列并行(Column Parallelism)(权重矩阵按列分割)。

AI并行计算方式详解:从数据并行到专家并行 数据并行  流水线并行 张量并行 专家并行 第12张

每个节点处理切分后的子张量,最后通过集合通信操作(如All-Gather或All-Reduce)合并结果。

AI并行计算方式详解:从数据并行到专家并行 数据并行  流水线并行 张量并行 专家并行 第13张

张量并行的优点在于适合处理单个张量过大的情况,能显著减少单个节点的内存占用。

张量并行的缺点是当切分维度较多时,通信开销较大。而且实现过程较为复杂,需要精心设计切分方式和通信策略。

下面用一张图简要对比数据并行、流水线并行和张量并行:

AI并行计算方式详解:从数据并行到专家并行 数据并行  流水线并行 张量并行 专家并行 第14张

专家并行

2025年初DeepSeek爆红时,有一个术语也随之火热,那就是MoE(Mixture of Experts,混合专家模型)。

MoE模型的核心是“多个专家层+路由网络(门控网络)”。

AI并行计算方式详解:从数据并行到专家并行 数据并行  流水线并行 张量并行 专家并行 第15张

专家层的每个专家负责处理特定类型的token(如语法、语义相关)。路由网络根据输入token的特征,选择少数专家处理该token,其他专家则不激活。

MoE实现了任务分工和按需分配算力,从而大幅提升模型效率。

专家并行(Expert Parallelism)是MoE(混合专家模型)中的一种并行计算策略。它将专家(子模型)分配到不同的GPU上,实现计算负载的分布式处理,提高计算效率。

专家并行与前述所有并行方式的最大不同在于,输入数据需要通过动态路由选择机制分发给相应专家,这涉及所有节点上的数据重分配动作。

在所有专家处理完成后,又需要将分散在不同节点上的数据按原始次序整合起来。

这种跨片通信模式称为All-to-All。(再次参考:到底什么是All-Reduce、All-to-All?)

专家并行可能存在负载不均衡的问题。如果某个专家接收到的输入数据超过其处理能力,可能导致Tokens不被处理或延迟,成为瓶颈。

因此,设计合理的门控机制和专家选择策略,是部署专家并行的关键。

混合并行

在实际应用中,尤其是训练万亿参数级别的超大模型时,几乎不会只使用单一的并行策略,而是采用多维度的混合并行(结合多种并行策略)。

例如:

数据并行+张量并行:数据并行处理批量样本,张量并行处理单样本的大矩阵计算。

流水线并行+专家并行:流水线并行划分模型层,专家并行划分层内专家模块。

更高级的是3D并行,通过“数据并行+张量并行+流水线并行”实现三重拆分,是超大模型训练的主流方案。

AI并行计算方式详解:从数据并行到专家并行 数据并行  流水线并行 张量并行 专家并行 第16张

3D并行

最后的话

好啦,以上就是关于DP、PP、TP、EP等并行训练方式的介绍。大家都看懂了没?

AI并行计算方式详解:从数据并行到专家并行 数据并行  流水线并行 张量并行 专家并行 第17张

并行计算方式其实非常复杂,刚才我们只是做了最简单的介绍。但在真实工作中,开发者无需了解具体的实现细节,因为业界提供了例如DeepSpeed(微软开源,支持3D并行+ZeRO内存优化)、Megatron-LM(NVIDIA开源,3D并行的标杆)、FSDP等开源软件,能够让开发者直接进行大语言模型训练。

小枣君之所以要专门介绍并行训练方式,其实更多是为了帮助大家深入地理解算力集群架构和网络的设计。

大家可以看到,不同的并行训练方式,有着不同的通信流量特点。算力集群整体架构和网络设计,需要尽量去适配这些并行计算方式的流量特点,才能满足模型训推任务的要求,实现更高的工作效率。

比如说,数据并行,由于需要频繁同步梯度信息,对网络带宽要求较高,需要确保网络带宽能够满足大量梯度数据快速传输的需求,避免因带宽不足导致通信延迟,影响训练效率。

流水线并行,大模型的每一段,在不同的服务器上以流水线的方式逐步计算,涉及到多个服务器“串起来”,就建议部署在比较靠近的服务器上(尽量部署在叶脊网络的同一个leaf叶下)。

张量并行,通信数据量大,就建议部署在一台服务器的多个GPU上进行计算。

专家并行中,不同专家分配在不同GPU上,GPU间需要交换中间计算结果等信息,其通信流量特点取决于专家的数量以及数据交互的频率等,也需要合理规划GPU间的连接方式和通信路径。

AI并行计算方式详解:从数据并行到专家并行 数据并行  流水线并行 张量并行 专家并行 第18张

总之,在GPU算卡性能越来越难以提升的背景下,深入研究并行计算的设计,从架构和网络上挖掘潜力,是业界的必然选择。

随着AI浪潮的继续发展,以后是否还会出现其它的并行训练方式呢?让我们拭目以待吧!