分布式机器学习:系统、工程与实战
上QQ阅读APP看本书,新人免费读10天
设备和账号都新为新人

1.5 从模型角度看如何并行

目前很多超大模型都是通过对小模型进行加宽、加深、拼接得到的,因此我们也可以反其道而行之,看看如何切分模型的各个维度,然后针对这些维度做反向操作。“分而治之”是分布式机器学习的核心思想,具体来说就是把拥有大规模参数的机器学习模型进行切分,分配给不同机器进行分布式计算,对计算资源进行合理调配,对各个功能模型进行协调,直到完成训练,获得良好的收敛结果,从而在训练速度和模型精度之间达到一个良好平衡。我们接下来就看如何“分而治之”,即从模型网络角度出发看如何并行。

1.5.1 并行方式

机器学习中的每个计算都可以建模为一个有向无环图(DAG)。DAG的顶点是计算指令集合,DAG的边是数据依赖(或者数据流)。在这样的计算图中,计算并行性可以用两个主要参数来表征其计算复杂度:图的工作量W(对应于图节点的总数)和图的深度D(对应于DAG中任意最长路径上的节点数目)。例如,假设我们每一个时间单位处理一个运算,则在单个处理器上处理图的时间是T1=W,在无穷数量处理器上处理图的时间是T∞=D。这样,计算的平均并行度就是W/D,这是执行计算图所需要的进程(处理器)的最佳数目。使用p个处理器处理一个DAG所需要时间的计算公式如下。[2]

min{W/p,D}≤TpO (W/p+D)

在前向传播和反向传播的时候,我们可以依据对小批量(mini-batch,就是数据并行切分后的批量)的使用、宽度(∞W)和深度(∞D)这三个维度来把训练分发(Scatter)到并行的处理器之上,按照此分发方式看,深度训练的并行机制主要有三种划分方式:按照输入样本划分(数据并行),按照网络结构划分(模型并行),按照网络层进行划分(流水线并行),具体如下:

• 第一种是数据并行机制(对于输入数据样本进行分区,在不同节点上运行数据样本的不同子集),其往往意味着计算性能的可扩展。大多数场景下的模型规模其实都不大,在一张GPU上就可以容纳,但是训练数据量会比较大,这时候适合采用数据并行机制,即在多节点之上并行分割数据和计算,每个节点只处理一部分数据。

• 第二种是模型并行机制(对于模型按照网络结构划分,在不同节点上运行模型同一层的不同部分),其往往意味着内存使用的可扩展。当一个节点无法存下整个模型时,就需要对图进行拆分,这样不同的机器就可以计算模型的不同部分,从而将单层的计算负载和内存负载拆分到多个设备上。模型并行也叫作层内模型并行或者张量级别的模型并行。

• 第三种是流水线并行机制(对于模型按照层来分区,在不同节点上运行不同层)。因为神经网络具有串行执行的特性,所以我们可以将网络按照执行顺序切分,将不同层放到不同设备上计算。比如一个模型网络有六层,可以把前三层放到一个设备上,后三层放到另一个设备上。流水线并行也叫作层间模型并行。

并行机制划分如图1-3(神经网络并行方案)所示(见彩插)。

图1-3所示的三种并行维度是两两正交的,DistBelief分布式深度学习系统就结合了这三种并行策略。训练同时在复制的多个模型副本上进行,每个模型副本在不同的样本上训练(数据并行),每个副本上依据同一层的神经元(模型并行)和不同层(流水线并行)上划分任务,进行分布训练。在实际训练过程中,小模型可能仅数据并行就足够,大模型因为参数多、计算量大,由一个GPU难以完成,所以要将显存和计算拆解到不同GPU上,就是模型并行。有时候数据并行和模型并行会同时发生。一些常见的拆解思路如下。

图1-3

图片来源:论文Demystifying Parallel and Distributed Deep Learning:An In-Depth Concurrency Analysis

• 对于与数据相关的模型,我们可以通过对数据的切分来控制切分模型的方式。这类模型的典型例子为矩阵分解,其模型参数为键-值对(key-value pair)格式。

• 有些模型不直接与数据相关(如LR、神经网络等),这时要分别对数据和模型做各自的切分。

1.5.2 数据并行

数据并行如图1-4所示(见彩插),其目的是解决计算墙问题,将计算负载切分到多张卡上,特点如下。

• 将输入数据集进行分区,分区的数量等于Worker(工作者/计算节点/计算任务/训练服务器)的数量。目的是先将每个批量的输入数据平均划分成多个小批量,然后把这些小批量分配到系统的各个Worker,每个Worker获取到一个小批量数据,这样每个Worker只处理训练数据的一个子集。

图1-4

图片来源:论文Communication-Efficient Distributed Deep Learning:A Comprehensive Survey

