大规模分布式 AI 模型训练系列——数据并行

生成式人工智能（GenAI）的进展速度，大家有目共睹。一个最直观的变化就是，模型能处理的序列长度正以前所未有的速度增长——现在，一些模型已经可以支持上百万甚至上千万个Token的序列了。这可不是小数字，它直接把分布式训练里的一个老问题推到了聚光灯下：如何有效并行处理这么长的序列？

过去几个月，学术界和工业界围绕“序列并行”（Sequence Parallelism）涌现了不少新成果。与此同时，我们手头也确实有长序列落地的实际需求。这就逼着我们得好好捋一捋那些序列并行的论文。可一动手才发现，各种分布式并行策略之间的关系错综复杂，互相影响。干脆，从头开始，把每一种并行方案都掰开了揉碎了讲清楚。

这篇文章作为第一站，咱们先从最基础的“数据并行”（Data Parallelism）说起。把数据并行的核心概念和技术梳理清楚，后面再聊更复杂的并行策略时，就有一个扎实的底子了。至于微软的Zero DP和Meta的FSDP，这次先放一放，留着以后专题聊。

二、引言

2.1 模型训练过程

要理解数据并行，得先明白模型训练到底是怎么跑起来的。整个过程包括前向传播（Forward Pass）、反向传播（Backward Pass）和权重更新（Weight Update）。咱们拿一个序列模型来举例，就是那种一层叠一层的形式：y = LayerN(LayerN-1(...(Layer1(x))))。

下面这张图展示了关键的Forward和Backward。每个模块从左边的输入开始，往右边输出结果。里面的Cache存的是中间结果，也就是Activation。图中wrt是“with respect to”的缩写，比如∇Loss wrt Weights of LayerN，意思就是LayerN里关于权重的梯度。

• Forward阶段：LayerN的前向计算需要两样东西——LayerN-1输出的Activation，以及LayerN自己的权重。

• Backward阶段：LayerN的反向计算需要三样东西——LayerN-1输出的Activation（也就是LayerN的输入）、LayerN的权重，以及LayerN输出层对应的Error。

再往下看，反向传播其实包含两个部分：

• 链式法则：反向传播算法的核心，负责把损失函数的梯度从输出层一路传回输入层。每一层根据链式法则算出损失函数对当前层输入的偏导数，再往前传。

• 权重梯度：每一层都要算损失函数对权重的偏导数，这个梯度就是用来指导权重怎么调整，让损失函数最小化。

把刚才的Forward和Backward拼接起来，就是完整的训练过程，下面这张动图展示得很清楚：

权重更新这一步，有了每一层的梯度，就可以用权重更新算法来更新模型的权重了。它需要三样输入：该层的优化器状态、该层的权重、以及该层权重的梯度。

2.2 梯度下降

梯度下降是机器学习里最常用的优化方法，目标是让损失函数最小化。根据每次更新用多少数据，可以分成几种变体：

• 全量梯度下降（Batch Gradient Descent）
  ：每次更新权重，都用整个训练集来计算梯度。假设有N个样本，每个权重都会得到N个梯度，先取平均值，然后用这个平均值去更新。
  • 优点：梯度估计非常准，收敛过程稳。
  • 缺点：数据集一大，计算梯度就特别慢，内存需求也高，不适合大规模数据。

• 随机梯度下降（SGD）
  ：每次只用一个训练样本来算梯度。每个权重只有一个梯度，直接用它更新。
  • 优点：计算快，适合大规模数据集，内存占用小。
  • 缺点：梯度估计不准，更新过程噪声大，收敛不稳定，容易在最优解附近振荡。

• Mini-batch梯度下降
  ：折中方案。每次用一小批数据（mini-batch）来计算梯度。
  • 优点：计算速度比全量快，比SGD稳定，减少了噪声。
  • 缺点：如果mini-batch样本不够多，梯度估计仍然可能有偏差。

2.3 AllReduce

AllReduce是集合通信里最常用的操作之一，专门用在多个设备（比如多台服务器或多个GPU）之间聚合数据。可以求和、求最小值、最大值等。拿AllReduce Sum来说，假设有K个设备，每个设备上有N个数据，那么AllReduce之后，每个设备上的输出就等于所有设备输入的累加。具体可以参考NCCL的文档。

下图展示了一种基于Ring的AllReduce实现，它其实可以拆成两个阶段：ReduceScatter和AllGather。

ReduceScatter的流程是这样：每个节点把一部分数据发给下一个节点，同时接收上一个节点的数据并累加，重复N-1步。完成后，每个设备都只拥有一部分数据的总和。

AllGather则反过来：每个节点把自己手里的部分结果发给下一个节点，同时接收上一个节点的部分结果，逐步拼出完整的数据。AllGather结束后，每个设备都有了全部数据。

2.4 AllReduce带宽

