MLSys 2021论文分析5

本文转自:知乎
作者:金雪锋

大家好,今天给大家带来MLSys 2021的一篇论文,来自Facebook AI Research的《VALUE LEARNING FOR THROUGHPUT OPTIMIZATION OF DEEP NEURAL NETWORKS》。这是本系列分享关于Tuning的第三篇论文,足以见得在深度学习领域算子Tuning的重要性。

如今机器学习的使用无处不在,神经网络的执行效率在很多场景下变得尤为重要。AI算子框架如Halide、TVM等利用算子调度(Schedule)将深度学习模型拆分成众多的算子。然而为算子寻找出优秀的调度,是一件困难的事情。自动调整(Auto-tuning)任务,希望在合法的算子空间中找出最佳的算子调度并完成实测。这篇论文通过强化学习(Value Iteration)+ CostModel(BiLSTM)来解决auto-tuning问题。

论文的主要的贡献点在于:

  1. 提出一种可以生成合法候选调度(schedule)的方式;
  2. 提出一种高精度的预测调度运行时间的Cost Model;
  3. 提出一种基于Cost Model的迭代式Value function,用于贪心地从空间中寻找最佳候选调度,比以往的方法速度提升2到3个数量级;
  4. 在这一系列的技术帮助下,优化后的神经网络比Halide提升1.5x,比TVM提升2.6x。

背景与基础

论文以Halide为基础,支持schedule的split、reorder、vectorize、parrallel、compute\_at、store\_at等操作。原始的DSL可以通过这一系列的操作转化成目标高性能schedule。与Halide相仿,论文同样将一个算子生成任务称作pipeline,而同一个pipeline的不同阶段称为stage。

图1. 支持的基本原语

图2. 由原始DSL到简单的schedule,再变换成高性能的schedule

空间生成

论文将空间生成分成5个部分:

  1. 在每一个stage中,将每一个变量v拆分成3重子循环(split);
  2. 在1 的基础上增加重排(reorder);
  3. 尝试对每个loop添加向量化选项(vectorize);
  4. 在满足生产者-消费者模型的条件下,增加临时buffer和inline(compute\_at、store\_at);
  5. 选择loop进行并行化(parallelize)。

经过这5个阶段后,空间将会变得庞大起来。论文通过一定的方式去减少空间大小,如:向量化时只考虑因子是向量本身的整数倍的情况。在常见网络中,平均每个stage有大约30万个选择。

图3 常见网络的空间大小

寻找最佳调度

论文采取基于Value function搜索最佳调度,采用了马尔科夫决策过程,MDP。

【注】关于马尔可夫决策过程、强化学习的相关学习可以参考https://www.cnblogs.com/pinard/p/9426283.html

通过一个Value Function V(s)来评估当前pipeline的“最低运行时间”。整体搜索流程如下:(1)输入一个pipeline(包含n个stages,按照拓扑序排列)、输入估价函数V(s)、设置初始化状态state0;(2)按照stage的顺序寻找当前stage最佳的调度决策,首先生成一系列的候选调度决策;(3)然后对于每一个调度决策s,计算其V(s),保留最大值vi以及对应决策si;(4)返回n个stages的决策s1...sn。N-stage的pipeline,每个stage平均调度决策为M的情况下,需要计算的总次数是M*N,而不是M^N。这样可以使得我们能以很快的速度计算完全部的结果。

图4,基于Value function搜索最佳调度

这种设计下,auto-tune的效果很大程度上取决于Cost Model的准确程度。论文提供了一种基于BiLSTM+特征工程的思路来解决估价问题。

论文中包含了三类特征:(1)与调度无关的指令特征,例如浮点数加减法、index下标计算、内存访问特征等等;(2)调度相关特征,例如向量化的数量、循环轴相关特征、inline相关特征、CPU利用率、影响性能的辅助消耗特征等等;(3)派生类型特征,如算术强度等。

论文构建了一个双向LSTM用作预测运行时间,结构如图所示。N个stages对应N个LSTM cell,最后汇总N个结果,得到整个pipeline的总结果。Cost Model的训练使用了13352个piplines,每个pipeline平均包含1184个调度。Loss的设计包含了预测值P、平均实际运行值M、重要程度C、减重程度D。

图5,Cost Model的Loss设计

【注】文章中的Cost Model是指用于预测经过完整调度的pipeline时间模型(BiLSTM),Value Function是已经应用部分调度后的pipeline可能达到的最佳时间预估函数。

图6, BiLSTM结构图

迭代式Value Function训练

为了训练出更加精确的模型V(s),论文采用迭代式优化Value Function:(1)对于一系列的pipelines以及上一个Value Function Vi-1,通过Algorithm1获得对应的最佳调度序列s0…sn;(2)对于任意一个调度sj,通过BeamSearch的方式(即Beam形式的Algorithm1),获得后面的调度tj+1…tn,同时获得s0..sj,tj+1..tn调度序列的BenchMark结果r,更新Vi(sj)。

