大规模深度神经网络训练系统浅谈：系统架构与性能优化

大规模深度神经网络训练系统浅谈：

系统架构与性能优化

孙鹏 I 文

张怀政 I 技术审稿

近几年来，深度神经网络（DNN）在自动驾驶、智慧医疗、机器翻译等众多领域取得了巨大的成功，并已经广泛地应用在了我们的日常生活中。从算法角度来看，DNN能比传统的机器学习模型在更多的场景上提供更高的准确率，其成功源于大数据与大模型的发展。

大数据：在DNN的训练过程中，我们需要通过迭代计算在网络的函数空间中寻找一个最优参数，从而能够更好地拟合训练数据。在大数据时代，DNN可以在前所未有的大规模训练数据集上，充分挖掘模型的表达能力。例如，在ImageNet-1K数据集上，DNN可以在120多万图片上训练一个高精度的分类模型。

大模型：与传统的机器学习模型相比，DNN为了获取了更强的特征表达能力，通常具有更多的模型参数和更高的计算量。例如，2012年提出的具有8层神经单元Alexnet具有近61M的参数，可以在ImageNet-1K数据集上获取39%的Top-1错误率，一次迭代计算需要14亿次浮点运算（1.4GFLOP）的计算量。2016年提出的具有152层神经单元的Resnet-152也具有60M左右的参数，能够在ImageNet-1K数据集上获取21.4%的Top-1错误率，但一次迭代需要高达22.6GFLOP的计算量，是Alexnet的16.1倍。

大数据与大模型已被证明十分有效，但同时也产生了对算力的巨大需求。通常来讲，DNN训练的总计算量为：训练样本数 * 模型一次迭代的计算量 * Epoch数（注：1个Epoch=所有训练样本被处理过一次）。如果用一个单核 CPU 在ImageNet-1K数据集上训练ResNet-50，可能需要花费几十年；如果在NVIDIA M40 GPU 上训练，需要花费14天 [1]。可以想象，如果DNN模型的每次训练都需要几周的时间，那么漫长的模型训练与调优过程将严重阻挡DNN的发展与落地。从系统角度来看，现代数据中心中的计算和通信能力的显著提升是保证DNN成功的重要因素。

分布式DNN训练系统架构介绍

由于资源限制，单机或者单计算单元（如GPU）不能满足对大数据和大模型处理要求，那么就需要使用多机多计算单元来扩展计算规模，满足DNN训练对算力的需求。参数服务器（ParameterServer, PS）和AllReduce是最常用的两类分布式DNN训练架构。

PS架构：在PS架构中，数据和计算负载都分布在工作节点（worker node）上；服务节点（server node）以向量、矩阵或者张量的形式保存的参数信息。在模型训练的每次迭代中，工作节点需要从服务节点拉取（Pull）最新的参数信息，在分配的训练样本上计算出梯度信息，并将其推送（Push）到服务节点。服务节点会汇集多个计算节点的梯度信息来更新对应的模型参数。Pull-Push是PS架构的基本通信方式。一般来讲，PS架构具有更强的容错性，一个或多个工作节点的崩溃并不会导致训练任务的终止。

P2P架构：在P2P架构中，所有的节点均为工作节点。在模型训练的每次迭代中，每个工作节点在其分配的训练样本上计算出梯度信息，用AllReduce通信方式汇聚并分发其他节点的梯度信息。在针对分布式DNN训练的AllReduce实现中，基于Ring AllReduce通信算法的NCCL2库最为高效，其理论通信负载不会随着GPU卡的增加而显著增加。

从任务分发的角度来看，分布式DNN训练系统可以采用数据并行和模型并行的工作方式。PS架构和P2P架构都支持数据与模型的并行工作方式。

· 数据并行：数据并行使分布式DNN训练系统支撑大数据的训练。训练样本会被分配到工作节点中，每个工作节点会有全部模型参数的拷贝。在迭代训练过程中，每个工作节点会用分配到的训练样本对每个模型参数计算出对应的梯度。

· 模型并行：模型并行使分布式DNN训练系统支撑大模型的训练。DNN模型会被切割并分配到工作节点中。这样，每个工作节点会有一部分模型参数的拷贝。在迭代训练过程中，每个工作节点会用训练样本对分配到的模型参数计算出对应的梯度。

