欢迎关注我的公众号 [极智视界]，获取我的更多经验分享
大家好，我是极智视界，本文来谈谈 AI训练框架的发展，是谈谈AI发展系列的第一篇。
时光斗转星移，AI 飞速发展。
满打满算，我真正进入到 AI 这个行业差不多五年的样子，这段时间其实是正处于 AI框架 和 AI算法本身快速发展的黄金阶段。也很有幸见证了一些 “古早” 的框架、算法，以及从 “古早” 到 “现代” 快速进化的过程，感慨良多。我决定在这几篇文章中谈谈我这个阶段对于 AI 的一些基础设施，如框架、算力、算法等的思考、感受 和 想法，作为一种输出，也作为一种记录，可能几年后回过头再来看，会别有一番滋味。
对于训练框架来说，记得我最开始是用 caffe、用 darknet、用 mxnet、用 keras 来进行训练，后面慢慢地还是和大多的 炼丹师 一样，逐渐收敛到了 pytorch。
把深度学习训练戏称为炼丹，把深度学习算法工程师戏称为炼丹师，其实是十分的形象。大家知道，炼丹术，是我国古代的一种传统技艺，是一种将大量灵材妙药使用丹炉火炼成丹的技术。而深度学习的训练也是一个类似的过程，需要喂大量高质量的数据 (灵丹妙药)、高算力的计算设备 (好的炼丹炉)，通过不断调整模型超参 (不同药材的搭配)，使模型更好地收敛。
不同框架的目的其实都是一样的，都是为了炼出一个在某一领域/甚至是通用领域具有很强能力的模型，所以这也导致了同一种算法可能会有多种框架的实现，这从整体的社会资源来说，肯定是属于浪费，属于重复造轮子了。比如拿 yolov7 算法来说，我接触过的就有 darknet版的、pytorch版的、tensorflow版的，这只是我接触过的，当然肯定还会有其他版本的，比如 paddle版的。这看起来就让人比较烦，不管是对于训练的人还是对于算法部署的人来说，应该都会觉得麻烦。对于训练的人来说，可能你本身已经实现了 pytorch 版本的 yolov7 了，那么你可能哪天还会想，是不是 darknet 对 yolo 的支持更加好，或者听说 paddle 性能不错，要不再试试 paddle 版的 … 而对于做算法部署的人，特别的拿做推理框架的人来说，一个很难的问题就是模型的前端解析，专业点可能叫 Parser。推理框架的前端可能像万花筒那样多，但是不管怎么样，你一般都得支持吧，那就必须要开发好几个路径的模型转换。虽然早就有了 onnx 这个 “美好的理想“，但 onnx 并不是特别的可靠。而更加可怕的是，因为有了 onnx，你可能还需要多出来一个执行 onnx 到你的 IR 之间的模型转换。而且往往这是一个多级路线，中间可能会经过多次的模型转换，比如 pytorch -> onnx -> IR、darknet -> onnx -> IR、paddle -> onnx -> IR 等。另外，有些情况下，你可能还会遇到像这种的模型转换：pytorch -> caffe、darknet -> caffe，甚至是 pytorch -> onnx -> caffe 等等。这里都转到 caffe，是因为 caffe 对于算法部署友好，比如像海思、像昇腾这种，就优先支持 caffe，意味着你需要从其他框架转到 caffe 来。而说到底，这一切的锅，可能都是来自于训练框架太多了。
那么训练框架为什么这么多呢，为什么就不能统一呢 (用户应该都希望统一吧)，来说说我的理解。 AI 框架很多时候其实就是一个公司的产品，甚至是核心产品。拿一流科技来说，他们最核心的产品就是他们的 AI 框架 –oneflow，那么他们肯定会不遗余力地去推广他们的产品。所以本身有 AI 框架的公司，或者说专门做 AI 框架的公司，本身肯定主要是用他们自己的框架，而且肯定也会自带一个生态圈。这部分人中的，觉得自研框架更牛、支持国产、情怀等，都是表层原因，更多的还是会回归到商业利益。这种现象就是导致 AI 框架泛滥的主要原因。
幸好，现在 AI 训练框架其实也有收敛的大趋势，现在大部分新的算法其实都已经是在用 pytorch 来实现。而现在来说，pytorch 其实已经成为了真正意义上最主流的训练框架了。而这要过去，要说主流的训练框架，肯定还是会提到 tensorflow 和 pytorch，而且 tensorflow 一般都放在前面。
说了这么多，再来一览一下 AI 训练框架的发展历程，我画了个图展示。
估计用过 13 年以前框架的人都 35 了吧，用过 caffe、darknet 的也都 30 多了吧。深度学习训练框架之所以从 2012 年开始发展这么快速，很大的原因来自于 2012 年诞生了 AlexNet，它是第一个在 ImageNet 上取得优异成绩的卷积神经网络，而后的训练框架的设计初衷都是围绕卷积计算。
2020 年其实很有意思，众多的国产训练框架百花齐放，跟国产芯片热遥相呼应。国产训练框架其实真的还算可以，特别是经过了几年的沉淀后。而这其中最有代表性的应该是百度的飞桨 和 一流的 oneflow 了，飞桨的算法仓十分丰富，特别像 ocr、姿态之类的算法做的很好。而 oneflow 在分布式训练方面优势明显，而这对于 AIGC 大模型来说就是刚需。而对于 mindspore，可能是属于那种华为奋力在推、奋力在建生态，但是不温不火的类型，这跟华为的一系列 AI 产品类似，如 昇腾卡、CANN 工具链 (后面我会专门说说卡 以及 推理相关的)。
训练框架的设计肯定是和算法的计算方式息息相关的。前几年大家都在 “卷” 卷积神经网络 (大家说的卷，可能就来自于卷积神经网络)，而现在的 AI 已经变了天，Transformer 已经不仅仅是 “大一统” 了，更有主导 AI 发展的趋势。而从算法本身的计算方式来说，现在流行的 Transformer 跟 “之前” 流行的卷积，差别很大。来看看卷积的计算方式，直观来说他就是一个滑窗的操作，然后实际计算会转换为卷积乘进行。
在传统的卷积神经网络，卷积的计算量会占网络整体计算量的大部分，一般这个占比会达到 90% 以上，甚至会达到 98% 以上，比如咱们熟知的 ResNet50，它的卷积计算量就占总计算量的 98% 左右；而对于另外一个常用的 backbone 网络 VGG16，则更加离谱，它的占比达到 99%。基于这种 “数据科学”，只要你对卷积计算优化的足够好，那么你的框架就不至于太差，一直以来都是如此。换句话说，因为卷积实际中会转换为矩阵乘，所以也可以说，只要你对稠密的矩阵乘优化的足够好，那么你的框架就不至于太差。这个理论既适用于训练框架也适用于推理框架，另外对于AI算力芯片亦是如此，而且对于芯片来说，这个遗留病会更加痛苦 (后面文章再说)。
毫无疑问，这种理论符合一切以卷积神经网络为基础的 AI 时代，但问题是，现在的 AI 还是 卷积神经网络的 AI 吗，这里可能需要打个问号。有同学可能会疑问，难道 Transformer 就没有很多矩阵乘了吗。Transformer 确实还是存在很多的矩阵计算，比如注意力机制中 Q、K、V 之间的矩阵乘、比如 Transformer block 中 mlp 中的矩阵乘等。有数据统计，在 Transformer 中，输入序列长度为 128 时，矩阵计算的占比约为 75%。可以看到，矩阵计算的占比虽然相比卷积神经网络的情况有些下降，但仍然还是主要的。那么问题在哪呢？问题就在于卷积神经网络是稠密的矩阵计算，而 Transformer 中穿插着很多的访存算子，让稠密不再稠密，如 reshape 之类的，可能做的事情重新排布一下数据 Layout，并没有真正需要执行计算。这就会导致在 Transformer 中，虽然矩阵计算的计算量还是占主要的，但是又因为本身对矩阵乘已经优化的足够好了，导致矩阵计算的实际执行开销可能还不如计算量占少部分的访存型算子的计算开销，这就是一个很大的问题。
从这个角度来说，当卷积神经网络不再是主导，当 Transformer 展露头角的情况下，AI 框架肯定需要做出一些相应的调整，这种调整在当下最直接来说就是对于 Transformer 的计算优化。从只需要 “照顾好” 计算密集型网络到要 “兼顾好” 访存密集型网络的过渡，肯定是个痛苦的过程。这个过程要是没做好，肯定就会加速被时代淘汰。可以预见的是，比如不再维护的 caffe、比如 darknet，在训练领域应该会首先就被淘汰了。对于caffe 来说，不维护肯定是主要原因了。另外还有一个重要的原因是，caffe 的设计对于 AIGC 大模型存在天然的缺陷，caffe 采用 google 的 protobuf 作为其模型 IR 的技术栈，而 protobuf 有一个限制是不能超过 2G，这是因为 protobuf 中有许多实现采用 32 位有符号数据表达。这个情况同样可以扩展到 onnx，onnx 的架构最初是由 caffe 的作者贾扬清设计，同样沿用了 caffe 的技术，也就是也用了 protobuf 来设计 onnx 的 IR，虽然现在有办法可以让 onnx 绕开 2 GB 的限制，但是天生的缺陷始终是存在的。对于 darknet 来说，darknet 的流行主要是因为 yolo，如果未来 transformer-base 的目标检测算法能够在精度、速度上完胜 yolo 系列，那么 darknet 就没有存在的必要了，更何况现在很多的新 yolo 算法其实都已经是基于 pytorch 来进行开放了。
从只需要 “照顾好” 计算密集型网络到要 “兼顾好” 访存密集型网络的过渡，对于经营了几年的成熟的训练框架来说，不太可能从架构上推倒重来。那么要去适应这种时代发展，可能会有几种做法，基本就是在原有大框架的基础上做一些针对性的优化，当然很明显，这些优化肯定目的在削减访存的开销。
- 算子转换 => 比如视觉 Transformer 中的图片打 patch，从数学上来说 或者 说从python单纯的计算角度来说，就是一个索引切片的过程。而咱们一般用框架去做的话，就会把这个过程用一个卷积核和步长相等的卷积来做，这样简单、明了 又 高效；
- 算子融合 => 把访存型算子和计算型算子通过数学计算先融合在一起，转化为矩阵计算。这个技巧在英伟达提供的 FasterTransformer 中，还有如 地平线的开工开物工具链中 广泛使用；
- KV Cache 优化 => 通过空间换时间思想，提高推理性能。KV Cache 优化的思想就是缓存当前轮可重复利用的计算结果，下一轮计算时直接读取缓存结果。这个技术已经在 pytorch 中使用；
- 等等；
可以看到，还是有挺多的方法可以让现代的训练框架来适应 Transformer 的计算特性。但这种 “适应” 和 当年专门为卷积神经网络来打造，待遇方面不可同日而语。没办法，历史原因让你没办法推倒重建，在原始大框架的基础上去适应新特性、去增加新特性是目前不得已 却又 现实的方案。可以预见的未来，训练框架会因为 “新” 的网络，而变得越来越臃肿。就跟 C++ 一样 …
好了，以上分享了 谈谈AI发展第一篇：AI训练框架。希望我的分享能对你的学习有一点帮助。
后面几篇的预告：
- 谈谈AI发展第二篇：AI推理框架
- 谈谈AI发展第三篇：AI训练算力
- 谈谈AI发展第四篇：AI推理算力
- 谈谈AI发展第五篇：AI编译框架
