文章标题：在GPU上优化循环神经网络的性能
作者/机构：Jeremy Appleyard (NVIDIA), Tomáš Kociský (牛津大学, Google DeepMind), Phil Blunsom (牛津大学, Google DeepMind)

A1 主要贡献

随着循环神经网络（RNN）变得越来越大和越来越深，单个网络的训练时间上升到数周甚至数月，因此提高这些网络的性能和可扩展性具有重要意义。尽管GPU已成为训练和部署循环模型的首选硬件，但所采用的实现通常只利用了针对这些架构的基本优化。

核心问题：随着循环网络变得更深（如【10】）且其核心单元结构更复杂（如【11】，【12】），要最大限度地利用最新一代GPU的计算能力变得越来越困难。虽然已有研究优化CNN在GPU上的实现（如【13】，【14】），但针对RNN运行时的优化工作相对较少。

研究目标与贡献：本文提出了一系列超越简单RNN GPU实现的优化方案，旨在为常见的网络架构实现接近峰值的计算吞吐量。通过暴露网络内部操作间的并行性，相比朴素实现，可以在各种网络规模上实现一个数量级的加速。文章详细描述了已被整合到NVIDIA cuDNN第五个版本中的三个优化阶段：
1. 优化单个RNN单元（cell）。
2. 优化单个RNN层（layer）。
3. 优化整个多层网络。

这些增强功能已在NVIDIA cuDNN v5库中针对Simple RNN、GRU和LSTM架构实现。

A2 方法细节

本节将以长短期记忆网络（LSTM）【11】（一种不带窥视孔连接的四门标准LSTM）的前向和后向传播过程为例来探讨性能优化。这些对LSTM性能有益的策略大多可以轻松迁移到其他类型的RNN中。以下是计算LSTM前向传播中时间步t输出的方程：

2.1 朴素实现

初始实现性能不佳。有多种方式可以朴素地实现循环神经网络的单步传播。我们以一种将每个独立操作（如矩阵乘法、sigmoid、逐点加法等）实现为单独内核（kernel）的方式作为起点。虽然GPU会并行执行每个内核内部的操作，但这些内核是按顺序依次执行的。这种实现的前向传播性能很差，对于一个隐藏状态大小为512、minibatch为64的测试案例，其性能约为0.4 TFLOPS，不到硬件峰值性能（基本时钟频率下约5.8 TFLOPS）的10%。

矩阵运算分组优化。一个广泛使用的优化是将共享相同输入的矩阵运算合并为单个更大的矩阵运算。在前向传播中，标准的LSTM公式会产生八次矩阵乘法：四次作用于循环输入（$R*h_{t-1}$），四次作用于前一层的输入（$W*x_t$）。在这两组各四次的乘法中，输入是共享的，但权重不同。因此，可以将一组四个矩阵乘法重构为一个四倍大小的单一矩阵乘法。由于更大的矩阵运算并行性更高（因此效率更高），在上述测试案例中，这种方法将前向传播吞吐量大致翻倍至0.8 TFLOPS。这种重构在大多数深度学习框架中非常容易实现，因此得到了广泛应用。GRU单元也可以进行类似的优化，将两组三个矩阵进行分组。单门RNN则无法从这种优化中受益，因为它们在每个输入端只有一个矩阵乘法。反向传播同样受益于此优化，因为四个输入被转换为一个输出。此实现的伪代码见清单1。


清单1: 伪代码，展示了前向传播优化的起点。

后续优化的普及性。虽然接下来将要讨论的一些优化之前已被实现过，但这些实现远非通用或标准实践。

2.2 单个单元优化

2.2.1 流式矩阵运算

RNN矩阵乘法并行性不足。RNN执行的矩阵乘法通常并行性不足，无法在GPU上获得最佳性能。当前最先进的GEMM（通用矩阵乘法）内核实现中，每个CUDA块计算输出矩阵的一个矩形瓦片（tile）。这些瓦片的维度通常在32到128之间。对于一个隐藏状态大小为512、minibatch为64的LSTM，其前向传播所需的矩阵乘法若以128x64的瓦片进行划分，总共只能产生16个CUDA块。由于CUDA块驻留在单个GPU流式多处理器（SM）上，而现代顶级GPU（例如我们使用的M40）当前拥有24个SM，这意味着该矩阵乘法最多只能使用GPU三分之二的可用性能。考虑到为了最大化延迟隐藏，每个SM上最好有多个块，因此很明显，要获得更好的性能，需要增加并行性。

