AI框架里的并行技术

AI框架里的并行技术

  • 发布时间:2022-02-14 07:00
  • 访问量:

【概要描述】近年来,深度学习网络及数据集的规模越来越大,AI计算任务也越来越重。如何实现计算任务在多节点上的分摊,以及节点之间的信息交互及同步,对于分布式计算集群的性能起到至关重要的作用。本文将分析AI框架里的并行技术及一些优化策略。

AI框架里的并行技术

【概要描述】近年来,深度学习网络及数据集的规模越来越大,AI计算任务也越来越重。如何实现计算任务在多节点上的分摊,以及节点之间的信息交互及同步,对于分布式计算集群的性能起到至关重要的作用。本文将分析AI框架里的并行技术及一些优化策略。

  • 分类:研究院
  • 作者:
  • 来源:
  • 发布时间:2022-02-14 07:00
  • 访问量:
详情

摘要

 

近些年来,为了取得更高的精度及更广泛的通用性,AI的深度学习网络及数据集的规模越来越大,这导致了AI计算的任务(主要是训练)也越来越重。在实际的操作中,这些规模庞大的计算任务,不可能在单个硬件节点上完成,而需要由许多计算节点组成的分布式计算集群完成。如何完成计算任务在多节点上的分摊,以及节点之间的信息交互及同步,对于分布式计算集群的性能起到至关重要的作用,而这些都是由AI框架提供支持。本文将结合常见的AI框架,分析其中的并行技术,以及一些优化的方向及策略。

 

分布式AI并行策略

 

分布式的AI集群实现计算加速的最重要一方面是采用并行化的计算策略。目前,AI框架常采用的并行策略包括:数据并行,模型并行,流水线并行等。在部署的过程中,需要根据不同的AI模型,及硬件资源,采用不同或者混合的并行策略,以达到最佳的加速比。

 

数据并行

 

图1 数据并行的示意图[1]

 

数据并行(Data Parallelism)的示意图,如图1所示,其主要特点包括:

l每个节点都包含完整的模型,以及模型的参数(weight,parameter),输入数据则根据模型的并行度进行拆分,分别被每个节点读取;

l每个节点读取相应的输入数据后,分别独自处理模型的前向和反向传播,并得到Gradients,归并所有的梯度并更新梯度;

l所有节点之间的通信,主要包括前向传播的Loss归并以及反向传播的gradient归并以及更新,这些通信则是通过相应的通信原语(gather/reduce/broadcast)操作。

数据并行模式,有一些非常明显的优势,包括:部署简单工作量小,每个节点内的计算效率高等。这使得其在规模小,计算密集度高的模型中,有着广泛的应用,受益于其部署简单,因此成为最先采用的并行方式,基本所有分布式AI软件引擎都支持数据并行部署。但数据并行由于需要在每个节点复制所有模型参数,所以其显存重复度高利用率低,并不适合大模型的部署。

 

模型并行

 

图2 模型并行示意图

 

模型并行(Tensor Model Parallelism)的示意图,如图2所示。其主要包括:

l模型被切分成多个层次,每个节点负责其中一个切分的计算,也即每个节点只保留了一部分的模型参数;

l节点间的数据有比较强的依赖关系,比如图2中GPU0的输入为模型的输入,输出则作为GPU1的输入,因而GPU1的计算需要在GPU0算完后才能开始;

l节点间的信息传递,除了数据并行的Loss及gradient外,还有激活信息的传递。

模型并行的方式,解决了数据并行显存利用率低的问题,其通过对模型的切分,每个节点只需要放置一部分的模型参数,从而使得其可以部署更大的模型。但如果不采用流水并行的方式,单纯的模型并行有很强的串行关系,即每个节点的计算开始都依赖于上一层节点的计算结束,这导致每一次都只有一个worker在运算,计算效率极低。

 

流水线并行

 

图3 流水线并行示意图

 

流水线并行(Pipeline Model Parallelism)的示意图,如图3所示。其主要包括:

l单独的模型并行计算流程,如图中上部分所示,每个时刻都只有一个worker在运行,计算效率低下;

l流水线并行在模型并行的基础上,进一步将数据划分成更小的micro-batch,每次micro-batch算完即将数据传递给下一层,从而使得计算流水线起来,其计算图如图中下部分所示;

l流水线的节点间信息传递,内容和模型并行一致,但激活数据的传递被分成了多份micro-batch,进行流水线传递。

流水线并行的方式通过引入micro-batch流水的方式,解决了模型并行计算效率低下的问题,在现今的大规模模型训练中有着广泛的应用。但流水线并行的方式更复杂,并且micro-batch的方式减少了单节点计算密集度,增加了节点间的信息传递频率,使得取得一个好的加速比成为一个难题,如何部署一个高效流水线模式,也是业界非常有挑战的问题。

 

AI分布式并行技术实现

 

AI框架里的并行技术对于一个特定计算任务的神经网络,如何实现高性能的分布式计算的加速效率,主要取决于两个方面:一方面,采用何种并行计算方式;另一方面,则是如何布局底层硬件的拓扑结构实现节点之间的信息同步,主要有PS(Parameter Server)和All-Reduce两种。AI分布式软件引擎的主要内容是支持上述并行分布式计算方式的实现,以及支持相对应的多节点之间信息同步的实现。

 

软件API

 

当前主流的AI框架,其提供的API函数接口也主要包含并行方式及硬件信息等方面,以PyTorch为例[1],其主要包括以下的一些概念:

n进程组(group):默认下,只有一个组代表一个分布式任务;

n进程数(world size):表示分布式进程的数量,如:GPU节点数;

n进程序号(rank):表示进程序号,其大小表示进程优先级(rank=0代表主节点),用于进程之间的通讯;

n本地进程序号(local rank):每个进程的本地序号,如:GPU编号。

在PyTorch 分布式编程中,其对环境的初始化接口如图4所示。其中,“backend”代表硬件的通讯库,nccl[2]是英伟达开发的一个通讯库,是一个可以实现多个GPU节点之间通过PCIe,NvLink, InfiniBand等硬件后端进行聚合通信的库,其他支持的通讯后端还包括mpi和gloo。“init_method”包含了分布式的硬件信息等,如:节点IP,节点包含的GPU数量等。“rank”和”word_size”即为上述所描述的,可以通过运行时指定进程数量以及当前节点的进程序号。

 

图4 PyTorch分布式初始化接口

 

对于数据并行的编程,其主要是对batch进行拆分,分布到不同的节点上,并通过节点之间的信息同步以及聚合,完成分布式的计算。如下图5所示,每个节点的local_batch_size由原来的batch_size平均分布到不同的节点得到,每个节点的模型(model)都是和原来的模型一样(Net),而在每个节点计算自身batch的loss后进行反向传播得到梯度,再通过barrier同步,再由all_reduce进行反向传播的梯度聚合,最后得到整个batch的梯度。

 

图5 PyTorch数据并行伪代码

 

对于模型并行的编程,需要显性指定模型的存放节点,并显性指定节点之间的数据传递。如图6所示,ToyModel在初始化时,其net1显性指定放在了cuda:0上,net2则显性指定放在了cuda:1上,而在具体计算的forward函数里,其显性指定了模型的数据在每个节点之间的传递。

 

图6 PyTorch模型并行伪代码

 

对于流水线并行的编程,其在模型并行的基础上,需要指定micro_batch的size,并指定在每个micro_batch的数据传递。如下图7所示,通过初始化函数里指定split_size,将原来的batch分割成split_size份的micro_batch,并在具体的计算forward函数里,显式表示出每个micro_batch的数据在每个节点之间的数据传递。

 

图7 PyTorch流水线并行伪代码

 

分布式编译

 

一个AI计算任务,通过框架提供的API进行相应的分布式编程改造后,会由软件引擎编译成更底层的计算及通讯节点。例如图8所示的OneFlow编译的SBP模型[3],其编译的时候,将上层应用接口包含对数据及计算拆分的并行方式,编译下降成S (Split),B(Broadcast),P(Patial-sum)三种行为节点,使得原始的计算图,下降成一张更具体的分布式排布图。其SBP三种行为,进一步会被编译成相应的tensor数据拆分,聚合及计算行为,再通过传统的AI编译器编译成可以被调用的kernel函数。

 

图8 OneFlow编译SBP模型[3]

 

另一方面,OneFlow通过定义了一个actor model的runtime模型,为每个行为节点提供了数据依赖关系的建模,从而可以高效地在主机runtime侧,排布每个行为节点的执行顺序,实现高效的计算及显存利用,完成整个计算任务的编译。

 

优化策略

 