上面提到，AllReduce分为ReduceScatter和AllGather两个阶段。假设GPU数量是K，每个阶段通信量是(K-1) * T * sizeof(dtype)。如果每个GPU的总线带宽是busBW，那么AllReduce的理论通信时延就可以算出来。

不过实际的总线带宽往往到不了理论值。举个例子，4块V100通过NVLink互联（没有NVSwitch），每两个GPU之间有两个NVLink连接，理论双向带宽是100GB/s，实测能到97GB/s。如果禁用NVLink，带宽就掉到16GB/s了。

AllReduce最关心的是总线带宽。在4*V100 NVLink互联的场景下，NCCL可以创建3个双向环（每个颜色代表一个双向环）。每个环的单向通信带宽是4*25GB/s，所以通信带宽理论上限是6*(4*25GB/s)=600GB/s（也可以用12个NVLink乘以50GB/s来算），平均每个GPU的busBW就是150GB/s。

当然，实际总线带宽和理论值之间还有差距，这跟数据量、GPU连接方式、通信算法、NCCL版本都有关系。下图为实际测试的AllReduce总线带宽，数据小于1MB时，busBW还不到30GB/s；数据量到128MB时，busBW能达到130GB/s，接近极限的150GB/s。没有NVLink的话，实测busBW只有10GB/s左右。

NVIDIA官方在V100上的测试结果与我们的实测性能相当。

三、Pytorch DDP

3.1 DP（Data Parallelism）

先看看单个设备用Mini-batch梯度下降是怎么训练的：设备加载完整的模型，每次塞进一个mini-batch（比如6个样本），先做Forward，再算Loss，然后Backward拿到梯度。

数据并行就简单了：在N个设备上存一模一样的模型参数，把mini-batch按设备数拆开（2个设备，每个设备分3个样本），然后各自做Forward、Loss计算和Backward。

但问题是，每个设备上只有自己那部分样本的梯度，要做Mini-batch梯度下降，得把整个mini-batch的梯度聚合起来。这时候AllReduce Sum就派上用场了，它可以在每个设备上都拿到全部梯度的平均值，然后再去更新权重。

3.2 DDP（DistributedDataParallel）

PyTorch提供了两种数据并行模式：DataParallel (DP) 和 DistributedDataParallel (DDP)。它们实现方式和应用场景差别不小。

• DataParallel (DP)
  ：简单模式，利用多GPU在单节点上并行，单进程多线程。流程是：把输入数据拆成多个mini-batch分到各GPU，同时跑Forward和Backward，然后在主GPU（默认GPU 0）上收集并汇总所有梯度，在GPU 0上更新参数，再同步给其他GPU。DP适合小规模模型，尤其单节点多GPU场景。但性能瓶颈很明显——主GPU负载过重，GPU间数据传输开销大，扩展性差。

• DistributedDataParallel (DDP)
  ：更高效的数据并行方式，多进程，支持多节点和多GPU。每个GPU上有一个模型副本，梯度独立计算，然后通过NCCL或Gloo等通信协议在GPU间同步梯度。DDP避免了DP的主GPU瓶颈，通信和计算可以重叠，大大减少通信开销，训练效率和扩展性都更好。

下面是一个PyTorch DDP的简单示例（注意代码中并没有初始化Process Group，完整例子可以参考官方文档）。第20行的loss.backward()里既做Backward又做AllReduce梯度聚合；第23行的opt.step()则用聚合后的梯度更新权重。

有一点要特别说明：数据并行会保证不同设备的初始状态相同，每次迭代都用聚合后相同的梯度来更新，所以每个设备上的模型参数在每个Iter的Forward和Backward计算时始终一致。

四、优化

4.1 梯度分桶

前面我们提到，AllReduce的总线带宽和消息大小关系密切。当消息小于1MB时，实际带宽可能还不到理论值的十分之一。怎么解决？把多个小梯度合并成一个稍微大点的消息再去通信，也就是梯度分桶（bucketing）。

在PyTorch DDP里，可以通过bucket_cap_mb参数控制Bucket的大小，默认是25MB。当然，Bucket也不是越大越好，太大反而可能影响性能。

4.2 通信与计算重叠

数据并行额外引入的梯度聚合（AllReduce）本质上是通信操作，可以跟Backward的计算重叠起来，这样能减少GPU在通信时的闲置。具体来说，某一层的Backward开始后，就可以着手做这一层的梯度聚合了，而且能跟下一层的Backward同时进行——PyTorch的loss.backward()里已经实现了这个机制。

通过Overlap，每个Iter的时延能有效降低。

不过，这里有个细节：所有的梯度聚合都完成之后，才会开始权重更新。实际上，权重更新在对应的梯度聚合完成后就可以立刻开始，但PyTorch DDP没有这么做。原因可能比较复杂，比如权重更新会跟Backward抢资源，而Optimizer都是单独实现的，进一步Overlap会大幅增加实现和管理计算图的复杂度。

