手把手教你从系统层面优化深度学习计算

深度学习在近几年里取得了巨大的进步，它已经或者是有望成功地被应用在我们许多生活场景中，比如自动驾驶、安防、翻译、医疗等等。可以说，计算机的计算和通信能力的大幅提升是促使深度学习成功的重要因素。

深度学习为什么依赖于超大的计算能力？

首先，深度学习本质上是基于统计的科学，所以大规模的样本数据对于深度学习的效果是至关重要的。其次，更大规模和更复杂的神经网络模型已经被证明非常有效，并在产品中有广泛的使用，这同时也产生了对计算能力的更大要求和消耗。举个例子，具有8层神经元的AlexNet网络2012年在ImageNet数据集上取得16%的错误率，该网络的一次迭代运行大约需要1.4 GFLOP的计算量。而微软提出的使用152层神经元的残差网络（ResNet）于2015年在该数据集上取得3.5%的错误率，其一次迭代的计算量大约是22.6GFLOP，是AlexNet的16倍。在当今的生产环境中，图像、语音以及自然语言处理相关的模型，例如人脸识别、语音转文字、机器翻译等，即使给予相当多的计算资源，很多仍需要几周的时间才能完成训练。

再次，深度学习模型是迅速迭代的。在AI领域，每年学术界和工业界都会提出大量的新模型。对每一个实际的问题，开发者需要不断尝试不同的模型和算法，甚至对于同一种模型算法，也需要去反复调试超参数以获得最好的预测效果。可想而知，如果模型的每次训练都要几周的时间，那么寻找最优模型的过程会非常漫长和痛苦。

另外，模型的线上推理具有更加极致的性能要求。线上的服务具有硬性的服务等级协议（SLA），所以在实际部署大型模型时，需要手工重新优化在深度学习框架（如TensorFlow）上已经训练好的模型，导致大量额外工程开销的产生。

由此可见，进一步优化深度学习计算对于深度学习的快速发展和成功应用起着至关重要的作用。

深度学习计算优化的挑战和机会

目前，优化深度学习的计算存在以下几个主要的挑战：

如何优化神经网络的计算使得它能够把单个硬件计算单元的效率发挥到极致

深度学习开发者和研究人员通常只想关注神经网络模型和算法本身，并不想被复杂的优化问题分散精力。这意味着深度学习框架这样的系统软件最好能够实现自动优化，而对模型开发者透明。那么，如何对特定的优化做合理的抽象使其更加灵活通用、更加容易地集成在系统框架中便是需要认真考虑的问题。

事实上，任何方面的优化问题都可以从模型算法和系统两个角度来看待。一方面，我们可以通过改变模型和算法来优化其对计算资源的使用效率从而改进其运行速度。这样的优化对特定的算法往往非常有效，但却不容易扩展应用到其它算法中。而另一方面，也就是微软亚洲研究院异构计算组正在进行的研究，则是在系统中实施模型算法无关的优化，这样的优化，通常可以为更多的应用带来性能的好处，同时也符合我们在前文提到的透明性的要求。

以系统优化助力深度学习计算

为了能够更好地理解系统这一层面的优化，我们先来简单介绍一下深度学习框架系统的背景知识。当今工业界流行的深度学习系统（包括TensorFlow、PyTorch、CNTK、MxNet、Caffe等）大都采用分层的体系结构设计。在前端提供高级语言（例如Python）的接口抽象，允许用户方便地描述神经网络结构，也就是深度学习的模型。描述好的模型在被系统运行前，首先会被转换成数据流图（Data-flow Graph）。在这个数据流图中，节点是特定的矩阵操作（也就是Operator，如Sigmoid、Matrix Multiplication等），而连接不同节点的边则是操作节点的输入和输出矩阵。这个数据流图也可以被看成是深度学习计算的中间表达。然后，深度学习系统的后端将这个数据流图映射到实际硬件上进行高效地执行，而大部分系统层面的优化就是在这个阶段完成的。

1. 加速分布式深度学习训练