为了提高分布式AI软件引擎对AI模型的加速以及扩展更大的模型,学术界和工业界也研究了许多优化策略,主要体现在增加节点的计算,减少数据在节点间的搬运,以及减少显存的占用等方面。本章节主要介绍两种目前分布式AI软件引擎支持的高效优化策略。

(1)1F1B [4]

对于pipeline并行方式的训练过程中,每个micro-batch进行前向传播后得到的activation数据,会在其反向传播的过程用到,这意味着需要大量的显存保存前向传播的中间激活数据。如图9所示,对于一个切割数为8的pipeline并行的训练,其流程图如(a)所示,前向传播得到的激活数据都需要保存,这意味着每个worker都需要保存8个mini-batch产生的中间激活数据。而如果能够每次都可以将前向传播产生的激活数据,立马被后向传播所消耗掉,那么就可以释放掉该激活数据,极大节约显存,如图(b)所示,这也就是1F1B(one forward one backward)的原理。但事实上,如果完全的1F1B策略,即图(b)所示,其worker的空置率会很高,计算效率非常低,因而可以通过对流水线进行一定的编排,实现一定程度的1F1B,如(c)图所示,这样整体上每个worker的计算效率和原来基本一致,但显存的占用可以明显下降。

 

图9 1F1B优化示意图

 

(2)Checkpointing[5]

对于需要进行分布式计算的AI模型中,其往往参数及数据的规模非常大,显存的大小对模型扩展的约束比较大,因而减少显存的占用是重要的研究方向。在图10所示计算中,对于通常的训练过程(图中中间部分),其每一层的前向计算的中间激活数据,都需要保存在显存中,直到相应层的反向计算利用完后才可以被释放。如果将前向计算完后的所有中间激活数据直接释放掉,而在做反向传播时,先重做一遍前向运算得到相应的激活数据,再去做反向计算,那么由于不用缓存中间激活数据,可以极大节约显存资源,这也是利用计算资源去换显存资源的方法。但由于计算每一层的反向传播时,都需要从input开始到相应的层,做第二次的前向计算得到激活数据,这导致了计算量会增大很多。为了解决计算量和显存的矛盾,于是引入了checkpointing机制,即缓存部分层的前向计算结果,如图10的最右边子图部分所示,存储了红色虚线框所定义层的数据为checkpoints,当进行第二次前向计算其后的层的激活数据时,可以用checkpoints作为开始点,而不用从input开始,这样可以减少一定的计算量,而又可以节约一定的显存占用。

 

图10 Checkpointing技术原理示意图[5]

 

总结与思考

 

本文主要介绍了分布式计算常用的并行模式,以及常见的AI框架中用于支持分布式计算的API及编译的实现,同时也简要介绍了编译器采用1F1B及Checkpointing的方式对显存优化的实现,探讨了AI分布式的优化策略方式。本人也希望通过本文,可以让读者对AI框架里并行技术及其实现过程和优化策略更加清晰一些,从中得到一点启发。

 

由于水平有限,文中存在不足的地方请各位读者批评指正,也欢迎大家一起参与我们的讨论。

 

 

参考文献:

[1]https://pytorch.org/tutorials/intermediate/model_parallel_tutorial.html.

[2]https://developer.nvidia.com/nccl

[3]Deepak Narayanan, etc. PipeDream: Generalized Pipeline Parallelism for DNN Training, 2019.

[4]Jinhui Yuan, etc. OneFlow: Redesign the Distributed Deep Learning Framework from scratch, 2021.

[5]Tianqi Chen, etc. Training Deep Nets with Sublinear Memory Cost, 2016.

 

近期文章

由简入繁探究机器视觉中的数据增强(上)

由简入繁探究机器视觉中的数据增强(上)

本文探讨图像数据增强的两大主流方法各自的优劣势,针对目前图像数据增强面临的挑战及未来思考可研究的技术方向。
2022-05-09
Kubric:高效地合成视觉数据集

Kubric:高效地合成视觉数据集

随着机器学习的发展,数据集的重要性越来越凸显。然而收集大量带标注的图像或者视频数据成本高昂。而使用合成的数据集则具有成本更低,且由于收集方便,标注信息准确,高度可控,使得其成为一种训练大规模深度学习模型的有效手段。
2022-05-02

联系我们

这是描述信息

招贤纳士

公众号二维码
这是描述信息

 版权所有   ©   上海壁仞科技有限公司   |    沪ICP备19047354号