越来越多的应用程序是由在不同加速器集上训练的大型复杂神经网络驱动的。这个过程由 ML 编译器促进，将高级计算图映射到低级、特定于设备的可执行文件。为此，ML 编译器需要解决许多优化问题，包括图形重写、设备上的操作分配、操作融合、张量的布局和平铺以及调度。例如，在设备放置问题中，编译器需要确定计算图中的操作到目标物理设备之间的映射，以便可以最小化目标函数，例如训练步骤时间。放置性能由复杂因素的混合决定，包括设备间网络带宽、峰值设备内存、协同定位约束等，这使得启发式算法或基于搜索的算法具有挑战性，这些算法通常满足于快速但低于-最佳，解决方案。此外，启发式很难开发和维护，尤其是在出现新的模型架构时。
最近使用基于学习的方法的尝试已经证明了有希望的结果，但它们有许多限制，使其无法在实践中部署。首先，这些方法不容易推广到看不见的图，尤其是那些由较新的模型架构产生的图，其次，它们的样本效率很差，导致训练过程中的资源消耗很高。最后，它们只能解决单个优化任务，因此无法捕获编译堆栈中紧耦合优化问题之间的依赖关系。
在最近在NeurIPS 2020 上作为口头论文发表的“Transferable Graph Optimizers for ML Compilers ”中，我们提出了一种用于计算图优化 (GO) 的端到端、可转移的深度强化学习方法，该方法克服了上述所有限制。与TensorFlow默认优化相比，我们在三个图优化任务上展示了 33%-60% 的加速。在包括Inception-v3、Transformer-XL和WaveNet在内的由多达 80,000 个节点组成的不同代表性图集上，GO 比专家优化平均提高了 21%，比现有技术提高了 18%，提高了 15 倍更快的收敛。
ML 编译器中的图优化问题 ML 编译器
中经常出现三个耦合优化任务，我们将它们表述为可以使用学习策略解决的决策问题。每个任务的决策问题可以重新定义为为计算图中的每个节点做出决策。
第一个优化任务是设备放置，其目标是确定如何最好地将图的节点分配给它运行的物理设备，从而最大限度地减少端到端的运行时间。
第二个优化任务是操作调度。计算图中的操作已准备就绪当它的传入张量存在于设备内存中时运行。一个常用的调度策略是为每个设备维护一个就绪的操作队列，并按先进先出的顺序调度操作。但是，这种调度策略没有考虑到其他设备上可能被某个操作阻塞的下游操作，并且经常导致调度未充分利用的设备。为了找到可以跟踪此类跨设备依赖关系的调度，我们的方法使用基于优先级的调度算法，该算法根据每个操作的优先级调度就绪队列中的操作。类似于设备放置，操作调度可以被表述为学习策略的问题，该策略为图中的每个节点分配优先级，以根据运行时间最大化奖励。
第三个优化任务是操作融合。为简洁起见，我们在这里省略了对这个问题的详细讨论，而只是注意类似于基于优先级的调度，操作融合也可以使用基于优先级的算法来决定融合哪些节点。在这种情况下，策略网络的目标再次是为图中的每个节点分配优先级。
最后，重要的是要认识到在三个优化问题中的每一个中做出的决定都会影响其他问题的最佳决策。例如，将两个节点放在两个不同的设备上会有效地禁用融合并引入可能影响调度的通信延迟。
RL 策略网络架构
我们的研究提出了 GO，这是一种深度 RL 框架，可适用于单独或联合解决上述每个优化问题。所提议的架构有三个关键方面：
首先，我们使用图神经网络（特别是GraphSAGE）来捕获计算图中编码的拓扑信息。GraphSAGE 的归纳网络利用节点属性信息泛化到以前看不见的图，这使得对看不见的数据做出决策而不会产生大量的训练成本。
其次，许多模型的计算图通常包含超过 1 万个节点。在如此大规模上有效地解决优化问题需要网络能够捕获节点之间的远程依赖关系。GO 的架构包括一个可扩展的注意力网络，该网络使用段级递归来捕获这种远程节点依赖关系。
第三，ML 编译器需要解决来自不同应用领域的各种图形的优化问题。用异构图训练共享策略网络的幼稚策略不太可能捕捉到特定类别图的特质。为了克服这个问题，GO 使用了一种特征调制机制，该机制允许网络专门针对特定的图类型而无需增加参数数量。
为了共同解决多个相关优化任务，GO 能够为每个任务添加额外的循环注意层，并在不同任务之间共享参数。具有剩余动作连接的循环注意层能够跟踪任务间的依赖关系。
结果
接下来，我们展示了基于真实硬件测量的设备放置任务的单任务加速评估结果，对具有不同 GO 变体的不可见图的泛化，以及联合优化操作融合、设备放置和调度的多任务性能。
加速：
为了评估该架构的性能，我们将 GO 应用于基于真实硬件评估的设备放置问题，我们首先在每个工作负载上分别训练模型。这种称为GO-one 的方法始终优于专家手动放置 (HP)、TensorFlow METIS放置和分层设备放置(HDP)——当前最先进的基于强化学习的设备放置。重要的是，凭借高效的端到端单次布局，GO-one 的布局网络在 HDP 上的收敛时间加快了 15 倍。
我们的实证结果表明，GO-one始终优于专家布局、TensorFlow METIS 布局和分层设备布局(HDP)。由于 GO 的设计方式可以扩展到由超过 80,000 个节点组成的超大图，例如 8 层Google 神经机器翻译(GNMT) 模型，因此它优于以前的方法，包括 HDP、REGAL和Placeto。GO 为 GNMT 等大型图实现了优化的图运行时间，分别比 HP 和 HDP 快 21.7% 和 36.5%。总体而言，GO-one 平均减少了 20.5% 和 18.2% 的运行时间分别与 HP 和 HDP 相比，包含 14 个不同的图形集。重要的是，凭借高效的端到端单次布局，GO-one的布局网络在 HDP 上的收敛时间加快了15 倍。
泛化：
GO 使用离线预训练对未见过的图进行泛化，然后对未见过的图进行微调。在预训练期间，我们在训练集中的图的异构子集上训练 GO。在切换到下一个之前，我们在每个这样的批次图上训练 GO 1000 步。然后在保留图上对这个预训练模型进行微调（GO-generalization+finetune），步骤少于 50 步，通常不到一分钟。GO-generalization+finetune for hold-out graphs 在所有数据集上始终优于专家放置和 HDP，并且平均匹配GO-one。
我们还直接在预训练模型上运行推理，而不对目标保持图进行任何微调，并将其命名为GO-generalization-zeroshot。这种未调整模型的性能仅比GO-generalization+finetune差一点，而略好于专家布局和 HDP。这表明图嵌入和学习到的策略都有效地转移，允许模型泛化到看不见的数据。
协同优化布局、调度和融合（pl+sch+fu）：
同时优化布局、调度和融合与未优化的单 GPU 情况相比提供 30%-73% 的加速，与 TensorFlow 相比提供 33%-60% 的加速默认放置、调度和融合。与单独优化每个任务相比，多任务 GO ( pl+sch+fu ) 比单任务 GO ( p | sch | fu )——一次优化所有任务——平均高 7.8%。此外，对于所有工作负载，与优化任意两个任务并为第三个任务使用默认策略相比，共同优化所有三个任务可提供更快的运行时间。
pl+sch：多任务 GO 使用默认融合协同优化布局和调度。sch+fu：多任务 GO 协同优化调度和融合与人类安置。pl | sch | fu：GO 分别优化布局、调度和融合。pl+sch+fu：多任务 GO 协同优化布局、调度和融合。
结论
硬件加速器日益复杂和多样化，使得开发健壮且适应性强的 ML 框架变得既繁琐又耗时，通常需要数百名工程师的多年努力。在本文中，我们证明了此类框架中的许多优化问题都可以使用精心设计的学习方法有效且最优地解决。
