[论文解读] Placeto: Learning Generalizable Device Placement Algorithms for Distributed Machine Learning

Ravichandra Addanki, Shaileshh Bojja Venkatakrishnan+3 more|arXiv (Cornell University)|Jun 20, 2019

Ferroelectric and Negative Capacitance Devices参考文献 23被引用 30

一句话总结

Placeto 提出了一种强化学习框架，通过将放置过程建模为迭代改进，并利用图神经网络编码结构信息，学习适用于分布式深度学习的可泛化设备放置策略。该方法在无需微调的情况下即可在未见过的图上实现泛化，训练速度最高提升 6.1 倍，且在速度和放置质量方面均优于先前基于 RNN 的方法。

ABSTRACT

We present Placeto, a reinforcement learning (RL) approach to efficiently find device placements for distributed neural network training. Unlike prior approaches that only find a device placement for a specific computation graph, Placeto can learn generalizable device placement policies that can be applied to any graph. We propose two key ideas in our approach: (1) we represent the policy as performing iterative placement improvements, rather than outputting a placement in one shot; (2) we use graph embeddings to capture relevant information about the structure of the computation graph, without relying on node labels for indexing. These ideas allow Placeto to train efficiently and generalize to unseen graphs. Our experiments show that Placeto requires up to 6.1x fewer training steps to find placements that are on par with or better than the best placements found by prior approaches. Moreover, Placeto is able to learn a generalizable placement policy for any given family of graphs, which can then be used without any retraining to predict optimized placements for unseen graphs from the same family. This eliminates the large overhead incurred by prior RL approaches whose lack of generalizability necessitates re-training from scratch every time a new graph is to be placed.

研究动机与目标

• 解决现有基于强化学习的设备放置方法训练开销高的问题，这些方法在面对新计算图时需要从头开始重新训练。
• 实现学习到的放置策略在同一家族但未见过的图上无需重新训练即可泛化。
• 减少在分布式训练中寻找高质量设备放置方案所需的放置评估次数。
• 设计一种对任意节点索引和遍历顺序不变的策略，以提升鲁棒性和可迁移性。
• 开发一种可扩展、高效且可泛化的设备放置方法，其在速度和性能方面均优于基于 RNN 的基线方法。

提出的方法

• Placeto 将设备放置建模为一系列迭代改进步骤，策略按顺序为每个节点预测一个设备，基于当前部分放置状态。
• 它使用图神经网络（GNN）编码计算图的结构，通过消息传递捕捉父子节点和并行节点之间的关系，且不依赖于节点顺序或标签。
• 策略网络通过专用子网络聚合节点的局部邻域特征——包括父节点、子节点和并行节点——后再做出设备决策。
• 基于 GNN 的架构使策略能够跨不同节点索引或遍历顺序的图实现泛化，而基于 RNN 的方法则依赖于序列处理，难以实现此类泛化。
• Placeto 采用强化学习目标，并引入中间奖励以引导策略训练，实现更快收敛。
• 该框架为每个图家族学习一个单一的通用策略，可直接应用于该家族中的任意新图，而无需重新训练。

实验结果

研究问题

• RQ1 能否训练一个强化学习策略，使其在分布式深度学习的未见过的计算图上实现泛化？
• RQ2 将放置过程建模为迭代改进是否能降低训练复杂度并提升策略泛化能力？
• RQ3 图神经网络是否能有效编码计算图中的结构信息，而无需依赖节点标签或遍历顺序？
• RQ4 Placeto 的泛化能力与基于 RNN 的方法相比，在放置质量和训练效率方面表现如何？
• RQ5 策略架构的选择在多大程度上影响其在不同图家族上的泛化性能？

主要发现

• 与基于 RNN 的方法相比，Placeto 将放置评估次数最多减少 6.1 倍，同时获得的放置结果与最优先前结果相当或更优。
• 在 NMT 模型上，Placeto 找到的放置方案比基于 RNN 的方法快 16.5%。
• Placeto Zero-Shot（使用无需微调的通用策略）的性能几乎与 Placeto Optimized（需数千次评估）完全一致。
• 基于 RNN 的方法在零样本设置下表现较差，RNN Zero-Shot 几乎与随机基线性能相当，凸显其缺乏泛化能力。
• Placeto 的策略在不同节点遍历顺序下表现出稳健的泛化能力，平均运行时间波动仅为 5%，而基于 RNN 的方法波动高达 30%。
• 消融实验表明，对父节点、子节点和并行节点的局部邻域聚合与池化机制对泛化至关重要；若移除这些组件，性能下降 13% 至 20%。