大模型原理--分布式训练策略之流水线并行、张量并行和专家并行

1.概述

GPT-2 揭示了规模效应（Scaling Effect）：随着模型参数与数据规模的增长，模型性能在多种任务上持续提升。基于此，大语言模型（Large Language Model, LLM）的参数量和训练数据规模必然持续增长，分布式训练已成为现代大模型训练体系中的核心技术。

本文主要介绍三个主流的分布式训练策略：流水线并行，张量并行和专家并行。

2. 流水线并行

流水线并行的核心思想是将大语言模型按层纵向切分，把不同层分配到多个设备上，每个设备负责一个模型分段。可以降低单设备的显存占用，使超大规模模型能够在单设备硬件资源有限的情况下，通过增加设备硬件的数量完成训练。

训练过程中，前向计算从第一个设备开始，逐级将中间激活值传递给下一个设备，直到完成整个模型的前向传播；反向传播则以相反顺序依次进行，梯度逐级回传和更新参数。这种按层接力的执行方式会导致部分设备在等待上下游计算完成时处于空闲状态，造成计算资源浪费，如下图所示：

Worker代表硬件设备，F代表前向传播，F下标编号代表分层的编号，B代表反向传播，B下标编号代表分层的编号。U代表参数更新，U下标编号代表分层的编号。

为提高不同硬件设备资源的利用率，主流方法是引入微批次（Micro-batches）与流水线调度策略，使相同时间段内不同硬件设备都有计算任务可执行，从而提高流水线并行的硬件设备资源的利用率。具体做法是：将原始批次进一步划分为多个更小的微批次。流水线中每个设备在处理完某个微批次（()内的编号是微批的编号）后，可立即切换到下一个微批次的计算。如下图所示：

总结：流水线并行通过将整体模型按照层，纵向切分，将模型分段的计算分摊给其他的GPU硬件，从而降低了对单个GPU的显存要求。

3. 张量并行

张量并行的核心思想是将模型单层内部的张量计算拆分到多个设备上并行执行，每个设备仅持有部分权重和中间结果。其目标仍是突破单设备显存和计算能力的限制，使超大模型的单层也能在多设备上协同完成。

在大模型计算中，计算瓶颈通常出现在激活矩阵与大权重矩阵的乘法操作上。由于矩阵乘法可分解为多个独立或可累加的点积运算，因此可以将权重矩阵均匀切分，每个设备仅持有其中一部分，然后将激活数据广播到所有设备，各设备使用本地分片权重独立完成部分矩阵乘法，再通过设备间通信聚合计算结果，得到完整的矩阵乘积。

总结：张量并行在流水线并行的基础上，对模型的单层进行横向切分，将计算分摊给其他GPU硬件，从而降低了对单个GPU的显存和计算要求。

5. 专家并行

专家并行的核心思想是将 Mixture-of-Experts（MoE）模型中的多个专家（Expert）分配到不同的设备上，每个设备仅存储和计算部分专家。MoE 结构天然适合并行化的原因是每个专家本质上是一个独立的前馈网络（FFN），彼此之间没有共享参数或依赖关系，因此可直接将不同专家切分并分布到多个设备上。

总结：通过各种形式的切分，将模型以及数据分配到其他GPU硬件，从而突破单设备在模型容量上的限制，使训练超大规模模型成为可能。