这个思考其实源于周五参加WAIC上组织的AI编译相关的闭门讨论的内容，观点有不少来源于现场讨论的朋友们，因为对这个主题感兴趣，我又结合自己的理解做了一些梳理。
对这次闭门讨论感兴趣的话，可以参考岚松（https://zhuanlan.zhihu.com/p/388415328）和孝强（https://zhuanlan.zhihu.com/p/388452164）两位的分享，蛮有趣的insights。以及这里的一个内容汇总：https://gitee.com/MondayYuan/WAIC-DLCompiler-MeetingMinutes
之所以关注AI框架算子层级的问题，是因为自己最近在关注AI硬件行业，也就在关注为一款新硬件提供软件支持可能面临的问题。
1. AI硬件公司，能成事非常不容易。首先是硬件要流片成功，然后是软件栈建设过得去，再之后还需要商业策略不掉链子（包括市场定位、切入策略、客户解决方案团队建设、量产方案等等）。这非常不像典型的互联网创业团队的玩法，小快灵就可以move起来，fail fast。硬件行业，有其很强的物理规律不能简单靠加人加班就超越，哪怕是天才型选手的加入，也只能确保这个过程犯错尽可能少，并不能确保大幅缩短执行周期。
2. 我的背景偏软件，至少希望确保软件这一环不会掉链子。所以围绕一款AI新硬件上面，软件栈的建设，就会有更多的思考。首先来看看一款AI硬件的软件栈应该长什么样子，下图是我的一个理解。在这套技术栈里面。有几个细节可以稍微展开一点：
1). 新硬件的device compiler和programming model的建设往往需要时间(新硬件上软件系统的建设本质上是一个bottom-up的工作，所以抽象层次越高的工作其实现复杂度往往更大，交付时间会靠后)，但上层AI应用的验证需求又是客观存在的，所以就会存在一个从Framework plugins/Libs甚至Programming Model直接打到assembler这一层的short-cut路径。因为assembler通常做的是一个字面汇编描述到二进制码的翻译，外加一些可轻可重的优化，更容易迭代式交付给上层使用。当然了，如果有可能的话，把device compiler和programming model也拆解成可以阶段式交付的milestone，尽可能让上层开发可以提前借助于高级语言来解放生产力会更有帮助（一个非常小的例子，在汇编码里做pass by value/pass by reference的区分，和在cpp里，前者还是需要消耗一些额外的心智负担，也会影响一些开发调试的效率），这属于实际操作中的考虑了。
2). 这里对接AI框架的层级我定义为"Framework plugins"，背后的思考是我在分析一些硬件公司的existing方案的时候，能够看到几乎所有新硬件厂商都有一个倾向，就是assume AI框架能够导出一个完整的静态图（包括训练和推理），用作自己软件栈的输入，来完成训练和推理能力的支持。这种作法存在的原因其实比较容易理解，新硬件为了获得相较于NV更好的性能表现，通常会添加DSA的支持，而DSA会加剧计算能力和访存以及计算kernel dispatch到加速器的gap，所以如果能全图都offload到硬件加速器上，从充分发挥新硬件性能是最理想的了。不过在eager mode广为用户接受的情况下，目前主流的AI框架想做到所有模型都导出一个干净的静态图我觉得并不容易，所以我其实prefer的是新硬件把对接AI框架的模块做成一个插件，可以集成进不同框架的后端，一方面可以更native地配合AI框架的eager mode，在python host端和加速器端做灵活的交互，另一方面，当用户在eager mode下完成交互开发，想追求更高性能的时候，可以再通过一些扩展的API来将模型可静态导出的部分尽可能导出成一张或几张静态图(相关的工作比较早也有人进行过尝试，比如MxNet Gluon)，交给加速器来执行，至于一些不易导出静态图的部分，仍然通过python与加速器交互的方式进行fallback执行。这样可以在性能和通用性上获得更好的trade-off。当然，这里有一个细节，如果引入了过多的python侧与加速器的交互，很可能会把加速器的速度明显拖慢，甚至相比NV不再有优势。所以eventually，从加速器角度，会期望不断扩大用户模型能够offload到加速器上的比例(以及通过类似cuda-Graph或pipeline执行的方式来隐藏这种开销)，但通过这种plugins的方式至少保证对接建模需求的可用性，不会因为用户模型一调整，就跑挂了，因为模型训练环节，对于交互式开发调试效率是非常看重的。
3. 在Framework plugins这个层面，想在对接主流AI框架时做到既保证在经典代表性模型上的极致性能(比如MLPerf里的模型)，同时兼顾对算法用户灵活开发的模型变种的适配程度，就需要考虑AI框架算子coverage的问题了。因为算法用户通常是使用AI框架提供的算子来完成模型搭建(以及少量的自定义算子)。TensorFlow现在有多少个算子呢，在最新的maser code base里已经接近2000个了。PyTorch也在相近的数量规模。这里更麻烦的是算子的数量会随着建模需求的演化还在不断地扩散增加。再加上AI框架版本的快速演化，这就让新硬件对接适配的时候可能出现疲于奔命的现象(这里不讨论业务场景里AI框架版本碎片化的问题，那也是一个考虑Framework plugins方案的原因)。
4. 我们在对接AI框架的时候，一定需要对接到算子层级么？我觉得未必。原理上，存在两条对接AI框架的路径。一条是直接对接到AI框架算子粒度，另一条是对接到AI框架算子粒度之下一个更细的原子IR粒度（比如HLO/Relay/ONNX/TorchScript，其中HLO只有200条左右指令)。从硬件厂商角度来说，最希望的还是能够有一套具备“封闭且可完备描述模型行为的”原子IR粒度可供对接。从用户使用角度，则更习惯于在算子层级这个粗粒度来进行模型描述，而不是使用原子IR描述模型。通过原子IR来装配组装出上层的AI框架算子，来bridge上层用户易用性和下层硬件backend对接性能/工作量的平衡。类似下图所示。
内容中包含的图片若涉及版权问题，请及时与我们联系删除
评论
沙发等你来抢