人工智能领域的许多最新进展都围绕着大规模神经网络展开,但训练大规模神经网络是一项艰巨的工程和研究挑战,需要协调GPU集群来执行单个同步计算。随着集群数和模型规模的增长,机器学习从业者开发了多项技术,让机器学习模型能在多个GPU上进行并行模型训练。

乍一看,这些并行技术令人生畏,但只需对计算结构进行一些假设,这些技术就会变得清晰:从某些角度来看,这也只是从 A 到 B 传递并不透明的位,就像数据包在网络交换机之间传递一样。

Fig 1. 各种 Parallel 模型

不同的并行技术将训练过程划分为不同的维度,包括:

  • 数据并行(Data Parallelism)在不同的GPU上运行同一批数据的不同子集
    • DP(Data Parallel)
    • DDP(Distributed Data Parallel)
    • FSDP(Fully Shared Data Parallel)
  • 流水并行(Pipeline Parallelism)在不同的GPU上运行模型的不同层
  • 模型并行(Tensor Parallelism)将单个数学运算(如矩阵乘法)拆分到不同的GPU上运行
  • 专家混合(Mixture-of-Experts)只用模型每一层中的一小部分来处理数据。

Note:Tensor Parallelism 翻译成模型并行可能并不是非常的恰当🤣

1. 并行模型

1.1 数据并行 (Data Parallesim)

数据并行是指将相同的参数复制到多个工作节点上,并为每个工作节点分配不同的数据子集同时进行处理。每个工作节点拥有完整的神经网络模型,每次训练仅将一批数据输入模型,进行前向传播、计算误差、反向传播,最后进行参数的更新。储了参数的更新,其余的操作都是互相独立的,所以可以在多个节点上进行并发的执行。

每个工作节点都有自己的模型和输入,当属于自己的模型参数推理完成后(产生了梯度参数),所有的工作节点会把参数(梯度参数)发给一个 Master,这个 Master 会把所有节点传进来的参数做融合,通过这些梯度参数更新生成新的模型参数,然后把这个模型的参数再发送给每个工作节点,拱他们进行下一轮的计算。

在 Pytorch 中提供了DP(Data Parallel)、DDP(Distributed Data Parallel)、FSDP(Fully Shared Data Parallel)三种不同的数据并行方法。

1.1.1 DP(Data Parallel)Parameter Server

DP 使用了 Parameter Server(PS) 作为理论依据。PS 结构是李沐老师提出来的方法,由server节点和worker节点组成。

[点击折叠] 李沐老师的 PS 讲解

Server 节点的主要功能是初始化和保存模型参数、接受worker节点计算出的局部梯度、汇总计算全局梯度,并更新模型参数。

Fig 2. Parameter Server

Worker 节点的主要功能是各自保存部分训练数据,初始化模型,从 Server 节点拉取(Pull)最新的模型参数,再读取参数,根据训练数据计算局部梯度,上传(Push)给 Server 节点。ps:李沐老师说这个 Pull 和 Push 叫法来源于 Git。

在这个模式下,会有负载不均衡的情况出现。因为在 Server 需要大量的显存来保存从 Woker 到来的梯度参数,而 Server 还需要将更新后的参数广播到每一个 Worker 上,那么 Server 的网络带宽就变得是非重要了。随着 Worker 数目的增多,网络的需求会更加大。并且 Server 作为最终处理所有 Worker 传来的梯度参数的机器,其计算消耗时间也是制约其他 Worker 不间断工作的瓶颈。

1.1.2 DDP(Distributed Data Parallel)Ring-All-Reduce

DDP 主要基于 Ring-All-Reduce 算法,该算法主要分两步:

  1. share-reduce:会逐步交换彼此的梯度并融合,最后每个 GPU 都会包含完整融合梯度的一部分。
  2. share-only:GPU 会逐步交换彼此不完整的融合梯度,最后所有 GPU 都会得到完整的融合梯度
Fig 3. Ring-All-Reduce 示意图

在 share-reduce 阶段,每个进程 p 将数据发送给进程 (p+1) % p,其中 % 是模运算符。所以进程A会向进程 B 发送。这就是创建类似于环形队列状态的的原因。此外,长度为 n 的数组被划分成 p 块,这些块中的每一个都将以 i 为索引。它看起来会是这样的:

Fig 4. Ring-All-Reduce 1

