数据并行在不同设备间复制模型并处理不同的数据批次时,当模型本身过大,无法载入单个加速器(GPU)的内存时,这种方法就无能为力了。针对此类情况,我们需要模型并行,它将一个模型划分到多个设备上。张量模型并行(TMP)是一种模型并行方式,它将单个层或层内的特定操作拆分到多个设备上。这种方法对于那些在特定层中参数量巨大的模型尤其适用,例如现代Transformer架构中常见的大型嵌入表或前馈网络(FFN)层。基本原理:张量拆分张量模型并行的基本原理是有序地将层的权重张量(有时也包括激活张量)拆分到多个GPU上。然后,计算在每个GPU上部分地执行,接着通过通信步骤来同步或合并结果,最终得到与原始未拆分层相同的输出。我们以一个标准线性层为例,其操作定义为 $Y = XA + b$,其中 $X$ 是输入激活,$A$ 是权重矩阵,$b$ 是偏置,$Y$ 是输出激活。张量模型并行提供不同的策略来并行化此操作。列并行在列并行中,权重矩阵 $A$ 按列拆分到 $N$ 个GPU上。设 $A = [A_1, A_2, ..., A_N]$,其中每个 $A_i$ 位于不同的GPU上。输入广播:输入 $X$ 广播(或已可用)到模型并行组中的所有GPU。并行计算:每个GPU $i$ 计算一个部分结果 $Y_i = X A_i$。收集/拼接:部分结果 $Y_i$ 沿列维度收集并拼接,形成最终输出 $Y = [Y_1, Y_2, ..., Y_N]$。偏置 $b$ 也可以按列拆分,$b = [b_1, b_2, ..., b_N]$,并在收集步骤前在每个GPU上独立添加。digraph G { rankdir=LR; node [shape=box, style=filled, fontname="Arial", fontsize=10]; edge [fontname="Arial", fontsize=10]; subgraph cluster_input { label = "输入"; style=invis; X [label="输入 X", fillcolor="#e9ecef"]; } subgraph cluster_gpus { label = "GPU (N个设备)"; style=dotted; rank=same; node [fillcolor="#a5d8ff"]; GPU1 [label="X * A1 + b1"]; GPU2 [label="X * A2 + b2"]; GPUn [label="..."]; } subgraph cluster_output { label = "输出"; style=invis; Y [label="输出 Y = [Y1, Y2, ..., Yn]", fillcolor="#e9ecef"]; } X -> GPU1 [label="广播"]; X -> GPU2 [label="广播"]; X -> GPUn [label="广播"]; GPU1 -> Y [label="拼接"]; GPU2 -> Y [label="拼接"]; GPUn -> Y [label="拼接"]; }列并行应用于线性层。输入 $X$ 是共享的,权重 $A$ 和偏置 $b$ 按列拆分,部分输出被拼接。这种方法在并行计算后需要通信(收集或全收集操作),用于组装完整的输出张量 $Y$。行并行或者,我们可以按行拆分权重矩阵 $A$:$A = \begin{bmatrix} A_1 \ A_2 \ \vdots \ A_N \end{bmatrix}$。输入拆分/分散:输入 $X$ 沿其最后一个维度(特征维度)在GPU间拆分:$X = [X_1, X_2, ..., X_N]$。注意,这与列并行不同,在列并行中,每个GPU都需要完整的 $X$。并行计算:每个GPU $i$ 使用其部分的输入和权重计算一个部分结果:$X_i A_i$。全归约:部分结果使用全归约操作在所有GPU上求和,生成最终输出 $Y = \sum_i X_i A_i$。偏置 $b$ 通常只在一个GPU上(例如,rank 0)在全归约求和之后添加。digraph G { rankdir=LR; node [shape=box, style=filled, fontname="Arial", fontsize=10]; edge [fontname="Arial", fontsize=10]; subgraph cluster_input { label = "输入"; style=invis; X [label="输入 X", fillcolor="#e9ecef"]; } subgraph cluster_gpus { label = "GPU (N个设备)"; style=dotted; rank=same; node [fillcolor="#ffd8a8"]; GPU1 [label="X1 * A1"]; GPU2 [label="X2 * A2"]; GPUn [label="..."]; } subgraph cluster_output { label = "输出"; style=invis; Y [label="输出 Y = Sum(Xi * Ai) + b", fillcolor="#e9ecef"]; } X -> GPU1 [label="拆分"]; X -> GPU2 [label="拆分"]; X -> GPUn [label="拆分"]; GPU1 -> Y [label="全归约"]; GPU2 -> Y [label="全归约"]; GPUn -> Y [label="全归约"]; }行并行应用于线性层。输入 $X$ 被拆分,权重 $A$ 按行拆分,计算部分结果,然后通过全归约求和。偏置 $b$ 在归约后添加。行并行在矩阵乘法之后需要通信(一个全归约操作)。一个重要好处是,输出激活 $Y$ 在所有参与的GPU上都得到复制,这可能是后续层(例如层归一化或另一个行并行线性层)所需的输入格式。Transformer中列并行与行并行的结合Transformer模型通常结合使用这些技术。例如,在一个由两个线性层组成的标准前馈网络(FFN)块中:第一个线性层可能使用列并行。这将较大的中间维度拆分到多个GPU上。输出激活是分布的(分片的)。第二个线性层可能使用行并行。它将分片激活作为输入(如果拆分维度匹配则无需通信),并生成通过全归约求和的输出,使FFN块的最终输出在GPU之间复制,为Transformer的下一部分(如残差连接或层归一化)做好准备。这种策略性组合有助于最大限度地减少通信开销,通过在FFN块内的两个线性层之间保持激活分片。嵌入并行化对于词汇量非常大的模型,嵌入表可能成为一个显著的内存瓶颈。在这里,张量模型并行可以通过将嵌入表按行(沿词汇维度)在GPU之间拆分来应用。当查找输入token ID的嵌入时:每个GPU接收完整的输入ID序列。每个GPU仅查找其所持有的词汇部分对应的嵌入。如果某个ID对应于另一个GPU持有的行,它会产生一个零向量。全归约操作会汇总所有GPU上的部分嵌入向量。由于对于任何给定的token ID,只有一个GPU生成了实际的嵌入向量(其他GPU生成了零),因此总和有效地收集了正确的嵌入。实现细节与通信手动实现张量模型并行需要仔细处理张量分片、计算和同步,使用 torch.distributed 包中的原语:torch.distributed.broadcast:将张量从一个进程发送到所有其他进程。torch.distributed.scatter:将张量块分散到各个进程。torch.distributed.gather:将张量从所有进程收集到一个进程。torch.distributed.all_gather:将张量从所有进程收集到所有进程。torch.distributed.reduce_scatter:在进程间执行操作(如求和)并分散结果。torch.distributed.all_reduce:在进程间执行操作(如求和)并使结果在所有进程上可用。通常会创建专用函数或包装器来封装这些操作,以适用于特定层类型(例如 ColumnParallelLinear、RowParallelLinear)。像NVIDIA的Megatron-LM这样的库率先使用了许多这些技术,并且其中一部分功能正在通过诸如 torch.distributed.tensor.parallel 等模块集成到PyTorch核心中,这些模块提供了更高级的API来简化这些实现。考虑一个使用辅助函数的列并行线性层示例:import torch import torch.nn as nn import torch.distributed as dist # 假设这些辅助函数用于管理并行组和通信 from .parallel_utils import ( get_tensor_model_parallel_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, copy_to_tensor_model_parallel_region, # 处理输入广播/拆分 gather_from_tensor_model_parallel_region # 处理输出收集/归约 ) class ColumnParallelLinear(nn.Module): def __init__(self, input_size, output_size, bias=True, **kwargs): super().__init__() world_size = get_tensor_model_parallel_world_size() # 确保 output_size 可以被 world_size 整除 assert output_size % world_size == 0 self.output_size_per_partition = output_size // world_size self.input_size = input_size # 权重矩阵沿输出维度(列)拆分 self.weight = nn.Parameter(torch.empty( self.output_size_per_partition, self.input_size, **kwargs )) # 初始化权重...(例如,使用 init.kaiming_uniform_) if bias: # 偏置也沿输出维度拆分 self.bias = nn.Parameter(torch.empty( self.output_size_per_partition, **kwargs )) # 初始化偏置...(例如,使用 init.zeros_) else: self.register_parameter('bias', None) def forward(self, input_): # 如果前一层输出已复制(例如 LayerNorm), # 输入可能需要广播或已可用。 # 此函数处理必要的通信。 parallel_input = copy_to_tensor_model_parallel_region(input_) # 执行局部矩阵乘法 output_parallel = nn.functional.linear(parallel_input, self.weight, self.bias) # 从张量并行组中的所有GPU收集结果 # 沿列维度拼接。 output_ = gather_from_tensor_model_parallel_region(output_parallel) return output_ 列并行线性层的实现草图。请注意权重/偏置的显式分片以及通信包装器的使用。权衡与考虑复杂性:手动实现张量模型并行需要对模型架构代码进行大量修改,这与通常作为包装器运行的DDP不同。调试并行逻辑可能具有挑战性。通信开销:张量模型并行在单个层的前向和反向传播中引入了通信。全收集或全归约等操作可能成为瓶颈,特别是当GPU之间的互连带宽(例如NVLink与PCIe)受限时。粒度:对于参数量大或激活值相对于执行的计算量大的层,它最有效。组合性:张量模型并行通常与数据并行(DDP)和流水线并行结合使用,以训练真正大型的模型,形成混合并行策略。总而言之,张量模型并行是一种不可或缺的技术,当单个模型层超出单个设备的内存限制时。通过在多个设备上拆分层内的权重和计算,它使得训练比以往可能更大的模型成为可能,尽管代价是增加了实现复杂性和通信开销。它是当前扩展大型神经网络的基础组成部分。