PyTorch 中的 graph mode 在性能方面示意更为杰出，本文介绍 Torch.FX 这个弱小工具，能够捕获和优化 PyTorch 程序 graph。
一、简介
PyTorch 反对两种执行模式：eager mode 和 graph mode。
eager mode 中，模型中的运算符在读取时会立刻执行，它易于应用，对机器学习从业者更敌对，因而被设置为默认的执行模式。
graph mode 中，运算符先被合成一个 graph，而后作为一个整体进行编译和执行，它的性能更高，因而在理论生产中大量应用。
具体来说，graph mode 反对算子交融，两个算子通过合并，能够升高或本地化内存读取以及内核启动总开销。
交融能够是横向 (horizontal) 的： 采取利用于多个 operand 的繁多操作（如 BatchNorm），并将这些 operand 合并到一个数组中。
交融也能够是纵向 (vertical) 的： 将一个内核与另一个内核合并，后者须要应用第一个内核的输入（如 ReLU 后接卷积）。
Torch.FX（缩写为 FX）是一个公开可用的工具包，作为 PyTorch 软件包的一部分，反对 graph mode 的执行。它能够：
- 从 PyTorch 程序中获取 graph
- 容许开发者在获取的 graph 上编写 transformation**
Meta 外部先前曾经在用 FX 来优化生产模型 (production model) 的训练吞吐量 (training throughput)。本文将通过介绍 Meta 开发的基于 FX 的优化，来展现利用图构造转换 (graph transformation) 优化 PyTorch 部署模型性能的办法
二、背景
embedding table 宽泛存在于举荐零碎中， 本节将介绍 FX 和 embedding table 的背景常识。
2.1. FX
图 1 是一个简略示例，演示了如何用 FX 转换 PyTorch 程序，它蕴含三个步骤：
- 从程序中获取 graph
- 批改 graph（在本例中，咱们用 GELU 代替 RELU）
- 从批改后的 graph 中生成一个新程序
图 1：在 PyTorch 模块中用 GELU 取代 RELU 的 FX
FX API 为检查和转换 PyTorch 程序 graph 还提供了许多其余性能。
2.2. embedding table
图 2：批尺寸 =1 的稠密特色 embedding table 示意图
在举荐零碎中， 稠密特色（例如，User ID，Story ID）由 embedding table 示意。
embedding table E 是一个 HxD 矩阵，其中 H 是哈希大小，D 是嵌入向量维度。E 的每一行都是一个浮点数向量。
feature hashing 的作用是将一个稠密特色映射到 E 的索引列表中，例如 [S1,S2,…,Sk]，其中 0≤Si<H。它的输入值计算为 f(E[S1], E[S2],…,E[Sk])，其中 E[Si] 是 Si 行的向量，f 是池化函数，通常是 sum，average，max 三个函数之一。
为了充分利用 GPU，稠密特色通常为批处理。 批处理中的每个实体都有本人的索引列表。如果一个批次有 B 个实体，能够简略了解为一个表征有 B 个索引列表。
更为谨严的示意办法是将 B 个索引列表合并成一个索引列表，并增加一个索引长度的列表（该批中的每个实体都有一个长度 length）。
例如，如果一批蕴含 3 个实体，其索引列表如下：
- Entity 1: indices = [10, 20]
- Entity 2: indices = [5, 9, 77, 81]
- Entity 3: indices = [15, 20, 45]
则残缺批尺寸的 indice 和 length 将是：
- Indices = [10, 20, 5, 9, 77, 81, 15, 20, 45]
- Lengths = [2, 4, 3]
而整个 batch 的 embedding table 查问，输入为是一个 BxD 矩阵。
三、3 种 FX Transformation
PyTorch 更新了 3 个 FX transformation，以减速对 embedding table 的拜访，本节将逐个介绍。
下文 3.1 对于将多个小输出张量联合成一个大张量的转换；3.2 对于将多个并行计算链交融成一个计算链的转换；3.3 对于将通信与计算重叠的转换。
3.1 联合输出稠密特色
batch 中的每个输出稠密特色，都能够示意为两个列表：一个索引列表和一个 B length 列表，其中 B 示意批尺寸。
在 PyTorch 中，这两个列表都能够以张量的模式存在。 当 PyTorch 模型在 GPU 上运行时，embedding table 通常存储在 GPU 内存中（它更靠近 GPU，读写带宽比 CPU 内存更高）。
须要应用输出稠密特色时，两个张量都要先从 CPU 复制到 GPU。然而每个主机到设施的内存复制都须要启动内核，这对于理论的数据传输来说，会更加消耗工夫。
如果一个模型应用了多个输出稠密特色，这种复制可能成为性能瓶颈（例如，1000 个输出稠密特色将须要从主机到设施复制 2000 个张量）。
一个缩小主机到设施 memcpy 数量的优化办法，就是在多个输出稠密特色发送到设施之前，先将其进行组合。
例如，给定以下三个输出特色：
- Feature_A: indices = [106, 211, 7], lengths = [2, 1]
- Feature_B: indices = [52, 498, 616, 870, 1013], lengths = [3, 2]
- Feature_C: indices = [2011, 19, 351, 790], lengths = [1, 3]
组合后的模式为：
Features_A_B_C: indices = [106, 211, 7, 52, 498, 616, 870, 1013, 2011, 19, 351, 790], lengths = [2, 1, 3, 2, 1, 3]
所以不须要从主机到设施复制 3×2=6 个张量，只须要复制 2 个张量。
图 3(b) 形容了这种优化的实现，它蕴含两个组件：
- CPU 端：输出 pipeline 被批改为将所有稠密特色的 indices 组合成一个张量，所有 length 组合成另一个张量。而后将这两个张量复制到 GPU 上。
- GPU 端：应用 FX，在模型 graph 中插入一个 Permute_and_Split 算子，从合并的张量中复原单个特色 indices 和 length 张量，并将其发送至上游的相应节点。
优化前：两个张量都要从 CPU 复制到 GPU
优化后：将输出稠密特色进行组合
3.2 从拜访 embedding table 开始的计算链横向交融
在一个生产模型中，每个 GPU 上有 10 个 embedding table 很常见。出于性能方面的思考， 对这些 table 的查问被分到一组，这样它们的输入就被串联在一个大张量中 （见图 4（a）中的红色局部）。
为了对单个特色输入进行计算， 应用 Split 算子将大张量分成 N 个小张量 （其中 N 为特色的数量），而后将所需的计算利用于每个张量。
如图 4(a) 所示，利用于每个特色输入 O 的计算是 Tanh(LayerNorm(O))。所有的计算结果都被串联成一个大的张量，而后传递给上游的算子（图 4(a) 中的 Op1）。
这里次要的 runtime cost 是 GPU 内核启动的开销。 例如，图 4(a) 中的 GPU 内核的启动次数为 2*N+3（图中的每个椭圆都示意一个 GPU 内核）。这会影响性能，因为 LayerNorm 和 Tanh 在 GPU 上的执行工夫，与它们的内核启动工夫相比很短。
此外，Split 算子可能会创立一个额定的嵌入向量输入张量的正本，耗费额定的 GPU 内存。
用 FX 来实现一种叫做横向交融 (horizontal fusion) 的优化，能够大大减少 GPU 内核的启动次数 （在这个例子中，优化后的 GPU 内核启动次数为 5，见图 4(b)）。
应用 Add_middle_dim 算子代替显式 Split，将 shape 为 (B, NxD) 的 2D 嵌入张量重塑为 shape 为 (B, N, D) 的 3D 张量。接下来将一个繁多的 LayerNorm 利用到它的最初一维。对 LayerNorm 的后果利用一个 Tanh。最初，用 Remove_middle_dim 算子将 Tanh 的后果复原成 2D 张量。
因为 Add_middle_dim 和 Remove_middle_dim 只是重塑张量， 并没有创立额定的正本，所以也能够缩小 GPU 内存的耗费。
优化前：所有输入被串联到一个大张量中
进行横向交融优化后
3.3 计算与通信间的重叠 (overlap)
面向投产的举荐模型的训练，通常是在分布式 GPU 零碎上实现的。 因为每个 GPU 的设施内存容量不足以包容模型中的所有 embedding table，因而须要将其散布在多个 GPU 上。
在训练步骤中，GPU 须要从其余 GPU 上的 embedding table 中读取 / 写入特征值。 这被称为 all-to-all 通信，可能是影响性能的重要起因。
通过 FX 实现一个 transformation，能够将计算与 all-to-all 通信重叠。 图 5(a) 显示了一个具备嵌入向量 table 拜访 (EmbeddingAllToAll) 及其他算子的模型 graph 实例。如图 5(b) 所示，在没有任何优化的状况下，它们会在一个 GPU 流上程序执行。
应用 FX 将 EmbeddingAllToAll 分成 EmbeddingAllToAll_Request 和 EmbeddingAllToAll_Wait，并在它们之间安顿独立的算子。
图 5：计算与通信的重叠
3.4 总结
表 1：本节探讨的优化及解决的相应性能瓶颈
为了发现哪些模型会从这些 transformation 中受害，开发人员对 MAIProf 收集的运行在 Meta 数据中心的模型的性能数据进行剖析。 得出与 eager mode 相比，这些 transformation 在一组生产模型上实现了 2-3 倍的速度晋升。
四、结语
从性能角度考量，PyTorch 中的 graph mode 比生产环境中应用的 eager mode 更受欢迎。FX 是一个弱小的工具，能够捕获和优化 PyTorch 程序 graph。本文展现了三种 FX transformation，用于优化 Meta 外部的生产举荐模型。
最初心愿更多 PyTorch 开发者能够应用 graph transformation 来晋升模型的性能。
—— 完 ——