清单2: 伪代码，展示了前向传播单个单元的优化。

使用CUDA流并发执行矩阵乘法。一个增加单个RNN单元并行性的简单方法是让GPU并发执行两个矩阵乘法（即$R*h_{t-1}$和$W*x_t$）。通过使用CUDA流，我们可以告知硬件这两个矩阵乘法是相互独立的。这使得GPU可用的并行性加倍，对于小型矩阵乘法，性能可提升高达2倍。对于较大的矩阵乘法，流式处理仍然有用，因为它有助于最小化所谓的“尾部效应”（tail effect）。如果启动到GPU的块数量只够填满其SM几次，这些块可以被看作是分批次（waves）通过GPU。第一批次的所有块大约在同一时间完成，然后第二批次开始。这个过程一直持续到没有更多工作为止。如果批次数很少，最后一批通常比其他批次有更少的工作要做，从而产生一个低性能的“尾部”。通过增加并行性，这个尾部可以与另一个操作重叠，从而减少性能损失。

2.2.2 逐点运算的融合

融合逐点运算以提高效率。尽管逐点运算（point-wise operations）天然具有并行性，但我们发现它们的执行效率很低。这有两个原因：首先，向GPU启动一个内核本身存在开销；其次，将一个逐点运算的输出完全写出到GPU主存，然后在片刻之后再为下一个运算读入，这个过程是低效的。由于逐点运算本质上是独立的，因此可以将所有逐点内核融合成一个更大的内核。

2.3 单层优化

跨时间步输入批处理。一个循环层由许多单元组成，每个单元的循环输入依赖于前一个单元的输出。然而，来自前一层（lower layer）的输入可能没有这种依赖关系，因此通常可以将多个时间步的输入连接起来，形成一个更大、更高效的矩阵乘法。选择连接的步数并非易事：更多的步数会导致更高效的矩阵乘法，但更少的步数可以减少循环操作可能等待的时间。确切的步数不仅取决于超参数，还取决于目标硬件。

权重矩阵重排。当将一个层作为一个整体来考虑时，另一个可能的操作是重新排序权重矩阵的布局。由于相同的权重矩阵在一个层的计算过程中被重复使用，重排的成本通常远小于操作这些矩阵的成本。在我们的测试中发现，预先转置权重矩阵能带来显著的性能提升。需要注意的是，因为权重矩阵的转置在反向传播中使用，所以这个预转置操作必须在每次网络传播中都执行。

清单3: 伪代码，展示了跨越一个层的前向传播优化。

2.4 多层优化

跨层并行计算。RNN中采用多层堆叠结构变得越来越普遍，即每个循环单元将其输出直接送入下一层的循环单元。在这种情况下，可以利用循环层之间的并行性：一个循环单元的完成不仅解决了当前层下一次迭代的依赖关系，也解决了下一层当前迭代的依赖关系。这允许多个层并行计算，从而在任何给定时间点都极大地增加了GPU的工作量。


清单4: 伪代码，展示了最终优化的前向传播。

2.4.1 调度

对角线波式调度。由于向GPU启动工作需要少量但不可忽略的时间，因此考虑内核启动的顺序非常重要。例如，如果GPU资源可用，几乎总是优先启动所有依赖项已解决的内核，而不是一个可能需要等待一段时间才能清除其依赖项的内核。通过这种方式，可以尽可能多地暴露并行性。为了实现这一点，我们选择了一个简单的调度规则：下一个被调度的工作是那些到达“第一个”循环单元需要遍历的边最少的工作。如果将一个循环网络视为一个二维单元网格，这会产生一个从第一个单元开始传播的对角线“波”式启动。

2.6 权重更新

高效的权重更新。上述优化仅适用于前向和后向传播步骤。通过在开始权重更新之前完成梯度传播，权重更新过程可以变得非常高效。可以使用单个大型矩阵乘法来更新每个矩阵，并且没有依赖关系，这通常能达到接近峰值的性能。与更新矩阵相比，更新偏置权重的成本非常低。

A3 cuDNN 实现细节

