AI框架里的并行技术
摘要
近些年来，为了取得更高的精度及更广泛的通用性，AI的深度学习网络及数据集的规模越来越大，这导致了AI计算的任务（主要是训练）也越来越重。在实际的操作中，这些规模庞大的计算任务，不可能在单个硬件节点上完成，而需要由许多计算节点组成的分布式计算集群完成。如何完成计算任务在多节点上的分摊，以及节点之间的信息交互及同步，对于分布式计算集群的性能起到至关重要的作用，而这些都是由AI框架提供支持。本文将结合常见的AI框架，分析其中的并行技术，以及一些优化的方向及策略。
分布式的AI集群实现计算加速的最重要一方面是采用并行化的计算策略。目前，AI框架常采用的并行策略包括：数据并行，模型并行，流水线并行等。在部署的过程中，需要根据不同的AI模型，及硬件资源，采用不同或者混合的并行策略，以达到最佳的加速比。
数据并行（Data Parallelism）的示意图，如图1所示，其主要特点包括：
每个节点都包含完整的模型，以及模型的参数（weight，parameter），输入数据则根据模型的并行度进行拆分，分别被每个节点读取；
每个节点读取相应的输入数据后，分别独自处理模型的前向和反向传播，并得到Gradients，归并所有的梯度并更新梯度；
所有节点之间的通信，主要包括前向传播的Loss归并以及反向传播的gradient归并以及更新，这些通信则是通过相应的通信原语（gather/reduce/broadcast）操作。
数据并行模式，有一些非常明显的优势，包括：部署简单工作量小，每个节点内的计算效率高等。这使得其在规模小，计算密集度高的模型中，有着广泛的应用，受益于其部署简单，因此成为最先采用的并行方式，基本所有分布式AI软件引擎都支持数据并行部署。但数据并行由于需要在每个节点复制所有模型参数，所以其显存重复度高利用率低，并不适合大模型的部署。
模型并行（Tensor Model Parallelism）的示意图，如图2所示。其主要包括：
模型被切分成多个层次，每个节点负责其中一个切分的计算，也即每个节点只保留了一部分的模型参数；
节点间的数据有比较强的依赖关系，比如图2中GPU0的输入为模型的输入，输出则作为GPU1的输入，因而GPU1的计算需要在GPU0算完后才能开始；
节点间的信息传递，除了数据并行的Loss及gradient外，还有激活信息的传递。
模型并行的方式，解决了数据并行显存利用率低的问题，其通过对模型的切分，每个节点只需要放置一部分的模型参数，从而使得其可以部署更大的模型。但如果不采用流水并行的方式，单纯的模型并行有很强的串行关系，即每个节点的计算开始都依赖于上一层节点的计算结束，这导致每一次都只有一个worker在运算，计算效率极低。
图3 流水线并行示意图
单独的模型并行计算流程，如图中上部分所示，每个时刻都只有一个worker在运行，计算效率低下；
流水线并行在模型并行的基础上，进一步将数据划分成更小的micro-batch，每次micro-batch算完即将数据传递给下一层，从而使得计算流水线起来，其计算图如图中下部分所示；
流水线的节点间信息传递，内容和模型并行一致，但激活数据的传递被分成了多份micro-batch，进行流水线传递。
流水线并行的方式通过引入micro-batch流水的方式，解决了模型并行计算效率低下的问题，在现今的大规模模型训练中有着广泛的应用。但流水线并行的方式更复杂，并且micro-batch的方式减少了单节点计算密集度，增加了节点间的信息传递频率，使得取得一个好的加速比成为一个难题，如何部署一个高效流水线模式，也是业界非常有挑战的问题。
AI框架里的并行技术对于一个特定计算任务的神经网络，如何实现高性能的分布式计算的加速效率，主要取决于两个方面：一方面，采用何种并行计算方式；另一方面，则是如何布局底层硬件的拓扑结构实现节点之间的信息同步，主要有PS（Parameter Server）和All-Reduce两种。AI分布式软件引擎的主要内容是支持上述并行分布式计算方式的实现，以及支持相对应的多节点之间信息同步的实现。
当前主流的AI框架，其提供的API函数接口也主要包含并行方式及硬件信息等方面，以PyTorch为例[1]，其主要包括以下的一些概念：
进程组（group）：默认下，只有一个组代表一个分布式任务；
进程数（world size）：表示分布式进程的数量，如：GPU节点数；
进程序号（rank）：表示进程序号，其大小表示进程优先级（rank=0代表主节点），用于进程之间的通讯；
本地进程序号（local rank）：每个进程的本地序号，如：GPU编号。
在PyTorch 分布式编程中，其对环境的初始化接口如图4所示。其中，“backend”代表硬件的通讯库，nccl[2]是英伟达开发的一个通讯库，是一个可以实现多个GPU节点之间通过PCIe，NvLink, InfiniBand等硬件后端进行聚合通信的库，其他支持的通讯后端还包括mpi和gloo。“init_method”包含了分布式的硬件信息等，如：节点IP，节点包含的GPU数量等。“rank”和”word_size”即为上述所描述的，可以通过运行时指定进程数量以及当前节点的进程序号。
对于数据并行的编程，其主要是对batch进行拆分，分布到不同的节点上，并通过节点之间的信息同步以及聚合，完成分布式的计算。如下图5所示，每个节点的local_batch_size由原来的batch_size平均分布到不同的节点得到，每个节点的模型（model）都是和原来的模型一样（Net），而在每个节点计算自身batch的loss后进行反向传播得到梯度，再通过barrier同步，再由all_reduce进行反向传播的梯度聚合，最后得到整个batch的梯度。
对于模型并行的编程，需要显性指定模型的存放节点，并显性指定节点之间的数据传递。如图6所示，ToyModel在初始化时，其net1显性指定放在了cuda:0上，net2则显性指定放在了cuda:1上，而在具体计算的forward函数里，其显性指定了模型的数据在每个节点之间的传递。
对于流水线并行的编程，其在模型并行的基础上，需要指定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函数。
另一方面，OneFlow通过定义了一个actor model的runtime模型，为每个行为节点提供了数据依赖关系的建模，从而可以高效地在主机runtime侧，排布每个行为节点的执行顺序，实现高效的计算及显存利用，完成整个计算任务的编译。
优化策略
为了提高分布式AI软件引擎对AI模型的加速以及扩展更大的模型，学术界和工业界也研究了许多优化策略，主要体现在增加节点的计算，减少数据在节点间的搬运，以及减少显存的占用等方面。本章节主要介绍两种目前分布式AI软件引擎支持的高效优化策略。
对于pipeline并行方式的训练过程中，每个micro-batch进行前向传播后得到的activation数据，会在其反向传播的过程用到，这意味着需要大量的显存保存前向传播的中间激活数据。如图9所示，对于一个切割数为8的pipeline并行的训练，其流程图如(a)所示，前向传播得到的激活数据都需要保存，这意味着每个worker都需要保存8个mini-batch产生的中间激活数据。而如果能够每次都可以将前向传播产生的激活数据，立马被后向传播所消耗掉，那么就可以释放掉该激活数据，极大节约显存，如图(b)所示，这也就是1F1B(one forward one backward)的原理。但事实上，如果完全的1F1B策略，即图(b)所示，其worker的空置率会很高，计算效率非常低，因而可以通过对流水线进行一定的编排，实现一定程度的1F1B，如(c)图所示，这样整体上每个worker的计算效率和原来基本一致，但显存的占用可以明显下降。
（2）Checkpointing[5]
对于需要进行分布式计算的AI模型中，其往往参数及数据的规模非常大，显存的大小对模型扩展的约束比较大，因而减少显存的占用是重要的研究方向。在图10所示计算中，对于通常的训练过程（图中中间部分），其每一层的前向计算的中间激活数据，都需要保存在显存中，直到相应层的反向计算利用完后才可以被释放。如果将前向计算完后的所有中间激活数据直接释放掉，而在做反向传播时，先重做一遍前向运算得到相应的激活数据，再去做反向计算，那么由于不用缓存中间激活数据，可以极大节约显存资源，这也是利用计算资源去换显存资源的方法。但由于计算每一层的反向传播时，都需要从input开始到相应的层，做第二次的前向计算得到激活数据，这导致了计算量会增大很多。为了解决计算量和显存的矛盾，于是引入了checkpointing机制，即缓存部分层的前向计算结果，如图10的最右边子图部分所示，存储了红色虚线框所定义层的数据为checkpoints，当进行第二次前向计算其后的层的激活数据时，可以用checkpoints作为开始点，而不用从input开始，这样可以减少一定的计算量，而又可以节约一定的显存占用。
图10 Checkpointing技术原理示意图[5]
本文主要介绍了分布式计算常用的并行模式，以及常见的AI框架中用于支持分布式计算的API及编译的实现，同时也简要介绍了编译器采用1F1B及Checkpointing的方式对显存优化的实现，探讨了AI分布式的优化策略方式。本人也希望通过本文，可以让读者对AI框架里并行技术及其实现过程和优化策略更加清晰一些，从中得到一点启发。
由于水平有限，文中存在不足的地方请各位读者批评指正，也欢迎大家一起参与我们的讨论。
[1] Austin Derrow-Pinion, Jennifer She, David Wong, et al. ETA Predictionwith Graph Neural Networks in Google Maps. 2021
[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]JinhuiYuan, etc. OneFlow: Redesign the Distributed Deep Learning Framework from scratch, 2021.
[5]TianqiChen, etc. Training Deep Nets with Sublinear Memory Cost, 2016.
往期推荐
壁仞科技研究院作为壁仞科技的前沿研究部门，旨在研究新型智能计算系统的关键技术，重点关注新型架构，先进编译技术和设计方法学，并将逐渐拓展研究方向，探索未来智能系统的各种可能。壁仞科技研究院秉持开放的原则，将积极投入各类产学研合作并参与开源社区的建设，为相关领域的技术进步做出自己的贡献。