分布式DNN训练系统性能优化介绍

分布式DNN训练系统的性能优化主要从如何提升多机加速比和提升单机运算效率来考虑。

分布式DNN训练系统的多机优化。分布式训练的主要瓶颈在于横向扩展能力和多机之间的通信开销。

· 增大系统横向扩展能力。以用NVIDIA GPU训练Imagenet数据集为例，在每次迭代计算过程中，每个GPU需要处理M张图片，当有N个GPU时，训练任务的整体batch size为N * M。由于硬件结构原因，在一定范围内，M越大，GPU的矩阵运算库会提高更高的运算效率。以Alexnet为例，在单卡训练中，M通常设为256，在多机多卡训练中，M通常不小于64。假设M=64，N=2014，训练任务的整体batch size高达64K。然而，使用传统的随机梯度下降算法（SGD），大batch size将会带来巨大的模型精度损失。为保证模型精度和单卡的高工作效率，分布式DNN训练系统往往只能限制M数值，即限制了系统横向扩展能力。为了解决这一问题，业界目前的方法是采用层级对应的适应率缩放（LARS）算法，针对DNN模型的每一层，动态调整学习率。目前领先的分布式DNN系统可以支持使用2048块GPU进行64K Imagenet的训练。

· 降低系统通信开销。降低分布式DNN训练系统的通信开销，可以显著增加系统的加速比，避免计算单元的浪费。例如，在一个没有通信优化的分布式训练系统中，16 GPUs卡的计算加速比甚至达不到8，也就是说该系统中计算单元经常处于等待通信完成，50%的计算能力得不到利用。为了降低系统的通信开销，分布式DNN通常采用RDMA方式。跟使用基于TCP/IP的通信方式相比， RDMA接口传输消息可以有2到10倍的性能提升。此外，NTU CAP组研发的ParameterFlow (PF) 通信技术，通过稀疏通信方式，可以大幅度降低网络通信开销。目前，领先的分布式DNN系统可以在1024卡达到920的加速比。

分布式DNN训练系统的单卡优化。提升单计算单元的运算效率对分布式DNN也十分重要。领先的DNN训练系统通常通过内核融合和内存优化的方法提升单计算单元的效率。

· 内核融合。最近几年提出的DNN，如Resnet-50，通常包含非常多的的小计算量的操作。由于GPU的硬件设计和GPU Kernel函数的启动销毁开销，大量的小运算操作往往不能充分利用硬件中的计算资源。解决这一问题的方法就是内核融合（Kernel Fusion）：将DNN训练过程中的用到的多个Kernel函数融合为一个函数进行调用。比如NVIDIA的RNN库函数，把整个循环神经网络（RNN）实现成一个Kernel函数，从而获得了非常好的性能。但是其缺点是手工融合好的Kernel函数不具有灵活性。目前在学术界和工业界开始研究使用编译的方法生成融合的内核代码。

· 内存优化。设备内存大小往往限制可以处理的DNN模型规模。例如，在有12G显存的K80 GPU上使用Caffe训练Resnet-50，单卡bath size最多设为32。为了解决这一问题，某些DNN训练系统会采用内存复用和计算换内存的方法降低内存需求。内存复用是指训练过程中的多个不同阶段的临时数据共享同一内存空间。计算换内存是指训练过程中不保留一些临时数据，需要时，通过其他数据再计算得出。

· 半精度训练。最新的NVIDIA GPU （如V100）带的TensorCore支持F16计算。F16计算的好处是一方面可以减少计算量，原来计算32位的单元处理FP16的时候，理论上可以达到两倍的速度，另一方面是模型训练需要的空间减少，可以使GPU有更大的空间缓存训练数据。但其挑战是如何保证训练出的模型的精度。

为应对大数据与大模型的发展，大规模深度神经网络训练系统的发展方兴未艾。本文简述了分布式DNN训练系统及其优化，仅代表个人意见，希望能给读者们带来启发。

[1] You, Yang, Zhao Zhang, Cho-Jui Hsieh, James Demmel, and Kurt Keutzer. "ImageNet training in minutes." In Proceedings of the 47th International Conference on Parallel Processing, p. 1. ACM, 2018.

本文版权归作者所有。

新加坡南洋理工CAP组