华为开源的深度学习框架 MindSpore

点击蓝字关注我们

扫码关注我们

公众号 : 计算机视觉战队

扫码回复：华为开源，获取下载链接

作者：杨军
来源：知乎

纯技术讨论，不涉及其他，部分我拿不准的地方，会直接以(?)标识出来，欢迎菊花厂同学来指正解惑。

华为的运营同学蛮专业的，在回答里介绍了一些比较重要的技术细节，哪怕不看code，对于做这个方向的同学，也大概也能捕捉到里面的一些core concepts，YongCHN同学也对MindSpore的auto parallel部分有一个小调研，也可以供感兴趣的同学了解。

最近关于AI框架开源的讨论比较多，包括Jittor和MegEngine以及MindSpore，所以我想还是尽量能够提供一些有额外信息量的输入。

1.从框架的设计原则上，个人认为MindSpore还是蛮中规中矩的，能够看到明显从TF，PyTorch里分别借鉴了两家的经验（不要忽略MxNet的Gluon，其实也是把静态图和动态图结合在了一起，可惜这两年日渐势微，其实也从一个侧面反映出引擎技术核心以外因素对AI框架推广的重要性，参见这里的一个回答，更不要忽略了Theano这样的已经deprecated框架里早就有的静态图的渊源以及当年还小引领过风骚的Chainer）。都是做这个方向的，就不客套了，官宣里提到的将python的模型描述JIT编译成计算图的思想，在Google的JAX/auto-graph以及PyTorch的TorchScript相关的工作里比较早就已经touch了，不过有决心去从头build并落实，这个还是值得点赞。

2. 如YongCHN同学关于auto-parallel的分析里所说，在TF里，为了在python层加入auto-parallel的支持，同时还兼顾TF 2.0的eager mode，在python代码里引入了大量的 if (eager_mode) 的判断，让整个代码的可维护性受到了不小的影响，而TF为了加入Distribution Strategy的支持，对其现有的python构图代码也做了大量的修改和插桩，其实让整个框架的python层实现复杂了不少。而MindSpore的整个code base，在auto-parallel的实现上能够感觉到还是清晰了不少，引入了ANF这个IR层（关于ANF的设计理念之前自己并不了解，感谢@叶子豪 同学的输入，稍微补了些课，相关的background材料可以参见这里和这里 ，也对照着重新看了一下ir/anf.h里的定义，认为设计动机确实如ANF的背景材料所说是为了简化source-level的变换复杂性，关于细节，这方面不是我的expertise，如果有了解的同行能够share更精准的理解就太好了），把分布式策略的工作有不少沉到了以ANF IR这一层，避免把python API层改得太惨，这个设计认为是更合理的分拆，也确实在没有历史包袱的框架里更容易make这样的design choice。站在前人的肩膀上总是能够也应该有更好的创新。

3.在我的理解中，MindSpore的架构层次大体上可以拆分为

• ME(Mind Expression)

• Graph Compiler

• 算子库(GPU/CPU平台上的比较常规，Ascend上是对CANN算子库的封装，只不过在代码库里只看到了少量的CPU算子和GPU算子的实现，Ascend的算子实现则没有看到源码，没有理解错的话，是通过拼装JSON串作为输入再通过一层bridge获取到对应算子的执行码，对应的具体执行码目前并没有开源出来，说得不对请花厂同学指正)

• Runtime

这几个大层次。

从git repo的组织来看，ME、算子库(部分算子库目前以闭源.so的方式提供)基本都放在MindSpore的主体repo里，Graph Compiler和Runtime放到了GE(Graph Engine)的repo里。按官网的doc似乎主要的fusion工作是在GE里完成的。但是呢，一些涉及到算子fusion的逻辑又散在了MindSpore主repo里，同时也在GE的repo里看到了一些fusion的pass。这让我稍微有些curious为什么是这样组织。从架构设计上，会感觉fusion相关的逻辑应该统一放到GE的repo里才自然。

目前我能够推测出的一个原因是GE主要target Ascend硬件，对于GPU相关的fusion不属于GE的范畴，所以在 MindSpore的repo里针对GPU的fusion做了一些比较直接的工作(?)。

4. 整个执行flow从代码结构来看是这样的，

• 用户完成python层的模型构图

• 调用python层的train API(mindspore/train/model.py)以后，会通过pybind11的接口触发 ExecutorPy::Compile()，在这个函数的实现里，会将用户python模型描述的AST解析成ANF的格式(了解JAX和auto-graph的同学就能看到相似之处了)

