正文
=================== ===================
| 0 | 1 | 2 | 3 | | 4 | 5 | 6 | 7 |
=================== ===================
gpu0 gpu1
我们将其垂直切分为2部分,将第0-3层放在GPU0上,将第4-7层放在GPU1上。
当数据从第0层传递到第1层,第1层到第2层,以及第2层到第3层时,这就像普通模型一样。但是当数据需要从第3层传递到第4层时,它需要从GPU0传输到GPU1,这会引入通信开销。如果参与的GPU位于同一计算节点(例如同一物理机)上,这种复制速度相当快,但如果GPU位于不同的计算节点(例如多台机器)上,通信开销可能会显著增加。
然后第4层到第5层到第6层到第7层的运行就像普通模型一样,当第7层完成时,我们通常需要将数据发送回第0层(标签所在的位置),或者将标签发送到最后一层。现在可以计算损失,优化器可以开始工作。
问题:
-
主要缺陷(也是为什么称之为"朴素"MP的原因)是在任何时刻只有一个GPU在工作,其他GPU都处于空闲状态。因此,如果使用4个GPU,这几乎等同于将单个GPU的内存量增加4倍,而忽略了其余的硬件。此外还有设备间数据复制的开销。所以4个6GB显卡使用朴素MP可以容纳与1个24GB显卡相同大小的模型,但后者会更快完成训练,因为它没有数据复制开销。但是,比如说,如果你有40GB的显卡,需要容纳一个45GB的模型,你可以用4个40GB的显卡(但由于梯度和优化器状态的存在,勉强可以)
-
流水线并行
流水线并行(PP)与朴素MP几乎相同,但它通过将输入批次分块成微批次并人为创建流水线来解决GPU空闲问题,这使得不同的GPU可以同时参与计算过程。
下面来自GPipe论文(https://ai.googleblog.com/2019/03/introducing-gpipe-open-source-library.html)的插图展示了朴素MP(上图)和PP(下图):
从中图可以很容易看出PP如何减少了GPU空闲的死区。这些空闲部分被称为"气泡"。
图中的两部分都展示了pp=4的并行性。也就是说有4个GPU参与流水线。因此有4个管道阶段的前向路径F0、F1、F2和F3,然后是反向顺序的反向路径B3、B2、B1和B0。
PP引入了一个新的超参数
chunks
来调优,它定义了通过同一管道阶段按顺序发送多少块数据。例如,在图中可以看到
chunks=4
。GPU0对块0、1、2和3执行相同的前向路径(F0,0、F0,1、F0,2、F0,3),然后等待其他GPU完成它们的工作,只有当它们的工作开始完成时,GPU0才会再次工作,对块3、2、1和0执行反向路径(B0,3、B0,2、B0,1、B0,0)。
注意,从概念上讲,这与梯度累积步骤(GAS)是相同的概念。PyTorch使用
chunks
,而DeepSpeed将相同的超参数称为GAS。
由于分块,PP引入了微批次(MBS)的概念。DP将全局数据批次大小分成小批次,因此如果DP度为4,全局批次大小1024会被分成4个每个256的小批次(1024/4)。如果
chunks
(或GAS)数量为32,我们最终得到微批次大小为8(256/32)。每个流水线阶段一次处理一个微批次。
要计算DP + PP设置的全局批次大小,我们执行:
mbs*chunks*dp_degree
(
8*32*4=1024
)。
让我们回到这个图。
当
chunks=1
时,你最终会得到朴素MP,这是非常低效的。而当
chunks
值很大时,你会得到非常小的微批次大小,这也可能不太高效。因此,需要通过实验来找到能够实现GPU最高效利用率的值。
虽然图中显示了一个无法并行化的"死亡"时间气泡(因为最后的
forward
阶段必须等待
backward
完成管道),但寻找最佳
chunks
值的目的是实现所有参与GPU的高并发利用率,这意味着要最小化气泡的大小。
调度的选择对高效性能至关重要,按发明顺序排列的最常见调度方式包括:
-
顺序 Gpipe: 使用流水线并行实现巨型神经网络的高效训练(https://arxiv.org/abs/1811.06965)
-
交错 1F1B Pipedream: 快速高效的流水线并行DNN训练(https://arxiv.org/abs/1806.03377)
-
循环、深度优先的高效大规模语言模型在GPU集群上的训练使用Megatron-LM(https://arxiv.org/abs/2104.04473)
-
广度优先的流水线并行(https://arxiv.org/abs/2211.05953)
这里是一个交错流水线的例子:
parallelism-sagemaker-interleaved-pipeline
在这里,气泡(空闲时间)通过优先处理反向传播进一步最小化。
DeepSpeed、Varuna和SageMaker等都使用了这种方式。
Varuna通过使用模拟来发现最有效的调度方式,从而进一步改进调度。
PP解决方案有两类 - 传统的Pipeline API和更现代的解决方案,后者通过帮助部分或完全自动化流程,使最终用户使用起来更加容易:
-
-
传统Pipeline API解决方案的问题:
-
必须对模型进行大量修改,因为Pipeline要求将模块的正常流程重写为相同模块的
nn.Sequential
序列,这可能需要更改模型的设计。
-
目前Pipeline API非常受限。如果在Pipeline的第一阶段有一堆Python变量需要传递,你必须找到解决方法。目前,pipeline接口只接受单个Tensor或Tensor元组作为唯一的输入和输出。这些张量的第一个维度必须是批次大小,因为pipeline会将mini batch分成micro-batches。可能的改进正在这里讨论(https://github.com/pytorch/pytorch/pull/50693)
-
在pipe阶段级别的条件控制流是不可能的 - 例如,像T5这样的编码器-解码器模型需要特殊的变通方法来处理条件编码器阶段。
-
必须安排每一层,使一个模型的输出成为另一个模型的输入。
我还没有尝试过Varuna和SageMaker,但根据他们的论文报告,他们已经克服了上述问题列表,并且对用户的模型只需要很小的改动。
实现:
-
Pytorch(https://pytorch.org/docs/stable/pipeline.html) (在pytorch-1.8中初步支持,并在1.9和1.10中逐步改进)。一些示例(https://github.com/pytorch/pytorch/blob/master/benchmarks/distributed/pipeline/pipe.py)
-
FairScale(https://fairscale.readthedocs.io/en/latest/tutorials/pipe.html)
-
DeepSpeed(https://www.deepspeed.ai/tutorials/pipeline/)
-
Megatron-LM(https://github.com/NVIDIA/Megatron-LM)有内部实现 - 没有API。
-
Varuna(https://github.com/microsoft/varuna)
-
SageMaker(https://arxiv.org/abs/2111.05972) - 这是一个只能在AWS上使用的专有解决方案。
-
OSLO(https://github.com/eleutherAI/Oslo) - 这是基于Hugging Face Transformers实现的。
-
PiPPy(https://github.com/pytorch/pippy) - 通过
torch.fx
自动PP
-
nanotron(https://github.com/huggingface/nanotron)
张量并行
在张量并行中,每个GPU只处理张量的一个切片,只在需要完整张量的操作时才聚合完整的张量。
在本节中,我们使用来自Megatron-LM(https://github.com/NVIDIA/Megatron-LM)论文的概念和图表:在GPU集群上高效训练大规模语言模型(https://arxiv.org/abs/2104.04473)。
任何transformer的主要构建块都是一个全连接层
nn.Linear
,后面跟着一个非线性激活函数
GeLU
。
按照Megatron论文的符号,我们可以将点积部分写为
Y = GeLU(XA)
,其中
X
和
Y
是输入和输出向量,
A
是权重矩阵。
如果我们以矩阵形式查看计算,很容易看出矩阵乘法如何在多个GPU之间拆分:
Parallel GEMM
如果我们将权重矩阵
A
按列分割到
N
个GPU上,并行执行矩阵乘法
XA_1
到
XA_n
,那么我们将得到
N
个输出向量
Y_1, Y_2, ..., Y_n
,它们可以独立地输入到
GeLU
中:
