转载于：知乎
作者： 金雪峰来自Google的论文，动机是将动态图执行模式和JIT编译有效的结合起来，论文地址：
https://arxiv.org/abs/2102.13267arxiv.org
概述
命令式顺序程序执行(在 ML 上下文中通俗地称为“Eager执行”或“按定义运行”)，易于理解、表达和调试，这就是为什么它被最广泛采用的原因，如Pytorch和Numpy。
另一方面，特定领域的编译器 (DSC)大大提高了机器学习模型的性能。 此外，这些编译器有时是使能特定领域加速器（例如Cloud TPU）的唯一方法。缺点是用户的程序必须以特定编译器的中间表示 (IR)形式呈现给这些DSC。由于这些IR 专注于特定领域，它们的表达力通常不如通用编程语言。虽然已经开发了许多通用编程语言的库来构建这些IR，但它们都存在语言子集问题，即为了能够构建IR，在用户程序中牺牲了部分表达能力。
在本文中，我们介绍了LazyTensor，这是一种将Eager Execution的优点与DSC 相结合的新方法。该技术允许在用户程序中使用所有主机编程语言功能，避免语言子集问题。
本文的主要介绍了：
1. 一种将Eager编程模型与领域特定编译器相结合的技术，该技术不会限制用户编程语言的表达能力，可以应用于任何按定义运行的机器学习框架。
2. LazyTensor在两种不同编程语言的两个不同机器学习框架中的实现：PyTorch和Swift for TensorFlow。
3. 通过跨多种语言、张量类型和加速器（GPU 和 TPU）进行验证。
相关背景
图模式
当前许多深度学习框架以数据流图结构来表示模型，比如TensorFlow v1.x；TensorFlow使用Python代码来构建计算图，这个图传递给用C++实现数据流执行引擎进行执行。 这个图可以序列化为被称作为GraphDef的缓存，因此它可以独立于宿主编程语言进行执行，我们将其称为GraphDef编程语言，将数据流执行引擎称为GraphDef的解释器。GraphDef编程语言可以转换为DSC的IR表示，但是由于其是对Tensor的优化，因此它不支持很多Python的语法。
Eager模式及异步执行
与图模式相比，在eager模式中，用户可以使用一种通用编程语言来进行模型的开发，因此编写和执行神经网络模型被认为更容易调试和更灵活。通过异步执行将计算硬件（CPU 以及 GPU 等加速器）得到有效利用。例如当用户的程序计算两个Tensor之间的矩阵乘法时，计算操作被分派到计算设备，控制权立即返回给用户的程序。只有当用户的程序试图查看其中的内容时，运行时才会进行阻塞，直到计算得到具体的Tensor值。因为Tensor计算从未在数据结构中具体化，所以不能转换为DSC的IR表示。
Eager模式和DSCs（特定领域编译器）
当前有很多机制尝试将eager模式和DSC结合：
- Tracing
通过一组对Tensor操作的方法，在计算过程中将这些方法的信息进行记录。这种跟踪过程类似于Autograd的tape方式，在正向执行过程中记录所有的执行过程，像 JAX和TensorFlow 2.x提供了跟踪装饰器，在装饰器中将捕获Tensor上的操作以供DSC优化。控制流和所有非Tensor操作的代码等对跟踪来说是“不可见的”，因此仅在跟踪时执行的代码，或者依赖于运行时张量值，则根本无法跟踪。
- 语言虚拟化
可以通过源到源转换来增强跟踪，该转换将无法跟踪的语言功能（例如控制流）重写为可跟踪的语言功能。一旦通过虚拟化进行增强，跟踪系统就能够跟踪 Python列表等内置类型的操作，或者需要判断运行时Tensor值的控制流。但是虚拟化可能无法覆盖所有语言功能，例如一个特别困难的案例是具有异常处理的控制流。
- 直接编译
另一种结合Eager代码和DSC的方法是实现一个编译器前端。 TorchScript实现了词法分析器、解析器、类型推理引擎和优化器，而Julia的 XLA 支持利用 Julia 编译器对其中的一个子集执行相同操作语。当嵌入到 Python 和 Julia 等动态语言中时，可以使用类似于跟踪的用户界面（例如函数装饰器）即时调用编译器。与跟踪类似，基于编译器前端的方法通常要求函数中的所有代码可以被静态编译，或者可以被DSC转换成目标IR。而这种限制是病毒式的：编译一个用户函数需要编译它调用的函数，因此一旦函数或库中有不受支持的行为（例如，调用物理模拟器或环境模型等外部函数）意味着所有传递调用者也不能被编译。
综上，这些现有技术的一个关键共同缺点是它们仅支持编程语言的一个子集。 此外，这些方法在可以用DSC编译的代码和不能用DSC编译的代码混合在一起的场景下，都无法正常工作。
LazyTensor方法及系统
LazyTensor方法介绍
LazyTensor的方法保持着eager执行的假象，因为该方法不会侵入通用编程语言。特别是支持了完整的通用语言特性，包括任意的Tensor和非Tensor的计算。
该方法建立在类PyTorch的异步执行之上，如果不需要查看Tensor的具体内容，用户无法感知在Tensor上的操作是否执行完成。相对于单独分发每个操作以异步执行，该方法会缓存操作序列，同时会判断该操作是否可以被转换成DSC的IR表达，最终将保存的操作序列转换成IR。IR图的构建过程总是从Tensor操作开始。
所有的Tensor API被分为两个部分，一种是可以被表达成IR图的API以及另一种无法被表达成IR的API。任意一个返回一个或多个Tensor的API，都是可以被转换成IR图，而返回非Tensor类型的API则无法被转换成IR图。
所有的API在用户程序中都可以被使用，包含以下优点：
- 函数抽象
不需要区分函数与IR图是否兼容，从而避免函数着色问题（不兼容的函数及其调用者都不能被转换成IR）。 任何函数都可以调用任何其他函数，无论它是否仅由IR兼容操作组成。调用IR不兼容操作的函数只是在运行时强制对IR图进行编译和执行，然后启动新的IR图的记录。
- 控制流
所有语言控制流操作，包括异常或虚函数调用等复杂情况在LazyTensor中都是可用的。
- 数据结构
程序可以将张量值嵌入到任意数据结构中。
因此，与之前的暂存系统相比，该方法支持所有的编程语言的特性。通过有效地构建一个复杂版本的运行时来实现这一点，该版本依赖于在幕后记录IR图，而不是将其暴露给用户。
LazyTensor系统
LazyTensor系统构建在 (a)底层Eager运行时和 (b)特定领域编译器之上。
LazyTensor有 3个主要组件：（1）自定义Tensor类型，其API与现有Tensor类型相同，（2）从Tensor操作到XLA HLO（DSC）序列的映射，以及（3）将张量操作序列转换成XLA HLO IR，并编排结果程序的编译和执行的运行时。因为编译通常非常耗时，所以LazyTensor系统会缓存和重用程序IR 图。
LazyTensor实现包括一个额外的barrier API（PyTorch 中的mark step ()，Swift for TensorFlow中的 LazyTensorBarrier()）。该API用于结束当前正在进行的IR图构建，并将其分发到运行时进行编译和执行。Barrier API采用一个布尔参数来控制调用是否应该阻塞，直到IR图执行完成并且所有张量数据都已在内存中具体化。IR不兼容操作的实现在继续执行之前调用barrier API 并设置阻塞位。未来可以从公共接口中删除该API。
- IR图
LazyTensor IR图将计算记录为有向无环图，其中叶子是输入，根是基于给定输入计算得到的结果。 图1包含从 PyTorch 构建的IR图的两个表示，用于在形状 [2, 4] 的浮点张量上进行简单的 x * y + z 计算。输入表示为 xla::device\_data。乘法和加法由节点 %6 和 %7 表示。在PyTorch中，加法操作允许为第二个参数指定一个缩放因子，它由附加节点%3表示。缩放因子是扩展到所需形状的常数1。缩放操作由底层 XLA 编译器后端优化掉，生成的本机代码不会实现常量，也不会执行乘以1。
与原生 PyTorch 类似，设备之间没有隐式传输。这反映在上面的IR图形表示中：类型为 xla::device 数据的输入与设备（在本例中为 CPU:0）相关联。 所有操作都要求所有输入都在同一设备地址空间上，随后它们的输出保持在同一设备上。乍一看，这种选择似乎给用户带来了很大的负担，但框架提供了通过一两行代码将整个模型传输到设备的方法。此外，此选择可防止难以调试与用户未请求的设备地址空间之间的隐式传输相关的性能问题。
- 控制流
目前，在LazyTensor IR图中没有控制流的表示。对于条件语句和循环语句，都会捕获执行生成的执行路径。 给定一个带有循环的简单程序（图2a），会生成一个展开的线性IR图（见图 2b）。
这反映了IR和主机语言（在例子中是 Python 或 Swift）之间是分离的。此外，该实现在实践中效果很好。例如，在少量模型配置之间进行选择的条件语句或从重复层构建模型的循环都可以很好地满足这种选择。另一方面，在需要判断Tensor值的条件语句中只会导致跟踪中断，而不是使跟踪变得不可能。
- In-place操作
Swift for TensorFlow 和 PyTorch 都允许在语法上就地更新：s += x 将s的值更新为 s + x。 此类操作是训练期间权重更新的基础，必须得到有效支持。这两个框架在支持此类操作的方式上有所不同。Swift的运算符重载会自动将就地版本转换为简单的赋值 s = s + x，因此可以重用常规加法实现。 XLA 编译器将利用不再需要s的旧值，并且可以有效地重用内存。
另一方面，Python不提供这种重写，因此PyTorch需要实现额外的、In-place版本的算术（以及其他）运算符。 然而，LazyTensor IR图没有mutation的概念，所有IR图都代表纯函数。为了实现mutation的语义，LazyTensor系统将mutation特性实现为将与目标相关联的底层计算替换为与就地操作右侧的表达式相关联的计算。 这实现了与 Swift 中的内置机制相同的效果：生成的 IR 图对于常规和in-place操作看起来相同。虽然可以以纯函数方式实现mutation语义，但在训练机器学习模型时，内存使用和性能至关重要。由于需要在替换目标之前存储右侧，因此该模型的简单实现将使用两倍的最佳内存量进行更新。 幸运的是，我们选择的纯函数表示与 XLA 模型相匹配，它允许指定输入和输出缓冲区之间的混叠。
- 未实现的Tensor操作
机器学习框架在其Tensor API中提供了数千种操作。 尽管DSC通常比预编译内核库支持更灵活的线性代数运算，但DSC通常并不支持所有的Tensor运算。为了提供一个替代的 Tensor API，所有没有映射到DSC的Tensor操作都使用以下模式实现：
1. 编译和执行所有操作输入的IR图。
2. 在输入上调用底层的Eager实现。
3. 使用Eager操作执行的结果开始一个新的 IR 图。
因此，当使用LazyTensor系统执行时，每个Tensor操作都将产生正确的程序，尽管有一些潜在的性能影响。使用调试器的断点工具，用户可以快速确定他们的代码调用操作的位置。
实验结果
1. 代码重用
LazyTensor 系统的核心已用于XLA，并支持Swift for TensorFlow和PyTorch集成。 表 1 记录了 Swift for Tensorflow 实现的各个文件夹中的源代码行 (SLoC)，并注释了它们是否在两个前端之间共享。这证明了这种技术（以及大部分实现）可以跨编程语言和底层 Eager 运行时重用。
2. 在Cloud TPUs上训练Transformers
Transformer是当今广泛用于自然语言处理领域的深度学习架构，它在从语言解析、机器翻译到问答等各种指标上都取得了最先进的性能。
通过PyTorch LazyTensor实现，我们启用了流行的HuggingFace Transformer库以使用 XLA 在云TPU上运行。与大致相当的GPU硬件相比，带有LazyTensor的PyTorch能够在TPU上展示出显著的性能提升（表2）。
3. ResNet-50在TPU集群上的性能
我们通过使用Swift for TensorFlow在TPU Pods上训练 ResNet-50来评估LazyTensor系统的缩放特性。Swift for TensorFlow在TPU上的性能是通过使用TPUv3-16、TPUv3-32和TPUv3-128集群（不使用 Cloud）在ImageNet 2012数据集上训练ResNet-50图像分类网络来衡量的，如图所示在表3中。该模型训练了90个epochs，并且记录了所需的时间以及每秒示例中的预热后吞吐量。 每个加速器的吞吐量基本保持不变，这表明 LazyTensor 技术可以扩展到大型TPU超级计算机。
4. 限制
不幸的是，并非所有模型在LazyTensor 系统中都比 Eager-system有更高性能。 例如程序运行时间不够长，或者编译优化后执行的性能收益没有超过JIT编译本身的开销。 因此，这种方法通常只在长时间运行的迭代计算（例如神经网络训练或批量推理）中才有意义。
最痛苦的限制之一不是来自技术本身，而是来自底层的DSC：XLA的静态形状限制。 所有张量形状都必须在IR图编译时已知，因为它们用于 XLA 编译器中的静态内存规划和其他优化。尽管系统缓存了基于LazyTensor IR图，但一些ML模型不仅仅是“形状稳定”的 IR 图集。 例如，在COCO数据集上训练 MaskRCNN显示加速器利用率很低，因为XLA在反复编译新的形状对应的程序。该应用程序验证了我们选择实现与原始 Eager Execution 相同的 API，因为它使用户能够在任何给定时间为其应用程序选择最有效的执行策略。
相关工作及演进
尽管 LazyTensor 系统已在多种应用中得到有效应用，但仍有许多改进方向可以拓宽其适用性。
尽管LazyTensor图构建开销不会影响大多数训练的收敛时间，但更新IR图会影响在小模型、小张量大小或低批量大小上运行的开销，这在推理中经常发生。虽然在第二次遇到相同的IR图时会跳过耗时的重新编译，但此类开销也可以从减少图构建开销中受益。
一种可能的方向是解决IR图的表示。传统的即时编译器（例如 WebKit 的 FTL JIT 9）中一种成功的策略是减少IR节点大小以及使用优秀数据结构和算法。可以在系统中借用此类技术来降低IR图的开销。
另一个方向将为用户提供一种方法来保证张量的底层计算将在各个步骤中保持不变。 在这样做时，在第一次迭代期间构建它之后，后续迭代可以跳过IR图的构建。
最后，自动截断、重新滚动循环和异步分发IR图片段的技术可以消除用户代码中对 LazyTensorBarrier() 的需求，并减少由于循环的可变上限而导致的重新编译。
总结
在本文中，我们介绍了 LazyTensor，这是一种将 Eager Execution 与特定领域编译器相结合的通用技术，它不会限制用户编程语言的表达能力。已经在两种机器学习框架的两种编程语言中成功实现了这项技术：PyTorch和Swift for TensorFlow，并且成功地重用了大部分实现。实验表明，对于只能通过 DSC (XLA) 访问的硬件，可以显著提高性能，正如在Cloud TPU 上的 HuggingFace Transformers 库所衡量的那样。 此外，还展示了这种方法如何扩展到大型分布式加速器集群。
点评
LazyTensor的基本做法：通过缓存API操作序列（而并不是真正执行），在到达一个barrier API时，进行统一的编译优化和执行。
这种方案，和现在的开发人员手工标记修饰符来确定静态图执行的范围有相似之处，LazyTensor可能更加自动一些，修饰符的方式有更好的使用范围。
优点：
将动态图和静态图的优点结合，既有灵活性又可结合DSC进行JIT加速。
挑战：
barrier API当前还无法自动生成与识别；
本身DSC的编译优化开销比程序运行时间或者编译优化后的性能提升还大的场景下，该方法为负优化；目前看，编译优化的开销大概率比执行时间长，所以比较好的方式是缓存，但是缓存的生命周期管理是个挑战。
动态shape的场景下，编译耗时远远大于编译后得到的性能提升。
控制流目前还需要依托Python解释器，性能上还是存在挑战
END
推荐阅读
更多嵌入式AI技术干货请关注嵌入式AI专栏。
