【AI大模型学习】Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM

Introduction

背景：

• 训练LLM的挑战：

1. GPU 内存容量有限，即使在多 GPU 服务器上也无法适应大型模型
2. 所需的计算操作数量可能会导致训练时间过长。

• 新的并行方法：

1. Tensor parallelism （张量并行，TP）
2. Pipline parallelism （管道并行 或 流水线并行，PP）
3. Data parallelism （数据并行，DP）

当然，你直接去使用这些并行方法会导致数千个 GPU 出现扩展问题。因为随着GPU和CPU的数量增多，通信量的增长却不是线性的。
所以本文通过组合张量、流水线和数据并行，实现扩展到数千个 GPU的大模型训练，提出了一种新颖的交错流水线调度，可以将吞吐量提高 10% 以上，并且内存占用与现有方法相当。
能够在 3072 个 GPU 上以 502 petaFLOP/s 的速度对具有 1 万亿个参数的模型执行训练迭代（每 GPU 吞吐量为理论峰值的 52%）。

以上是简述的abstract翻译。回到正文本身。

我们提到并行度的时候，很难避开两个问题（限制）：

• 横向扩展超过某个点，每个 GPU 的batch size 变得太小，降低了GPU利用率并增加了通信成本；
• 由于可以使用的设备的最大数量是batch size ，这限制了可用于训练的加速器数量。

分层的张量并行，一定程度上是可以解开以上这些限制，：

假设在一台八卡的NVIDIA DGX A100 服务器上（ 8个80GB-A100 GPU），张量并行能很好的解决问题，但我们谈到更大的模型，那就需要更多服务器，也就是多机（多节点）的情况，这个时候张量并行就存在新的问题：

• 张量并行所需的all-reduce必须经过服务器间的连接，通信会消耗大量时间，因为多机通讯明显比单机内GPU的通信更慢（尤其是单机内还有nvlink这样的东西）
• 高度的模型并行化可能会产生小矩阵乘法（GEMM），可能降低 GPU 利用率。

所以本文的目的是提出一个DP，PP和TP相结合的并行方法（简称为 PTD-P），提升训练吞吐量，同时保留严格的优化器语义。

MODES OF PARALLELISM

我们跳过繁杂的定义，直接来看第一张图（Figure2），这是一个张量并行和流水线并行的组合。
左边和右边分别是流水线并行，在每一块中间又按照张量（Tensor）并行分块。

流水线并行：

管道并行化：我们将模型的层被分配到多个设备上。对于重复的transformer块，每个设备可以分配相同数量的transformer层。（当然文章没有考虑更不对称的模型架构，其中将层分配给管道阶段更加困难，这会在后面提到）

micro batches 划分：一个batch被划分为更小的micro batch；随后micro batch之间可以进行流水线执行。流水线方案需要确保输入在前向和后向传播中看到一致的权重版本。

可以再来看Figure3。文章把切分好的batch放到设备之间进行流水线调度，但是这就使得某些设备存在空闲——也就是灰色的块。这些部分被定义为气泡（pipeline bubble）。

由于每一次batch计算的开始和结束时，设备都会处于空闲状态。而且流水线是严格且随时刷新的。在一个简单的流水线中，后向传递(BP)时需要更新权重，而在前向传递却不需要。从Figure2下方这条流水线图可以看出。

G-Pipe：

我们可以用一个公式来计算气泡率：
$$\frac{t_{pb}}{t_{id}}=\frac{(p-1)}{m}$$

GPipe 管道调度，其中所有micro batch（用数字表示，）前向传播（蓝色），反向传播（绿色）。灰色区域代表气泡。为简单起见，我们假设反向传播所需的时间是前向传播的两倍。管道调度的效率不取决于这个因素。本例中的每个批次由 8 个micro batch组成，每个蓝色或绿色框中的数字是对应微批次的唯一标识符（特别是，第一批次由micro batch 1 – 8 组成，第二批次由micro batch 9 – 16 组成） ， 等等）。优化器是步进式的，权重参数在管道刷新时更新，以确保严格的优化器语义，从而导致空闲设备和流水线气泡。

未完待续施工中