简单来说，分布式训练是指将一个大型的机器学习模型的训练任务，拆分到多个计算设备（如多个GPU、多个服务器节点）上并行执行，以加速训练过程、解决单设备内存不足的问题

1. 模型越来越大（如GPT-3有1750亿参数），单个GPU的内存（显存）无法容纳整个模型和其优化状态。
2. 数据量越来越大，在单个设备上顺序训练需要花费数周甚至数月时间,难以快速迭代。

分布式训练的核心思想与目标

• 核心思想：分工协作，将计算负载和存储负载分摊到多个设备上，通过设备间的通信来协调工作,共同完成训练任务。
• 主要目标：
  • 加速训练：通过并行计算，将训练时间从“月/周”缩短到“天/小时”。
  • 扩大模型规模：训练在单卡上无法装载的巨型模型。
  • 提高资源利用率：充分利用集群中的计算资源。

分布式训练的主要并行范式

根据如何“拆分”任务，主要有以下几种并行策略,通常会组合使用：

数据并行

这是最常用、最直观的分布式训练方式。

• 怎么做：
  • 将训练数据集划分为多个子集（Batch）。
  • 将完整的模型副本复制到每个计算设备（如每个GPU）上。
  • 每个设备用自己分配到的一个数据子集（一个Mini-Batch）进行前向传播和反向传播,计算得到本地梯度。
  • 关键步骤：在所有设备间同步梯度（求所有设备梯度的平均值）。
  • 每个设备用同步后的梯度更新自己副本的模型参数，由于初始参数相同,更新后的参数也保持一致。
• 优点：实现相对简单，通信模式固定（主要是梯度同步）,对于模型可以放进单卡的情况非常有效。
• 缺点：每个设备都需要存储完整的模型和优化器状态，对于超大规模模型,单卡显存仍然不够。
• 通信核心：All-Reduce 操作（特别是对梯度进行求和或求平均）。

模型并行

当模型太大，单个设备放不下时,就需要拆分模型本身。

• 怎么做：
  • 将模型的网络结构在层间或层内进行切分,不同的部分放置在不同的计算设备上。
  • 数据（一个Mini-Batch）需要依次流过这些设备来完成前向传播,反向传播时梯度则反向传递。
• 优点：可以训练巨大的模型。
• 缺点：会引入设备间的通信依赖和等待空闲，设备2必须等待设备1完成计算后才能开始工作，导致设备利用率不足（“气泡”问题）。
• 典型技术：流水线并行 是模型并行的一种优化形式。

流水线并行

这是对朴素模型并行的改进,旨在解决设备闲置问题。

• 怎么做：
  • 同样将模型按层切分到不同设备上。
  • 将训练数据分成更小的 “微批次”。
  • 像工厂流水线一样，当第一个微批次在设备2上计算时，第二个微批次可以立刻在设备1上开始计算，从而让多个微批次在不同设备上同时处理,提高设备利用率。
• 核心挑战：需要精心调度微批次，处理流水线启动和排空时的“气泡”。

张量并行（层内并行）

一种更细粒度的模型并行，将单个层内部的运算（如一个大型矩阵乘法）拆分到多个设备上。

• 怎么做：将一个大的权重矩阵按行或列切分，每个设备持有矩阵的一部分,共同协作完成该层的计算。
• 代表：NVIDIA的 Megatron-LM 就大量使用了张量并行来训练百亿、千亿级参数的Transformer模型。
• 通信特点：通信通常发生在层内计算期间,通信模式比数据并行复杂。

混合并行

在实际训练超大规模模型（如百亿、万亿参数）时，几乎没有单一并行策略,工程师会组合使用上述所有策略：

• 数据并行 在不同 组 的设备间使用。
• 张量并行 在 组内 设备间使用,以放下一个巨大的层。
• 流水线并行 在 组间 使用,以放下整个模型。

使用1024个GPU训练一个大模型时，可能会先分成8组流水线（流水线并行），每组128个GPU；这128个GPU内部再通过张量并行紧密协作；如果有更多数据需要并行,可以为整个系统再配置多份副本进行数据并行。

分布式训练的关键技术组件

1. 通信后端：负责设备间的数据交换。

   • NCCL：NVIDIA的集合通信库，针对GPU间高速互联（NVLink, InfiniBand）进行了极致优化,是当前主流。
   • gRPC：更通用的TCP/IP网络通信,常用于CPU集群或作为备用方案。
   • MPI：高性能计算领域的传统通信标准。

2. 同步 vs 异步更新：

   • 同步（Synchronous）：等待所有设备的梯度都计算完成后，同步（平均）梯度，再更新，主流方法，训练稳定,但速度受限于最慢的设备。
   • 异步（Asynchronous）：设备计算完梯度后立即更新中心参数，不等待其他设备，训练速度快但可能不稳定,现已较少使用。

3. 参数服务器架构 vs 集合通信架构：

   • 参数服务器：有中心节点（Server）存储和管理全局模型参数，工作节点（Worker）从Server拉取参数，计算梯度后再推送回去，早期常用,易成为通信瓶颈。
   • 集合通信：无中心节点，所有设备通过 All-Reduce 等操作直接对等通信，是现代分布式训练（尤其是数据并行）的标准架构,通信效率高。

4. ZeRO（Zero Redundancy Optimizer）：

   • 由微软DeepSpeed提出的一种革命性的优化器状态分区技术。
   • 核心思想：消除数据并行中的内存冗余，它将优化器状态、梯度、甚至模型参数本身分区存储在所有设备上,每个设备只负责更新自己那部分。
   • 在通信时，按需从其他设备收集所需的参数。ZeRO-Offload 和 ZeRO-Infinity 还能进一步将部分数据卸载到CPU内存甚至NVMe硬盘,实现了在有限GPU资源下训练超大模型的能力。

分布式训练是当代大规模深度学习训练的基石，其演进路径是： 单卡训练 -> 数据并行（加速）-> 模型/流水线并行（放大模型）-> 张量并行（更细粒度放大）-> 混合并行 + ZeRO（训练千亿/万亿模型）。

对于入门者，从理解数据并行和梯度同步开始最为重要，而对于希望训练或使用大模型的从业者，需要深入理解混合并行策略以及像 PyTorch DDP, FSDP（PyTorch对ZeRO的实现）, DeepSpeed 这样的框架和工具。

标签： 分布式训练 并行计算