分布式训练的主要瓶颈在于多机之间的通信开销。如今计算机网络的硬件技术已经有了很大的发展，InfiniBand的RDMA网卡（Remote Direct Memory Access，这是一种硬件的网络技术，它使得计算机访问远程的内存时无需远程机器上CPU的干预）已经可以提供50~100Gbps的网络带宽和微秒级的传输延迟。目前许多以深度学习为目标应用的GPU机群都部署了这样的网络。然而深度学习系统如何才能充分利用好硬件提供的通信能力使分布式的训练获得更大的性能提升呢？另外，使用RDMA的软件接口进行通信能够绕过TCP/IP协议栈，减少了操作系统内核态的运行开销。在这样的网络通信技术的支持下，任何与通信相关的计算处理的开销都会变得非常显著，而这正是许多原先基于TCP/IP而设计的网络通信机制中所存在的问题。

RPC（Remote Procedure Call，远程过程调用）是一个被广泛使用的多机之间的通信抽象原语，它的主要设计目标是通用性。在没有考虑RDMA的情况下，很多深度学习框架都会采用RPC的机制（例如gRPC）来实现多机之间的通信。然而，RPC需要维护一个内部的私有缓存，从而不得不引入用户数据的存储空间和内部缓存之间的数据拷贝。这种内存拷贝的开销在使用RDMA网络的情况下会变得非常明显。我们通过micro-benchmark观察到，跟使用基于TCP/IP的gRPC相比，直接通过RDMA的接口传输消息（对不同的消息大小）可以有2到10倍的性能提升。

那么针对深度学习的应用负载，如何才能更好地利用RDMA硬件的能力？首先，我们来分析一下深度学习应用的几个特点：

Tensor是深度学习计算中最主要的数据结构，大量的计算开销都是花在对Tensor的处理上

基于以上几个特点，我们可以对数据流图进行分析，找到那些可以静态决定shape信息的Tensor，以便在运行前，在接收端预先为其分配RDMA可访问的内存空间，并将其相应的可远程访问的地址传送给发送端。这样一来，在运行时，发送端可以通过单边的RDMA请求将Tensor的数据直接传输到接收端，从而完全避免了没有必要的额外内存拷贝，达到零拷贝的通信过程。我们将这种机制在TensorFlow上进行实验， 和基于TCP/IP的gRPC相比，这一方法在一系列典型模型上均取得了多倍的性能改进。甚至和针对RDMA优化过的gRPC相比，我们的方法仍然能够取得超过50%的性能提升。

另外，我们在分布式深度学习方向上关注的另一个问题是如何自动地对资源无关的数据流图做优化的分布式执行，也就是自动划分数据流图中的计算任务并为其分配相应的计算资源，以使计算效率最优化。Google的Jeff Dean团队在这个方向上已经做了很好的先驱性工作。但局限于模型并行和单机多卡的运行环境，目前这仍然是一个非常重要并且大有可为的方向，需要结合数据并行，分布式及异构环境来综合考虑。

2. 提升单个计算单元的运算效率

前面提到过，使用深度学习框架来实现的模型算法，在运行时前会被转换成数据流图。不少具有实际应用价值的模型都非常复杂，由它们所转换出来的数据流图通常是由成千上万的操作节点构成，其中包含了很多运算量非常小的节点，也就是说它们的输入矩阵的大小很小，或者是其计算逻辑的复杂度相对于对输入数据访问的复杂度来说很低。大量这样的操作节点会引入以下一些运行时开销，并且这样的开销会非常显著。

深度学习系统运行时需要根据数据流图中节点的依赖关系来调度节点的执行

计算量小的操作节点往往难以挖掘出足够的数据并行性，因此不能充分利用处理器硬件中的计算资源

解决这一问题的主要思路是内核融合（Kernel Fusion）。一些手工的优化方法就运用了这一思想，比如NVIDIA基于CuDNN的RNN库函数。它把整个循环神经网络实现成一个GPU的内核函数，因此获得了非常好的性能。然而它的缺点也非常明显，那就是不够灵活和通用，无法应用在其它网络或一些变种的循环神经网络中。而我们更加关注的是如何在深度学习的系统中自动地对任意的网络模型实施优化。

