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

文 / 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/
  

机器学习编译器中的图优化问题

在机器学习编译器中经常出现三个耦合优化任务，我们用数学模型将它构造成一个解决策略可学习的决策问题。这些任务的决策问题都可以重构为对计算图中的每个节点做决策。

第一个优化任务是 训练设备配置，目标是确定如何以最佳方式将计算图的节点分配到其运行的物理设备上，使端到端运行时间最小化。

第二个优化任务是 运算调度。当设备内存中存在传入张量时，计算图中的运算 即可 运行。常用的调度策略是为每个设备维护就绪的运算队列，并以先进先出的顺序调度运算。然而，这种调度策略并未考虑配置于其他设备的下游运算。这些运算可能会被某个运算所阻塞，经常导致调度中设备利用率不足。为了找到可以跟踪此类跨设备依赖关系的调度，我们的方法使用了基于优先级的调度算法，基于每个优先级在就绪队列中调度运算。与训练设备配置问题类似，运算调度随后可以被数学建模为一个策略学习问题，该策略为计算中的每个节点分配一个优先级，使基于运行时间的奖励最大化。

第三个优化任务是 运算融合。篇幅有限，我们在此省略对该问题的详细讨论，仅指出运算融合与基于优先级的调度类似，也可使用基于优先级的算法确定待融合的节点。在这种情况下，策略网络的目标同样是为计算图中的每个节点分配优先级。

最后，还要注意在三个优化问题中，每个问题所采取的决策都会影响其他问题的最优决策。例如，将两个节点放在两个不同的设备上，会有效地禁用融合，并引入可能影响调度的通信延迟。

RL 策略网络架构

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

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

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

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

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

GO 概览：结合了计算图嵌入向量和顺序关注的端到端计算图策略网络

为了同时解决多个依赖性优化任务，GO 可以为每个任务添加额外的递归注意力机制层(Recurrent Attention Layers)，并在不同任务之间共享参数。具有动作残留连接的递归关注层可以实现任务间依存关系的跟踪。

多任务策略网络，扩展了 GO 的策略网络，为每个任务和残留连接增加了递归关注层。GE：计算图嵌入向量，FC：全连接层，Nxf：融合动作维度，Fxd：配置动作维度，Nxs：调度动作维度

实验结果

接下来我们将介绍评估结果，分别取自基于真实硬件测量的设备配置任务的单任务速度提升、向具有不同 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。这说明计算图嵌入向量和习得策略均可有效迁移，使模型可以泛化到未见过的数据。

跨异构工作负载计算图的泛化。图中显示了使用 6 种工作负载（Inception-v3、AmoebaNet、递归神经网络语言模型 (RNNLM)、Google 神经机器翻译 (GNMT)、Transformer-XL (TRFXL)、WaveNet）中的 5 种（除了保留计算图）的计算图进行训练时，GO 两种不同泛化策略的比较，并在保留工作负载上进行了评估（x 轴）

协同优化配置、调度和融合 (pl+sch+fu)

同时优化配置、调度和融合比单 GPU 未优化的情况提速 30%-73%，比 TensorFlow 默认配置、调度和融合提速 33%-60%。与单独优化每个任务相比，多任务 GO (pl+sch+fu) 比单任务GO (p | sch | fu)（一次优化所有任务）平均提高了 7.8%。此外，对于所有工作负载，协同优化所有三个任务，比优化其中任意两个任务并对第三个任务使用默认策略具有更快的运行时间。

多任务优化的各种工作负载的运行时间。TF-default：TF GPU 默认配置、融合和调度。hp-only：仅包含默认调度和融合的人工配置。pl-only：仅包含默认调度和融合的 GO 配置。pl | sch：GO 通过默认融合分别优化配置和调度。pl+sch：多任务 GO 通过默认融合协同优化配置和调度。sch+fu：多任务 GO 通过人工配置协同优化调度和融合。pl | sch | fu：GO 分别优化配置、调度和融合。pl+sch+fu：多任务 GO 协同优化配置、调度和融合

结论

硬件加速器的复杂性和多样性日益增加，使得稳健且适应性强的 ML 框架的开发既繁琐又耗时，往往需要数百名工程师的多年努力。在本文中，我们证明了使用精心设计的学习方法可以更快、更好地解决此类框架中的许多优化问题。

致谢

本研究由 Daniel Wong、Amirali Abdolrashidi、Peter Ma、Qiumin Xu、Hanxiao Liu、Mangpo Phitchaya Phothilimthana、Shen Wang、Anna Goldie、Azalia Mirhoseini 和 James Laudon 共同完成。

更多 AI 相关阅读：

• 智能滚动：让转录后的文本编辑、共享和搜索更容易

• Vision Transformer：用于大规模图像识别的 Transformer

• Coral：让终端人工智能更加触手可及

• 推出 LIT：NLP 模型的交互式探索和分析

• MediaPipe Holisitic：实现端侧人脸手势姿态同时估计