Tensor Parallelism

张量并行是一种先进的分布式训练技术，用于在多个硬件加速器（如多个GPU或 TPU）之间拆分大型数学结构，即张量。在训练大型深度学习模型时，参数量很容易超过单个设备的内存容量。张量并行不会将整个神经网络层放在一个 GPU 上，而是将权重矩阵分片，并把矩阵乘法等数学运算拆分到集群中的多个设备上。这样，模型就能利用整个硬件配置的合并内存和计算能力，按照单程序多数据（SPMD）范式执行并行计算，同时通过NVIDIA NVLink等高速互连同步结果。

张量并行的工作原理#

矩阵乘法是神经网络的核心运算。张量并行会沿行或列拆分矩阵，以分配这些运算。例如，在全连接层或 Transformer注意力机制中，一个 GPU 可以计算矩阵的左半部分，另一个 GPU 则计算右半部分。并行计算完成后，设备会进行通信，通常使用快速的全归约集合通信操作汇总部分结果，再将完整张量传递给下一层。2025 年的最新学术进展进一步优化了这一过程，采用部分同步激活来减少通常会限制大型计算集群性能的通信开销。

区分相关并行技术#

要了解张量并行在更广泛的分布式计算领域中的定位，需要将其与其他常见策略区分开来：

• 张量并行与模型并行：张量并行是模型并行的一个非常具体的子类别。模型并行通常指以任意方式将模型拆分到多个设备上，而张量并行严格指在单个层内部对各个张量进行分片。
• 张量并行与流水线并行：流水线并行是另一种模型并行形式，它按网络深度划分模型：前几层放在 GPU 0 上，接下来的层放在 GPU 1 上，依此类推。这会产生称为流水线气泡的序列依赖。张量并行则拆分各层本身，使它们能够同时执行而无需顺序等待，但需要更高的网络带宽。
• 张量并行与数据并行：在数据并行中，完整模型会在每个 GPU 上完整复制，只有训练数据集会在各设备之间拆分。对于Ultralytics YOLO26这类经过高度优化、可轻松运行于现代 GPU 的架构，通过 PyTorch 的 DistributedDataParallel 实现数据并行是默认方法。通常只有当单层参数超出硬件 VRAM 容量、导致内存不足（OOM）错误时，才需要使用张量并行。

实际应用#

张量并行在现代 AI 基础设施中不可或缺，尤其适用于需要海量计算资源的最先进架构：

• 训练大型语言模型（LLM）： Meta 的 Llama 3 等大型基础模型依赖张量并行，通常通过 NVIDIA Megatron-LM等框架实现。不过，DeepSeek V3是一个显著例外，它依靠流水线并行和专家并行进行训练，没有使用张量并行。由于这些模型的隐藏维度和注意力头规模都很大，必须将其拆分到一台 8-GPU 节点上，才能高效训练并在实时推理期间保持低延迟。
• 大型视觉模型（LVM）与 3D 生成： 随着计算机视觉朝着大规模多模态推理系统发展，研究人员会在 AWS SageMaker 等服务上结合使用张量并行和流水线并行来训练大型视觉 Transformer（ViT）。这项技术能够处理需要海量连续内存块的高分辨率图像和视频生成任务。

在 PyTorch 中实现张量并行#

过去，工程师必须编写复杂的自定义分布式逻辑来对张量分片。近期，PyTorch推出了 DTensor（分布式张量），原生简化了这一工作流。下面的示例展示了如何使用官方 PyTorch Distributed Tensor API创建按行分片的张量：

import torch
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor import Shard, distribute_tensor

# Initialize a 1D device mesh across 2 GPUs
mesh = init_device_mesh("cuda", (2,))

# Create a standard PyTorch tensor (e.g., representing a layer's weights)
local_tensor = torch.randn(1024, 1024)

# Distribute the tensor across the GPUs by sharding along the first dimension (row-wise)
# Each GPU now holds a (512, 1024) chunk of the original tensor
distributed_tensor = distribute_tensor(local_tensor, mesh, [Shard(0)])

print(f"Global shape: {distributed_tensor.shape}, Local shape: {distributed_tensor.to_local().shape}")

对于边缘优化的视觉任务和快速模型部署，开发者通常依赖 Ultralytics Platform自动优化硬件利用率。数十亿参数的基础模型需要手动配置张量并行；而对于 YOLO26 等模型，你可以直接使用简单的 CLI 命令高效扩展训练。这能让系统无缝结合原生数据并行技术和可靠的模型训练技巧，从而实现最大吞吐量。