在深度学习训练中,当单块 GPU 的计算能力和显存不足以满足训练需求时,可以通过增加 GPU 数量扩大每一步能够处理的样本规模。数据并行(Data Parallelism)的基本思路并不复杂:每个 GPU 都保存一份完整的模型参数,不同 GPU 同时处理不同的训练样本,然后把各自计算得到的梯度进行聚合,再按照统一的结果更新模型。这样,多块 GPU 就能够共同完成一次训练迭代。
以最常见的同步数据并行(Synchronous Data Parallelism)为例,假设有 8 块 GPU,每块 GPU 都保存相同的模型副本。一个 global batch 被划分成 8 个 local batch,每块 GPU 独立完成前向传播和反向传播。由于不同 GPU 看到的是不同样本,它们计算出的梯度自然不会完全相同。此时需要通过集合通信把这些局部梯度进行聚合,使各个 GPU 最终得到相同的梯度结果。
如果第一个 GPU 得到梯度 G1,第二个 GPU 得到 G2,其他 GPU 分别得到 G3 到 G8,那么同步数据并行通常需要计算这些梯度的总和或平均值。具体采用求和还是平均,与训练框架的实现以及学习率和 batch size 的定义有关,并不能简单地说 AllReduce 永远直接产生“平均梯度”。完成聚合以后,各 GPU 使用相同的梯度执行 optimizer.step(),模型参数因此继续保持一致。
这就是梯度同步最重要的意义。它并不是为了保证模型一定收敛到所谓“全局最优解”。深度神经网络通常是非凸优化问题,训练过程并不存在这样简单的保证。更准确地说,同步梯度使多 GPU 训练在数学上接近于使用更大 batch 的单机训练:所有 GPU 的局部梯度共同构成当前 global batch 对模型参数的梯度估计。在学习率、batch size、随机性、混合精度以及优化器状态等条件适当匹配的情况下,多 GPU 可以在保持训练算法基本一致的同时提高吞吐量。
实际系统中最重要的通信操作之一就是 AllReduce。NVIDIA NCCL 对 AllReduce 的定义是:各个 rank 提供自己的数据,经过指定的归约操作后,每个 rank 都得到相同的结果。对于梯度同步来说,归约操作通常是求和,然后根据训练实现进行相应缩放。
AllReduce并不意味着某一块 GPU 必须先收集所有梯度,再把结果广播给其他 GPU。现代高性能实现会利用 Ring、Tree、CollNet、NVLS、NVLSTree 等不同通信算法,根据 GPU、NVLink、PCIe、网络交换机和节点拓扑选择合适的方案。当前 NCCL 文档明确列出了多种 AllReduce 算法,并且默认情况下由 NCCL 根据硬件拓扑和架构自动选择,而不是固定使用“多阶段多树”一种算法。
以 Ring AllReduce 为例,每个 GPU 与通信环中的邻居交换数据,同时进行局部归约。经过多个通信阶段后,每个 GPU 都能够得到完整的归约结果。Ring 算法的一个重要特点是,它能够避免由单一中心节点承担全部通信压力。在理想条件下,每个 GPU 的通信量与梯度总大小相关,而不是随着 GPU 数量呈指数级增长。因此,“GPU 数量增加后梯度同步的网络传输量呈指数增长”是不准确的。真正的问题是,当 GPU 数量继续扩大以后,通信链路、跨节点网络、同步等待以及拓扑不均衡等因素会越来越容易限制整体扩展效率。
多 GPU 训练的另一个关键技术是计算与通信重叠。现代框架通常不会等整个反向传播完全结束后才一次性同步全部梯度,而是把参数梯度划分为若干 bucket。当某个 bucket 中的梯度已经准备完成,就可以启动相应的 AllReduce,同时反向传播继续计算其他层的梯度。PyTorch DistributedDataParallel 就采用了这种梯度同步机制,并提供了梯度 bucket、通信 hook 等相关功能。
因此,理想情况下,GPU 在计算后面网络通信的同时,前面的通信也可以进行。如果通信时间能够被后续计算隐藏,GPU 就不需要长时间等待 AllReduce 完成;如果网络带宽不足、通信延迟过高,或者模型本身计算量较小,通信就可能暴露出来,形成典型的通信瓶颈。
高速互联技术在这里非常重要。单机多 GPU 系统可以利用 NVLink 等 GPU 互联技术减少数据跨越 PCIe 或主机内存路径的开销;多节点训练则通常需要依靠高速网络,例如 InfiniBand 或高速以太网,并通过 NCCL 等通信库完成跨 GPU、跨服务器的集合通信。NCCL本身并不是完整的分布式训练框架,而是专门负责 GPU 集合通信和点对点通信的库。
混合精度也可以降低通信压力。梯度和通信缓冲区采用较低精度的数据类型时,传输的数据量可能减少,从而降低网络带宽需求,但具体能否直接减少通信量以及采用什么精度,要取决于训练框架、优化器和数值稳定性设计。与此同时,梯度压缩、梯度累积、通信融合以及不同的 bucket 配置,也可以用于进一步改善计算与通信之间的比例。
需要特别区分的是,数据并行、模型并行和流水线并行解决的是不同问题。数据并行主要复制模型并切分数据;模型并行则把模型计算或参数分布到不同 GPU;流水线并行通常把模型按层划分成多个阶段,让不同 GPU 处理不同的微批次。大型模型训练中,这些方法往往会组合使用,而不是简单选择其中一种。
至于异步梯度更新,它并不是同步数据并行中的一个普通“通信优化开关”。异步训练允许不同 GPU 或训练工作者在不等待所有参与者的情况下继续更新参数,会改变优化过程中的梯度一致性和参数陈旧性特征,因此需要重新考虑收敛行为和系统稳定性。在现代大规模训练中,同步数据并行仍然是非常重要的基础方式,而更复杂的并行体系通常是在数据并行基础上进一步组合张量并行、流水线并行、专家并行以及参数或梯度分片等技术。
多 GPU 训练真正困难的地方,因此并不是简单地“把几块 GPU 连起来”。GPU 数量增加以后,计算、显存、通信带宽、通信延迟、拓扑结构、梯度同步以及负载均衡会共同决定扩展效率。理想状态下,增加 GPU 可以让 global batch 更快处理完成;现实系统中,如果通信时间逐渐接近甚至超过计算时间,继续增加 GPU 的收益就会下降。
数据并行的核心思想可以归纳为:复制模型、切分数据、分别计算梯度、通过集合通信聚合梯度,再让所有模型副本按照一致的更新结果继续训练。AllReduce 是其中最重要的通信机制之一,而真正决定大规模训练效率的,则是通信算法、GPU互联、网络拓扑、梯度分桶以及计算与通信重叠之间的整体配合。随着训练规模从单机多 GPU 扩展到数百甚至数千 GPU,如何让通信尽可能隐藏在计算之后,已经成为分布式训练系统设计中的重要问题。
数据并行是怎样工作的,多 GPU 训练为什么需要同步梯度
图片说明:示意图 图片来源:Public Domain(公有领域)
喜欢这篇报道?
使用下面的功能,方便以后继续阅读和分享 MNewsTV
关于文章收藏
收藏不需要注册帐号,收藏信息仅保存在当前浏览器中。
删除收藏请进入「我的收藏」进行管理。
捐助(Paypal): https://www.paypal.me/observeccp 订阅中国观察电报 Telegram : https://t.me/s/ObserveCCP