IT思维

论文作者 | Paul Barham，Michael Isard

编译 | Maglish

编辑 | 陈思

AI 前线导读：数值计算系统的性能和可编程性正陷入低谷。系统研究人员出色的工作使机器学习的基准在过去 5 年不断提升，但是探索创新的机器学习研究想法变得越来越难。在本文中，Google Brain 的研究人员解释了硬件加速器的进化如何有利于编译器后端优化巨大的单内核，介绍了这种对高性能但不灵活的内核的依赖如何增强了编程模型的主导风格。作者认为这些抽象编程缺乏表现性、可维护性和模块性，而这些都阻碍了研究的进展。最后，作者指出了该领域的发展方向，并提倡在现代加速器上推进高性能通用数值计算系统的发展。

更多优质内容请关注微信公众号“AI 前线”（ID：ai-front）

在试图改进胶囊网络的实现，以扩大到更大的数据集时，研究团队有了这篇论文的初步想法。胶囊网络是一个令人兴奋的机器学习研究思想，其中标量值的“神经元”被小矩阵取代，使它们能够捕捉更复杂的关系。胶囊或许不是机器学习中的“下一个大事件”，但它算是创新性机器学习研究理念的一个代表性例子。

虽然卷积胶囊模型需要的浮点运算（FLOPs）比卷积神经网络小 4 倍，并且训练参数少 16 倍，但是用 TensorFlow 和 PyTorch 实现的胶囊模型都比卷积神经网络模型慢得多，并且很小的模型就已耗尽了内存。到底是为什么呢？

胶囊网络内部循环的简化形式类似于传统的 CNN 层中的计算，但是是对 4×4 的矩阵进行操作，而不是标量。

当前机器学习框架的一个基本组成部分是 2D 卷积。大多数框架提供一个原语操作，输入大小为 H×W 的 N 个图像，其中每个像素的“深度”为 Ci 个通道。对于“核大小”k＝3 和“步长”s＝2，conv2D 计算以每个（x，y）坐标为中心的重叠的 3×3 像素块的加权和，然后产生像素深度为 Co 的 N 个较小的图像（如图 1 所示）。数学上可以表达如下：

其中的点乘表示标量乘法，O,I 和 K 都是 4 维的标量矩阵。最后的代码只是围绕一个乘法累加操作的 7 个嵌套循环，但矩阵布局、矢量化、并行和缓存对最后的表现都很重要。

卷积胶囊计算 3×3 卷积块中“姿态”矩阵的加权和，形成“投票”：

其中的点乘表示矩阵乘法，V,P 和 W 是 4×4 矩阵的 4 维数组，或等价于 6 维标量矩阵。

下面将解释为什么机器学习框架难以有效地运行胶囊计算。

卷积胶囊的原语可以在 CPU 上合理地实现（表 1），但是在加速器（例如 GPU 和 TPU）上就会出现问题。在加速器上的表现至关重要，因为几乎当前所有的机器学习研究，以及大多数生产模型的训练，都需要使用加速器。使用加速器执行特定的 ML 训练或大量推理工作的边际成本比使用 CPU 要低得多。

加速器非常适用于机器学习，因为训练任务中计算量较大的部分是密集的多维数组的线性代数运算。密集的线性代数与 CPU 所承担的工作负载相比是比较规则的，并且相对容易并行化。因此，人们为常规并行计算设计了越来越复杂的加速器。常见的加速器特性包括“warp”、块和线程网格、非常宽的矢量计算单元（ALU）和脉动阵列乘法器（MXUs）。但即使这些常规计算也很难获得良好的加速器性能。虽然常见的计算方式受到关注，并得到了很好的优化，但是非标准计算的性能会受影响，如卷积胶囊。

很难从常规计算中获得良好性能的一个主要原因是编译器必须考虑加速器的内存系统以及 ALU。为了防止数据瓶颈，加速器的并行能力已经与内存系统紧密耦合。GPU 上的 ALU 性能峰值需要“联合负载”，一个 warp 的所有 32 个线程同时访问同一缓存线中的不同值；程序实现必须适应内存的大小和步长；并且高效的程序必须利用单次内存访问中加载的所有值。

通常，加速器代码必须通过内存结构执行显式调度，而不是依赖透明的多级缓存。内存访问粒度一般要求线程协同加载彼此的值，然后交换它们；这样，代码还包含跨循环迭代的复杂指令调度。虽然匹配并行 ALU 的内存访问能带来良好的硬件利用率，但任何不匹配都会导致性能下降一个数量级。为了避免这种情况，需要为每一代加速器调整内核参数，例如 padding、步长和维度布局。

对于像卷积这样的“模版计算”，输入值需要被重叠的计算窗口重复使用，调度负载和存储以优化存储器带宽是非常具有挑战性的任务。而卷积胶囊中的数据重用模式具有几个额外的复杂性维度。

