如风暴般探索，如雨滴般利用：基于坐标下降微调内核调度器的益处
机器学习
2024-07-16 v2 编程语言
摘要
机器学习模型由内核组成，内核是对张量（由自然数线性组合索引的数据）执行操作的算法。内核的示例包括卷积、转置和矢量积。实现内核有多种方式，这些实现构成了内核的优化空间。内核调度是在给定目标函数（通常是执行速度）下寻找最佳实现的问题。诸如 Ansor、Halide 和 AutoTVM 等内核优化器通过搜索启发式方法解决此问题，该方法结合了两个阶段：探索与利用。第一阶段评估许多不同的内核优化空间；后者则试图通过在同一空间内考察内核来改进最佳实现。例如，Ansor 通过草图生成进行探索，并利用进化算法来利用最佳草图。在本工作中，我们展示了通过将 AutoTVM 算法 Droplet Search 引入 Ansor 的探索阶段，在提高内核质量的同时减少 Ansor 搜索时间的潜力。该方法包括限制 Ansor 探索的样本数量，选出最佳样本，并使用坐标下降算法对其进行利用。将此方法应用于 Ansor 生成的前 300 个内核时，我们通常比让 Ansor 分析 10,000 个内核能以更短时间获得更好的内核。该结果已在运行于四种架构上的 20 个著名深度学习模型（AlexNet、ResNet、VGG、DenseNet 等）中复现：AMD Ryzen 7 (x86)、NVIDIA A100 tensor core、NVIDIA RTX 3080 GPU 和 ARM A64FX。包含此组合方法的补丁已于 2024 年 2 月在 Ansor 中获得批准。作为此搜索方法通用性的证据，2024 年 6 月向 TVM 的 MetaSchedule 提交了类似的补丁，并取得了同样良好的结果。
引用
@article{arxiv.2406.20037,
title = {Explore as a Storm, Exploit as a Raindrop: On the Benefit of Fine-Tuning Kernel Schedulers with Coordinate Descent},
author = {Michael Canesche and Gaurav Verma and Fernando Magno Quintao Pereira},
journal= {arXiv preprint arXiv:2406.20037},
year = {2024}
}
备注
22 pages, 19 figures, original work