目前在学术界和工业界已经存在一些系统采用编译的方法生成融合的内核代码，比如TVM、Halide和Taco等。这些系统使用Tensor Algebra作为前端表示方法，每个Tensor Algebra表达式进而可以被编译成相应的内核代码。而Tensor Algebra可以作为更低一层的中间表达被集成到深度学习系统中，也就是说高层的数据流图可以先转换成由Tensor Algebra表达式组成的代码块，再被编译成可执行的代码。然而，这些系统对于可以进行融合的操作节点有很多限制，不能很好地融合多个非pointwise的操作，例如多个矩阵乘操作。然而，我们发现如果打破这一限制从而融合更多操作节点是可以带来更多显著的性能提升的。

在GPU的运行环境下融合多个非pointwise的操作具有一定的挑战性，因为非pointwise的操作中输入矩阵的每个元素都可能依赖于前一个操作的输出矩阵中的许多不同位置的元素值，所以在这两个操作之间需要插入Barrier同步原语。而在GPU中实现Barrier需要保证该内核的所有线程块在运行时都是保持活动状态的，这意味着我们必须要求融合后的内核采用有限个数的线程块，但同时又能够处理远超过线程块数量的数据块。

为了解决这一问题，我们尝试采用persistent-thread的线程块模型，也就是说在融合后的内核的整个生命周期启动固定数目的线程块并让它们保持活动状态。我们的优化系统在产生融合的内核代码的过程中类似于解决一个装箱（bin-pack）问题，即把待融合的子数据流图中的每一个操作节点所要处理的数据块分派给适当的活动线程块，从而使得每个线程块的负载尽可能均衡，并且保持操作节点的运算在原数据流图中的并行性。

为了生成优化的GPU内核函数，一个重要的考虑因素是线程块和数据块的合理划分。然而这又依赖于一些非常复杂的因素，比如操作节点运算中计算和访存复杂度的比率、GPU的shared memory的大小、寄存器文件的大小及分配方法等等。因此一个最优的选择是很难通过静态的方法决定的。幸运的是，深度学习的迭代性以及需要相当多的迭代才能收敛的特性使得我们可以利用早期的迭代过程来收集运行时的动态信息以帮助优化系统做更明智的决定。

3. 克服设备内存资源限制

设备内存的大小往往限制了可以处理的模型规模，解决这一问题的一个思路是对模型进行压缩和量化。如今学术界和工业界已经有大量的研究工作提出不同的压缩和量化的方法，然而，在实际的应用场景中使用压缩和量化仍然是个繁琐的迭代过程。在这个过程中，用户可能会进行以下几个方面的尝试。

不同的压缩方法

不同的压缩程度

显然，这样一个繁琐的过程需要一个好的工具来使之变得方便。这也是我们组正在关注的一个问题。我们正在尝试扩展TensorFlow的API来使用户可以在模型脚本中直接控制量化和压缩的方法、对象、程度和过程。

压缩和量化通常是用来解决模型部署时的性能和内存资源不足的问题，而解决模型训练时内存不够的问题的思路之一是用计算来换内存。比如，如果数据流图中某一个操作节点的计算量很小，但是输出的中间结果数据量很大，一个更好的处理方式是不在内存中保存这个中间结果，而在后面需要用到它的时候再重新执行这个操作节点的计算。当然，重新计算还是引入了一定的额外开销。

事实上，还存在另外一种解决这个问题的思路，就是将大的输入数据就保存在CPU端的主存里，并将操作节点实现成流式的处理，将大的输入数据分段拷贝进GPU的设备内存，并通过异步的拷贝使得对每一分段的计算时间和下一分段的拷贝时间能够重叠起来，从而掩盖住数据拷贝的开销。对于矩阵乘法这样的操作，由于计算复杂度相对于访存复杂度较高，当分段较大的时候，计算时间和拷贝时间是可以达到完美重叠的。然而，如果所要进行的操作不是矩阵乘法，而是一些简单的pointwise操作，计算的复杂度就没有办法和内存拷贝的开销相抵消。所以这种做法还需要跟内核融合相结合。比如将矩阵乘法和后续的pointwise操作相融合，每一个分段的计算都会把该分段的矩阵乘和pointwise操作都做完，然后再处理下一个分段。

原文发布时间为：2018-05-26

本文来自云栖社区合作伙伴“数据派THU”，了解相关信息可以关注“数据派THU”。