趋近智
数据并行(DP)通过复制模型和处理不同的数据块来有效使用多个设备,但这要求每个设备承载整个模型。对于规模很大的模型,即使是单个副本也可能超出单个加速器的内存容量。此外,模型中的一些操作,例如前馈网络或注意力机制 (attention mechanism)中的大型线性变换,可能成为计算瓶颈。在这种情况下,张量并行(TP),有时也称为层内模型并行,是必不可少的。
张量并行不是分割数据或层序列,而是在特定操作内部,将实际张量(权重 (weight)、激活、梯度)拆分到多个设备上。这使得对这些大张量的计算能够并行进行,分散单个层的内存占用和计算负担。
张量并行最直接的应用是在大型线性层中,它们是Transformer中多层感知机(MLP)和注意力块的重要组成部分。假设一个线性变换 ,其中是输入激活,是权重 (weight)矩阵。我们可以将这种矩阵乘法并行化,例如在两个设备上,主要有两种方式:列并行和行并行。
在列并行中,权重矩阵垂直(沿其列)拆分到设备上。如果我们有两个设备,我们将拆分为 。输入通常会被广播或提供给两个设备。每个设备随后计算一部分输出:
最终输出通过沿列维度拼接结果获得:。
的列并行。权重矩阵被拆分为和。输入在不同的GPU上与每个部分相乘。结果和被拼接起来形成最终输出。
从数学角度看,前向传播涉及 。反向传播 (backpropagation)需要计算相对于和的梯度。梯度 是自然地划分的(,)。然而,计算梯度 需要汇总来自两条路径的贡献:。这种求和通常通过在持有和的设备之间使用all-reduce通信操作来完成。
在行并行中,权重矩阵水平(沿其行)拆分。对于两个设备,。输入也被认为是沿其列划分的(通常因为它是前一个列并行层的输出)。为了在这种独立的视角下简化,我们假设输入在两个设备上都可用。每个设备根据其权重切片计算部分结果:
与列并行不同,最终输出是部分结果的和:。
的行并行。权重矩阵被拆分为和(行方向)。部分结果和在不同的GPU上计算。一个all-reduce操作将这些结果求和以生成最终输出。
从数学角度看,如果被视为一个单独的块,并非标准的矩阵乘法表示。在这种情况下(例如,紧随一个列拆分层之后)的实际操作是,其中。前向传播需要all-reduce操作以在设备之间执行此求和。反向传播中的计算涉及基于已分区的分区梯度:和。梯度直接在每个设备上为其对应的权重切片计算。
张量并行通常有策略地应用于Transformer块内部,以平衡计算并减少通信开销。Megatron-LM等框架推广的一种常见模式是在MLP块中结合了列并行和行并行,并对注意力机制 (attention mechanism)应用了类似的分区。
MLP块: 标准的Transformer MLP块计算 (包含残差连接)。张量并行按以下方式应用:
这种组合巧妙地安排了并行,使得前向传播所需的通信(第二线性层后的all-reduce)和反向传播所需的通信(第一线性层后的all-reduce)不会不必要地重叠,从而优化了流程。
注意力块: 张量并行也可应用于自注意力 (self-attention)机制。
注意力块内部的实现细节可能比较复杂,包含优化的核函数和通信模式,以高效处理序列长度和头维度。
张量并行一个主要的缺点是与数据并行相比,它增加了通信开销。虽然数据并行通常每个训练步骤(对于梯度)涉及一次all-reduce,但张量并行在每个Transformer块的前向和反向传播 (backpropagation)内部引入了通信。
这些通信操作(如all-reduce)涉及在所有参与张量并行组的设备之间同步和交换数据。交换的数据量取决于激活或梯度的大小。这种通信成本随张量并行组中设备数量的增加而变化,如果设备间互连带宽(例如NVLink、InfiniBand)相对于计算速度不足,可能成为性能瓶颈。
然而,请注意以下几点:
这是一个PyTorch风格的代码片段,展示了为列并行拆分权重 (weight)矩阵的思路(这高度简化,省略了通信和梯度处理):
import torch
import torch.nn as nn
# 假设 world_size 是用于张量并行的 GPU 数量
# 假设 rank 是当前 GPU 在张量并行组中的排名
class ColumnParallelLinear(nn.Module):
def __init__(self, input_size, output_size, world_size, rank):
super().__init__()
self.input_size = input_size
# 每个 GPU 处理输出特征的一个切片
self.output_size_per_partition = output_size // world_size
self.world_size = world_size
self.rank = rank
# 只初始化当前 rank 对应的权重矩阵部分
self.weight = nn.Parameter(
torch.randn(self.output_size_per_partition, self.input_size)
)
# 偏置项也进行分区
self.bias = nn.Parameter(
torch.randn(self.output_size_per_partition)
)
def forward(self, x):
# 假设输入 x 在所有 GPU 上都可用(已广播或来自
# all-reduce 的结果)
# 矩阵乘法只在权重的本地分区上进行
# Output_partition = X * A_partition^T + b_partition
# 注意:PyTorch 线性层期望权重为
# (out_features, in_features)
output_partition = nn.functional.linear(x, self.weight, self.bias)
# 在实际实现中,如果下一层不是行并行,
# 这个 output_partition 需要在 GPU 之间进行收集。
# 对于 Megatron-LM 的 MLP 模式(列并行->行并行),
# 此处不需要前向通信。
# 反向传播通信(grad_X 的 all-reduce)由
# Megatron-LM 等库中的自定义自动梯度函数处理。
return output_partition # 只返回此 GPU 计算的切片
总结来说,张量并行是一种有效的技术,用于将单个大层的计算和内存分散到多个设备上。虽然它带来了较大的通信开销,但它常常是一个重要的构成部分,与数据并行和流水线并行一起,用于训练突破单加速器能力限制的先进大型语言模型。Megatron-LM和DeepSpeed等框架大大简化了复杂性,提供了优化的构建模块,以有效实现张量并行。
这部分内容有帮助吗?
© 2026 ApX Machine LearningAI伦理与透明度•