Tensor Parallelism
了解张量并行如何将权重矩阵拆分到多个 GPU 上以训练大规模模型。探索它与 Ultralytics 数据并行的区别。
张量并行 (Tensor Parallelism) 是一种高级 分布式训练 技术,用于机器学习中将庞大的单个数学结构(即张量)划分到多个硬件加速器(如 GPU 或 TPU)上。在训练大规模 深度学习 模型时,参数量往往会轻易超过单个设备的内存容量。张量并行不再将整个 神经网络 层放在一个 GPU 上,而是对权重矩阵进行切分,并将矩阵乘法等数学运算分配到集群中的多个设备上。这使得模型能够利用整个硬件架构的组合内存和计算能力,在 单程序多数据 (SPMD) 范式下执行并行计算,同时通过 NVIDIA NVLink 等高速互联技术同步结果。
Link to this section张量并行是如何工作的#
神经网络的核心是矩阵乘法。张量并行通过按行或按列拆分矩阵来分配这些运算。例如,在全连接层或 Transformer 注意力机制中,一个 GPU 可能计算矩阵的左半部分,而另一个 GPU 计算右半部分。并行计算完成后,设备会进行通信——通常使用快速的 全归约集合操作 (All-Reduce collective operations)——在将完整张量传递到下一层之前聚合部分结果。2025 年的最新学术进展通过引入部分同步激活机制,进一步优化了这一过程,以减少通常会成为大规模计算集群瓶颈的通信开销。
Link to this section区分相关并行技术#
了解张量并行如何融入分布式计算的广阔领域,需要将其与其他常见策略区分开来:
- 张量并行 vs 模型并行: 张量并行是模型并行的一个高度具体的子类。通用模型并行是指以任何方式将模型跨设备拆分,而张量并行严格指的是对单个层 内 的各个张量进行分片。
- 张量并行 vs 流水线并行: 流水线并行是模型并行的另一种形式,它通过深度划分网络——将前几层放在 GPU 0 上,接下来的放在 GPU 1 上,依此类推。这会产生所谓的流水线气泡 (pipeline bubbles) 导致的顺序依赖。张量并行则拆分层本身,在没有顺序延迟的情况下同时执行它们,但这需要高得多的网络带宽。
- 张量并行 vs 数据并行: 在数据并行中,整个模型被完全复制到每个 GPU 上,只有训练数据集被拆分到各个设备中。对于像 Ultralytics YOLO26 这样可以轻松放入现代 GPU 的高度优化架构,通过 PyTorch 的
DistributedDataParallel进行数据并行是默认方法。张量并行通常仅在单个层的参数超过硬件显存并导致内存不足 (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(分布式张量),从原生层面简化了这一工作流。以下是使用 官方 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.






