Pruner: A Draft-then-Verify Exploration Mechanism to Accelerate Tensor Program Tuning
Fuente:
arXiv
Guardado en:
| Autores principales: | , , , , , , , , , , , , |
|---|---|
| Formato: | Preprint |
| Publicado: |
2024
|
| Materias: | |
| Acceso en línea: | |
| Etiquetas: |
Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
|
| _version_ | 1866912317830070272 |
|---|---|
| author | Qiao, Liang Shi, Jun Hao, Xiaoyu Fang, Xi Zhang, Sen Zhao, Minfan Zhu, Ziqi Chen, Junshi An, Hong Tang, Xulong Li, Bing Yuan, Honghui Wang, Xinyang |
| author_facet | Qiao, Liang Shi, Jun Hao, Xiaoyu Fang, Xi Zhang, Sen Zhao, Minfan Zhu, Ziqi Chen, Junshi An, Hong Tang, Xulong Li, Bing Yuan, Honghui Wang, Xinyang |
| contents | Tensor program tuning is essential for the efficient deployment of deep neural networks. Search-based approaches have demonstrated scalability and effectiveness in automatically finding high-performance programs for specific hardware. However, the search process is often inefficient, taking hours or even days to discover optimal programs due to the exploration mechanisms guided by an accurate but slow-learned cost model. Meanwhile, the learned cost model trained on one platform cannot seamlessly adapt online to another, which we call cross-platform online unawareness.
In this work, we propose Pruner and MoA-Pruner. Pruner is a "Draft-then-Verify" exploration mechanism that accelerates the schedule search process. Instead of applying the complex learned cost model to all explored candidates, Pruner drafts small-scale potential candidates by introducing a naive Symbol-based Analyzer (draft model), then identifies the best candidates by the learned cost model. MoA-Pruner introduces a Momentum online Adaptation strategy to address the cross-platform online unawareness.
We incorporate Pruner into the TVM and conduct extensive experiments on three GPU-based platforms. Results show considerable speedup in schedule search time. In online tuning scenarios, Pruner and MoA-Pruner achieve an average speedup of $2.6 \times$ and $4.82 \times$ compared to Ansor. In offline tuning scenarios, Pruner achieves an average speedup of $4.75 \times$ and $4.05\times$ compared to TenSet and TLP, respectively. Furthermore, Pruner achieves an average speedup of $4.08 \times$ compared to MetaSchedule on TensorCore. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2402_02361 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | Pruner: A Draft-then-Verify Exploration Mechanism to Accelerate Tensor Program Tuning Qiao, Liang Shi, Jun Hao, Xiaoyu Fang, Xi Zhang, Sen Zhao, Minfan Zhu, Ziqi Chen, Junshi An, Hong Tang, Xulong Li, Bing Yuan, Honghui Wang, Xinyang Machine Learning Tensor program tuning is essential for the efficient deployment of deep neural networks. Search-based approaches have demonstrated scalability and effectiveness in automatically finding high-performance programs for specific hardware. However, the search process is often inefficient, taking hours or even days to discover optimal programs due to the exploration mechanisms guided by an accurate but slow-learned cost model. Meanwhile, the learned cost model trained on one platform cannot seamlessly adapt online to another, which we call cross-platform online unawareness. In this work, we propose Pruner and MoA-Pruner. Pruner is a "Draft-then-Verify" exploration mechanism that accelerates the schedule search process. Instead of applying the complex learned cost model to all explored candidates, Pruner drafts small-scale potential candidates by introducing a naive Symbol-based Analyzer (draft model), then identifies the best candidates by the learned cost model. MoA-Pruner introduces a Momentum online Adaptation strategy to address the cross-platform online unawareness. We incorporate Pruner into the TVM and conduct extensive experiments on three GPU-based platforms. Results show considerable speedup in schedule search time. In online tuning scenarios, Pruner and MoA-Pruner achieve an average speedup of $2.6 \times$ and $4.82 \times$ compared to Ansor. In offline tuning scenarios, Pruner achieves an average speedup of $4.75 \times$ and $4.05\times$ compared to TenSet and TLP, respectively. Furthermore, Pruner achieves an average speedup of $4.08 \times$ compared to MetaSchedule on TensorCore. |
| title | Pruner: A Draft-then-Verify Exploration Mechanism to Accelerate Tensor Program Tuning |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2402.02361 |