Distributed Training

分布式训练将一个模型的训练工作分配到多个处理器上，通常是在一台机器中使用多个 GPU或 TPU，也可以跨云端或本地集群进行，从而缩短训练时间，或让原本无法装入单个设备内存的模型得以训练。设备会使用 NVIDIA NCCL 等通信库，通过高速互连交换梯度、激活值或参数分片。现代基础模型和大型语言模型都采用这种方式训练，它也能缩短多 GPU 工作站上的日常目标检测训练时间。

分布式训练的工作原理#

数据并行是最常见的策略。每个 GPU 都保存一份完整的模型副本，处理各自分配的训练数据批次切片，并参与全归约操作来计算梯度平均值，使每个副本应用相同的梯度下降更新。PyTorch 的 DistributedDataParallel (DDP)会将通信与反向传播重叠进行。全局批次大小等于每个 GPU 的批次大小乘以 GPU 数量，因此学习率通常也需要相应缩放：Goyal 等人使用线性缩放规则和预热阶段，在 256 个 GPU 上以 8,192 张图像的批次大小，在一小时内完成了 ImageNet上的 ResNet-50 训练；从 8 个 GPU 扩展到 256 个 GPU 时，扩展效率约为 90%。

分片数据并行可降低这些模型副本的内存开销。ZeRO和 PyTorch 的 Fully Sharded Data Parallel (FSDP)会将参数、梯度和优化器状态分散到多个 GPU 上，并且只在某一层进行计算时才汇集该层的完整权重，因此随着 GPU 数量增加，每个 GPU 所需的内存会减少。

模型并行会拆分神经网络本身。张量并行将单个矩阵乘法分配到多个 GPU 上，流水线并行将连续层分配给不同 GPU，并让微批次依次流经这些 GPU；上下文并行拆分长序列；专家并行则将混合专家模型中的专家分配到不同 GPU 上。Megatron-LM等框架会将这些方式与数据并行结合起来，训练拥有数千亿参数的模型。

多 GPU 训练实践#

通信是主要开销：全归约流量会随模型规模增长，因此节点内通过 NVLink 连接的 GPU 比通过较慢网络连接的节点更容易扩展。除非将批归一化层转换为同步批归一化，否则统计量会按 GPU 分别计算；当每块 GPU 分到的批次较小时，这一点很重要。Ultralytics YOLO26 通过 DDP 运行多 GPU 训练：传入多个设备 ID 时，会通过 torch.distributed.run 为每块 GPU 启动一个进程；CUDA 上使用 NCCL 后端，并将 batch 值平均分配到各块 GPU。

from ultralytics import YOLO

# 加载预训练的 YOLO26 模型
model = YOLO("yolo26n.pt")

# 使用 DDP 在 GPU 0 和 1 上训练；batch=32 表示每个 GPU 每步处理 16 张图像
results = model.train(data="coco8.yaml", epochs=100, imgsz=640, batch=32, device=[0, 1])

实际应用#

• 大型语言模型：Meta 使用最多 16K 个 H100 GPU 训练了 Llama 3 405B，并结合张量、流水线、上下文和数据并行；其作者将这种方式称为 4D 并行。
• 自动驾驶汽车：感知团队会使用车辆车队采集的大量摄像头、激光雷达和雷达录制数据来训练检测与分割模型，而这只有大型 GPU 集群才能做到。
• 医学图像分析：高分辨率 3D CT 和 MRI 影像会耗尽单个 GPU 的内存，将批次或模型分片分散到多个 GPU 上，就能训练更大的影像和模型。

分布式训练与相关概念的区别#

• 联邦学习：联邦学习也会跨多个设备训练，但为了保护隐私，数据保留在手机、医院或其他客户端上，只有模型更新会传到中央服务器。分布式训练则将数据集中在一个集群中，以最大限度提高速度。
• 拆分学习：拆分学习会在客户端和服务器之间拆分模型，确保原始数据不会离开客户端，以额外的通信开销换取隐私保护。
• 梯度累积：梯度累积会先对多个小批次的梯度求和，再执行一次更新，从而在单个 GPU 上实现与大批次相同的计算结果，但需要更多时间；数据并行则会在多个 GPU 上同时计算这些批次。