AI框架的未来将何去何从，跟着MindSpore一起探讨其趋势

自MindSpore在3月底开源至今已有6个月，一直想找个时间帮助大家更好的理解和应用MindSpore，本篇文章主要是分析AI框架的发展趋势以及介绍一下MindSpore的构想，希望有助于大家去了解MindSpore，后续计划开设AI专栏，更新一些列文章，初步主题规划为下图所示，欢迎大家一起探讨。

AI框架的发展个人认为可以大致分为3个阶段：

第一阶段的代表是torch、theano以及Caffe，奠定了基于Python、自动微分、计算图等基本设计思路。

第二阶段的代表是TensorFlow、PyTorch，一个通过分布式训练、多样的部署能力在工业界广泛使用，另外一个是提供动态图能力，在灵活性上吸引了大量的研究者和算法工程师。

第三阶段的方向是什么？目前看还没有定型，我们看到Google也有多条技术路径在探索，包括TF2.0，JAX、MLIR、Swift for TF等；有些在做动态图和静态图的统一，有些在构建统一的IR基础设施，还有些在探索新的表达方式等，总之还是百花齐放的状态。

不管AI框架未来的发展方向是什么，AI框架的驱动力还是相对清晰，我们认为AI框架的驱动力主要是5个-“ABCDE”：

Application+Bigdata：AI框架的Application就是AI的模型和算法，Bigdata就是AI模型需要的数据

Chip：AI的芯片，代表AI的算法发展

Developer：AI的算法研究者和算法工程师

Enterprise：AI部署和AI责任

下面，我们想通过对AI框架的驱动因素的分析来讨论一下AI框架的发展方向。

· 模型和数据的规模和复杂度持续提升

今年5月，OpenAI发布GPT-3模型，千亿参数量，数据集（处理前）45T，训练一次的成本号称接近500万美金，超大模型不仅仅是算法和算力成本的挑战，同时对AI框架的挑战也是巨大，主要体现在三点，性能墙（内存、通信、计算利用率）、效率墙以及精度墙：

性能墙：

大模型下，GPU和NPU的单卡(内存一般只有32G)肯定装不下整个模型，传统的数据并行是不够的，未来内存复用、混合并行（数据并行/模型并行/pipeline并行）等技术的使用将会是常态；

同时，混合并行的切分策略寻优是很困难的，不同的切分产生的通信量差异巨大，计算利用率也很不一样，比如pipeline并行往往在计算利用率存在较大的挑战，算子切分的并行在通信量的挑战会很大，这些都需要AI框架来支持；

最后在大规模并行训练系统中，当性能要求越来越高时，数据预处理也会成为一个瓶颈，比如我们在ResNet50性能调优时发现，当单Step跑到17~18ms的时候，host数据处理的时间就跟不上了。

效率墙：

混合并行策略如果让算法工程师来手工确定的，这个门槛就会很高，既要懂算法，又要懂系统，怎么做到自动并行是关键。

精度墙：

大规模模型训练天生就是一个大Batchsize的训练，怎么更好的收敛，达到精度要求，都是一个大的挑战。

除了上面讲的三点外，超大规模的训练还面临其他的挑战，比如大规模集群的可用性、尾部时延、动态调度等等。

· 框架的负载从单一的深度学习模型向通用的张量可微计算演进

目前主要看到三个方向：

1）DNN与传统机器学习结合，比如深度概率学习、图神经网络等，这一块基本上业界的AI框架都已经支持。

2）AI与科学计算结合，看到业界在探索三个方向。

方向一：AI建模替代传统的计算模型，这个方向刚刚起步，进展还不是很多；

方向二：AI求解，模型还是传统的科学计算模型，但是使用神经网络的方法来求解，这个方向已经有一定的进展，目前看到不少基础的科学计算方程已经有对应的AI求解方法，比如PINNs、PINN-Net等，当然现在挑战还很大，特别是在精度收敛方面，如果要在AI框架上AI求解科学计算模型，最大的挑战主要在前端表达和高性能的高阶微分；