由于对参数调整分析很困难，加速器的高性能后端在一组计算“内核”（通常是孤立的循环嵌套）上花费了大量的开发工作，例如 2D 卷积和批矩阵乘法，它们决定了基准的表现性能。对于这些内核中的每一个，后端维护人员需要花费数小时或数天为一小组具有代表性的操作数形状（矩阵维度的基数）寻找最佳算法和参数设置，然后在运行时使用启发式或自动调整选择这些预调整的实现。

对于机器学习论文来说，提出新原语，却无法用现有内核有效计算的情况是非常常见的。研究人员开发了像张量理解（Tensor Comprehensions，TC）和 PlaidML 这样的编译器，允许终端用户编写自定义内核，它们都为领域特定语言提供了与数学公式相似的简洁语法，例如图 2 中胶囊原语的 TC 实现与公式 2 的对比。

但现实情况是，这些工具只适用于编译小代码片段。编译时间通常很长，并且代码质量往往不能达到峰值性能。

图 3a 说明了当前编译器框架为传统 2D 卷积生成 GPU 代码的困难，其性能需要与 cuDNN 中经过精心调整的库竞争。在每一种情况下，作者使用框架中可用的最低级别原语实现了 conv2d（公式 1）。TC 使用遗传搜索算法来优化参数，性能仅为 cuDNN 的八分之一，但是仅用了一个小时的搜索时间。最终性能也非常依赖于输入操作数（NCHW vs NHWC）的内存布局。TVM 具有类似的自动调整卷积模板，在 30 分钟搜索之后亦无法与 cuDNN 的性能相抗衡（图 3b）。PlaidML 获得最佳性能，比 cuDNN 慢 4 倍，编译时间快。TVM 也为这些操作数形状提供了手工调度的 conv2d 内核，但是它比 cuDNN 慢 19 倍。

回到胶囊示例中，作者接下来尝试为核心的胶囊原语编写自定义内核（公式 2）。用 gcc 围绕 4×4 矩阵乘法（matmul）函数编译明显的 C++ 循环嵌套，可以得到了高质量的矢量化代码，将其作为基线。它在一个 x86 内核上运行约 60ms，用 OpenMP 在 6 个内核并行化时达到 11.7ms。自己编写的 CUDA 实现运行了 1.9ms，但花费了两天时间进行手动调整。

虽然 PlaidML 在 gcc 上编译得很快，但内核执行要慢得多。TC 需要近 3 分钟来找到一个优于 CPU 的内核，但最终发现了运行时间少于 1.8ms 的调度（见表 1 和图 3C）。

作者对这些结果给出如下解释：当前框架在工作负载上较为先进，手动调整由特定模型或模型簇使用的小计算集是合理的。不幸的是，该框架并不适合研究，因为对非常规计算做实验时存在性能悬崖。虽然在生产部署前花费几个小时的搜索时间是可以接受的，但期望研究人员忍受这样的编译时间是不现实的（要记住这只是一个内核，而需要编译的可能是一个巨大的整体计算）；即使优化的内核定期缓存在本地，它也是传播研究的主要障碍。任何下载模型源代码的人都必须花费数小时或几天的时间在硬件上进行编译，然后才能进行实验。

在像 TensorFlow 和 PyTorch 这样的 ML 框架中使用自定义计算并不直接。因此，在 TensorFlow 和 PyTorch 中实现卷积胶囊的最简单和最有效的方法是使用已经由这些框架支持的高级操作。

不幸的是，两个框架都不能将最终的缩减融合到批处理矩阵乘法（预优化内核）中，这大大增加了所需的存储器带宽以及中间存储要求。为了计算两个相对较小的量，API 迫使程序将比真正所需数据高出两个数量级的数据复制、重新排列和实现到内存。正是上述问题让研究人员无法找到任何高性能的卷积胶囊的实现。

即使对于像 ResNet 那样简单的 ML 模型，操作数也有不同的形状，因此不同的卷积调用可能具有不同的最优参数。如果中间值的生成操作和消费操作之间的布局不同，那么该值必须被“转置”（转换为不同的布局）。

ML 计算与许多计算不同，因为它们通常涉及近似的浮点数或定点数。通过改变精度或非均匀量化的选择，可以或多或少地获得目标度量，例如测试集预测精度。精度也会影响计算 / 通信瓶颈。

常见子表达式消除（CSE）：CSE 对于 ML 框架非常重要，因为它采用了反向传播计算结构。标准 CSE 倾向于将激活值物化到内存，但是加速器内存的数量有限以及 ALU 的缓慢，意味着人们更倾向于重新计算激活。因此，CSE 为框架引入了另一个组合搜索问题：选择哪些值应该被物化。

分布式执行： 由于数据并行结构，机器学习计算通常可以有效地分布在多个加速器上。在实践中，编译器和机器学习框架通常利用与设备中使用的机制不相交的方法来实现分布式并行性。据我们所知，没有框架试图共同优化选择计算片段应该在哪个设备上运行，以及如何选择子计算在设备内的结构。

