ai深度学习编译器工程师需要哪些技术栈?
大概两年半前入了这个坑,就一直在思考这个问题。
深度学习编译器需要了解的知识面大的恐怖,理想的情况是:
像算法工程师一样了解各个模型的结构
像深度学习框架工程师一样了解各个功能模块的实现
像编译器工程师一样了解编译原理
像高性能工程师一样了解算子编写及优化方法
像硬件工程师一样了解体系结构
还有基础的编程语言,数据结构算法,操作系统原理....
这显然不现实,你不可能在每个子领域都和专注这个子领域的人懂得一样多,这也是困扰我挺长一段时间的问题,学的再多,似乎也只是沧海一粟,这个清单要仔细地列出来,远比楼上的高赞回答长得多。
于是我只能退一步想,技术栈的目的是为了解决问题,那么问题就可以变成,如何分配有限的时间,最大化积累的解决问题能力。找到最有学习价值的内容,也许比学习本身更重要。知道这么是暂时不需要了解的,比知道什么是需要了解的更宝贵。具体来说,我把我的方法大致总结为“抓住重点,构建框架,灵活深入”。抓住重点,指的是每个领域都有一些入门的,必须掌握的,不懂就完全看不了其他的知识。构建框架,指的是在了解了基础之上,对知识有框架性的认识,知道这个领域要解决哪些问题,哪些问题大致有哪些方法,哪些方法可以去哪里看到。灵活深入,以前两点作为基础,也是最后要达到的效果,既然我们难以做到直接充分掌握各个领域的技术细节,那我们就应该追求,当工作中分配到某个任务的时候,基于任务的需求,快速地获得这项任务需要的知识。
先挖个坑,回头再写各个领域具体哪些是入门重点,框架如何,快速深入的方法是什么。
这里插播一个声明,我说的这些主要是针对真正的具体工作需要积累哪些知识和能力。至于面试会问啥,还是多搜搜面经吧,毕竟国内技术面试脱离实践这个问题也不是一两天了。
先填一下算法层的坑吧。算法工程师要解决的问题是,端到端地解决具体的问题,具体来说就是从用户给出的原始数据到用户需要拿到的最终数据,比如如果是一个对图像的分割任务,那就是从原始的图片数据集到最后标号分割框的输出图片。而对于深度学习编译器来说,要解决的问题是对于一个已知结构的模型,找到最优的计算实现,这个链路,是比算法工程师要短很多的。更进一步来说,深度学习编译器只需要理解一个模型长什么样。至于模型有什么用,模型为什么长这样,怎么训练出这个模型,怎么获得这个模型想要的输入,并不重要。
举例来说,算法工程师要知道怎么处理训练数据不平衡的问题,这跟深度学习编译器就没什么关系。又比如NLP里面,输入数据要怎么做embedding,也不重要,又比如,bert在最后一层加上不同的尾处理,就可以拿来解决不同的问题,这也不重要。
那理解一个模型需要长什么样,具体需要懂啥呢?最基本的当然是各个算子的语义,这个没什么好说的,应该都能掌握,对自己要求高一点的话,可以对着ONNX文档一个个看过去。至于模型层面,模型的数量是浩如烟海的,但是常见的模块来来回回就是那些,深度学习编译器的很多优化都是基于这些模块来的。所以我会建议大家把重点放在积累理解常见的模型模块,比如conv-bias-batchNorm,比如transformer的encoder/decoder blocker。对于整个模型来说,不需要看太多了,看一下最经典的几个网络,也就是MLPerf跑的那几个就好了,resnet,yolov3,bert。
在工作的时候,你完全可能会遇到一个相对陌生的模型,所以你需要具备快速把握一个模型的能力,这里推荐大家积累的具体能力有两个。一个是通过netron阅读模型的onnx能力。在你掌握上一段说的理论内容之后,你可以下载一个模型的onnx文件试着打开看看,刚开始也许还是蛮痛苦的,看过的应该都懂,但你应该积累出的能力是,打开一个陌生模型的onnx文件,很快能够看出这个模型大致包括哪些基本结构模块,进而理解优化要点有哪些。第二个能力是阅读模型代码能力,比如NLP方面的HuggingFace, CV方面的OpenMMlab,有些时候你需要了解一些模型细节,直接进到代码里面看是更方便的。
今天来填填体系结构的坑,目前还不熟悉NV生态之外的硬件,就以NV作为代表来说了,以及,我会把对CUDA的掌握放在这个部分。这个技术点的重要性取决于你是否希望走性能优化这条路,要理解GPU的性能,就必须深刻地理解体系结构。计算机里面讲的体系结构,我理解就是,对于软件设计有帮助的,不涉及具体电路设计的所有硬件细节。在NV的生态里,CUDA其实就是GPU体系结构的软件抽象。所以,了解CUDA是了解GPU的窗口,而进一步了解GPU的体系结构,其实就是在更深刻地理解CUDA。进一步说,认知建立的步骤应该是,先从功能性的角度过一遍CUDA,有一个整体性的,框架性的认知;然后学习GPU的体系结构,能否充分地发挥硬件性能,本质上在于能否充分利用硬件的并发性和局部性,所以对GPU体系结构的掌握,需要知道每一个可以提高并发性和局部性的硬件特性;最后一个步骤是带着对硬件的理解,再回到CUDA,理解如何通过CUDA来利用各种硬件特性,或者说规避各种性能陷阱。这里推荐三个学习资源,分别对应上述说的三个步骤,第一个是NV官方的CUDA Sample, 第二个是景乃峰,柯晶和梁晓峣的《通用图形处理器设计》,第三个是官方的CUDA Programming Guide。如果还想继续深入的话,可以继续学习PTX,如果还想学习指令集(SASS),也许可以看一下开源的逆向工程做的汇编器,比如商汤最近开源的CUAssmbler。
今天来填一下高性能计算库的坑,所谓高性能计算库,指的是面向一类计算问题的kernel库,比如cublas, cutlass, cutensor, cub(这里我没有提cudnn,cudnn最近几个版本的演化已经有点从算子库向一个处理局部图的框架走的感觉了)。开发高性能算子库大概需要三层能力。
第一层是如何针对特定的问题写出高效的kernel,以GEMM为例,不同的problem size(M,N,K),数据类型,layout在不同的卡上遇到的瓶颈,和对应的最优实现方式都是不一样的。所以你首先得掌握使用各种工具(比如nsight compute)进行性能分析,找到性能瓶颈的方法。然后你得知道算子实现的各种“玩法”,比如各种流水编排,各种优化访存pattern的手段。最后整个流程积累下来,你是应该对这个算子有性能建模的能力的。
第二层是让算子库变得模块化和工程化,一个算子库是面向一类计算问题的软件,不是简单的kernel的集合,作为一个成熟的软件工程,高性能算子库需要尽可能将各种“玩法”变成像积木一样可以自由拼接的模块,最大程度地实现代码复用,同时也方便对整个代码库的管理。除此之外,各种activation操作,add bias,sclase(可以统称epilogue)也需要模块化,并允许被fused到gemm/conv的kernel里面,这对上层的图优化是一个非常重要的支撑。再进一步,代码是死的,团队是活的,从算子库团队而言,应该是一个高性能算子的高效的工厂,新的需求转化成新的功能,或者新的kernel,新的kernel投入测试,解决回归测试中发现的性能问题,每一环的效率都有很大的讲究。
第三层是帮助用户解决算子选择的问题,这个问题有点像第二层做完之后又回到第一层的问题,就是面对一个具体的用户case,怎么找到那个最优的实现。这一层其实已经算是高性能算子库和框架/DL编译器的交界问题了,有些算子库,比如cublas,会有heuristic来一定程度上解决这个问题,而cutlass相比之下就是一个更干净的模板库,并不打算解决这个问题。
让我们回到深度学习编译器研发工程师的视角。深度学习编译器和高性能算子库有两种关系,
第一种关系是深度学习编译器本身就是高性能更算子库,深度学习编译器往往会有自动生成高性能算子的模块(codegen),或者有一些自己写的高性能kernel,如果你是负责这部分的工程师,那你应该就得完全具备高性能算子库工程师的能力了,不过目前深度学习编译器的codegen做得比较好的基本上是pointwise这类实现起来相对简单的算子。
第二种关系是深度学习编译器会调用外部的高性能算子库,这也主要是面对gemm/conv这些计算密集型,玩法比较多,难度比较大的算子。负责管理对外部库调用的工程师就可以更有针对性地积累第二层和第三层的技术了,对于第一层来说,可以暂时只需要了解有哪些“玩法”,每种“玩法”是为了解决什么问题,对于具体的实现方法,优先级可以稍稍往后放。