图7, 迭代优化Value Function流程

【注】关于Value Iteration过程,可以参考https://zhuanlan.zhihu.com/p/33229439

结果对比

  1. Cost Model。Cost Model的对比项主要是平均预测误差、最大预测误差、决定系数。由实验结果可以看出,论文的结果均远超于Halide、TVM。
  2. 整体的结果对比:加速效果与搜索时间对比。

图8,Cost Model实验对比

图9,加速实验、搜索时间实验

总结

下面总结一下已经分析的三篇论文对tuning问题的思考与处理方法(TVM-Ansor[3]、Google-TPU[2]、MIT-TIRAMISU[4]、Facebook-本篇[1])差异。在tuning的基础流程上,各个研究机构是相似的,都是遵循空间生成、Cost Model训练与预测、空间中搜索寻找最优解。而在关键问题上(Cost Model和搜索),各家又各有不同。

首先在Cost Model上,从模型选择上来说,基于XLA的TPU-tune选择了GNN、TIRAMISU和Facebook使用LSTM、Ansor则采用一种特殊的packed-XGBoost。不同的机构的AI框架与设计都大有不同,各个论文给出的speedup很难从数值上直接断定“哪个模型更好”。具体模型落在具体场景的性能优劣,可能是一个工程问题。但是从设计思路上来说,它们不约而同的采取“子模块建模+求和”的方式,将一个子图转化成若干子模块来建模:TPU-tune的opcode、TIRAMISU基于loop的computation、Facebook的stages和Ansor的stages。这些子模块都可以被单独建模(如一个LSTM cell或者单个xgboost的预测值等),汇总之后就可以用来表达整个子图的最终结果。想要更精确的Cost model,则重点关注建模子模块(features)即可。从这点上来看,是跟子图整体式建模有区别的。

在搜索算法上,TPU-tune的论文侧重点在于Cost Model,理论上来说任意搜索算法都适用;TIRAMUSA采用的是树搜索Beam search或MCTS,Beam search属于一种贪心但是有一定的容忍度,而MCTS则为搜索加入了反馈,以希望修正误差;Facebook采用的是另一种强化学习的思路Value Iteration;而Ansor则采用的是Guided random search中的Evolutionary algorithm。从搜索算法而言,越来越多的机构采取各种强化学习(RL)的方式处理搜索,RL确实是一种更理想的方式,毕竟“调option、调切分”这件事某种程度上来说跟“agent在environment上探索然后根据reward去调整”实在太像。但是就目前而言,tuning框架上似乎还是没有像AlphaGo一样的模型能直接证明RL比传统的搜索方式有优势。传统的搜索方式没有模型引导[3],但是更方便研究人员加入各式各样的策略去解决搜索方向问题;RL则通过模型去学习策略[1],而能不能学到复杂有效的策略又成一个难题。

总体而言,AI框架处于一个百花齐放的时代,而其下的算子(子图)框架、tuning框架也多种多样。从学术研究上来说,很多方面都可以有进一步的尝试:优化空间生成、更精确的建模、搜索的速度与方向引导等等;从工程上来说,让tuning优美的落地不同层级的AI场景,实现更快的搜索生成更高性能的算子,同样需要工程师们做出更多的努力。

参考文献

[1]Steiner, B., Cummins, C., He, H., and Leather, H., “Value Function Based Performance Optimization of Deep Learning Workloads”, arXiv e-prints, 2020.

[2] Kaufman S J, Phothilimthana P M, Zhou Y, et al. A Learned Performance Model for Tensor Processing Units[J]. arXiv preprint arXiv:2008.01040, 2020.

[3] Zheng L, Jia C, Sun M, et al. Ansor: Generating high-performance tensor programs for deep learning[C]//14th {USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 20). 2020: 863-879.

[4] Merouani M, Leghettas M H, Pr R B D T A, et al. A Deep Learning Based Cost Model for Automatic Code Optimization in Tiramisu[D]. Master’s thesis, École nationale supérieure d’informatique, Algiers, Algeria, 2020.

推荐阅读

更多嵌入式AI技术相关内容请关注嵌入式AI专栏。
3 阅读 370
推荐阅读
0 条评论
关注数
14116
内容数
421
嵌入式端AI,包括AI算法在推理框架Tengine,MNN,NCNN,PaddlePaddle及相关芯片上的实现。欢迎加入微信交流群,微信号:gg15319381845(备注:嵌入式)
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
Arm中国学堂公众号
关注Arm中国学堂
实时获取免费 Arm 教学资源信息
Arm中国招聘公众号
关注Arm中国招聘
实时获取 Arm 中国职位信息