Tensor Parallelism

学习张量并行如何跨 GPU 切分权重矩阵以训练大型模型。探索它与 Ultralytics 数据并行之间的差异。

张量并行 (Tensor Parallelism) 是一种用于机器学习的高级分布式训练技术，用于将大型数学结构（即张量）拆分到多个硬件加速器（如 GPUs 或 TPUs）上。在训练海量深度学习模型时，参数量往往会轻易超过单个设备的内存容量。与其将完整的神经网络层放置在一个 GPU 上，张量并行会将权重矩阵进行分片，并将矩阵乘法等数学运算拆分到集群中的多个设备上。这使得模型能够利用整个硬件设置的综合内存和计算能力，在单程序多数据 (SPMD) 范式下执行并行计算，同时通过 NVIDIA NVLink 等高速互连方式同步结果。

Link to this section张量并行是如何工作的#

神经网络的核心是矩阵乘法。张量并行通过按行或按列拆分矩阵来分配这些运算。例如，在全连接层或 Transformer 注意力机制中，一个 GPU 可以计算矩阵的左半部分，而另一个 GPU 计算右半部分。并行计算完成后，设备之间会进行通信（通常使用快速的 All-Reduce 集合操作），在将完整的张量传递到下一层之前聚合部分结果。2025 年的最新学术进展通过引入部分同步激活进一步优化了这一过程，从而减少了通常会导致大型计算集群瓶颈的通信开销。

Link to this section区分相关并行技术#

要理解张量并行如何适应更广泛的分布式计算领域，需要将其与其他常见策略区分开来：

张量并行与模型并行： 张量并行是模型并行的一个高度具体的子类别。虽然通用的模型并行是指以任何方式将模型拆分到设备上，但张量并行特指对单个层内部的单个张量进行分片。
张量并行与流水线并行： 流水线并行是模型并行的另一种形式，它按深度划分网络——将前几层放在 GPU 0 上，接下来放在 GPU 1 上，依此类推。这会产生称为流水线气泡的顺序依赖。张量并行拆分层本身，在没有顺序延迟的情况下同步执行它们，但需要高得多的网络带宽。
张量并行与数据并行： 在数据并行中，整个模型被完全复制到每个 GPU 上，只有训练数据集被拆分到各个设备上。对于像 Ultralytics YOLO26 这样可以轻松装入现代 GPU 的高度优化架构，通过 PyTorch 的 DistributedDataParallel 进行数据并行是默认方法。张量并行通常仅在单个层的参数超过硬件 VRAM 从而导致内存不足 (OOM) 错误时才必要。

Link to this section实际应用#

张量并行在现代 AI 基础设施中不可或缺，特别是对于需要大规模计算能力的尖端架构：

训练大型语言模型 (LLMs)：像 Meta 的 Llama 3 和 DeepSeek V3 这样的海量基础模型使用 NVIDIA Megatron-LM 等框架来实现张量并行。由于这些模型的隐藏维度和注意力头非常大，必须将其拆分到 8-GPU 节点上才能高效训练并在实时推理过程中保持低延迟。
大型视觉模型 (LVMs) 和 3D 生成：随着计算机视觉扩展到海量多模态推理系统，研究人员在 AWS SageMaker 等服务上将张量并行与流水线并行结合使用，以训练巨大的视觉 Transformer (ViTs)。该技术允许处理需要巨大连续内存块的高分辨率图像和视频生成。

Link to this section在 PyTorch 中实现张量并行#

过去，工程师必须编写复杂的自定义分布式逻辑来对张量进行分片。最近，PyTorch 引入了 DTensor (Distributed Tensor)，原生简化了此工作流程。以下是使用官方 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}")

For edge-optimized vision tasks and rapid model deployment, developers typically rely on the Ultralytics Platform to automatically handle optimal hardware utilization. While multi-billion parameter foundation models require manual tensor parallelism configurations, you can efficiently scale training for models like YOLO26 using simple CLI commands out-of-the-box. This ensures maximum throughput by seamlessly utilizing native data parallelism techniques alongside robust model training tips.