方向三：使用框架来加速方程的求解，就是科学计算的模型和方法都不变，但是与深度学习使用同一个框架来求解，其实就是把框架看成面向张量计算的分布式框架。

3）计算图形相关的，类似Taichi这样的语言和框架，提供可微物理引擎、可微渲染引擎等。

AI芯片主要是两种形态，一种是基于SIMT的GPU，另外一种是类SIMD的NPU，随着NVIDIA A100的发布，我们看到两种芯片架构在相互借鉴，相互融合：

大算力还是要靠SIMD，SIMT的灵活性虽然强，但是芯片的面积/功耗挑战太大，Tensor Core的规模越来越大；

片内片间高速互联；

多硅片，大内存封装，x倍体积（特别是推理芯片）。

AI芯片的持续演进对AI框架也提出了许多关键的挑战：

优化与硬件耦合度提升，图算编译一体化：图层融合优化已经趋于收敛，需要和算子层联动优化，子图级和算子级的界限打破，基于Graph tuning的整体优化是目前的热点;

多样的执行方式：图下沉模式和单算子调用模式混合，不同的情况下，可能采用不同的方式；

可编程性挑战变大：由于有了大量的SIMD加速指令，不在是单纯的SIMT编程模型，异构编程的挑战变大。

· 新的AI编程语言尝试突破Python的限制

目前看到有代表性主要是Julia和Swift for TensorFlow

Julia的历史较早，目前在科学计算领域已经有不错的应用基础，当前结合这些领域的积累，逐步进入AI和深度学习领域，宣称的特点是动态语言的特征，静态语言的性能，和Python比，比较有特色的地方包括：

1）类MATLAB的张量原生表达

2）多态的能力+动态类型推导/特化

3）IR的反射机制

所谓的IR反射机制是指，Julia的IR可以让开发者做二次加工，以Julia的机器学习库Flux+Zygote为例：

Flux是一个Julia上的扩展库，定义基本的算子，如conv，pooling，bn等等，方便开发者基于Flux去定义机器学习模型；

Zygote基于Julia提供的IR反射机制实现了源到源的自动微分，在这个过程中，是不要改变Julia本身的。

与Julia相比，Swift for TensorFlow则完全是另外一套思路，他试图从工业级开发的角度来找到差异化，包括静态类型、高性能、与端侧结合的易部署等。

从目前看尽管Julia和Swift都有些特色，但是短期内还很难撼动Python在AI领域的地位。

· AI编译器成为框架的竞争焦点

AI编译器目前看到有三个方向在发力。

第一类是致力于动静态图统一的AI JIT能力，比如TorchScript、JAX等；

第二类是偏向于面向芯片的优化，如XLA、TVM；

第三类是想做AI编译器的基础设施，MLIR希望是提供MetaIR，成为构建AI编译器的基础，Relay/TVM则是想把编译器接口开放出去支撑第三方框架。不过我们看到这三个方向都还存在比较大的挑战。

AI JIT：

Python的灵活性和动态性太强，动态shape好搞定，但是动态Type就难了，更不用说Python里面有大量非常灵活的数据结构的用法，一个动态图想无缝的转到静态图的确不是一件很容易的事情。

编译加速：

目前主要还是基于Pattern的算子融合和基于模板的算子生成技术，这里的瓶颈在于算子生成技术，因为算子融合的组合条件太多，基于模板无法做到完全枚举，要提升整个编译加速的泛化能力，下一个需要攻克的技术就是不需要模板的自动算子生成技术。

AI编译器的基础设施：

MLIR在设计理念上确实是先进和宏大的，MLIR目标通过Dialect的扩展来支持各种领域编译器，包括AI编译器（图层和算子层）。

从进展看，MLIR在TF Lite的应用最快，主要是用作模型转换工具。我们需要思考的是，MLIR带来的好处究竟是什么？我想MLIR本身并不会打来框架性能和易用性的提升，主要还是重用和归一，如LLVM的基础结构CFG+BB、大量的优化等，问题是这些重用是否有利于AI图层和算子层的编译。