机器学习训练算法特别依赖于自动优化策略，因为大多数计算通常在由源代码中出现的“前向传递”中合成的梯度算子中执行。由于计算梯度的代码不是程序员编写的，因此程序员来引导优化的机会有限。即使在没有自动生成梯度的情况下，也很难写出有手动优化标注的模块化代码。

最近的研究对自动全程序优化技术的兴趣越来越大，但是方法还在起步阶段，并且通常集中每次只优化程序的一个方面。毫无疑问，多维整体程序优化是一项艰巨的任务，但或许可以从最近混合搜索、学习方法的成功中获得一些希望，如 AlphaGo，它表明在巨大的组合搜索空间中有希望找到好的解决方案。似乎有必要在搭建机器学习框架时考虑自动优化器，才可能充分利用全程序优化。

作者首先观察到数值计算得益于传统编程语言中不存在的特性。自动微分就是一个这样的特性，并且数值程序也是不寻常的，因为它们使用在其参数秩（维数）上具有多态性的函数进行编写。再考虑标准卷积表达式（公式 1）。对批（n）和输出通道（Co）维度的每个元素的计算是独立的，并且表达卷积的自然方式是以三维输入 I 和 K 所写的子计算。

当输出具有更多维度时，语言可以自动地“提升”函数，跨越批和输出通道维度。早在 APL 语言中，数字语言就已经包含了提升秩多态性，但是关于如何将这种多态性与现代模块化类型和语言结合起来，还存在大量的开放性研究问题。

回想一下，后端是围绕对大型单内核的调用而构建的。作者认为这种后端设计方法正在阻碍编程模型的可维护性、可调试性和表达性方面的进展。更糟糕的是，由此带来的对语言创新的阻碍本身就减少了后端开发人员对当前情况进行改进的动力。

单后端内核的一个结果是前端选择内核或“操作符”作为抽象点。在流行的框架中，如 TensorFlow 和 PyTorch，用户程序是在 Python 中编写的，并调用以后端特定的语言和库（如 C++、MKL、CUDA 和 cuDNN）编写的运算符，或者有时是较低级别但可移植的特定领域语言，例如 Tile 和张量理解。当现有的运算符不足以执行任务时，用户必须使用较低级别的语言来编写新的运算符，通常还必须手动编写其梯度，这是一个困难且容易出错的过程。

有的框架，例如 Julia，它名义上使用相同的语言来表示操作符图及其实现，但是后端设计抵消了前端的有效性。

正如先前所建议的，语言层面的另一种选择是提供低维 conv2d 函数，并在必要时将其提升到更高的维度。作者认为这将大大提高可读性和可维护性。单后端内核是这种方法的一个障碍，因为语言必须自动发现在某些维度排序下等价于单内核的模式，然后在调用内核之前自动重新排序维度。这种模式匹配如果出现任何失败，都会使性能严重下降。但是，如果有个编译器可以为通用提升函数编译出高效代码，那么进程会快得多。

人们经常提倡“命名维度”，即让矩阵的维数与文本名称相关联，或者与其数字位置相关联。命名维度能够提高可读性，但作者认为，在提升代码模块性方面它们的影响可能会更大，因为命名维度可以使语言远离固定维度顺序，从而让函数提升更方便。例如在 APL 中，维度有严格的顺序，在调用函数前参数必须经过转置才能确保正确的顺序。

为了有效地实现具有命名或无序维度的程序，有必要重新考虑后端的设计。有几个项目已经在这个方向上采取了初步尝试。他们在一定程度上解相关源代码的维数顺序与低级代码的维度顺序，但是仍然要求前端指定每个阵列的有序维度。

目前针对上述的这些问题，还没有很好的解决方案。一方面是工程实现上存在难度，一方面是没有足够的动机去研究这些问题的解决方案。面向机器学习的端到端工具链的研究需要许多不同领域的专业知识。很遗憾的是，尽管研究人员在一些子问题上已经做出了效果卓越的工作或者一些具有前瞻性的努力，但是仍无法提出一个能基本解决端到端问题的综合性思路。本着逐个击破的想法，作者提出了一些探索性的思路：

• 语言设计，包括自动化微分以及使用语言的语法表达纯命名维度和内核。
• 定义布局不可知的通用循环嵌套图的后端 IR(接口需求，Interface Requirements)。
• 转化上述 IR，将其解构为具体的 CSE 策略。
• 在上面提到的降低 IR 的情况下生成加速器代码的编译过程，快速产生足够的代码并接近最好的表现。

对于现有的机器学习工具链，这篇论文并不是轻视他们的思想和工程实现，并且显而易见的是它们对于大部分领域相关人员来说是非常具有价值的。作者的考虑主要来自于语言和后端的不灵活性会成为阻碍创新型研究的真正障碍，这会使得这个非常活跃的领域进展缓慢。因此在设计加速器、工具链时，作者希望大家能够更多的考虑这一点，尤其是对于基准测试的研究。

查看论文原文：

https://dl.acm.org/citation.cfm?id=3321441