前言
曾在游戏世界挥洒创意,也曾在前端和后端的浪潮间穿梭,如今,而立的我仰望AI的璀璨星空,心潮澎湃,步履不停!愿你我皆乘风破浪,逐梦星辰!
Synchronous SGD(同步SGD)
首先介绍一个概念,SGD: 同步随机梯度下降(Synchronous Stochastic Gradient Descent, Sync-SGD)是一种分布式深度学习的训练方法,它在多个GPU上并行计算梯度,并在每次更新时同步所有计算结果(All-Reduce),以确保参数的一致性。
工作原理
-
数据并行:
- 数据集被划分成多个子集,每个 GPU 负责一个子集。
- 假设有 n 个 GPU,每个 GPU 处理 b 个样本,则总的批量大小为 nb。
-
前向传播 (Forward Pass):
- 每个 GPU 计算自己批次数据的前向传播,得到损失。
-
后向传播 (Backward Pass):
- 每个 GPU 计算自己批次数据的梯度(即参数的偏导数)。
-
梯度同步 (All-Reduce 操作):
- 计算出的梯度会通过全局同步(如 All-Reduce 操作)聚合:g=1/n(gi+g2+....+gi)
- 其中 gi 是第 i 个 GPU 计算出的梯度,g 是平均后的全局梯度。
-
参数更新:
- 每个 GPU 使用相同的全局梯度 g 来更新本地模型参数。
等价于单 GPU 运行更大批量
从数学上来看,同步 SGD 等价于在单 GPU 上运行批量大小为 nb 的SGD。
更通俗易懂的解释,我们来做一个计算时间对比
假设:
- 在 单 GPU 训练 时,每批数据(batch size = 128)训练一次需要 0.4 秒。
- 在 同步 SGD 训练(4 GPU) 时,每个 GPU 处理 batch size = 32,因此:
- 前向+后向计算时间 仍然是 0.1 秒/批次(每个 GPU 计算较少数据)
- 梯度同步(All-Reduce)时间 设为 0.02 秒/批次
- 总训练时间 = 计算时间 + 通信时间 = 0.1+0.02=0.120.1 + 0.02 = 0.120.1+0.02=0.12 秒/批次
方案 | GPU 数量 | 每 GPU 处理的批量大小 | 总批量大小 | 训练时间/批次 |
---|---|---|---|---|
单 GPU | 1 | 128 | 128 | 0.4s |
4 GPU 同步 SGD | 4 | 32 | 128 | 0.12s |
理想情况下,理论加速比:
0.4 / 0.12 ≈ 3.33
接近 4 倍加速,但因为通信开销(All-Reduce 需要时间),实际加速比 小于理想情况。
思考:分布式计算最大的问题是计算被收发数据的时间卡住了,假设t1(Compute)是单GPU计算b个样本梯度的时间,t2(Communication)是接收和发送的时间,那么每个批量的计算时间是max(t1, t2)。这怎么解决?
分布式训练的瓶颈
- 如果计算(t1)远大于通信(t2):
- 计算成为主要的瓶颈,通信开销影响较小。
- 这种情况一般发生在模型计算量非常大(如 Transformer、ResNet-152 这种计算密集型模型)。
- 如果通信(t2)远大于计算(t1):
- 训练被通信时间卡住,即“计算等数据”。
- 这种情况通常出现在:
- 批量大小较小(计算量小但通信成本不变)。
- 网络带宽较低(如 GPU 之间通信速率低,或分布式节点远离)。
- All-Reduce 效率低(通信拓扑、数据传输方式不优化)。
如何优化通信瓶颈?
如果通信时间成为瓶颈(t2≫t1),可以采用以下几种优化方案:
-
增加计算负载(提升 t1)
- 增大 batch size:让每个 GPU 处理更多数据,减少通信相对开销。
- 使用更复杂的模型:计算量大,通信占比自然下降。
-
降低通信开销(减少 t2)
- 梯度压缩(Gradient Compression):减少传输数据量,如量化梯度、剪枝。
- 拓扑优化:
- NVLink/InfiniBand:比传统 PCIe 速度更快,减少通信时间。
- Tensor Fusion(张量融合):将多个小梯度合并成一个大梯度,一次性传输,减少通信开销。
-
异步通信
- Overlap Compute & Communication(计算-通信重叠):
- 让计算和通信并行进行,尽可能利用 GPU 和网络资源。
- 例如,在计算下一个 batch 时,上一批的梯度仍在 All-Reduce。
- Overlap Compute & Communication(计算-通信重叠):
这里抛出一个有意思的思考:假设硬件不动,咱们采用增加batch的方式,虽然会减少每个epoch的训练时间,但是batch增加的副作用是精度会下降,然后咱们会打算增加训练的epoch轮数来达到所需的精度,但是增加epoch也耗时间呐,这会失去咱们本来分布式训练的目的:增快训练速度。怎么办?
如何权衡计算效率和精度?
方法 1:使用线性缩放学习率规则
- 规则:当 batch size 从 B 增加到 kB 时,学习率 η也增大 k 倍: η′=k×η
- 例如:
- Batch size = 32,学习率 η=0.01。
- Batch size = 256(增大 8 倍),则 新学习率 η′=0.08。
📌 理由:大 batch 计算出的梯度更稳定,所以可以用更大的学习率,使优化器跳得更远,加速收敛。
方法 2:使用 Warmup 预热
- 直接使用大批量 + 高学习率可能导致训练不稳定。
- Warmup 预热策略:前 几千个 batch 逐渐增加学习率,比如:
- 第 1 个 batch,学习率 0.010.010.01。
- 第 1000 个 batch,学习率 0.080.080.08(最终值)。
- 这样可以让模型逐渐适应大 batch 的训练,避免震荡过大导致收敛不稳定。
方法 3:梯度聚合(Gradient Accumulation)
- 问题:如果 GPU 显存不足,无法一次性处理超大 batch。
- 解决方案:
- 使用小 batch 计算多个 step 的梯度后,累积梯度,再进行一次更新。
- 例如:
- GPU 只能容纳 batch size = 64,但想要等效于 batch size = 256。
- 运行 4 次 batch size 64,累积梯度后再做一次梯度更新。
方法 4:使用 LARS 优化器
- 在 ResNet、BERT 等大规模模型上,LARS(Layer-wise Adaptive Rate Scaling) 可以帮助大 batch 训练稳定收敛。
- 它对不同层使用不同的学习率,避免大 batch 训练时梯度更新过平滑的问题。
结论
- 增大 batch size 可以减少一个 epoch 的训练时间,但可能导致模型泛化能力下降,甚至需要增加 epoch 来弥补精度损失,最终可能并没有真正节省时间。
- 最佳策略:增大 batch size 时,需要调整学习率(线性缩放)+ warmup 预热 + LARS 优化器,否则可能会影响模型的收敛效果。
- 在实际应用中,适当增大 batch size(比如 32 → 512),配合优化策略,可以加速训练;但如果 batch size 过大(比如 32 → 8192),可能导致训练精度下降,最终得不偿失。
🚀 最终目标:找到一个 batch size、学习率和通信效率之间的平衡点,让训练尽可能快,同时保证模型的准确率和泛化能力。