个人的观点，对于算子层来说MLIR+LLVM肯定是适用的，对于图层就未必，LLVM虽然统一了许多编程语言的编译器，但是优势领域还是集中在静态编译领域，在JIT和VM领域，LLVM的优势并不明显，CFG+BB这样的基础架构未必适合需要自动微分/JIT这样的AI图层编译器。

从AI部署看，我们看到三个趋势：

大模型在端侧的部署，特别是语言型的模型

端云协同的场景逐步开始应用，这里主要有两类，一类是云侧训练，端侧做在线finetuning的增量学习，第二类是联邦学习。

AI无处不在，在IoT那些资源受限的设备上进行AI模型的部署。

目前看AI框架主要是两个挑战：

AI框架在云侧和端侧能否做到架构上的统一，这里的统一主要是指IR格式，只有这样才能做到云侧训练出来的模型，在端侧可以平滑的进行增量学习，并方便地进行联邦学习。

AI框架能否做到可大可小，比如小到K级的底噪，能够内置到IoT的设备里。

AI的责任涉及的范围非常广，包括安全、隐私、公平、透明、可解释。

作为AI业务的承载，AI框架需要具备使能AI责任的能力，目前AI框架需要解决的几个挑战：

对于AI责任的各个方面，缺乏通用的分析方法和度量体系，缺乏场景感知的自动化度量方法。

AI模型鲁棒性、隐私保护技术、密态AI在实际场景下对模型性能影响较大。

AI的可解释性还缺乏理论和算法支撑，难以给出人类友好的推理结果解释。

MindSpore作为一个新兴的框架，大家经常问的一个问题是他的差异化在哪里？结合前面分析的AI框架的驱动力和挑战，我们希望MindSpore在5个方向上引领AI框架的演进：

· Beyond AI：从深度学习框架演进到通用张量可微计算框架

MindSpore会提供一个更通用的AI编译器，为支持更多的应用类型提供可能性

· 分布式并行原生（Scale out）：从手工并行演进到自动并行

MindSpore在支持大规模训练方面除了性能和扩展性好外，更重要的是想降低这一块的门槛，让分布式训练变得更加简单

· 图算深度融合（Scale up）：从图算分离优化演进到图算联合优化

MindSpore提供图算联合优化、自动算子生成、深度图优化等关键技术，充分发挥AI芯片的算力

· 全场景AI：从端云分离架构演进到端云统一架构

MindSpore的云侧框架和端侧框架做到统一架构，方便云端训练/端侧finetuing或者端云协同的联邦学习

· 企业级可信能力：从消费级AI演进到企业级AI

MindSpore会内置对抗性训练、差分隐私、密态AI、联邦学习以及可解释AI等能力。

当然，软件架构都是持续演进的，很少有技术能够做到独门绝技，MindSpore也希望和业界其他框架在正确的方向上一起进步。

同时，MindSpore社区也对外发布了10大课题，邀请开发者一起参与创新。具体参见链接：github.com/mindspore-ai

MindSpore主要分为四层：

MindSpore Extend：

这个是MindSpore的扩展包，现在的数量还比较少，希望未来有更多的开发者来一起贡献和构建

MindSpore的表达层：

MindExpress是基于Python的前端表达，未来我们也计划陆续提供C/C++、Java等不同的前端；

MindSpore也在考虑自研编程语言前端-仓颉，目前还处于预研阶段；同时，我们内部也在做与Julia这些第三方前端的对接工作，引入更多的第三方生态。

MindSpore的编译优化层：

MindCompiler是我们图层的核心编译器，主要基于端云统一的MindIR实现三大功能，包括硬件无关的优化（类型推导/自动微分/表达式化简等）、硬件相关优化（自动并行、内存优化、图算融合、流水线执行等）、部署推理相关的优化（量化/剪枝等）；

MindAKG是MindSpore的自动算子生成编译器，目前还在持续完善中。

MindSpore全场景运行时：

这里含云侧、端侧以及更小的IoT。

同时MindSpore也提供了面向AI责任的MindArmour，以及面向数据处理/可视化/可解释的MindData。