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