Pruner：一种加速张量程序调优的“先起草后验证”探索机制
机器学习
2025-04-10 v3
摘要
张量程序调优对于深度神经网络的高效部署至关重要。基于搜索的方法在为特定硬件自动寻找高性能程序方面展现出了可扩展性和有效性。然而，由于探索机制由准确但学习缓慢的代价模型引导，搜索过程通常效率低下，需要数小时甚至数天才能发现最优程序。同时，在一个平台上训练的代价模型无法在线无缝适应另一个平台，我们将其称为跨平台在线无感知。在本工作中，我们提出了 Pruner 和 MoA-Pruner。Pruner 是一种“先起草后验证”的探索机制，可加速调度搜索过程。Pruner 没有将复杂的已学习代价模型应用于所有探索的候选者，而是通过引入简单的基于符号的分析器（起草模型）来起草小规模的潜在候选者，然后通过已学习代价模型识别最佳候选者。MoA-Pruner 引入了动量在线适应策略以解决跨平台在线无感知问题。我们将 Pruner 集成到 TVM 中，并在三个基于 GPU 的平台上进行了广泛实验。结果显示调度搜索时间显著加速。在在线调优场景中，与 Ansor 相比，Pruner 和 MoA-Pruner 分别实现了 和 的平均加速。在离线调优场景中，与 TenSet 和 TLP 相比，Pruner 分别实现了 和 的平均加速。此外，在 TensorCore 上，与 MetaSchedule 相比，Pruner 实现了 的平均加速。
引用
@article{arxiv.2402.02361,
title = {Pruner: A Draft-then-Verify Exploration Mechanism to Accelerate Tensor Program Tuning},
author = {Liang Qiao and Jun Shi and Xiaoyu Hao and Xi Fang and Sen Zhang and Minfan Zhao and Ziqi Zhu and Junshi Chen and Hong An and Xulong Tang and Bing Li and Honghui Yuan and Xinyang Wang},
journal= {arXiv preprint arXiv:2402.02361},
year = {2025}
}