越来越多的应用是由在不同加速器上训练的大型复杂神经网络驱动的。ML编译器为这一过程提供了便利，它将高级计算图映射到低级的、特定于设备的可执行文件。在这样做的过程中，ML编译器需要解决许多优化问题，包括图的重写、设备上的操作分配、操作融合、时序器的布局和分层以及调度。例如，在设备放置问题中，编译器需要确定计算图中的操作到目标物理设备之间的映射，以便使目标函数（如训练步骤时间）最小化。安置性能由各种错综复杂的因素决定，包括设备间网络带宽、设备内存峰值、协同定位约束等，这对启发式或基于搜索的算法来说是个挑战，它们通常满足于快速但次优的解决方案。此外，启发式算法很难开发和维护，尤其是在更新的模型架构出现的时候。

最近使用基于学习的方法的尝试已经展示了有希望的结果，但它们有一些局限性，使得它们在实践中无法部署。首先，这些方法不容易推广到未见过的图形，特别是那些由较新的模型架构产生的图形；其次，它们的样本效率很差，导致训练过程中的资源消耗很高。最后，它们只能解决单一的优化任务，因此，不能捕捉到编译栈中各个紧密耦合的优化问题的依赖性。

在 &34;（最近作为口头论文发表在NeurIPS 2020上）中，我们提出了一种用于计算图优化（GO）的端到端、可转移的深度强化学习方法，该方法克服了上述所有限制。我们展示了与TensorFlow默认优化相比，在三个图优化任务上提速33%-60%。在由多达80000个节点组成的多样化的代表性图集上，包括Inception-v3、Transformer-XL和WaveNet，GO比专家优化平均提高了21%，比之前的技术状态提高了18%，收敛速度提高了15倍。

在ML编译器中经常会出现三个耦合优化任务，我们将其表述为可以使用学习策略解决的决策问题。每一个任务的决策问题都可以重构为为计算图中的每个节点做一个决策。

第一个优化任务是设备放置，目标是确定如何最好地将图中的节点分配到其运行的物理设备上，使端到端运行时间最小化。

第二个优化任务是操作调度。当计算图中的一个操作在设备内存中存在其传入的时序时，它就准备好运行了。一个经常使用的调度策略是为每个设备维护一个准备好的操作队列，并按照先进先出的顺序调度操作。但是，这种调度策略没有考虑到放在其他设备上的下行操作可能会被某项操作阻断，往往会导致调度中设备利用率不足。为了找到能够跟踪这种跨设备依赖关系的调度，我们的方法使用了一种基于优先级的调度算法，根据每个操作的优先级来调度就绪队列中的操作。与设备放置类似，操作调度就可以被表述为学习一个策略的问题，该策略为图中的每个节点分配一个优先级，以使基于运行时间的奖励最大化。

第三个优化任务是操作融合。为了简洁起见，我们在这里省略了对这个问题的详细讨论，而只是指出与基于优先级的调度类似，操作融合也可以使用基于优先级的算法来决定融合哪些节点。在这种情况下，策略网络的目标同样是为图中的每个节点分配一个优先级。

最后，重要的是要认识到，三个优化问题中每个问题所采取的决策都会影响其他问题的最优决策。例如，将两个节点放在两个不同的设备上，就会有效地禁用融合，并引入一个通信延迟，从而影响调度。

我们的研究提出了GO，这是一个深层的RL框架，可以适应解决上述每个优化问题--既可以单独解决，也可以联合解决。所提出的架构有三个关键方面。

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

其次，许多模型的计算图往往包含10k以上的节点。在如此大的尺度上有效地解决优化问题，需要网络能够捕捉节点之间的长程依赖关系。GO的架构中包含了一个可扩展的注意力网络，它使用段级递归来捕捉这种长程节点依赖关系。

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

GO的概述。一个结合了图嵌入和顺序关注的端到端图策略网络。

为了联合解决多个依赖性优化任务，GO能够为每个任务增加额外的递归注意层，不同任务之间的参数是共享的。循环关注层与动作的残余连接，可以跟踪任务间的依赖关系。

多任务策略网络，扩展了GO的策略网络，增加了每个任务的递归关注层和剩余连接。GE：Graph Embedding，FC：全连接层，Nxf：融合动作维度，Fxd：放置动作维度，Nxs：调度动作维度。

接下来，我们介绍了基于真实硬件测量的设备放置任务上的单任务提速、对不同GO变体的未见图的泛化以及联合优化操作融合、设备放置和调度的多任务性能的评估结果。

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

我们的实证结果表明，GO-one的性能始终优于专家放置、TensorFlow METIS放置和分层设备放置（HDP）。由于GO的设计方式可以扩展到由超过80,000个节点组成的超大型图，比如8层的谷歌神经机器翻译（GNMT）模型，因此它的性能优于之前的方法，包括HDP、REGAL和Placeto。GO实现了像GNMT这样的大型图的优化图运行时间，分别比HP和HDP快21.7%和36.5%。总的来说，与HP和HDP相比，GO-one在不同的14个图形集上平均实现了20.5%和18.2%的运行时间缩短。重要的是，通过高效的端到端单枪匹马的投放，GO-one在投放网络的收敛时间上比HDP加快了15倍。

GO使用离线预训练，然后在未见图上进行微调，将GO泛化到未见图上。在预训练期间，我们在训练集的异构图子集上训练GO。在切换到下一个图形之前，我们在每一批这样的图形上训练GO 1000步。然后对这个预训练的模型进行微调（GO-generalization+finetune），在hold-out图形上进行少于50步的训练，这通常需要不到一分钟的时间。在所有的数据集上，GO-generalization+finetune在hold-out graphs上的表现都一致地优于expert placement和HDP，平均来说与GO-one相当。

我们还直接只在预训练的模型上运行推理，而不对目标图进行任何微调，并将其命名为GO-generalization-zeroshot。这个微调模型的性能只比GO-generalization+finetune略差，而比专家放置和HDP略好。这说明图嵌入和学习到的策略都能有效转移，使得模型可以泛化到未见数据。

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

本站所发布的文字与图片素材为非商业目的改编或整理，版权归原作者所有，如侵权或涉及违法，请联系我们删除