4.3 梯度累加

GPU显存有限，模型规模却越来越大，有时候就算用了数据并行，Batch Size还是上不去。这时候可以用梯度累加（Gradient Accumulation）。每次用micro-batch的数据Forward和Backward，只在本设备累加梯度，不做全局聚合，也不更新权重。等到最后一个micro-batch，才进行梯度聚合和权重更新。

在PyTorch DDP中，用no_sync()就能实现梯度累加。

4.4 激活重计算

模型训练时需要保存大量中间状态（Activation），占不少显存。如果计算资源充裕但显存不够，可以用计算换内存——这叫激活检查点（Activation Checkpointing）或激活重计算（Activation Recomputing）。具体做法是：Forward过程中只保留部分Activation，比如每隔3层保存一个，相当于把模型分段。Backward时，再通过之前保存的Checkpoint重计算来获取需要的Activation。间隔越大，越省内存，但重计算代价越高；间隔越小，重计算代价小，省的内存也有限。

4.5 ZeroBubble

我们之前提到，Backward过程实际上包括两部分：对输入的梯度（链式法则）和对权重的梯度。不同权重的梯度计算之间没有任何依赖关系。在Zero Bubble Pipeline Parallelism论文中，作者把这个特性用到了流水线并行上，这里就不具体展开了。

五、附录

5.1 参数平均

数据并行训练除了用梯度平均，还可以用参数平均。每个训练节点独立训练一段时间后，把各节点的模型参数直接平均。两者各有优劣：

• 参数平均：实现简单，适合数据分布均匀且节点间通信成本高的情况。但每个节点本地跑多个迭代后，节点间模型差异可能很大，导致收敛不佳，速度也可能慢一些。

• 梯度平均：每个mini-batch后立刻平均梯度，模型各节点保持同步，有助于加快收敛。但对通信要求高，每个mini-batch后都需要通信，节点间同步要求也高。

5.2 AllReduce算法

AllReduce有不同的实现算法，常见的有Ring AllReduce和Tree AllReduce。

• 通信过程：先做Reduce-Scatter，再做AllGather，共2*(N-1)步。
• 优点：带宽利用率高，因为所有节点同时在发送和接收数据。
• 缺点：节点多时延迟高。

• 通信过程：从叶子向根做Reduction，根得到完整结果；然后从根向叶子Broadcast。步数约为2*log(N)。
• 优点：通信步骤少，大规模节点时延迟低。
• 缺点：每一步只有部分节点参与通信，带宽利用率相对低。

总结一下：Ring AllReduce更适合高带宽、低延迟的网络环境，尤其节点数不多时；Tree AllReduce更适合节点多或网络延迟高的情况，因为它的延迟随节点数增长更慢。

下面是Tree AllReduce的示意图，向上箭头是Reduction阶段，向下箭头是Broadcast阶段。

我们用nccl-test在NVLink+NVSwitch全互联的8*H100上测试了Ring和Tree的性能。Ring的配置是NCCL_ALGO=Ring ./build/all_reduce_perf -b 1M -e 1G -f 2 -g 8，Tree的配置类似。结果很明显：在这种全互联架构下，Ring AllReduce性能比Tree高出不少。

5.3 Adam优化器

Adam优化器结合了动量和自适应学习率（如RMSprop），是深度学习训练中最流行的选择。它的核心优势在于能根据参数更新历史自适应调整每个参数的学习率，从而加速收敛、提高稳定性。Adam维护了两个额外的状态：

• 一阶矩
  ：过去梯度的移动平均，提供平滑的梯度方向，增加稳定性。

• 二阶矩
  ：过去梯度平方的指数移动平均，代表梯度变化率，允许每个参数有自适应学习率。

Adam的关键公式和伪代码在原始论文中有详细描述。每个参数对应一个一阶矩m和二阶矩v，更新过程如下（伪代码）：

class AdamOptimizer:
    def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, eps=1e-8):
        self.learning_rate = learning_rate
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.m = None
        self.v = None
        self.t = 0

    def initialize_moments(self, params):
        self.m = {k: np.zeros_like(v) for k, v in params.items()}
        self.v = {k: np.zeros_like(v) for k, v in params.items()}

    def update_params(self, params, grads):
        if self.m is None or self.v is None:
            self.initialize_moments(params)
        self.t += 1
        params_new = {}
        for k in params.keys():
            self.m[k] = self.beta1 * self.m[k] + (1 - self.beta1) * grads[k]
            self.v[k] = self.beta2 * self.v[k] + (1 - self.beta2) * np.square(grads[k])
            m_new = self.m[k] / (1 - self.beta1 ** self.t)
            v_new = self.v[k] / (1 - self.beta2 ** self.t)
            params_new[k] = params[k] - self.learning_rate * m_new / (np.sqrt(v_new) + self.eps)
        return params_new

PyTorch的Adam优化器实现可以参考官方文档。