基于端到端可迁移深度强化学习的图优化

文 / Yanqi Zhou 和 Sudip Roy，研究员，Google Research

各式各样加速器上训练的大规模复杂神经网络驱动着越来越多的应用的发展。机器学习 (ML) 编译器为这一进程提供了便利，它将高级 计算图 映射到特定设备可执行的低级文件。为此，机器学习编译器需要解决许多优化问题，包括计算图重构、将算子分配给（不同的）计算设备、算子融合、张量的布局和平铺以及调度。例如，在训练 设备配置优化问题 (Device Placement problem) 中，编译器需要确定计算图中的运算到目标物理设备之间的映射，以使目标函数（如训练步时间）最小化。

• 计算图

  https://www.easy-tensorflow.com/tf-tutorials/basics/graph-and-session#GRAPH

• 设备配置优化问题

  http://proceedings.mlr.press/v70/mirhoseini17a/mirhoseini17a.pdf

配置性能由多种复杂因素共同决定，包括设备间网络带宽、峰值设备内存、共置约束等，这对启发式算法或基于搜索的算法来说具有挑战性，而启发式算法或基于搜索的算法通常被用作快速但次优的最佳解决方案。此外，启发式算法很难开发和维护，尤其在新模型架构不断涌现的当下。

最近基于学习方法的研究成果展现出很好的前景，但也有一定局限性，难以在实践中部署。首先，这些方法不易泛化到未见过的计算图，特别是那些由较新的模型架构产生的计算图；其次，它们的样本效率很差，导致训练期间的资源消耗很高。最后，它们只能解决单一的优化任务，无法捕获编译栈中紧密耦合的优化问题的依赖关系。

在最近发表于 NeurIPS 2020 会议的口头报告论文“ Transferable Graph Optimizers for ML Compilers ”中，我们提出了一种用于计算图优化 (Graph Optimization, GO) 的端到端可迁移深度强化学习方法，该方法克服了上述所有限制。对照 TensorFlow 的默认优化，我们在三个计算图优化任务上展现出 33%-60% 的速度提升。在由多达 80,000 个节点（包括 Inception-v3、Transformer-XL 和 WaveNet）组成的多样化典型计算图集上，GO 比专家优化平均提高了 21%，比先前的最先进技术提高了 18%，收敛速度快了 15 倍。

• Transferable Graph Optimizers for ML Compilers

  https://papers.nips.cc/paper/2020/file/9f29450d2eb58feb555078bdefe28aa5-Paper.pdf

• TensorFlow

  https://tensorflow.google.cn/

我们的研究提出了 GO，这是一个深度强化学习框架，可以解决上述每个优化问题 - 既可以逐个解决，也可以解决联合问题。提出的架构有三个关键方面：

首先，我们使用计算图神经网络（特别是 GraphSAGE ）来捕获计算图中编码的拓扑信息。GraphSAGE 的归纳网络利用节点属性信息泛化到先前未见过的计算图，从而对未见过的数据做出决策，并且不会产生大量的训练成本。

• GraphSAGE

  http://snap.stanford.edu/graphsage/

其次，许多模型的计算图通常包含超过 1 万个节点。要在如此大的规模上有效解决优化问题，网络必须能够捕获节点之间的长距离依赖关系。GO 的架构包括一个可扩展的注意力网络（Attention Network），该网络使用段级递归来捕获这种长距离节点依赖关系。

第三，机器学习编译器需要解决来自不同应用领域的各种图优化问题。用异构计算图训练共享策略网络的朴素策略不太可能捕获特定类计算图的特殊性。为了克服这个问题，GO 使用了一种特征调制机制，允许网络在不增加参数数量的情况下，针对特定的计算图类型进行特异化。

接下来我们将介绍评估结果，分别取自基于真实硬件测量的设备 配 置任务的单任务速度提升、向具有不同 GO 变体的未见过的计算图的泛化以及联合优化运算融合、设备 配 置和调度的多任务性能。

速度提升

为了评估该架构的性能，我们将 GO 应用到基于真实硬件评估的设备 配 置问题，首先在每个工作负载上分别训练模型。这种被称为 GO-one的方法始终优于专家人工配置 (HP)、TensorFlow METIS 配 置和 分层设备配置 (HDP) - 目前最先进的基于强化学习的设备 配 置。重要的是，通过高效的端到端单发 配 置，GO-one 比 HDP 在 配 置网络的收敛时间缩短 15 倍。

• 分层设备配置

  https://research.google/pubs/pub46646/

我们的实证结果表明， GO-one 的性能始终优于专家 配 置、TensorFlow METIS 配 置和分层设备 配 置 (HDP)。由于 GO 旨在扩展到像 8 层 Google 神经机器翻译 (GNMT) 模型这样由超过 8 万个节点组成的极大型计算图，因此它的性能优于之前的方法，包括 HDP、REGAL 和 Placeto。GO 为 GNMT 等大型计算图实现了优化的计算图运行时，分别比 HP 和 HDP 快 21.7% 和 36.5%。总体而言，与 HP 和 HDP 相比，GO-one 在 14 个多样化计算图中的平均 运行时间减少了 20.5% 和 18.2% 。重要的是，通过高效的端到端单发 配 置，GO-one 可以通过 HDP 将 配 置网络的收敛速度 增加 15 倍 。

• Google 神经机器翻译

  https://ai.googleblog.com/2016/09/a-neural-network-for-machine.html

泛化能力

GO 使用离线预训练泛化到未见过的计算图，然后在未见计算图上进行微调。预训练期间，我们在训练集的异构计算图子集上训练 GO。我们在每一批这样的计算图上对 GO 进行 1000 步训练，然后再切换到下一批。然后，在保留计算图上对预训练的模型进行不到 50 步的微调 ( GO-generalization+finetune )，这通常需要不到一分钟的时间。 GO-generalization+finetune （保留计算图）在所有数据集上都始终优于专家 配 置和 HDP，平均而言与 GO-one 相当。

我们还只在预训练的模型上直接运行推理，而不对目标保留计算图进行任何微调，并将其命名为 GO-generalization-zeroshot 。这种未微调模型的性能仅略逊于 GO-generalization+finetune ，仍略优于专家 配 置和 HDP。这说明计算图嵌入向量和习得策略均可有效迁移，使模型可以泛化到未见过的数据。