[论文解读] ProTuner: Tuning Programs with Monte Carlo Tree Search

Ameer Haj-Ali|arXiv (Cornell University)|May 27, 2020

Machine Learning and Algorithms参考文献 47被引用 7

一句话总结

ProTuner 提出使用蒙特卡洛树搜索（MCTS）自动调优 Halide 程序，以实现高性能的深度学习与图像处理工作负载。通过将调度问题建模为马尔可夫决策过程（MDP），并利用带有上限置信区间（UCB）的 MCTS 实现前瞻与长期奖励估计，ProTuner 在 16 个真实基准测试中性能优于束搜索（beam search），最高提升达 3.25×，且对成本模型噪声更具鲁棒性。

ABSTRACT

We explore applying the Monte Carlo Tree Search (MCTS) algorithm in a notoriously difficult task: tuning programs for high-performance deep learning and image processing. We build our framework on top of Halide and show that MCTS can outperform the state-of-the-art beam-search algorithm. Unlike beam search, which is guided by greedy intermediate performance comparisons between partial and less meaningful schedules, MCTS compares complete schedules and looks ahead before making any intermediate scheduling decision. We further explore modifications to the standard MCTS algorithm as well as combining real execution time measurements with the cost model. Our results show that MCTS can outperform beam search on a suite of 16 real benchmarks.

研究动机与目标

• 为解决深度学习与图像处理工作负载中自动程序调度的挑战，此类手动调优成本高昂，启发式方法常无法找到最优调度方案。
• 通过采用更具鲁棒性的前瞻搜索策略，克服束搜索的贪婪性及其对噪声中间性能估计的依赖。
• 探索将带有 UCB 的 MCTS 应用于程序调度问题作为马尔可夫决策过程（MDP）的可行性与有效性。
• 评估结合学习到的成本模型与真实执行时间测量对提升调度性能的影响。
• 识别并解决 MCTS 中的性能瓶颈，如高子节点枚举开销与对成本模型的过拟合问题。

提出的方法

• 将 Halide 调度问题建模为马尔可夫决策过程（MDP），其中状态表示中间调度方案，动作表示调度优化操作（如分块、向量化、融合等）。
• 应用带有上限置信区间（UCB）的蒙特卡洛树搜索（MCTS），在搜索树中平衡探索与利用，实现前瞻与长期奖励估计。
• 通过基于模拟的滚动回溯评估完整调度方案，避免依赖易出错的中间性能预测。
• 将真实执行时间测量与学习到的成本模型结合，提升搜索准确性，降低对模型不准确性的过拟合风险。
• 通过调优探索超参数 $C_p$ 并减少模拟过程中的冗余子节点枚举，优化 MCTS 性能。
• 在每次决策后以新根节点重建搜索树，实现自适应与动态的调度探索。

实验结果

研究问题

• RQ1 带有 UCB 的 MCTS 是否能在深度学习与图像处理工作负载中，优于束搜索找到高性能的 Halide 调度方案？
• RQ2 MCTS 中通过完整调度模拟实现的前瞻能力，是否能降低噪声或不准确成本模型的影响，相比束搜索更具鲁棒性？
• RQ3 将真实执行时间测量与学习到的成本模型结合，如何影响自动调优过程的性能与鲁棒性？
• RQ4 基于 MCTS 的自动调优中，主要性能瓶颈是什么？如何缓解？
• RQ5 在不同程序与搜索深度下动态调优探索参数 $C_p$，是否能提升 MCTS 的效率与有效性？

主要发现

• 在所有 16 个真实基准测试中，ProTuner 搭载 MCTS 的性能均优于束搜索，最高提升达 3.25×。
• MCTS 中使用完整调度模拟可降低噪声中间成本估计的影响，使搜索过程比束搜索更具鲁棒性。
• 将真实执行时间测量与学习到的成本模型结合，可进一步提升性能，证明了混合评估的价值。
• 当前 MCTS 实现存在 7.5% 的成本评估开销，其中 88% 的时间消耗在枚举未使用的子节点上，表明通过优化可实现最高 8 倍的性能提升。
• 当 MCTS 运行时间过长时，会出现对成本模型的过拟合，提示需要引入早期停止或自适应搜索终止机制。
• 在不同程序与搜索深度下动态调优 $C_p$ 可提升探索效率，尤其在运行时间较短、方差较低的程序中效果更显著。