一种在多芯片模块上划分机器学习模型的可移植方法
A Transferable Approach for Partitioning Machine Learning Models on Multi-Chip-Modules
摘要：
多芯片模块（MCM）降低了机器学习（ML）加速器的设计和制造成本，同时提供与单片大型芯片相当的性能和能效。然而，针对MCM的ML编译器需要最佳和有效地解决复杂的优化问题，以实现这种高性能。一个这样的问题是多芯片划分问题，其中编译器确定MCM中的小芯片上的张量计算图中的操作的最佳划分和放置。划分MCM的ML图特别困难，因为搜索空间随着可用小芯片的数量和神经网络中的节点数量呈指数增长。此外，由底层硬件施加的约束产生有效解极其稀疏的搜索空间。在本文中，我们提出了一种策略，使用深度强化学习（RL）框架发出一个可能无效的候选分区，然后由约束求解器纠正。使用约束求解器确保RL在稀疏空间中足够频繁地遇到有效的解决方案，以便与非学习策略相比以更少的样本收敛。我们为策略网络所做的架构选择允许我们在不同的ML图中进行泛化。我们的生产规模的模型，BERT，在真实的硬件上的评估表明，使用RL策略生成的分区实现6.11%和5.85%的吞吐量比随机搜索和模拟退火。此外，微调预训练的RL策略将搜索时间从3小时减少到仅9分钟，同时实现与从头开始训练RL策略相同的吞吐量。
英文摘要：
Multi-Chip-Modules (MCMs) reduce the design and fabrication cost of machine learning (ML) accelerators while delivering performance and energy efficiency on par with a monolithic large chip. However, ML compilers targeting MCMs need to solve complex optimization problems optimally and efficiently to achieve this high performance. One such problem is the multi-chip partitioning problem where compilers determine the optimal partitioning and placement of operations in tensor computation graphs on chiplets in MCMs. Partitioning ML graphs for MCMs is particularly hard as the search space grows exponentially with the number of chiplets available and the number of nodes in the neural network. Furthermore, the constraints imposed by the underlying hardware produce a search space where valid solutions are extremely sparse. In this paper, we present a strategy using a deep reinforcement learning (RL) framework to emit a possibly invalid candidate partition that is then corrected by a constraint solver. Using the constraint solver ensures that RL encounters valid solutions in the sparse space frequently enough to converge with fewer samples as compared to non-learned strategies. The architectural choices we make for the policy network allow us to generalize across different ML graphs. Our evaluation of a production-scale model, BERT, on real hardware reveals that the partitioning generated using RL policy achieves 6.11% and 5.85% higher throughput than random search and simulated annealing. In addition, fine-tuning the pre-trained RL policy reduces the search time from 3 hours to only 9 minutes, while achieving the same throughput as training RL policy from scratch.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
