Page 213 - 《软件学报》2026年第1期

210 软件学报 2026 年第 37 卷第 1 期

量分桶, 并以批处理方式进行梯度通信, 同时注册反向传播钩子. 梯度张量准备好后进行通信和数据传递, 从而实
现反向传播和梯度通信的重叠, 进一步提升性能.
在数据并行的异步通信中, 如果采用框架不可知的实现方式, 则需要控制平面来确保通信顺序以防止死锁.
Horovod 提出新的协调模式, 通过高效去中心化的编排策略, 以降低控制平面的开销. 框架钩子方式一般适合只使
用一种框架后端, 目前由于大部分的大语言模型基于 PyTorch 为后端, 所以目前许多大语言模型库的数据并行策
略后端默认采用 PyTorch DDP 的数据并行方案, 例如 Transformers, FairSeq [48] 和 TorchScale [49] 等.
3.2.1.2 张量并行
提高模型容量被认为是提升模型质量的有效手段, 但当模型尺寸超出单个加速器内存限制时, 需要开发特殊
算法以支持模型的训练. 在数据并行中, 要求将整个模型部署在单个 GPU 上, 当模型过大时无法实现. 为应对这一
挑战, 一些系统提出了模型并行方案, 通过对模型进行切片, 降低显存占用, 以便训练更大的模型. 模型并行方案包
括张量并行和流水线并行, 其中张量并行对单层进行切分, 而流水线并行在层间进行切分. Megatron-LM 使用了简
单高效的层内模型并行, 也被称作张量并行方法. 其实现方式不需要编译或更改库实现, 且与流水并行方案正交不
冲突. 张量并行通常有两种方案来切分通用矩阵乘中的权重矩阵, 以两层 MLP 为例, 第 1 层的通用矩阵乘之后还
要进行 GELU 激活函数运算, Y = GELU (XA).
行切分方式 [26] : 一种选择是通过行切分权重矩阵 A, 以及列切分输入矩阵 . 行切分权重矩阵方式为 X =
X
[ ]
A 1
[X 1 ,X 2 ],A = , 这种切分方式计算公式为 Y = GELU(X 1 A 1 + X 2 A 2 ). 由于 GELU 是一个非线性函数, GELU(X 1 A 1 +
A 2
X 2 A 2 ) , GELU (X 1 A 1 )+GELU(X 2 A 2 ), 最终这种方式需要一个在 GELU 激活函数之前的同步点.
列切分方式 [26] : 另一种方式是在列上切分矩阵 A = [A 1 ,A 2 ], 这种切分可以让 GELU 独立地作用于每个切片的
输出矩阵上 [Y 1 ,Y 2 ] = [GELU(XA 1 ),GELU(XA 2 )].
两种方案各有优劣. 行切方案如果后续结果需要完成 GELU 激活函数, 则需要在给 GELU 输入张量之前, 先
进行 All-Reduce 通信, 聚合矩阵乘计算产生的张量, 进而完成矩阵乘计算; 列切方案则生成两个独立的矩阵输出,
且无需通过 All-Reduce 完成后续激活操作, 但对于需要完整张量输入的后续操作, 则需使用 AllGather 聚合输出结
果的张量切片. Megatron-LM 根据 Transformer 的计算特点, 设计 MLP 第 1 层的通用矩阵乘使用列切权重矩阵方
式, 第 2 层使用行切的权重矩阵方式, 进而只需要在第 2 个通用矩阵乘的输出上跨 GPU 进行一次聚合. 相比于
Megatron-LM 依赖混合 Python 和优化好的 CUDA 内核的方式, MaxText 利用 TensorFlow、JAX 和 XLA 的静态
编译, 进而达到通过纯 Python 代码实现并自动优化内核, 其内部也支持数据并行, 全切片数据并行, 序列并行和张
量并行.
3.2.1.3 流水并行
张量并行方案通常是为特定结构设计的, 难以迁移到其他任务. 为了实现高效且与任务无关的模型并行性,
Google 提出了一个流水并行库 GPipe, 允许扩展任何可以表达为层结构的模型. 通过在单独的加速器上流水线化
不同的层序列, GPipe 提供了缩放和切片各种不同网络结构的灵活性, 将网络扩展到更大的规模. 在前向传播计算
中, 每个加速器只存储切片边界的输出激活. 在反向传播过程中第 k 个加速器重算复合前向函数 F k . 因此峰值激
( L N ) N
活内存需求降低为 O N + × , 其中 N 是批尺寸, K 是切片分区数量, L 是模型层数, M 是微批次的数量,
K M M
L
是微批尺寸, 是每个切片的层数. 与基线模型的内存需求 O(N × L) 相比, 这节省了缓存的激活张量内存. 由于
K
采用了流水并行, 每个加速器会产生一些空闲时间, 称作气泡 (bubble) 时间. 其中气泡时间占比的复杂度为
( )
K −1
O , 气泡时间可以被微批次的尺寸 M 所摊销, 越多的微批次尺寸, 气泡占比越低. GPipe 发现当
M + K −1
M ⩾ 4× K 的时候, 气泡开销几乎可以忽略, 原因在反向传播时候的重算可以更早地进行调度执行, 从而填充气泡
时间的闲置计算资源. 为了降低流水并行中的内存开销, Narayanan 等人 [50] 提出了“一轮前向传播紧接着一轮反向
传播 (one forward pass followed by one backward pass, 1F1B)”, 虽然这并没有降低气泡的复杂度, 但由于能够更短