Tensor Parallelism
学习张量并行如何跨 GPU 切分权重矩阵以训练大型模型。探索它与 Ultralytics 数据并行之间的差异。
张量并行是机器学习中一种先进的分布式训练技术，用于将大型的单个数学结构（即张量）划分到多个硬件加速器（如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和DeepSeek V3这样的大型基础模型利用NVIDIA Megatron-LM等框架来实现张量并行。由于这些模型的隐藏维度和注意力头非常巨大，因此必须将它们拆分到8-GPU节点上，以便高效训练并在实时推理期间保持低延迟。
- 大型视觉模型（LVM）与3D生成：随着计算机视觉向海量多模态推理系统发展，研究人员将张量并行与流水线并行结合在AWS SageMaker等服务上，以训练巨型的视觉Transformer（ViT）。这项技术允许处理需要巨大连续内存块的高分辨率图像和视频生成。
在 PyTorch 中实现张量并行#
历史上，工程师必须编写复杂的自定义分布式逻辑来对张量进行分片。最近，PyTorch引入了 DTensor（分布式张量），原生简化了这一工作流程。以下是使用官方PyTorch分布式张量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来自动处理最优的硬件利用率。虽然拥有数十亿参数的基础模型需要手动配置张量并行，但你可以使用简单的命令行界面（CLI）命令开箱即用地高效扩展像YOLO26这样的模型的训练。这通过无缝利用原生数据并行技术以及强大的模型训练技巧，确保了最大的吞吐量。