cuDNN v5中的实现与调优。第2节中描述的优化已在NVIDIA cuDNN库的第五个版本中为单门RNN、GRU和LSTM实现。该实现能够以比当前接口更低的层次与cuBLAS进行交互，并针对此用例调整用于确定cuBLAS操作模式的启发式算法。具体来说，如果cuBLAS检测到一次调用可能导致GPU未被充分利用，它通常会选择一个并行度更高但资源效率较低的路径。由于我们在比cuBLAS更高的层次上知道预期的并行量，因此在存在高流式并行性的情况下，覆盖此行为以偏向资源效率更高的路径，有时会带来整体的加速。希望未来的cuBLAS版本能包含一个允许这种手动调优的接口。

A4 实验

实验环境

  • 模型架构：实验主要使用LSTM网络，同时也支持Simple RNN和GRU。网络参数在不同实验中有所不同，例如：
    • 基准测试：隐藏状态大小512，minibatch 64。
    • 优化效果分析：4层LSTM，1000个时间步，隐藏状态大小512，minibatch 64。
    • 超参数影响分析：4层LSTM，隐藏状态大小和minibatch大小在很大范围内变化。
  • 硬件配置：NVIDIA M40 GPU，拥有24个流式多处理器（SMs），在固定基础时钟频率下，其峰值性能约为5.8 TFLOPS。
  • 软件配置：

实验结果

各项优化的累积性能影响

实验在一台M40 GPU上对一个包含1000个时间步、4层、隐藏状态大小为512、minibatch为64的LSTM网络进行了前向传播性能测试。如表1所示，与完全朴素的实现相比，所有优化叠加后实现了约11倍的加速；与仅使用标准GEMM分组优化的实现相比，也实现了约6倍的加速。

表1: LSTM前向传播性能。每项优化都在前一项的基础上应用。这些测量是在一个100次迭代、四层、隐藏状态大小为512、minibatch为64的LSTM上使用cuBLAS 7.5进行的。

超参数对优化效果的影响

图1展示了在4层LSTM网络上，不同隐藏状态大小和minibatch大小下各项优化的影响。
* 多层加速效果：在某些情况下，将层数从1层增加到4层可以使吞吐量翻倍（即用2倍的时间完成4倍的工作）。这种性能提升在低并行度的情况下尤其明显（例如，minibatch较小，隐藏状态大小较小，或两者都小）。
* cuBLAS启发式算法的影响：在minibatch为32的情况下，当隐藏状态较大时性能出现下降，这归因于cuBLAS因内部启发式算法而选择了不同的执行路径。
* 并行性与优化效果：minibatch为64时加速效果最显著。随着minibatch大小和隐藏层大小的增加，GPU已有的并行性也随之增加，因此专注于增加并行性的优化策略效果会减弱。尽管如此，即使在最大的基准测试问题（minibatch 256，隐藏状态大小 4096）中，层间重叠仍然能带来性能提升。
* 部分优化策略的权衡：“输入批处理”策略在除了minibatch为32之外的许多情况下被发现是有害的，这可能是由于第2.3节中讨论的权衡。总体而言，除了“输入批处理”，每项优化在大多数情况下都有帮助。


图1: 各项优化对一个四层LSTM网络前向传播的影响。所用M40 GPU在固定基础时钟下的峰值性能约为5.8 TFLOPS。

cuDNN v5 最终性能

图2展示了使用cuDNN v5候选版（RC）时，不同网络类型（Simple RNN, GRU, LSTM）的前向和后向传播性能。该实现通过与cuBLAS的底层交互和启发式调优，在各种网络配置下均达到了很高的计算吞吐量。


图2: 使用cuDNN v5 RC的不同网络类型的前向和后向传播性能。所用M40 GPU在固定基础时钟下的峰值性能约为5.8 TFLOPS。

A5 结论

本文提出了一种能够在GPU上高效执行循环神经网络的方法。与之前在GPU上实现良好加速的工作（如【19】，【20】）相比，本文提出的方法达到了更高的性能水平。
* 核心策略：主要策略是向GPU暴露尽可能多的并行性，以最大化硬件资源的利用率。
* 适用场景：这些方法在处理层数较多但规模较小（即并行性较低）的深度循环网络时尤其有效。

未来工作与挑战：
本文的一个未利用的性能优化点是参数在循环迭代间的重用。可以设想将这些参数存储在更低级别的GPU内存（如共享内存）中，并在迭代之间重用。在受带宽限制的情况下，这可能会极大地提高性能。然而，这种方法有几个缺点：
1. 存储限制：参数的存储空间有限，因此会对参数数量设置上限。
2. 实现风险：任何此类实现都必须做出在CUDA编程模型中无效的假设，因此容易出现意外故障。

A6 附录

本文没有附录。