第一个 share-reduce 操作进程 A 向进程 B 发送 a0,进程 B 将向进程 C 发送 b1,等等。类似这样:

Fig 5. Ring-All-Reduce 2

然后,当每个进程收到前一个进程的数据时,它就会应用 reduce operation,然后继续将其再次发送到环中的下一个进程。如果 reduce operation 是 sum。它看起来就会像这样:

Fig 6. Ring-All-Reduce 3

然后,再进一步:

Fig 7. Ring-All-Reduce 4

第二步,shared-only,这个操作非常的简单,只是以环形方式共享数据,而不应用 reduce operation。这就把每个进程中的每个分块的结果合并起来。

Fig 8. Ring-All-Reduce 5

Ring-All-Reduce 有这些优点:

  1. 负载分散在每个 Worker 上,通讯时间基本是一致的。并且不需要通过 Server 分发全模型的参数到每个 Worker 上。
  2. 使用 ring-all-reduce 的方式进行通讯,随着 Worker 数量 N 增加,总传输量恒定。也就是理论上,随着 Worker 数量的增加,ring all-reduce 有线性性能拓展能力。

DP 和 DDP 之间的区别

Data Parallel Distributed Data Parallel
更多的开销;模型在每次参数聚合后需要被复制,并在每次前向传递后被销毁(因为只需要梯度参数被传回就够了)。 模型只会被复制一次
只支持单个节点的 Master 可以拓展到多台机器上
单进程多线程的实现方式 采用多进程的方式

1.1.3 FSDP(Fully Shared Data Parallel)

DDP 用起来简单方便,但是要求整个模型能加载一个GPU上,这使得大模型的训练需要使用额外复杂的设置进行模型拆分。Pytorch 的 FSDP 从 DeepSpeed ZeRO [8] 以及 FairScale 的 FSDP 中获取灵感,打破模型分片的障碍(包括模型参数,梯度,优化器状态),同时仍然保持了数据并行的简单性。

FSDP 是对模型参数、优化器状态和梯度进行分片,可以选择将部分模型参数卸载到 CPU 中。FSDP 是在 DDP 的基础上,将 All-Reduce 操作分解为独立的 Reduce-Scatter 和 All-gather 的操作这种方式能够使得每个工作节点仅需存储一个参数和优化器状态的分片。FSDP 非常类似于 ZeRO-3

Fig 9. FSDP workflow [3]

如 FSDP 完全流程图所示,通过 All-gather 操作可以从其他工作节点获取模型权重来进行前向传播和反向传播。Reduce-Scatter 操作聚合局部梯度,并在各工作的节点上分片,还可以选择将梯度卸载到CPU 中,等需要时再从 CPU 中加载回来。

Fig 10. Reduce-Scatter and All-gather

原本的 All-Reduce 操作可以拆分成 Reduce-Scatter 和 All-gather 操作。在进行Reduce-Scatter 时,gradients 在相同块的各个 ranks 中被加起来,在进行 All-gather 时,每个 Worker 上可用的聚合梯度分片的部分可供所有 Worker 使用。

1.2 流水并行(Pipeline Parallesim)

通过流水线机制进行并行训练,我们将模型的参数分配到GPU上。每个GPU只持有一部分参数,因此,同一个模型在每个GPU上消耗的内存成比例减小。

将大模型分为若干份连续的layer很简单。然而,各层的输入和输出之间存在着顺序上的依赖性,所以一个简单的实现会导致大量的GPU空闲时间,因为 Worker 在等待前一个 Worker 的输出来用作其输入。这些等待时间块被称为 “气泡(bubble)”,而在这些等待的过程中,空闲的机器本可以继续进行计算。。

Fig 11. Naive way for Streaming [1]

为了减少气泡的开销,我们可以重用数据并行的思想,通过让每个 Worker 一次只处理一个数据元素的子集,使我们能够巧妙地将新的计算与等待时间重叠(overlapping)起来。核心思想是将一个批次分成多个微批(microbatches);每个微批的处理速度应该是成比例的,每个 Worker 在下一个微批可用时就开始工作,从而加速流的执行。有了足够的微批,Worker 可以在大部分时间内被利用,在进程的开始和结束时,气泡最小。梯度是跨微批的平均数,只有在所有微批完成后才会对参数进行更新。

模型被分割的节点数的数量通常被称为流水线深度(Pipline depth)。

