Paper Reading: Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM
Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM
姑且算作 Megatron LM v2,因为是晚一年发表的,粗略过一下,因为都是同样的 motivation 和 background 等等
作者 Jared 的视频介绍
其中和 ZeRO 的对比很有意思,吞吐量高了很多。周末看看 Gpipe 和 ZeRO
ABSTRACT
本文中,我们展示了如何将张量并行、流水线并行和数据并行结合起来,以扩展到数千个 GPU。
提出了一种新颖的交错流水线调度方法,可以在内存占用与现有方法相当的情况下,将吞吐量提高 10%以上。
INTRODUCTION
数据并行扩展通常效果良好,但存在两个局限性:
- beyond a point, the per-GPU batch size becomes too small, reducing GPU utilization and increasing communication cost
- the maximum number of devices that can be used is the batch size, limiting the number of accelerators that can be used for training.
模型并行技术、张量(层内)模型并行化
- 张量并行所需的 all-reduce 通信需要通过服务器间链路,这些链路比多 GPU 服务器内可用的高带宽 NVLink[9]慢;
- 高度的模型并行化可能会产生小的矩阵乘法(GEMMs),可能会降低 GPU 利用率。
Pipeline model parallelism 流水线模型并行化,一个批次被拆分为较小的微批次,执行过程在这些微批次之间进行流水线处理。
将流水线并行、张量并行和数据并行结合起来,这种技术我们称之为 PTD-P
还与 ZeRO[36]进行了比较,发现由于跨节点通信较少,我们的方法在 1750 亿和 5300 亿参数的模型上比 ZeRO-3 高出 70%。这些模型太大,无法容纳在多 GPU 服务器上。
MODES OF PARALLELISM
Data Parallelism
每个工作节点都拥有完整模型的副本,输入数据集被分片,工作节点定期聚合它们的梯度,以确保所有工作节点看到一致的权重版本
Pipeline Model Parallelism
在流水线并行化中,模型的各层被分配到多个设备上。当用于具有重复相同 Transformer 块的模型时,每个设备可以被分配相同数量的 Transformer 层。
一个批次被拆分为较小的微批次;然后在微批次之间进行流水线执行。
其实流水线和 CPU 的流水线调度很相似?
量化 GPipe 的流水线气泡大小(𝑡𝑝𝑏)。我们将批次中的微批次数量表示为 𝑚,流水线阶段的数量(用于流水线并行的设备数量)表示为 𝑝,每次迭代的理想时间表示为 𝑡𝑖𝑑(假设完美或理想的扩展),执行单个微批次的前向和后向传递的时间表示为 𝑡𝑓 和 𝑡𝑏。在此调度中,流水线气泡包括批次开始时的 𝑝−1 个前向传递和批次结束时的 𝑝−1 个后向传递。流水线气泡中花费的总时间为 𝑡𝑝𝑏 = (𝑝−1)·(𝑡𝑓 +𝑡𝑏)。批次的理想处理时间为 𝑡𝑖𝑑 = 𝑚·(𝑡𝑓 +𝑡𝑏)。因此,流水线气泡中花费的理想计算时间的比例为
$$ Bubble\ time\ fraction (pipeline\ bubble\ size) = \frac{t_{pb}}{t_{id}} = \frac{(p−1)\cdot (t_f + t_b)}{m \cdot(t_f + t_b)} = \frac{p−1}{m} $$
为了减少流水线气泡的影响,可以增加微批次的数量或减少流水线阶段的数量。
使流水线气泡时间占比(bubble time fraction)尽可能小,我们需要满足 $m » p$
Schedule with Interleaved Stages: 新调度将气泡时间减少了 v 倍,这种流水线气泡大小的减少并非没有代价:这种调度需要额外的通信。
Tensor Model Parallelism
- MLP parallelism
- multi-head attention parallelism
PERFORMANCE ANALYSIS OF PARALLELIZATI ON CONFIGURATIONS
Tensor and Pipeline Model Parallelism
$$ \frac{p - 1}{m} = \frac{n / t - 1}{m} $$
因此,我们可以看到,张量模型并行化增加了设备之间的通信量。因此,当 𝑡 大于单个节点中的 GPU 数量时,跨较慢的节点间链路执行张量模型并行化的开销可能是不切实际的。
𝑡 表示张量模型并行的规模,𝑝 表示流水线模型并行的规模,𝑑 表示数据并行的规模。
𝑏:微批次大小 𝑚:每个流水线中一个批次内的微批次数量 𝑛:GPU 的数量
Data and Model Parallelism
在使用张量模型并行化时,每个微批次都需要执行 all-reduce 通信。这在跨多 GPU 服务器时可能会非常昂贵
Microbatch Size
微批次大小 b 的选择也会影响模型训练的吞吐量。
Activation Recomputation
激活重计算(Activation Recomputation)[12, 18, 20, 21] 是一种可选技术,通过在后向传递之前再次运行前向传递(并仅保存给定流水线阶段的输入激活值,而不是整个中间激活值集,后者占用内存更大),以增加计算操作的数量为代价来减少内存占用。
IMPLEMENTATION
略
EVALUATION
-
PTD-P 的性能如何?它是否能实现实际的端到端训练时间?
-
流水线并行化在给定模型和批量大小下的扩展性如何?交错调度对性能有多大影响?
-
不同的并行化维度之间如何交互?微批次大小等超参数的影响是什么?
-
scatter-gather 通信优化的影响是什么?在大规模运行训练迭代时,我们对硬件施加了哪些限制?
End-to-End Performance
评估了系统在参数量从 10 亿到 1 万亿的 GPT 模型上的端到端性能
启用了 scatter/gather 优化 的交错流水线调度
如何计算 FLOPS?
Comparison to ZeRO-3
PTD-P 在这两个模型上的性能比 ZeRO-3 高出 70%,这得益于更少的跨节点通信。
Pipeline Parallelism
Weak Scaling. 较高的批量大小具有更好的扩展性,因为流水线气泡被分摊到更多的微批次上
Interleaved versus Non-Interleaved Schedule. 启用 scatter/gather 通信优化 的交错调度比非交错(默认)调度具有更高的计算性能。
Comparison of Parallel Configurations
Microbatch Size
Activation Recomputation
Scatter-Gather Optimization
Fused Operators
Inter-Node Communication Bandwidth
Checkpoint Loading and Saving
各种配置 evaluation 略