• 完成ANF graph的构建之后，剩下的事情就比较自然了，在这个graph上bla bla bla做一系列的transformation，直到生成一个编译好的byte string，并将这个编译好的结果序列化保存下来

• 再通过一个pybind11的接口触发ExecutorPy::Run()，对上面的编译结果进行实际运行。

这个流程其实蛮自然的，至少在JAX和TF auto-graph项目里都有类似的作法，在PyTorch里其实也有类似的尝试，核心的难点我认为是在于Python的语法太灵活了，于是在不太起眼的AST2XXX(XXX可能是TF GraphDef，可能是JAX里的HLO graph也可能是MindSpore里的ANF graph)这个步骤可能会出现大量的corner case。PyTorch通过python解释器和PyTorch核心交互调用的方式（算是一种trace的方式）来牺牲一定性能但根本性的避免了这个陷阱，而JAX也好，auto-graph也好，以及MindSpore也好，在想获取动静结合的组合优势的同时，也需要pay for对应的engineering cost。

5. 当前开源出的结果里，model zoo还是偏少，这其实不难理解。AI框架演化到现在这个阶段，我个人认为，哪怕是为一个新的硬件平台，从头开发一个AI框架的核心idea也都已经基本是显学了（扒扒TF/PyTorch的code，读读相关的system design的paper，有精力兴趣的话再去刨一下Chainer/Theano/MxNet相关工作的材料，主流AI框架的核心idea也就拿到了），难点其实在于如何结合具体的硬件平台以及业务场景将构成框架的各个pieces有效地组织起来，获得好的trade-off。这背后有着大量的系统设计和工程架构的要求。之前Jittor刚出来那会儿，有商汤的朋友在朋友圈里感叹说两三个可以做一个AI引擎，我对这个观点持很大的保留意见，如果说两三个精干的人搭一个能够跑通几个public model的AI引擎原型甚至在个别模型上有不错的性能表现我没有任何疑问，但从跑通几个模型到达到工业级生产水准，我认为还有着巨大的宏沟需要跨越(性能的鲁棒性，计算结果正确性，训练到推理链路的完整和流畅性，从云到端不同设备覆盖的广阔性，针对不同场景的架构设计的妥协等等，背后都对应着比较巨大的工程资源的投入)，这里先不展开了。

6. 关于材料中提到的整图下沉的技术点，确实需要承认，当能够自已掌握从硬件到软件全链路时，系统的整体优化空间变得就更为广阔，整图下沉就是一个例子了。在DaVinci的芯片，配备了ARM的cpu，就使得加速器上也具备了通用计算的处理能力，四个CPU core跑一个OS，四个CPU core跑一些不适合DaVinci的Scalar/Vector/Cube计算特点的逻辑（比如一些条件控制逻辑甚至IO预处理逻辑等）是一个听起来make sense的技术点。这里可能有一个小的陷阱是，在AI作业里，有些时候，对CPU的算力需求可能会比较大（比如IO或网络通信），甚至超过DaVinci的片上ARM CPU的算力capacity，这种情况下，MindSpore在DaVinci芯片上的性能表现就更值得关注，也是更考究系统设计鲁棒性的地方了。系统设计没有magic，每一个晶体管，每一个cycle，用到这个地方以后，就不能再用到另一个地方了。软件系统的设计也类似。当然了，花厂有很强的硬件自研能力，如果说未来的芯片会引入IO core或对AI芯片上的通信能力进行强化，进一步offload对片上通用CPU算力的需求，让整个AI作业的执行过程尽可能不要跳出片子，我倒是也不会意外。不过另一个设计维度也许是，不一定要求作业执行过程完全不能跳出片子，而是通过一些机制来保证当作业部分执行环节跳出片子，也能进行性能补位（类比GPU设计理念里通过高Occupancy来确保可并行调度的计算需求足够高，来缓解消隐访存latency开销的作法）也许是一个可以考虑的点。Again，这里的难点在于设计trade-off的抉择，比如说作业完全在片子内能cover性能是100分的话，如果不能cover，是期望对这类作业退化到90分呢，还是退化到80分但同时反过来把那10分的余量用于确保片内完全能cover作业的性能达到120分？这需要结合workload趋势理解，软件栈系统把握能力，以及硬件架构设计，用至最终的硬件工艺实现来进行一个全局的把握，是一个值得玩味也迷人的复杂系统设计问题。

7. 对于MindSpore背后的计算密集算子相关AI编译的工作，其实还有着比较强的兴趣想进一步了解，即对应于TBE(以及TIK(?))背后的工作，不知道未来是否也会开源出来。