张量并行是一种将单个神经网络层的计算分布到多个硬件设备上的技术。在这种方法中，层的权重矩阵被分割成多个分片，每个设备持有部分参数并执行相应的矩阵乘法部分。随后，部分结果通过集体通信操作（如全归约）进行合并，以生成该层的输出。这种方法属于模型并行的一种形式，与数据并行不同，后者中每个设备持有模型的完整副本并处理不同的数据样本。张量并行对于训练和运行超大规模模型（如大型语言模型）至关重要，这些模型超出了单个加速器的内存容量。
张量并行的主要动机源于现代深度学习模型对内存和计算的需求。例如，一个拥有数千亿参数的模型无法装入单个GPU或TPU的内存中，这些设备的内存通常为16到80吉字节。通过分割权重矩阵，张量并行使得模型能够分布在多个设备上，从而支持使用更大的模型并加快训练速度。它还减少了每个设备的内存占用，并允许更大的批处理大小。然而，张量并行引入了通信开销，因为设备必须同步部分结果，这可能成为瓶颈，尤其是在较慢的互连上。
张量并行通常与其他并行策略（如流水线并行和数据并行）结合使用，以在数千个设备上扩展训练。它是许多最先进模型训练基础设施的关键组成部分，包括由OpenAI、Anthropic和Google DeepMind开发的模型。该技术对于推理同样重要，通过将计算分布到多个加速器上，能够以低延迟服务大型模型。
历史背景
将矩阵运算拆分到多个处理器的概念可以追溯到早期的并行计算研究。在20世纪80年代和90年代，研究人员探索了神经网络的并行实现，但硬件和软件工具十分有限。现代形式的张量并行随着2010年代基于GPU的深度学习的兴起而出现。最早的重要应用之一是在大型卷积网络的训练中，多伦多大学及其他机构的研究人员尝试了模型并行。然而，正是基于Transformer的模型（如2017年提出的Transformer架构）的出现，使张量并行成为主流技术。Transformer的注意力层和前馈网络由大型矩阵乘法组成，这些运算天然适合分片。
2019年，NVIDIA及其他组织的研究人员发表了关于训练大型语言模型的张量并行的详细描述。NVIDIA开发的Megatron-LM框架展示了Transformer模型的高效张量并行，在多个GPU上实现了接近线性的扩展。这项工作为后续的大规模训练系统奠定了基础，例如谷歌的GShard和微软的DeepSpeed，它们都将张量并行作为核心功能。
张量并行的工作原理
在典型的神经网络层中，前向传播计算矩阵乘法：Y = XW，其中X是输入激活，W是权重矩阵，Y是输出。在张量并行中，权重矩阵W沿一个或多个维度进行划分。对于全连接层，常见的方法是按列拆分W（即沿输出维度）。每个设备持有W的列子集并计算部分输出。然后使用全归约操作将部分输出合并，以生成最终的Y。这被称为列并行分区。或者，W可以按行拆分（沿输入维度），这需要对输入执行全收集操作，并对输出执行全归约操作。
对于Transformer模型，张量并行应用于多头注意力和前馈网络。在多头注意力中，查询、键和值权重矩阵在设备间分片，注意力头被分布。输出投影随后合并。在前馈网络中，两个线性层通常被分区：第一层按列拆分，第二层按行拆分，这样中间激活被分片，输出被归约。这种设计通过仅在特定点要求全归约操作来最小化通信。
通信模式至关重要。每次全归约操作都需要在张量并行组中的所有设备之间交换数据。通信量正比于激活的大小，而非权重的大小，这通常更小。然而，对于非常大的模型，通信仍然可能显著。为缓解这一问题，张量并行通常在同一节点内使用高速互连（如NVIDIA的NVLink或AMD的Infinity Fabric），而跨节点则使用数据并行。
与其他并行技术的比较
张量并行是多种模型并行策略之一。例如，流水线并行按层拆分模型，每个设备负责一组连续的层。与张量并行相比，这减少了通信量，因为只需在层边界交换激活值。然而，流水线并行可能因流水线气泡而出现空闲时间。数据并行则相反，它在每个设备上复制模型并处理不同的数据样本，每一步后都需要进行梯度同步。张量并行与这两者互补：它可以与流水线并行结合以降低每台设备的内存需求，也可以与数据并行结合以提高吞吐量。
另一项相关技术是序列并行，它在Transformer模型中沿序列维度进行分片。这通常与张量并行结合使用，以进一步减少内存使用。专家并行用于混合专家模型，它将专家模块分片到不同设备上，这属于模型并行的一种形式，但并非严格意义上的张量并行。
应用与使用场景
张量并行广泛用于大型语言模型的训练和推理。例如，由OpenAI开发的GPT-3模型，在数千块NVIDIA GPU上结合了包括张量并行在内的多种模型并行策略进行训练。同样，Anthropic的Claude模型和Google DeepMind的Gemini模型也依赖张量并行来扩展到数千亿参数。在推理方面，张量并行使得GPT-4等模型能够通过将计算分布到多个GPU上来实现低延迟服务，正如NVIDIA的Triton推理服务器和Hugging Face的文本生成推理所做的那样。
除了语言模型，张量并行还用于其他领域，如计算机视觉和科学计算，其中大型模型或大批量需要分布式计算。例如，在高分辨率图像上训练大型视觉Transformer（ViT）可以从张量并行中受益。
挑战与局限
尽管张量并行有诸多优点，但它也面临几个挑战。主要问题是通信开销。All-reduce操作需要高带宽、低延迟的互连。在以太网等较慢网络的系统上，张量并行可能变得低效。因此，它通常用于单个服务器或紧密耦合的集群内。另一个挑战是负载均衡：不均匀的分片可能导致某些设备利用不足。此外，张量并行需要仔细的内存管理，因为每个设备必须存储其权重分片和中间激活值。最后，手动实现张量并行很复杂；它需要修改模型代码并处理通信原语。幸运的是，PyTorch和TensorFlow等框架提供了对张量并行的内置支持，抽象掉了大部分复杂性。
软件与框架支持
多个深度学习框架和库支持张量并行。PyTorch提供了torch.distributed模块和tensor_parallel API用于张量分片。TensorFlow提供了tf.distribute，支持模型并行。专门的库如Megatron-LM、DeepSpeed和Hugging Face的Transformers已为Transformer模型实现了张量并行。这些工具使研究人员和工程师能够以最少的代码更改应用张量并行，通常只需指定设备数量和分片策略。
未来方向
随着模型不断增大，张量并行仍将是一项关键技术。未来的发展可能包括更高效的通信算法，如分层all-reduce，以及更好地与异构硬件集成，包括AMD GPU、Intel加速器以及AWS Trainium和Groq的LPU等定制芯片。此外，自动并行发现（即系统确定最优分片策略）是一个活跃的研究领域。随着对大规模AI需求的日益增长，张量并行将继续发展，以应对规模和效率的挑战。