在前向传递期间,Worker 只需要将其 Layer 的输出(称为激活)发送给下一个 Worker;在反向传递期间,它只将这些激活的梯度发送给前一个 Worker 。如何安排这些传递以及如何在微批中聚集梯度,有很大的设计空间。 GPipe让每个工作者连续处理前向和反向的传递,然后在最后同步聚合来自多个微批的梯度。而 PipeDream 则安排每个 Worker 交替地处理前向和反向 Stream。

Fig 12. GPipe and PipeDream [1]

1.3 模型并行(Tensor Parallesim)

流水并行将一个模型按层“垂直”拆分。也可以在一个 layer 内 “水平” 分割某些操作,这通常被称为模型并行训练。对于许多现代模型(如Transformer)来说,计算的瓶颈是将激活值与大的权重矩阵相乘。矩阵乘法可以被认为是成对的行和列之间的点积:有可能在不同的GPU上计算独立的点积,或者在不同的GPU上计算每个点积的一部分,然后将结果相加。无论采用哪种策略,我们都可以将权重矩阵切成大小均匀的 “块(Shards)”,将每个块放在不同的GPU上。要得到完整矩阵的结果,需要进行通信将不同部分的结果进行整合。

一个例子是 Megatron-LM,它在 Transformer 的 自注意 和 MLP层内并行化矩阵乘法。 PTD-P 同时使用张量、数据和流水线并行;它的流水线并行将多个不连续的 layer 分配到单设备上运行,以更多网络通信为代价来减少气泡开销。

有时,网络的输入可以在一个维度上进行并行化,相对于交叉通信来说,并行计算的程度很高。 序列并行就是这样一个想法,一个输入序列在不同时间被分割成多个子集,通过在更细粒度的子集上进行计算,峰值内存消耗可以成比例地减少。

1.4 混合并行(Mixture-of-Experts)

混合专家(MoE)方法,对于任何一个输入,只有一部分网络被用来计算输出。在有许多套权重的情况下,网络可以在推理时通过门控机制选择使用哪一套。这样就可以在不增加计算成本的情况下增加许多参数。每组权重被称为 “专家(experts)",希望网络能够学会将专门的计算任务分配给每个专家。不同的专家可以托管在不同的GPU上,这也为扩大模型使用的GPU数量提供了一个明确的方法。

Illustration of a mixture-of-experts (MoE) layer. Only 2 out of the n experts are selected by the gating network. (Image adapted from: Shazeer et al., 2017)[1]
混合专家(MoE

Illustration of a mixture-of-experts (MoE) layer. Only 2 out of the n experts are selected by the gating network. (Image adapted from: Shazeer et al., 2017)[1]

GShard将 MoE Transformer 的规模扩大到6000亿个参数,其中MoE layers被拆分到多个TPU上,其他layers是完全重复的。 Switch Transformer 通过将一个输入只路由到一个专家,将模型规模扩展到数万亿的参数,甚至有更高的稀疏度。

难点:

  • 现在的设备,比如GPU,比较擅长做运算,不擅长做分支。
  • 大批量是训练模型的必须,但是这种方式下会导致每个小模型的样本数较少,无法训练得到好的模型。
  • 为了控制稀疏性,可能需要在loss上去做些改进,来保障小模型上的均衡负载。

Moe Github, code

2. 自动并行方法

2.1 Exploring Hidden Dimensions in Parallelizing Convolutional Neural Networks [9]

Get This Paper

这篇文章勾勒了自动并行的基本框架,很多解决自动并行的工作都是这样一个流程。


Reference:

[1] Techniques for training large neural networks (openai.com)

[2] 先进编译实验室-自动并行-并行划分 - 知乎 (zhihu.com)

[3] Getting Started with Fully Sharded Data Parallel(FSDP)

[4] chenzomi12/DeepLearningSystem: Deep Learning System core principles introduction. (github.com)

[5] Deep Learning in a Nutshell: Core Concepts | NVIDIA Technical Blog

[6] Visual intuition on ring-Allreduce for distributed Deep Learning | by Edir Garcia Lazo | Towards Data Science

[7] 数据并行Deep-dive: 从DP 到 Fully Sharded Data Parallel (FSDP)完全分片数据并行 - 知乎 (zhihu.com)

[8] DeepSpeed: Extreme-scale model training for everyone - Microsoft Research

[9] Exploring Hidden Dimensions in Parallelizing Convolutional Neural Networks