• 模型在多个Worker上复制,每个Worker维护和运行的都是完整的模型。虽然不同Worker的数据输入不同,但是运行的网络模型相同(也可以认为是模型参数共享)。

• 每个Worker在本地数据分区上进行独立训练(梯度下降)并生成一组本地梯度。

• 在每次迭代过程中,当反向传播之后需要进行通信时,将所有机器的计算结果(梯度)按照某种方式(集合通信或者参数服务器)进行归约(Reduce)(比如求平均),以获得相对于所有小批量的整体梯度,然后把整体梯度分发给所有Worker。每次聚集传递的数据量和模型大小成正比。

• 在权重更新阶段,每个Worker会用同样的整体梯度对本地模型参数进行更新,这样保证了下次迭代的时候所有Worker上的模型都完全相同。

由于是多个Worker并行获取/处理数据,因此在一个迭代过程中可以获取/处理比单个Worker更多的数据,这样大大提高了系统吞吐量。而通过增加计算设备,我们可以近似增加单次迭代的批量大小(batch size)(增加的倍数等于Worker数)。这样做的优势是:批量大小增大,模型可以用更大的步幅达到局部最小值(需要相应地调整学习率),从而加快优化速度,节省训练时间。

1.5.3 模型并行

模型并行如图1-5所示(见彩插),其目的是解决内存墙问题,通过修改层内计算方式,将单层的计算负载和显存负载切分到多张卡上,其原理如下。

• 将计算进行拆分。深度学习计算主要是矩阵计算,而矩阵乘法是并行的。如果矩阵非常大以至于无法放到显存中,只能把超大矩阵拆分到不同卡上进行计算。

• 将模型参数进行分布式存储。“基于图去拆分”会根据每一层的神经元特点(如CNN张量的通道数、图像高度或者宽度)把一张大图拆分成很多部分,每个部分会放置在一台或者多台设备上。

• 每个Worker仅仅对模型参数的一个子集进行评估和更新。

• 样本的小批量被复制到所有处理器,神经网络的不同部分在不同处理器上计算,这样可以节省存储空间,但是在每个层计算之后会引起额外的通信开销。

图1-5

图片来源:论文Communication-Efficient Distributed Deep Learning:A Comprehensive Survey

1.5.4 流水线并行

流水线并行如图1-6所示,其目的同样是解决内存墙问题,将整个网络分段,把不同层放到不同卡上,前后阶段分批工作,前一阶段的计算结果传递给下一阶段再进行计算,类似接力或者流水线,将计算负载和显存负载切分到多张卡上。流水线并行特点如下。

图1-6

• 在深度学习领域,流水线指可以重叠的计算,即在当前层和下一层(当数据准备就绪时)连续计算;或者利用神经网络串行执行的特性,根据深度划分深度神经网络(DNN),将不同层分配给不同的设备,从而达到切分计算负载和显存负载的目的。

• 流水线并行将一个数据小批量再划分为多个微批量(micro-batch),以使设备尽可能并行工作。

• Worker之间的通信被限制在相邻阶段之间,比如前向传播的激活和反向传播的梯度,因此通信量较少。当一个阶段完成一个微批量的前向传播时,激活将发送给流水线的下一个阶段。类似地,当下一阶段完成反向传播时将通过流水线把梯度反向传播回来。

• 流水线并行可以看作是数据并行的一种形式,由于样本是通过网络并行处理的,也可以看作模型并行。流水线长度往往由DNN结构来决定。

由于神经网络串行的特点使得朴素流水线并行机制在计算期间只有一个设备属于活跃状态,资源利用率低,因此流水线并行要完成的功能包括以下方面。

• 为了确保流水线的各个阶段能并行计算,必须同时计算多个微批量。目前已经有几种可以平衡内存和计算效率的实现方案,如PipeDream。

• 当一张卡训练完成后,要马上通知下一张卡进行训练,目的是让整个计算过程像流水线一样连贯,这样才能在大规模场景下提升计算效率,减少GPU的等待时间。

1.5.5 比对

下面我们来比对一下数据并行和模型并行(把流水线并行也归到此处)两者的特点。

• 同步开销:数据并行每次迭代需要同步N个模型的参数,这对带宽消耗非常大;模型并行通信量也大,因为其与整个计算图相关,因此更适合多GPU服务器;流水线并行只传输每两个阶段之间边缘层的激活和梯度,由于数据量较小,因此对带宽消耗较小。从减少通信数据量角度看,如果模型参数量较少但是中间激活较大,使用数据并行更适合。如果模型参数量较大但是中间激活较小,使用模型并行更适合。但是超大模型必须采用流水线并行模式。

• 负载均衡:模型并行通过模型迁移实现负载均衡;数据并行通过数据迁移实现负载均衡。调节负载均衡其实就是解决掉队者(Straggler)问题。