幻灯片#
- 11 Computational Graph
- 12 Compiler Frontend (
01:05:20)
计算图#
引言(04:44)#
我们准备上课。我们上一周主要是在讲计算图,然后已经从计算图开始再去讲计算图的构建。然后我们提到了控制流,是我们计算图当中比较复杂的一种特殊的操作。为了支持控制流,我们提供了各种各样不同的 AI 框架的一些解决方案。现在也想简单地稍微回顾一下,因为上次课讲的没有讲完,可能稍微总结一下上次课,然后再接着推荐我们这一次。
计算图构建策略的演进(05:16)#
好的,上一次我们也主要介绍了最早的一些探索,主要就是围绕控制流。然后有的我们会选择使用自己的一些语言接口 API 去支持控制流。这样的话就生成了以 TensorFlow 0.x、1.x 为代表的这种静态图构建的声明式编程,作为计算图构建的主要派别。而与之相对的可能就是 PyTorch 0.x、1.x 所代表的这一类,主要是支持 Python 自己前端的一些控制流,然后选择动态地根据你当前的类型参数去进行算图的构建。这里面我们也进行了一些对比,我们就会发现从使用便捷度上一定是这种动态的更为便捷,方便我们去进行调试等操作。但在优化层面,我们又发现静态的图就更容易优化一些。所以我们上节课就简单地介绍了这样两种。然后也提到了后续我们主要是想要做一些融合的计算图构建策略,那也是现在主流的 AI 框架都做的一个事情。
动静态图融合策略:以函数为单位(06:29)#
如果说想要做动态图和静态图的融合混合,那必然你也需要有一个粒度。比如说在某种程度上,你用的是静态图,然后在某些层面上你用的是动态图。那一个比较直观的符合大家认知的方式就出现了,就是我们整体依然是动态图这样的一个逻辑,但是我们把相应完整的一个子部分的一个子逻辑去把它构建成一个静态的子图。那么这样的一个子图该以什么样的东西为单位呢?目前现在 AI 框架普遍选择的都是我们基于 Python 的原生语言,我们以 Python 的这个函数为我们的一个单位,做一个静态子图的构建。那这样的话就是用户可以自由地去写它的一个 Python 函数,然后就可以调用我们框架提供的一些语言或者是一些 API,把这样的一个 Python 的函数把它转化成我们的一个静态图。我们上节课也主要讲了就是这样的一类函数式编程,我们认为它是属于声明式编程。也就是说你是先把你要做的事情作为一个函数写在那里了,然后我们再去把它转化成图等等。但是它相对于我们过去的静态图,就是 TensorFlow 那种基于自己的原语的一个静态图的构建,已经有了显著的不同。主要就是它基于 Python 的原语,然后它就能更好地去支持我们的用户。
函数式编程的两种主要思路(08:02)#
我们上节课也主要介绍了这种函数式编程的两种主要思路。一种是追踪式的思路,它可能更符合我们 Python 动态图的一些构建思维,就是我们顺着这个语言逐渐地去把我们的这个图构建出来,然后把这个图固化,就是一个追踪式的思路。另外一个思路就是语言转换,我们大概也讲了,语言转换你可以从 Python 的语言转换成你自己的一个构成的语言,你也可以去捕捉 Python 在编译的时候产生的一些语言,然后来去做我们的动态图的转换。
追踪式与源码转换的优劣与发展(08:41)#
好的,那上节课大概花了一段时间去展示了一下。就是现在以 PyTorch 为例,它的一些比如说基于追踪的进行图转换的最终后果,以及基于源码转换的然后固化成静态图的一些效果。然后在这些效果当中,我们没有详细地去讲说,到底 PyTorch 是怎么实现的 trace based,怎么实现的源码跟踪。但是通过一些举例,我们也能感受到,比如对于一些比较经典的例子,就是我们上次课一直在举例的那样的几个函数。就是我们的条件判断,就是有控制流的这样的一些函数,以及有可能有一些变化的需求的,比如说形状跟着输入在变化,或者是涉及到一些外部代码。那么我们就以这样的几个例子为测试用例。然后目标是希望大家可以感受一下,这样子的一种生成静态图的逻辑,它存在怎样的不足。以及我们逐渐的发展的。
动态图和静态图的融合:PyTorch Dynamo(09:49)#
我们认为追踪的,如果是直接追踪源码,它遇到的问题就是它很难实现各种各样的根据输入的一些变化,比如说根据输入改变形状等等。那一个部分就是说我们在追踪的时候,我们稍微滞后一点追踪,我们做一个 just-in-time 的追踪。那当你做 just-in-time 的追踪的时候,因为你已经获得了你要执行的当前的这个输入,那有一些跟输入相关的事情你就能够解决得比之前的那个纯追踪的那个就会就能解决得更好一些。但是我们依然意识到,如果你只有当前的输入,那对于控制流来说,有可能你不能完整地走完整个控制流,所以它还是有所不足。然后这就使得我们逐渐走到了源码转换,走到源码转换这边,好处就是它能更好地去支持控制流。然后如果说你做的是一个 Python 自己的代码转换,你依然面临着输入未知的情况下存在一些不足的问题。所以我们就只需要把这二者都进行一个结合,那都进行结合以后,你就到了我们要讲的这个 PyTorch 的 Dynamo,它的策略就是说我是进行源码转换,同时我还要等到你有输入了以后,我才开始进行转换,而且更具体地来说就是,等你有了数据输入以后。Python 你自己不也是要编译的吗?然后我就去捕捉你编译的这样的一些 bytecode,然后我去把这个 CPython bytecode 转换成我自己的一个计算图语言的构建。然后我们认为至此为止,我们的源码转换工作就已经做得比较好了。比如说控制流,我们可以按照我们的需求依次生成完整的所有的分支。然后比如说我们上次课可能没有讲到的是根据 shape、根据形状、根据输入的形状有变化的。那如果大家尝试这个 torch.compile 的话,我们也会发现它也能够准确地把这个发现跟输入一的形状相关的这件事情能够意识到,然后做一些相应有的一个转换。
动态形状处理与外部代码支持(12:01)#
这里比较有趣的事情就是我们看到一上来就是它的这个 reshape,它写的是一个固定的 和 的形式。这可能意味着比如说 Python 它的一个执行在这里,它编译出的字节码可能就直接把这个 shape 就复制到这里。就导致我们 PyTorch 在捕捉这个字节码的时候,就直接发现了我 shape 的时候要做的是一个固定的形状。它第一次就构建了这样的一个分支。那我们知道如果你后续再去用其他的 shape 去其他形状的 去调这个函数的时候,我们就知道当前这个分支就不再适用了。其实大家可以想象一下,就是对于我们的 Python 自己来说,Python 对于 Python 自己来说,Python 本身其实再一次执行到我们调的这样的一个 torch.compile 的语句的时候,它也会发现后面比如说你换了另外一个其他的 shape,然后它就需要再去重新生成一下字节码。那这时候对于这个 torch.dynamo 再去做原语转换的时候,我就会发现新来的这一段新的代码,新的这个字节码是我没有转换过的,就是它出现了一些变化。比如说它的输入有所变化。那这时候它就会意识到我要重新再去执行源码转化。然后再去根据我新的这个需求,新的 Python 的中间的这个字节码去生成一套新的静态图,然后来去完成后续的一些操作。
我们注意到就是如果你第二次调,比如说你放了一个 的一个形状,其实 Python 自己这边就已经会意识到,就是第二次出现这个函数,然后它的形状变掉了,它就已经会稍微聪明一点地就是告诉你这个东西是跟输入大小相关。那如果说我们读到了这样的一个字节码,对于我们的这个来说,它转换的时候也就相应地把它变成了一个输入放进来 放进来。然后有了这个 以后,它就作为我的应该做形变的这样的一个参数。然后我就能够实现更灵活的这样的一个随着我们的输入去改变它的 shape 的编译操作。我们也就发现就是说虽然你是根据 input 随便去改变形状。并且 PyTorch 第一次做的这个源码转换,它明显不能支持所有的 input。但是在第二次的时候,我就已经很聪明地扩展到了任何一个形状都能够支持,而且这只是编译层面的一个事情。对于用户来说就是无论是第一次还是后几次,你都能够很正确地拿到你想要的结果,而且事实上是在到后面几次的时候,它的速度会更快,因为我们的静态图就已经构建好,并且可以重用了。好的,这也是这个静态图构建的一些意义。
然后其实就是如果大家对这个事情感兴趣的话,就是 PyTorch 什么时候决定要根据一个新的字节码重新去生成一个函数的新的子图的时候,它其实是有一些源码的,就是每一次源码的一些比较的一些工作,然后它有一些守卫机制,如果发现有比较大的差异,然后它就会意识到这个时候我是需要重新去写我的代码的。好的,那就是我们刚才讲到的跟这个形状相关的一些使用 PyTorch 2.0 torch.compile 它能做到的一些构建子图的一些方式,我们发现它也比较优秀了。
接着我们关心的可能就是外部的一些代码。然后对于外部的代码,它也是非常有意思的。当然我们最后首先先看到,如果你用的是一个不正确的 shape,就是本身它逻辑上就应该是错误的话,那我们看到它也能够正确地返回,这里存在一个 shape 不相同的一个错误。然后除此之外,如果你是按照正常的逻辑去执行,然后我们也会发现在 PyTorch Dynamo 里面,因为我们其实有把这个外部的代码,就是本身是这个 NumPy 的代码,提供一个与它相对称的 PyTorch 版本。所以通过调用这个 PyTorch 对外部代码的支持,然后它就可以比较妥善地在这个子图当中去构建,就是通过原语的一些翻译,构建一个 PyTorch 支持的和外部代码逻辑一样的 PyTorch 子图的运行版本。所以它就是对于所有它能够支持的这种外部的语言,尤其是比如说我们经常会用到的 NumPy 的各种各样的计算的这一类语言。PyTorch 是可以提供这样自动地把外部代码转换回来的操作的。大家知道就是对于 NumPy 来说,你并不需要它的梯度,所以 PyTorch 在转换的时候就做了一个 re-trace,那这也是可以避免后面部分的一些微分,去浪费我们的一些计算资源,计算空间等等。所以也是有一定优化在里面的。大家能够感受到它可以把这样的一些事情做一个子图的转换。
好的,而且这样的一个准确的一个子图转换,其实在调用的时候我们也可以想象了。就是因为它本身是一个正确的代码逻辑,所以两次随机它必然能够做到不会是一模一样的一个结果。所以我们比较严格发现,就是它的运行结果就满足我们的预期,每一次随机都是不一样的一个随机数。
PyTorch Dynamo 与 AOT Autograd(17:38)#
好的,那我们就大概举了一下这个 PyTorch 在源码转换新的 Dynamo 解决方案。然后 Dynamo 它之所以很优秀,就是能够撑起 Python 2.0,其实我们觉得很重要的一个事情。一方面就是说它完美地解决了整个是一个动态图,然后局部是一些静态子图的这样的一个计算图构建的问题。如果说要完美解决这个静态图构建的问题的话,我们很值得一提的事情就是关于反向的一些事情。就是我们刚才举的更多的例子都是说这个 function 它执行变它的前向,你能够把它与之相对应的这个前向的计算图进行一个符合我们逻辑的,所有控制流的逻辑都正确的这样的一种构建。这是很必要的。但是作为一个 AI 的编程语言,更重要的一个事情就是你也需要能够让用户去调用梯度的计算以及反向的这种优化的所有功能。
所以对于我们的 Dynamo 来说,更重要的一个工作是它支持了这个子图以及其对应的反向微分计算的子图的同时构建,这里面就是我们的 AOT Autograd 这样的一个函数,支持了这样的一个事情。那这里的 AOT 其实它的意思是 ahead of time,也是我们在编译的一些工作当中经常会用到的一种语言。那 ahead of time 它其实意思就是说就是这个时候你还没有去主动地去调用。比如说你要算 gradient 等等,backward 这一类的,去触发反向的这样的一个操作。但是就是你调这个函数,你已经定义了当前这个前向的函数。那我再给这个当前前向函数做编译的时候,我就可以把它的这个反向的图也去做一个构建。因为可能我认为如果用户需要的话,这个东西在 AI 当中前向和反向都是非常重要的一个事情。好的,那 AOT Autograd 它一个核心思想就是说我想要去支持我们现在这种固化的一个子图,然后做一个子图的一个反向的构建。那做子图反向的构建的话,其实如果说我们很直觉地去想,是只要你有了这个前向的图,那反向的图就可以按照我们之前习题讲的这样的一个方式,然后去反向地去进行一个自动微分,然后构建一个子图。然后同样的,它一定能够满足 backward 的总的计算,那如果是按照这个逻辑的话,它就是我们的 eager mode 的一种方式。但是如果说就是你把它构成一个完整子图,就是你可以想象,其实你可以不用删掉这张图,而是可以把它留下来,并且可以做一些相应的优化。优化完了之后,然后你可以再把它分成两个图,分成一个反向的。那这种方式就是我们要讲的这个 ahead of time 的,增加了更多优化的一种方式。
好的,对,然后我们后面的许多信息都是参考的这样的一个链接,大家感兴趣的话也可以去看一下。然后后面我们会有一些更具体的例子,其实这个例子也是官方的这个 PyTorch 的这个 AOT ahead of time 的这个 tutorial 里面提供的一些例子。好的,那上节课就是讲到的就是说我们要用 Dynamo。
torch.compile 的使用与调试(21:07)#
额外要提一下的就是说每次在最开始使用的时候,它是需要一个 reset 的这样的一个功能来完成一些初始化的一些事情,所以这个是比较重要的。然后后面就是大家就可以去定义自己的函数。然后等我们需要有一个函数,能够把它固化成一个子图的时候,我们就给它调用这个 torch.compile。然后这边是这个函数,然后后面是各种各样的一个 backend。大家可以现在就认为我们的所有 backend 做的都是一些优化器,AI 优化器后端的一些操作。你也可以什么都不做,比如说我们就放一个自定义的一个优化器在这里,那这个优化器除了打印一下什么都不做。那这样的话,目前我们就用的是 torch.compile 一个不完整的版本,就等于是它只是把这个函数做一个静态子图的构建,然后接着就把它输出了,省略了后续的一些优化。但总之就是在我们讲构建的时候,这是比较够用的,并且我们希望把它打印出来来看一看。好的,那对于这样的一个设置,就是这样的一个 torch.compile 本身已经能给我们提供一个前向图的构建了。我们首先可能想的是一个 eager mode。所谓 eager mode 的就是说我既然已经有了这个子图,那我就调这个子图去执行各种各样的输入。然后就输入会给我一个输出,然后这个输出它本身也是一个 PyTorch 的 tensor,然后我就可以给它执行 backward 的。如果你执行 backward 的话,那相应的它会提出,它会计算出所有输入应该有一个怎样的 gradient。
AOT Autograd 的演示与图分割(22:43)#
最后我们的程序就是检查了一下通过 PyTorch 固化的子图的一个反向自动微分计算的 gradient 和我们完全不固化,完全是一个 eager mode。整个走这个流程的一个 loss 的 backward 带来的比如说 的自动微分。我们 check 了一下它俩是不是一样,那这边我们好像没有打印结果,但结果是一样的,大家可以放心。好的。我们更关心的事情是在这个 eager mode 下面,就是我们在这个 PyTorch 执行固化的时候,大家可以看到就是因为我们的这个后端提供了一个打印的功能。然后它能够把它构建的这样的一个图打印出来。然后相应地我们就能看到这图其实是没有什么优化的,等于每一个操作符都相应地在这里产生了对应的一个运算。然后这些运算其实都是这个构图的时候的一个节点,然后这些节点都分别地正确地执行,最后返回它的返回值,就是一个比较直观的一个前向图的一个构建。
好的,那我们就想尝试一下 ahead of time 的 AOT Autograd 到底做了一个什么样的一个事情。那我们就还是去测试我们之前提的这样的一个只是做一些加减乘除,做一些三角函数的这样的一个这个 PyTorch 是官方的一个 tutorial 的一个例子,然后这时候我们调的暂时调的不是 torch.compile。我们后续会提到,就是你可以调 torch.compile,然后通过其他的参数来完成一模一样的事情。但是现在我们就显示地去调了这个 AOT function,然后这个也是 AOT Autograd,其实它是 by default 实现的,主要就是实现了这个反向图前向图和反向图一同构建的这样的一个核心的函数。好的,那这个函数它跟 torch.compile 类似,它的第一个输入依然是我们 Python 定义的这样的一个 function,然后也就是 function 的一个指针。好的,那接着就是你也可以给他提供一些前向图的优化器和反向图的优化器,那当然你也可以不提供,它就用 default 的版本。那我们这里提供也只是为了像类似的,我们还想要打印一下,来看一下前向图和反向图都长成一个什么样子。然后这里就大家都已经注意到,就是这个 ahead of time function,它本身它的任务就是要同时去构建前向的图和反向的图。
好的,那我们可以稍微看一下,然后你就会发现它这个前向图,这是其实它前向图、反向图都已经构建了,只不过在你执行这个前向的函数的时候,它会先把前向的图给打出来。然后等你执行这个反向图的时候,它会把这个反向的图也会打出来。所以这边就是我们打印的时候写的是这个就是写成 the forward,但其实这是反向的,上面这个是前向的。那如果关注这样两个 ahead of time section 去编译的子图,你就会发现它跟前面的那个图是有着比较大的不同的。
当然我们前面那个图只有前向图,但是你只看前向图的话,你就会发现,首先它好像做了一些精简。除此之外最重要的事情,你会注意到它返回的值并不是我们程序本身唯一的一个返回值的一个出口。比如说程序最终要返回的是最后计算的这个 cosine ,然后它后面多了两个 和 这样两个冗余的值。而且如果你只计算前向的话,其实它就只有这个 cosine 被赋给了我们的返回。所以后面这两个其实就没有用。但是 AOT 它本身是同时提供前向和反向图,后面这两个我们马上就看到,它主要的作用体现在它是为了加速反向图的计算所构建而成的。你就会发现这个反向图它也会和一般的反向的一个计算图有所不同。它不仅是有过去的一些 gradient 作为输入,它还有 和 这样两个额外的新的冗余的输入。
好的,这里我们想说的就是其实整个 AOT function,它做的事情就是说我同时构建了前向图和反向图。然后接着我就会给前向图反向图进行一些优化。然后后续我还会要做一个分割,就是我认为哪些主要是前向图的部分,哪些是反向图的部分,然后做一个分别的保存。这就是所有在发生的事情,然后发生完了之后,我们自己看到的打印后的结果就存在了这样的一种图。分割之后留下的一些前向图像,反向图去传数据的这样的一些额外的信息。那根据我们的这个函数,其实大家也可以想象一下,就是有了这两个以后,对于我们的反向图来说,它很多的计算就可以不用再从 input 去获取,不用再重新算一遍,对它的优化是有比较大的一个一个优化不是的,好的,那这里就是为了展示一下这个 ahead of time 的这样的两个事情。
然后这里我们也想以这个例子,对我们这个 ahead of time autograd 去进行一个总结,也就是刚才说的那个事情。那就是说之所以叫 ahead of time,也就是说它在运行前就已经开始执行了。那它执行的时候,它就需要用一些假的 tensor。然后有了这些执行,我就可以去通过 trace 的方式去构建我前向的图以及反向的图。然后把这个联合的图去进行分割,并且把这些分割的图转换成一个可以被调用的 PyTorch 的这样的一个函数,这大概就是 ahead of time 它所做的所有事情。好的,其实这个东西就像我们刚才说的,就是我们举例是举的 AOT function,然后当然它也是可以在这个 torch.compile 里面被具体被实现的。
torch.compile 的超参与调试流程(28:48)#
好的,那 torch.compile 就是如果说你不加任何后续的这些参数,你只是 torch.compile 一个 function 的话,那它其实按照它的默认会做的事情是非常多的。首先就包括我们这个前向图的构建,以及如果是需要的话,它有时候还会有这个前向图和反向图的构建。然后同时也有各种各样的优化,那是我们后面要讲的一些优化。如果说你加一些限制的话,它就可以被裁剪成某一种功能,那这里想讲的是就是因为 torch.compile 它做的事情可能是非常多的,比如说它就会调用我们刚才提到的 AOT Autograd,然后去做整体的一些图。那这样的话对于大家来说可能就是比如说一个一个函数,然后直接调 torch.compile。然后如果它报了一些错,或者它没有报错,但是它提供的结果跟我们想的不一样,或者是它梯度计算的结果跟我们想的不一样。这都是有可能会给我们的这个编程带来一定的复杂性的一个情况。所以它当时也是考虑到了这一点,所以它给它的 torch.compile 提供了一些不同的超参的选项。然后他们自己也在他们的这个 FAQ 这样的一些网页上提供了一些建议,就是说建议大家在使用的时候,你可以先直接采用这个 torch.compile,然后 backend 选择 eager 的这样的一种模式。
如果你 backend 选成 eager 的这种模式,那 torch.compile 它运行的逻辑就非常接近于我们最开始讲的这个动态的一个动态编程,动态去构建动态图的这样的一种策略。它整体就会把我们的前向图去做一个根据输入的一个动态的一个构建,构建完了以后返回给这个函数。那如果比如说你跑这个 eager mode 能够成功跑了,这一般都应该是比较容易做到的一个事情。因为我们是比较熟悉这种 eager mode 下的这种动态图的构建的。
如果这一步成功了,然后大家就可以再去把这个 backend 调成这个 AOT eager。如果调成 AOT eager,它就意味着它会运行我们刚才的那个 AOT function 的那个逻辑。也就是说它既会构建前向,同时也会按照我们反向构图的逻辑,把这个反向的图像大家的习题一样去给它构建出来。但是它也不会做太多的优化,就直接把这个图裁剪了以后还给大家。这样的话如果说前向反向都能正常访问,这也是一种很好的一个测试手段,然后等你确定这一步也成功了,那最后一步我们就可以调 torch.compile 真正默认的这个 backend 的等于 Inductor 的这样的一个优化器的这样的一种使用方式。
那对于 Inductor 来说,它就会在构建了两个子图之后,其实构建一个完整的子图之后,它就会调用这个 PyTorch 自己的一些优化器,这个 TorchInductor 优化器去进行很多的基于计算图的优化。然后可能是我们今天后面马上就会讲的一些具体的内容,然后优化完了之后再把这些这个图进行拆分反馈给用户等等。所以这也可以是我们在调用 torch.compile 的一种调试的逻辑。好的,当然它也相应地把这个 torch.compile 它的一些功能也展示给我们。那我们认为 PyTorch 需要做的一些功能就是一些前向图构建、反向图构建以及计算图优化这样的一些事情。
如果大家感兴趣的话,这边我们提了一下,如果我们要的是 torch.compile 对应的一个函数,然后你可能你的 backend 本来应该是选 AOT Autograd 和 AOT eager。意思就是说就是如果执行这样的一个参数的话,最终其实在 PyTorch 实现的时候,它调用的就是 AOT Autograd 这样的一个函数,我们依然还是把这个函数稍微写了一下,写成一个可以去进行一个 print 的一个版本,就是在编译的同时做一个 print。然后这样的话我们会看到就是你可以直接调用 torch.compile,然后依然是我们之前的那个函数。然后如果说你调用的是 AOT Autograd,那么它 print 的两个子图就自动地做了。我们前面提到的就是前后图联合的一个分割的这样的一个操作,大概就是说这么一个事情。
动静态图混合编程的总结与 TensorFlow 2.0 的演进(33:28)#
好的,我们刚才就是以 PyTorch 为主大概讲了一下。现在我们都进入到了这个静态图和动态图混合的一个环节。然后比较重要的事情是在我们会希望这个编程语言能够首先支撑我们去做非常灵活的编译和去做非常灵活地写代码以及找错误。然后这种时候我们需要用到的都应该是一个全局的这样的一个动态的图。然后你就会发现如果你是以函数化的一个编程方式的话,那么对于一个 function 来说,它显然是可以直接被用 eager mode 去正常地调度,或者是求 gradient 检查等等。那它都可以利用到我们本身是一个动态的彩图这样的一个优势,等到如果说你确实已经完成了代码写作,然后完成了这个调试整体都运行没有什么逻辑问题了,然后我们就是想要提高一下效率。
这时候它的一些 function 就希望能够把它转而固化成一些静态图,那这里我们就会希望就是以 function 为单位,然后去提供一些本来是动态图的一些函数,换成静态图的这种图捕捉的这样的一些技术。然后 PyTorch 主要提供了基于追踪的图捕捉,基于原语翻译的图捕捉的技术。整体我们的思路都是说,无论你是基于追踪还是基于原语捕捉,我们都和最开始的静态图构图完全不一样了。
我们现在都是让用户可以去写 Python 原语的一些控制流,语言等等。然后我们把更多的事情交给我们的 AI 框架,由 AI 框架来去完成相应的无论是追踪的建图还是语言转换的建图等等。当然这也会带来很多的工作,同时带来一些优势和劣势。比如说对于我们基于追踪的这些思路,它可能就是对于 AI 框架来说更容易实现,但是它适用的场景更有限一些。主要是我们刚才也体会到了各种各样,比如说控制流支持不完整这些的可能性。然后对于我们原语转换来说,我会觉得它好像似乎已经完成了所有逻辑上的事情。能够广泛地去支持种种的一些控制流,或者是用户的一些代码。同时还可以支持高效的反向微分的计算等等。但是它的不足一方面就是它编程是非常困难的一件事情。另外一方面就是它因为有很多逻辑在里面,所以它的调试以及生成的代码本身的可理解性、可读性也都会变得差一些。然后灵活性可能会变得差一些。但是现在因为我们主要是一个面向应用,有很大的需求。所以目前主流的方式还是说按照原语的转换去进行计算图构建的这样的一种流派。
好的,我们刚才主要是以这个 PyTorch 为代表,然后我们也还想稍稍提一下,就是我已经看到了如果最开始就是动态图,那它走它自己的这个研究方向的路线。它当然可以引入一些把某些动态图捕捉固化到的这样的一种策略。然后走到了我们第三代这种动静态图混合的这样的一个新时代。但是对于如果过去,比如说 TensorFlow,他们最开始是以静态图为主流路线的。他们是怎么挪到这个比如说 TensorFlow 2.0 动静态结合的这样的一种新的策略下面的,那这里面我们要提到的事情就是 TensorFlow 1.0 的时候,最开始他们其实就是使用了一些自己的原语。然后他们就存在一些比如说他们自己的一些控制流的一些语言,有 TensorFlow.condition,TensorFlow.while_loop 等等,或者是一些调度这些基础的 API 的一些更高级的 API,比如说一些 case,这样的一些关键词等等。
在他们最开始有了自己的原语之后,然后他们很自然地就会有一些从这些原语去构建他们的计算图的这样的一种类似于编译的这样的一些功能的一些研发。但是后来当然它就是有这样的一些原语帮助他们做计算图做图优化。这是所有走静态图路线比较擅长的在图优化,发挥了很多研究,推出很多高效的优化策略的这样的一个研究路线。但是就是随着 PyTorch 占领了市场,然后我们就会发现用户更希望使用的不是这种很特异的一种专门的 API,更希望的是使用 Python 原生的一些控制流。那对于 TensorFlow 2.0 来说,其实它的转换也比较简单,我就主动地去把所有的这个 TensorFlow 基于这个原语的这样的一些 API 去建立一个像 Python 的原语的一个映射。然后整体我做一个类似于编译器的一个原语转换,我也能够完成 Python 的这个部分该怎么去构建静态图的这样的一个部分。所以其实对于静态图的一些研究的路线来说,他们是都能够比较快地转换成这种 function 的这样的一种构建子图的一种策略的。
然后除此之外就是 TensorFlow 2.0 最重要的一个事情是说,我们这样的话就是对于一个 function 构建成一个静态子图,我能做的比较好了。但我做的不太好的事情是在于我整个的一个编程环境。全局上我可能过去是静态图的思路,我现在缺乏一个动态图的思路。所以在 TensorFlow 2.0 的时候,很重要的一件事情是它提出了它自己的一个 eager mode。如何去就整个写了一遍动态的这样的一些思路的一套支撑的框架。好的,那有了这两套之后,它就可以在全局上去动态地支撑,然后在局部去调这些子图的结合,它就能够做得比较好了。
这是一件好事情,就是说它也就是很快地跟上了现在的一个需求。但是从历史上发展上来看,就是 TensorFlow 2.0 在推出的时候遇到的最大的阻碍,就是因为它整体全局变成了动态,跟过去它的 1.0 全局的静态有着很本质上的一些逻辑上的不同。然后 TensorFlow 2.0 在发出的时候就公开表示他们不会向下兼容 TensorFlow 1.0 的,就基本不会提供官方的兼容支持。然后这也是当时大家会很很有反对意见的。这个导致它进一步损失了更多的市场的一个历史性的因素。
AI 编译器发展历程与现状(40:09)#
好的,大概就是讲一些小故事,然后后面我们会讲一些 AI 框架的一些历史。但是在讲 AI 框架历史之前,我们现在想先讲一下 AI 编译器的一个发展的一个历程。之所以提到 AI 编译器,其实是 AI 编译器它应该是说从 AI 框架当中独立出来的一个小部分。就是我们发现有很多部分都跟我们传统的编译器很相关。然后 AI 已经需要了一个自己独立编译器的这样的一个到了这样的一个时代。所以到了 AI 编译器以后,我们也想关注一下,就是刚才我们提到的那些静态建图、动态建图以及混合建图,对于我们 AI 编译器的发展都会是怎样的一种一种对应。
好的,那我们认为我们 AI 编译器的就是从编译器的角度来看,发展其实也是分了三个阶段。第一个阶段其实主要就是静态图阶段。在静态图阶段,比如说 TensorFlow 0.x、1.x,它就已经要做很多编译器的一些构建的工作,比如说它要去利用编译器的方式构建计算图,然后把这个算子做一定程度的优化等等。这都是第一个阶段的 AI 编译器。我们认为就是在 PyTorch 1.0 这一类的这种动态图的构建当中,其实它是不存在任何编译器的 AI 编译器的功能的。主要是因为它并不存在一种从某个语言转换成另外一个语言,并进行一定优化的事情。但是随着到了后来我们到了动态图,全局动态图、局部静态图的 function 编程的一个时代之后,我们就认为我们到了 stage two。那 stage two 来说就是我们多种的 AI 框架,都目前正处于的这样的一种 AI 编译器发展的阶段,那在这样的一个发展阶段,我们希望前端支撑的是类似于 Python 或者是 Python 的比较让用户更熟悉的自然语言的一些表述。然后后端我们会做这个 AI 自己使用的一些原语构建的图和算子,并且去进行图的优化,算子的优化等等。然后我们觉得面向未来,其实大家更想做的事情是我们不仅想优化计算图,我们还想去优化算子。
这时候我们已经提到了 AI 编译器的一个关于优化的一些主要功能,相信优化计算图大家现在都非常能够理解。比如说就是有一些分支是类似的,你可以去进行合并等等这些计算图级别的一些优化。算子的优化其实大家也是有非常多的体会的,比如说我们之前学了这个卷积层的这样的一个基础的算子。那如果你是最或者说是矩阵乘那这样的一些算子,如果你是最浪费的精力,但是你是完全不加思索地去做各种各样的循环遍历的话,你可以实现一个算子版本,但是相应地你当然也可以去利用一些更优化的一些策略,然后就有更优化的算子。那这样在一个比如说 convolution 计算的这样的一个粒度来说,我们做的就是一些算子的优化。那这时候就是大家的工作相当于是做的是一个很专业的程序员,专门地给某一个功能做了一个很人工的优化。而 AI 编译器它想做的事情是说,我可能认为就是用户给我提供的是一个相对来说比较欠优化的一个算子的实现。然后我本身可能也不是一个专门面向某一个功能,但是我根据我的语言的调度,然后我去进行尽可能的一些做一些算子优化。这也是 AI 的编译器想要做的事情。好的,那我们后面在未来第三个时代,我们就会希望就有图的优化和算子的优化能够去进行更统一的一些优化。它就涉及到一些图算统一表达的事情,这里是一个代码,等我们到讲到后面计算图的时候,会更多的还会涉及到这样的一些内容,所以大家就会有更多的一些体会。好的,那我们认为当前就是 AI 编译器发展当前的一个现阶段的一个现状,是说我们目前绝大部分情况,下图和算是分开进行表达,分开进行优化的一个事情。然后接着就是我们算子的实现缺乏一些自动化的一些优化技术。所以这是未来发展大家如果感兴趣可以去进行研究的一些方向。
好的,以上就差不多是我们在目前混合式编程的一个总结了。就是全局我们是一个动态的图,然后在一些子图的范围内可能有一些标志标识,就代表着它是一个可以被固化的捕捉的一个子图。那 PyTorch 我们大概讲了 TorchScript 的一些标记,或者是一些 torch.compile 这样的一些函数。那其实 TensorFlow 对应的也是有这个 tf.function,以及这个 @tf.function 的标记等等。总之大家最终都选择了这样的一种混合框架的这样的一个概念。
AI 框架发展历程(45:22)#
好的,那后面其实想要简单介绍一下,就是我们认为 AI 框架它的一个发展的一个历程。AI 框架发展的历程,其实它是大于等于我们刚才提到的这个 AI 编译器发展的历程的。最开始的时候我们认为就是还不存在 AI 框架,那可能是在 2010 年以前的时候。然后主要的是一些计算库,比如说 NumPy、MATLAB 或者 SciPy 这样的一些科学计算库,他们来提供一些基于 library 的一些计算,然后具体来说就提供了一些前向以及跟前向相关的一些数学微分链式法则的一些方法。
然后主要的思路是使用这个库函数去进行一些编程。然后当时代表的一些工作主要是由 Caffe 这一类的工作。然后对于这应该也是 Caffe 的一些论文里面的一些图。然后当时它因为提供了一些这种 library,它就可以使得用户更容易地去调度。比如说去调用它的卷积计算,然后池化计算,然后 batch normalization,或者是一些激活函数的计算等等。然后大家就可以调用它们去构建一些子图等等。
然后这里面也会涉及到一些计算图的概念,就是这样的一个建图就会建立一个前向的图。然后有了前向的图就可以支持反向的一些微分计算,但是我们认为就是以 Caffe 为主的这样的一些过去的 library,它是不支持扩展的模式的这样的一种自动微分的形式。然后主要是一些按照层的这种形式的简单的一个网络的一些设计。所以说它是最开始神经网络发展的一个很重要的一个路径,但是最终它还是不太能够很高效地支持现在的一些微分计算等等的一些功能。好的,还是存在这个比较多的局限性的。比如说现在层出不穷的网络结构,然后现在需要一些特殊优化的前向和后向等等,它都是不能进行特别好的一些支撑这样的一些事情。
这里面也举了一些当时的一些 library 的一些例子。然后大家看到就是在 library 的时候,其实就是这个 library,它要给大家提供一个前向的网络架构以及反向的一些微分计算。再反向微分计算的时候,我们可以把它基本分成是按照比如说源于转换 ST source transformation 这样的这种原语转换的一些策略,或者是它是操作符重载 OO operator overloading 的这样的一种策略。我们就会发现当时也存在很多种编程语言,他们主要都是在比如说 2000 年、2006 年、2007 年,然后这样的一些 library,然后提供了各种选择了各种各样的一些实现的路径。然后这个 PyTorch 最早也是有一些这些方面的研究。比如说 PyTorch 是在 17 年做一些基于基于操作符重载的一些微分的这样一些云计算等等。好的,接着就是我们认为这个 AI 框架进入到第二个时代,其实就是以我们计算图为标记的这样的一个时代。
也就是说当 AI 框架选择去构建这个静态图,以这个静态图为主要的一个创建以及维护以及优化的这样的一个目标的时候,我们就进入了这样的一个 2.0 的时代。那么在 2.0 的时代的话,整体静态图就有了我们刚才提到的整体计算图就有了我们刚才提到的两种策略。第一种策略就是以 TensorFlow 为代表的静态图的构建,以及以 PyTorch 为代表的动态图的一个构建。然后其实之前每次教学讲到这里的时候,我们都会展示一下这个 Paper with Code 的一张当前主要是研究方向的这样的一些领域的文章。他们使用的底层的这个平台 code,使用的框架都是哪些,然后它有一个比较好的统计。然后我记得是在 2023 年的时候,当时 TensorFlow 还有一点点占比,然后等到 2024 年的时候讲这门课的时候,就是这张图 TensorFlow 已经很少的一定的占比了,然后主要部分都是 PyTorch 这样的一个占比。
但是在 2024 年的时候,很有趣的一个事情就是华为的昇腾,它的份额在有一个增长的一个趋势,不过今年我在更新这个课件的时候有发现这个 Paper with Code 的这个平台,它关闭服务器了。然后现在变成了 Hugging Face 的一些 daily papers。然后他们已经不再不再继续跟踪这个 framework 的一个具体的一个趋势了。但是其实大家也可以想象,就是现在的话应该 TensorFlow 可能会进一步地减少,但是 PyTorch 可能依然是为主流的一个状况,但是由于一些国际关系,有可能其他的一些包括昇腾在内的一些框架也依然有很大的比例,大概这样一个清晰。好的,我们认为目前已经发展到了第三代,就是图像融合的这样的一个,主要是我们的动静态图融合的这样的一个思路。而且我们动静态图融合的时候,我们主要构建的策略是基于我们源码转换的这样的一个策略。
大家可以回忆一下,就是当我们提到静态图的时候,我们当时说过,对于我们框架来说,尤其是跟 AI 编译器相关的一些事情来说,一个静态图它整个实现的功能是比较完整。比如说它有前端语言的一些转换,然后它有中间的图的构建,以及更重要的它有自己图的优化,以及面向内核代码的一些优化等等。我们认为它是比较完整。那到了后来我们进入了这种整体是动态图,然后局部构建静态图的这样的一个时代之后,我们会发现就是 AI 整个框架,它作为一个包含编译器在内的这样的一个结构,它其实做的事情就会更多一些,更加丰富了。那主要就是因为这时候前端我们就限定了是 Python 本身的这样的一些语言。然后后续我在这个静态图的构建就产生了我们刚才提到的非常类似编译器的这种源码。对源码转换的这样的一些构建的一些方法,以及后续非常可以借鉴许多传统编译器的一些 AI 编译的一些优化这样的一个事情。总之作为一个框架,我们就提供了编程的灵活性,同时也满足了计算的高效性。所以这是整个目前我们 AI 的框架所想要的支撑和缩小而支持的这样一系列事情。
好的,反正就是整体大概介绍一下 AI 框架的这样的一个情形,然后整体来说就是我们认为我们有混合的,有静态图的,有动态图的,然后有如果朝静态图方向更加不灵活,但是更加高效的话,可能是这些以 library 为主的这样的一种编程方式。而如果想要更加灵活,但是就是缺乏一些很高效的优化等等。它主要是 PyTorch 的动态图的版本,以及各种基于 Python 的各种各样的,比如说 NumPy 这样的一系列的 library 的一些情况。这里就是对各种各样的 AI 框架去进行一些总结是好的。
软硬件协同与 AI 编译器(53:17)#
那随着我们这个 AI 框架本身,它是可能是一个更偏向软件层面的一种发展。我们讲了第一代、第二代以及第三代等等,但是我们希望提及的事情是说,就是对于 AI 编程这件事情,一个很明显的特色就是软硬件是非常协同的一件事情。那在早期的时候,有可能我们对硬件主要就是有并行计算就已经很好了的这样的一种状态。那么到中间的时候,大家就已经发现这个算力处于不够的这样的一种状态。所以我们就不仅使用这个 GPU 为主,同时还使用一种类似于我们叫 NPU,也就是专门面向这个神经网络高效计算的这样的一系列计算单元。然后他们可能并不是通用 GPU 的类,但是是一些类似的一些架构。然后同时在第二代也已经产生了这种异构的这种多计算单元混合调用的这样的一种策略。
在然后到了我们目前的现在的这样的一个 AI 框架的一个时代,我们就需要更多的硬件去更好地管理它们,然后满足我们的高效计算的一些需求。所以也会涉及到我们的一些分布式的一些等等等等的一些内容。好的,这边主要是对 AI 的这个计算图的构建去进行一个总结。然后下节课我们来讲这个 AI 编译器以及图优化的一些内容。我们休息一下。
好的,这是我们在这个计算图构建的一个总结,相信大家现在已经非常熟悉了。就是声明式构建、命令式构建,然后函数式的混合构建,以及在过去 AI 框架的这三个阶段,之前的时候我们有基于操作符重载以及原语转换的一些自动微分的实现。当然现在这个对大家对这个表达式的就是我们之前也是讲过的,相信大家也都还是非常熟悉的,然后有各种各样的一些优缺点等等。
编译器前端#
AI 编译器前端与传统编译器的对比(01:05:32)#
好的,接下来我们想要进入的内容是这个 AI 编译器的一些部分了。具体来说,主要是 AI 编译器的前端这样的一个内容的一些介绍。
首先我们也想看一下 AI 编译器它的一些发展的一些背景。总的来说就是说,我们认为在 AI 框架当中,有很大一部分它都跟我们过去的编译器的功能是比较类似的一个事情,它主要的功能就是想隔绝我们的硬件层面和我们的这个编程层面的两边的一个复杂度。然后面向编程层面,我们希望能够提供以计算图为核心的一种自动计算图的构建以及优化等等。然后面向我们后端的硬件,我们又希望以计算图为我们的主要表达,去提供面向硬件的一些内核级的优化等等。相应地我们也就把它分成了根据它是属于面向计算图的还是面向计算硬件的,把它分成了属于 AI 编译器的前端和 AI 编译器的后端。我们在讲前端的时候,主要讲的就是图的优化以及一些调度等等。
好的,那我们也可以把这个编译器以及我们的 AI 编译器去进行一个纵向的一些对比。那在传统的时候,我们在做编译器的时候比如说大家可能比较了解的是 GCC 这类的编译的一些语言。然后这里我们来举例,以及后续常会提到的,主要是 LLVM 这样的一种编译的语言。然后我们后面还会再具体地提到它。
但这个我们主要想说它和比如说 GCC 这一类编译语言最主要的不同,就是它有自己的一个中间的一个表达。有了这样的一个专业表达 IR,它就能够更好地去隔离前端和后端,然后去允许我们编译优化的一些重用。对于传统的编译器,如果大家比较熟悉的话,它主要的功能模块就是说面向源码,我们做一些前端的优化,然后接着这个编译器做一些优化,然后再面向后端再做一些优化,完成这样的一些以优化为主的编译器的一些功能。
类似的就是我们发现,对于我们在过去 LLVM 这一类的这个基于 LLVM 的编译语言的一种优化器当中,我们就已经发现它遇到的问题。就是说我们前端的代码是多种多样的,然后我们后端的硬件也会是各种各样的一种硬件。所以这时候有一个统一的表达,就能有更好地推进我们编译器的一些发展,我们会发现这和我们现在的这个 AI 编译器它的需求也是非常有类似性的一个事情。那对于我们前端来说,就是 AI 编译器想要做的事情。我们也是要构建这样的一个比较通用的一个基于图的这样的一种表达。然后接着我们希望能够在图这个表达上有一堆 AI 的编译器,能够做一些高效的优化和加速等等。最后面向后端我再去执行一些它相应的一些面向硬件的一些优化的一个操作。整体来说就是因为这两个研究方向如此类似,然后我们就会发现有非常多的语言是这种专业词汇是共用的。
比如说一些我们之前已经提到的 JIT just-in-time 的这样的一些功能,以及这个提前的 AOT ahead of time 的这样的一些功能,以及 IR 就是我们常用的这种中间表达,以及我们的这个语法树 AST,所有的这就是很多都是比较类似的一些关键的一些用词。然后最后,就是作为一个编译器,传统的编译器它也经常会是做很多 pass,做很多遍。比如说第一遍做一些死代码的裁剪等等。就是所有的这些优化的 pass,其实也可以在我们 AI 编译器当中去沿用这个逻辑。也是类似的一遍一遍地去进行各种各样能做优化的一些事情,所以他们是有比较多的相似性的。
PyTorch 2.0 与 AI 编译器的发展(01:09:30)#
好的,那在后面讲这个 AI 编译器当中,其实我们会大量地以这个 PyTorch 2.0 为例。因为我们认为 PyTorch 2.0 它其实它的 torch.compile 很大一部分的力气就花在了这个 AI 编译器这样的一个事情的逐渐提升上,那作为 PyTorch 最开始是一个动态图,它可能在优化方面是比较不足的。然后到了 PyTorch 2.0 以后,它其实已经补充了很多相关的这样的一些内容。然后如果大家看一些这个 PyTorch 2.0 他们自己的一些推文,一些宣传的话,你会看到 PyTorch 2.0 它主要会表示我们在 2.0 做出了这样的一些新的一些改进,然后我们就会发现这里面的很多部分我们也已经讲到并且非常了解了。比如说 TorchDynamo,它主要研究的问题是如何转换成动静态图融合的这样的一种图转换的这样的一些事情。比如说我们提到的这个 ahead of time 的 autograd,它能够提供高效的带自动微分的反向图构建的一个整体的一个图优化的这样的一些功能。然后后续这两个是我们后面在讲到编译器的时候会提到的两点。一个是 PrimTorch,大家可以理解为就是 PyTorch 的原语的这样的一些优化。
这里面比较有趣的一个事情就是 PyTorch 它做了一些统计,就是说我们现在 AI 发展越来越迅速,然后出现了各种各样的算子,最开始可能是卷积算子,或者是有 transformer 的算子等等,然后到现在已经有两千多个 PyTorch 的这样的一个官方认为应该要支持的这样的一些算子。那你可以想象如果有两千多个算子的话,对于很多硬件如果它想要支撑这两千多个算子,就是非常复杂的一个事情,那对于编译器来说也是一个很复杂的一个很大的一块内容,所以 PrimTorch 它的一个主要研究策略是说,我想要给这两千多个算子去进行一个抽象。我就找到了他们之中最经典的,最必不可分的这样的很基础的 250 个算子。我就认为剩下这 2000 个算子都是可以由我这最基础的 250 多个算子经过一定的组合来完成的这一类 2000 个算子的一个功能。然后这样的一个抽象化其实是很重要的一个事情。当然了它是说这 250 个就是说你排列组合和组合成这 2000 个,可能它的效率并不是最优的。但是对于一些比如说小的硬件厂商,他们这时候就可以首先去支持这 250 个最经典的算子。然后这样通过这 250 个的排列组合,它也能够支撑 PyTorch 的全部的功能。当然它可能不是效率最优的一种方式,那随后它就可以花更多的力气去逐渐地扩充它的一个算子库。
然后这也是比较适合于我们整个生态,尤其是硬件多样化的这样的一种一系列的贡献。好的,那么除了这个原语的算子的一些分级以外,然后 PyTorch 2.0 还提到了这个 TorchInductor,其实就是它自己的一个编译器进行优化的一些操作。我们后续也要讲到一些以 PyTorch 为例,或者是以其他一些 AI 编译器为例的一些优化的操作。然后这里面也比较有趣的一点,就是说做一个编译器,它不仅要面对计算图去进行优化,它还要面对计算算子去进行优化。然后在算子优化的时候,它其实调用了很大部分是使用了这个 OpenAI 的 Triton 作为它的一个算子级别的一些中间的一些表达,我们后面再讲到更多中间表达的时候,还会再提到这一点。
其他 AI 编译器介绍(01:13:24)#
好的,那除了 PyTorch 2.0,然后我们课程也会介绍到其他的一些编译器作为一些例子,然后来讲一讲。其中值得一提的是这个 TVM Tensor Virtual Machine,它是一个开源的这样的一个编译器,然后是陈天奇团队开发的这样的一个深度学习的一些编译器。那对大家来说,可能陈天奇团队应该并不陌生。因为我们的课程很大程度地参考了陈天奇老师提供的一些 deep learning 相关的一些课程。然后大家感兴趣也可以去他的个人网站去看一看。
然后对于 TVM 来说,它主要的一个特色就是说它对后端的各种各样的硬件的支撑会更好一些,作为一个平台级的,它里面也大量地使用了比如说 LLVM 这样的一些传统的编译器的语言。然后它能够支持 CUDA,然后也能够支持这个 Metal 就是 Apple 的一些显卡的硬件的一些具体的语言等等。反正它对它后端的支持,硬件后端的支持是更为比 PyTorch 灵活一些的。
好的,那除了我们会提到 TVM 以外,作为一种开源的 AI 编译器。然后另外一个就是也是类似的 AI 编译器的功能的是这个 XLA。然后它也是最近对开源的这样的一种编译的机器学习编译的一种平台语言。
可能大家有人对他如果不陌生的话,其实他最开始是 TensorFlow 中的一部分,主要用来做 TensorFlow 的一些线性代数相关的一些计算等等。然后后来它自己就足够足够内容足够庞大,然后就逐渐地独立出来,并且不仅仅支持 TensorFlow,也会支持 JAX,支持 PyTorch 等等。然后它就变成了 OpenXLA。然后它在一定程度上它并不是完全开源,OpenXLA 是它的这个开源的部分,就是一定程度上去进行了一些相应的开源。然后它最大的特色主要是在这个线性代数相关的一些计算上,比如说张量算子,以及这个算子融合加速上有一些比较好的一些优势。但是它最开始从 TensorFlow 脱身出来,所以它其实核心的还有一个特色就是说它比较适合静态图,然后相对来说一致性会差一点。然后它的这个编译质量虽然比较高,但是有可能是比较这个算力局限的。
AI 编译器与硬件生态(01:16:01)#
好的,我们其实就是想大概介绍一下现有的一些 AI 编译器,他们的一些比较经典的一些研究的代表性。然后我们也想把它们纵向去进行一些横向去进行一些比较。我们刚才提到就是说最开始其实 AI 编译器,它是 AI 框架当中很重要的一个部分。一些经典的 AI 框架,比如说有 PyTorch,有 TensorFlow,还有 JAX 等等各种各样的编译框架,那么这些 AI 框架他们自己一般会有自己的编译器。比如说 PyTorch 就主要在研发它的 TorchInductor,然后 TensorFlow 主要是研发了 XLA,然后基于这个基础上开源了 OpenXLA,以及很多很多的编译语言。然后这个时候其实 TVM 在这里面比较特别的一个点就在于,它完全是一个开源的。然后希望能够普遍地去支撑各种 AI 框架以及各种底层硬件的这样的一个编译器。所以其实 TVM 它是能够支持比如说 PyTorch、TensorFlow 以及种种的这样的一些 AI 框架的一些使用的。
在 AI 编译层底端的话,主要就是各种各样的硬件。为了想跟这些硬件更好地去做具体硬件化的一些加速等等。其实往往我们会需要一个底层的跟硬件结合更紧密的一个算子库。这里面最为大家所熟悉的可能就是 CUDA 所有英伟达的显卡,他们统一都支撑的这个 CUDA 的一些算子库,当然 CUDA 它是一个 PyTorch 或者 XLA,大家都非常支持的这样的一种算子库。是现在 AI 发展的非常核心的一个底层的一个算子计算单元。但是我们也知道其他的显卡相应地也在发展,而且 CUDA 本身它并不是一个开源的一个生态,所以我们看到比如说英伟达他们主要用的,不好意思除了英伟达以外,比如说这个 AMD 他们主要用到的一个后端的算子是这个 ROCm。我们后面也会讲到一点点。然后 Apple 的话主要用到的是这个 Metal 的这样的一个库,以及对于其他的一些国产的一些 GPU 等等。如果你不是英伟达的 GPU,然后很多的话它也还是支持这个 Vulkan,我可以认为是 OpenGL 的一个升级版本,也可以进行一定的并行运算等等。所以我们会发现对于各种各样的硬件底层来说,它都需要一些比较高效的一些算子库。而我们的这个编译器也就需要和这些算子库打交道。
然后这里进行了一个比较,主要是想分析一下这些优化器他们在跟这个 CUDA 的关系是怎么样的。我们一般认为就是 XLA 就是 TensorFlow,它这一套编译器它基本上是非常高度地去依赖 CUDA 的。主要是因为 XLA 它是面向张量计算的这样的一种优化语言。然后它很多关注的是一些很特殊的硬件,比如说英伟达的 GPU,更或者是这个 TensorFlow 经常用的 TPU。所以他们是非常硬件和这种 CUDA 是很紧密地绑定的。
好的,那作为 PyTorch 来说,就是作为现在最主流的这样的一个 AI 框架,首先它肯定是也比较高地依赖 CUDA,但是整体来说它也在努力地去做一些市场的均衡。所以相对来说它也会比如说调用 OpenAI 的 Triton,然后使得它可以在底层的算子的结构上,就是利用一种不是 CUDA 的平台。它也可以做一些类似于 CUDA 的这样的一种算子的一些优化等等。所以它也就相应地能够支撑主要是英伟达。但是除了英伟达以外,比如说苹果或者是 AMD 一些显卡,它也能做到一些一些扩展。但是主要来说就是 PyTorch 自己可能会默认支持英伟达的 CUDA。然后记得 PyTorch 发过来一个版本是能够支持这个 AMD 的一些显卡的。但是这里面肯定是需要 AMD 和 PyTorch 去做很多的合作。
好的,最后就是我们刚才提到的 TVM,它作为一个想要做第三方的这样的一种中间的一些 AI 编译器。他很多时候就能够支持各种各样的一个后端的一个展示,比如说主流的有 CUDA,或者是你也可以把你的它给他转成这个 ROCm 的,也就是 AMD 的一些硬件,或者是苹果的等等。以及如果你是比如说是移动端,那有可能是比较低级的一些并行的一些语言,那有可能他会支持 Vulkan,然后如果说甚至连这个 GPU 都没有,然后他也会以这个 LLVM 的形式去支持一些 CPU 的这样的一些情况。
好的,那我们现在就大概看到就是看作为一个 AI 编译器来说,大家可能都根据自己的一个市场需求有了一些不同的选择。但是对于整个生态来说,我们是希望能有一个 AI 的编译器,它能够更加独立地去进行一些相应的发展。然后使得我们的整个市场就不至于过于垄断这样的一些事情。
AI 编译器的工作流程与优化策略(01:21:32)#
好的,我们大概讲了一些背景。然后接下来我们想看一下 AI 编译器整体的一个工作环境要做的事情有哪些。然后我们刚才讲了,我们主要关心的就是是一些编译器,然后我们想要参考过去 LLVM 的思路,是在某一种表达上的这种一遍一遍去做优化的这样的一种 AI 编译器。然后我们前面可能是各种语言,比如说 Python 作为我们的前端,然后后面是各种硬件去作为我们的后端,那具体来说我们就会关注到所有这些都是 AI 编译器所关心的事情。然后具体要执行的功能,比如说前端我就要做计算图的一些生成,去做自动微分的一些图的扩展。然后我可能要把各种各样的原语转换成我的计算图等等,这是我前端所要做的所有事情。
然后接着我在优化器的部分,尤其是跟前端紧密相接的,是我们后面要讲的就是各种各样的面向计算图的优化。比如说循环优化,内存分配等等。那等到后面面向后端的时候,同样还有这个编译器的 AI 的这个编译器后端可能要做一些图算融合,一些张量化算子的一些具体的一些优化的策略等等。等真正到了后端的时候,那么后端需要做的一些事情主要是在硬件上面。我要做一些比如寄存器的分配,然后一些的这样的一些策略等等这样的一些工作。如果到了我们的这个硬件之上,然后具体根据自己的每一个硬件不同的驱动,它本身也还要再走一下硬件自己的编译器,把我们的这个 AI 编译器生成的这样的一些命令,转化成硬件上能够去进行执行的这样的一些语言,那如果你在 CPU 之上,很可能你就会把你的命令又转换成了传统的编译器,比如说 LLVM 的也就一些中间的一些表达。那如果说你是一个比如说 NPU 或者 GPU,那很可能也会对应一些比如说 CUDA 或者是其他的一些 IR 的一些中间的表达。然后最后我们可以在这个等机器上去进行一些相应的执行等等。
所以我们这里想说的是,在我们整个 AI 编译器的最底层,我们会发现对于传统的这个编译器,我们可以去进行一个复用。而对于我们整体 AI 编译器的一个构架上,我们也会发现在很多概念上它会类似于传统的编译器,那具体我们一会儿会关注的,主要是我们一层一层的优化,那也是走好几个 pass,也是很类似于过去你可能会有向量组合代码消除。这些也都是我们 AI 编译器,尤其是 AI 编译器前端想做的一些工作。
好的,在这做一些对比的时候,我们还想先回顾一下传统的编译器,比如说 LLVM 它的一些中间表达者长什么样子。然后由此来看一下我们计算图,到了我们 AI 编译器,我们的中间表达又长什么样子。在过去就是 LLVM 它的一个中间的一些表达,比较常见的可能有一些这种对账机代码或者三地址的这种代码等等。比如说三地址代码的话,它具体就会是一些操作数和一些运算符合目标这样的一些组成的这样的一些中间的一些原语。有了这些中间的原语,它就是我们 LLVM 能够去进行优化的一个最最基础的这样的一个表达。好的,那我们来看一下,这是过去的以这样的一个三地址代码就能够支撑过去 LLVM 的各种各样的一些优化策略。但是到了我们这个 AI 的编程当中,我们就会发现,比如说最开始你要构建计算图,然后就后来你要做各种各样的图优化。然后到最后你可能还要再面向 LLVM 的表达去做一些优化。
那有一个问题就来了,就是我们在 AI 的编译器当中最核心的表达到底是什么?最开始大家肯定还是想到的是一些 graph IR,就是类似于我们计算图的这样的一种表达。最开始就是可能有一些面向应用的一些原语转换,然后可能会有一些抽象语法树,这些抽象语法树它也可以是一些类似于计算图的这样的一种抽象表达。然后这种抽象语法树它可以进一步地被优化。优化出来的就能够给我们提供服务各种各样的计算的一些计算图的一些表达。
这显然是 graph IR 是属于 AI 编译器,和传统翻译编译器不太一样的。我们独自新有的而且很重要的一系列的这种表达是我们的这个图,然后我们当然也已经在之前的课程当中,有了非常多的一些理解。然后我们认为这个图的这样的一个表达,对于我们 AI 编译器来说,它很明显比过去的那种,比如说三地址的那些代码,更适合我们的这样的一个面向 AI 的一些优化的一些逻辑。比如说你可能想要优化一些控制流等等的这样的一些事情。但是我们后续也还要支撑,比如说线性的一些 IR,所以其实就是我们通常在图当中可能就只能到达算子的级别,也就比如说矩阵乘是这里的一个算子等等。
然后等到更底层我们就会发现我们还需要一些其他的一些表达。好的,因为我们既需要计算图的表达,又需要一些更底层的表达。所以在一些编译器当中它就出现了一种使用混合表达的一些逻辑,它想起来也比较直观。
那具体的一些想法就是说对于整体宏观上,比如说控制流等等我都拿这个计算图的思路去进行表达。但是对于某一个步骤当中你具体做的事情,比如说你做一些赋值等等这样的一些事情。我就会采用一些更类似于以前的一些计算原语的一些中间表达,一些编译器中间表达的一些方式去进行表达。那对于这种混合的表达,它其实就是一个复杂的结构,可能就是说它整体看起来依然是一个计算图。但是在每个节点上我做的是一个根据算子原语的一些扩充,就等于是每个算子具体做的事情。中间可能比如说是三地址的这样的一些码,我就把它放在了这个中间的我的某一些叶子节点上面。这样的话我整个图就完完整整地包括了我 AI 编译器需要关注的所有信息,就是我的运算逻辑以及具体的执行的一些运算。所以它是一个混合表达,我们认为是一个完整的表达。
但是并不是所有的 AI 编译器都会采用这样的相对来说比较混合的一种表达。对于 PyTorch 来说的话,它其实是阶段性的,也就是说在 Python 里面你是能够找到这样的一些混合表达的那如果说你可以去给这个图去静态地追踪一下,然后给他去转换一下它的相应的一个图的一个打印。然后它的图你就能看到有一些是一些基础的一些操作,也有一些它其实涉及到的是一个一个比较大的一个循环这一类的代码逻辑。然后在这个代码逻辑里面,它可能就是一个 block 去执行某一些的操作。整体我们就能看到就是 PyTorch 它本身是存在这样的一些混合表达的。
好的,那我们这里就想说对于 AI 编译器来说,图表达很重要,底层的表达很重要,混合的表达也很重要,那这时候就出现这样的一些不同,比如说 PyTorch 它就选择第三种,中间表达都存在。尤其是图的表达,它主要是在构建图的过程当中产生的。然后接着为了优化,我逐渐把这个图扩充成一种混合式的表达。如果我已经优化完了,我把它切割部署,最后在部署的时候,最终它就只剩下了适合硬件的这样的一种 LLVM 的一种最底层的一种表达。所以 PyTorch 的话它整体存在同时存在三种和三种这一类的中间的表达,它的整体逻辑其实是比较复杂这一类的思想。对,然后我们这里面虽然提到了传统这个 LLVM 的中间的一些表达,但是它的计算也已经并被进行了进一步的扩展,尤其是我们现在经常需要一些比如说矩阵的一些乘这样的一些计算。所以我们也把过去的一些单值的一些加减乘除,扩展成了一些当量的加减乘除这样的一些运算。那我们可以看一下,如果说你确实是同时混合使用多种表达,就是 PyTorch 的这样的一个策略,那它其实是存在多种优点和缺点的,那它最大的优势其实符合 PyTorch 的一贯的主张的方向。就是我一定要满足最好的一个灵活性,可以在各个层面上去接各种各样的一种优化的手段。
你可以去优化图的,也可以专门优化算子的,总之我是灵活性是最好的。然后同时他也给用户提供了一种层次的这样的一个结构,用户就可以比如说我只关心图,或者说我只关心某一部分算子等等,这都是它的优势。但它最大的劣势就在于就是同时维护三种不同的表达,中间还要在适时地对这些表达去进行相应的转换。具体如果说你上来有一个定义的一个用户的一个 function,那这个 function 最开始就是要构建图的表达,然后图的表达接着就要有一些编译器的操作,把它转化成一个混合的表达。然后接着还有一些更底层的编译器的工作,把这个混合的表达迁移到我们的硬件上等等。你就会发现整个转换,它不仅是实现起来比较复杂,而且在层与层之间做转换的时候,它相应的也会比较耗时。
然后除此之外还有一些优化,就是如果在图表达的时候,你必然只能做跟图相关的优化。而你在扩充成算子表达的时候,这时候你相应的才能做一些算子相应的算子上面的一些优化。所以如何让你的这个复杂的表达去对应你的一个优化策略,也是这里面比较复杂的一些事情。总的来说就是 PyTorch 它为了灵活性,然后确实面临这样的一些一些挑战。然后我刚才跳掉的一些也有一些 AI 框架,比如说华为的昇腾,他们就全面推广的是支持统一的一种表达,其实就是混合的表达。
因为如果你只支持统一一种表达的话,那必然是一个是图和算子都相互融合,那它它的缺点肯定是就是这个图整体会是一个非常庞大的一个图,维护起来会比较困难。但是它的优化就会相应地会灵活一些,就是你可以随时随地地把这些算子不同的快去进行协同。这样的一些优化等等,然后也是昇腾他们在推的一个主要的一个思路。
AI 编译器中的图优化(01:33:07)#
好的,那我们大概就讲到了这个 AI 编译器当中的各种各样的表达,然后接下来我们就关心的是这个 AI 编译器当中需要做到的一些事情。然后关于构建的这些我们都已经了解过了。然后接下来我们主要关心的就是 AI 编译器作为一个优化器,它该怎么去做一些优化。然后就像我们说的,就是所有的这些都东西,我们在后续的课程当中可能会主要以 PyTorch,尤其是 PyTorch 2.0 来作为我们的一些例子来去做一些举例。然后我们想说 PyTorch 2.0,它也是加入了 Linux 的一个 foundation,所以现在也是一个开源的这样的一个状态。
然后我们前面也提到了一点,就是 PyTorch 2.0 一个很大的一个优势是在他提出了 PrimTorch 的这样的一个算子的一个分层的一个逻辑。把整体广泛的两千多个算子总结成其中必不可少的这 250 个最基础的一个算子,方便我们比如说有一些硬件,就可以最开始先去让这 250 个算子可以在这些硬件上面跑。然后对于我们编译器来说也是一样的,就是最开始可以先选择去支持这 250 个算子,它的一些算子优化策略等等。然后逐渐地再去向更多的一些算子去进行一个一个扩展。
AI 编程的软硬件协同与大作业(01:34:29)#
好的,那这里提到 PyTorch 的话就是我们在提到 AI 框架,就是整体我们可能对它有了一定的理解。然后有这样的一个宏观概念的话,其实我们会希望这个课程可以给大家提供一些就是当大家如果面对面向未来的一些 AI 应用,比如说你去设计了一个新的一个算子等等,就是希望大家能够了解到有哪些事情是可以做的。比如说你设计了一个新的算子,然后你希望它能够在 AI 的生态当中被广泛地利用,那就意味着你要给这个算子去打通它前面和后面的所有的事情。我们一般是从前向开始来的,比如说你这个算子天然的就能放在很多 Python 代码里面去跑。那对于后端,我们就可能希望有专属的一些硬件的核的计算能够支撑这个算子。那这样的话你要做的事情可能是,比如说在这个算子上实现它的一些算子级别的一些实现,以及这个算子级别的一些面对我当前硬件的一些特有的一些加速,然后当前这个算子它和其他的算子是不是能做一些优化,这可能是在 AI 编译器当中能做的一些事情。所以就是了解到这个全局,可以有助于我们如果未来做一些工作的话,了解到哪些是与它相关的这样的一些事情。
好的,另外还想提的一点就是关于我们课程的大作业,大家可以在这边再看一下大作业的逻辑。那等于是我们想说,对于我们大作业的实现来说,我们实现的更像是一个 PyTorch 1.0,那我们的输入就是各种各样的 Python 的一个原语。然后我们去进行动态的建图,然后通过建图,通过自动微分的构建,它自动地去调用底层的算子库。然后同学们也在前面的 lab 当中去手动的实现了各种各样的算子,以及算子的加速的一些模块。所以我们整个走的相当于是最右边的这样的一条路径。就是从 Python 的源头,然后走 runtime,最后走到我们的硬件。而我们 AI 的编译器,它其实是一条与之相并列的这样的一条路径。
对于编译器来说,他走的其实会是包括比如说图的一个构建。尤其是我们之前讲到的这个基于 function 的一些子图的全局动态,中间子图的一种构建等等。然后接着他会做一些尽可能做一些我们后面会讲到的各种各样的优化。然后会做做面向后端的硬件的一些部署和专有的优化等等。好的,那我们就认为所有 runtime 这些它都是更适合我们的这个程序员的 PyTorch 本身的 1.0 的这样就有的这样的一些逻辑。而这样的一条逻辑就是它这里面核心的依赖,就是依赖一个很好的一个算子库。就是在 runtime 你能做的优化是非常有限的。
唯一的优化都来自于我们的这个算子库。比如说 cuDNN 等等这一类的库,它能够提供的一些优化,或者是我们手动实现的一些优化。所以相对来说这个 PyTorch runtime 这一端的生态是比较稳定的,绑定在了英伟达 CUDA 的这样的一个生态当中。但是对于我们实际的产业界的应用来说,就是有很多事情。比如说大模型的训练或者是大模型的应用,对于这一类来说,它相应的它不是一个频繁迭代改变的一个需求。然后它很可能很多地方是可以通过 AI 的编译器去进行一些图的固化,去进行一些硬件的加速,然后起到一个很好的整体的一个调优的一个效果。所以对于左半边的这一派的这一条研究路径,其实是当然英伟达也占据很大的一部分市场空间。然后除此之外也是许多其他的 GPU 厂商非常极力去支持的,尤其是我们现在有非常多的国产 GPU。然后对于国产 GPU 来说,如果能够支持大量的,比如说大模型的一些哪怕是 inference 或者是它的,如果还能支持它的训练,就已经能够完成我们现在非常多的一些任务。所以在整体在左半边的这个研究路线上,我们会发现,当然英伟达也在推。然后各个公司也都有非常多的一些工作,包括华为昇腾在内,他们也是沿着这个路线做了非常多的一些开源的贡献等等。
Triton 与 AI 编译器工作流程(01:39:12)#
好的,我们刚才讲了一些 PyTorch 2.0 的内容,然后最后想提一下的是这个 PyTorch 2.0 的它的这个 TorchInductor,作为一个面向硬件的一个编译器,然后我们刚才提到了它不仅想要支持英伟达的显卡,然后也希望对 AMD 或者是或者是 Apple 的显卡有一定的兼容性。然后这时候它就需要使用一些面向硬件的一些编程语言,就是使用到了 OpenAI 的 Triton。其实这里面也比较有趣的一点就是 OpenAI 它作为一个主要是比较应用层面的一个 AI 领域的研究的一个公司。他们因为他们很巨大的一些算力的需求,所以他们研发了 Triton 这样的一些编译加速的这样的一些表达,然后这个表达也被 PyTorch 引用进来。就是如果大家想理解一下,就是 Triton 它做些什么事情,其实它的一个想法就是说我们会希望我们知道对于很多 GPU 来说,你需要编写很好的 CUDA 代码,它才能够比较高效地去用这样的一个一个硬件。但是对于一般的用户来说,有可能他们的编码能力、编程能力没有不能写出特别高效的一个 CUDA 代码。然后 Triton 的一个想法就是说它希望能够让用户写一些很简单的像 Python 一样的这种编程语言。但是它把自己通过一些编译器的转转码转换加速,给它转换成一种更高级的,更 CUDA,更适用于 GPU 并行计算的这样的一些语言转换的一些技术,一些技术。
所以这是他们主要研究的事情,然后他们也就被 PyTorch 发现,非常适合以这样的一种表达为基础去支撑除了 CUDA 以外的其他的硬件。所以这也是他们被引用到了 PyTorch 的编译器当中的比较重要的部分。好的,接下来我想讲一下的就是了解了这么多 AI 编译器要做的一些内容,然后也了解了前端想要做一些优化。然后我们也提到在前端的优化上,我们想要参考传统的编译器去做各种各样的 pass。我们就可以看一下比如说一些操作,就是某一个 pass 等等,它是在 AI 编译器当中是怎样被被触发的,是怎样的一个优化流程。
我们认为整体可以是走这样的一个流程图,最开始就是我们的这种用户提供的代码语言,然后我们要进行原语转换,转换成 AI 编译器所支持的中间表达 IR。然后最开始往往它应该是一些图,基于图的 IR。好的,有了这些图的 IR 首先第一步要做的是一些静态分析。这一步其实跟我们传统的编译器的静态分析是非常类似的。主要分析一下这个对象它的类型,它的大小等等这一类的事情。然后如果确定没有什么问题,可以安全的继续执行的话,那我们接下来就会做很多的基于跟优化相关的这样的一些 pass,一遍一遍地把我们这个中间的表达变得更加的高效。然后接着就是你可能做了一定的前向图的优化之后,我们就可以用 autograd 的这个方式去把我的反向图构建出来。构建出来反向图之后,就可以再去过一遍所有的这个优化器。最终我就会获得一个完整的前向图和反向图的一个表达。
最后这个图的表达想要做的事情就是面向各种各样的硬件。我去做它的一些面在硬件的各种派遣,以及执行,同时也是要利用用硬件的一些优势。所以最终我要把我整个图切割成各种各样的子图,子图其实借助的就是一些 subgraph 的一些重点表达的这样的一些形式。这整体就是我们 AI 编译器前端从源代码到后面子图的一些切割,整个要做的一些事情。
好的,最前面的这个静态分析和这个静态检查,都是类似于我们传统编译器的。然后它的目标就是说用户可以写得更简洁一些,然后我们编译器来及早地发现一些问题,比如说你这里面是不是有一些死循环,然后可能可以给用户去进行一些提示,又或者是我可以去做一些类型的检查,这里面会不会有一些类型的转换,然后出现了类型不能兼容的一些情形。又或者是有的时候你可能就是写得比较冗余。我可以把一些常用的,完全不变的一些常量去转化成固定的值,这样避免一些不必要的一些额外的计算等等。所以最开始的这个静态检查是比较传统的一些一些功能。
图优化策略与常量折叠(01:44:19)#
好的,那静态检查完了之后,主要是我们图的各种各样的优化,那这里面呢图的优化,我们也可以把它整体地大概地分为两大类。一大类是就是跟我们计算图更相关的,就是随着神经网络的需求出现了以后,我们需要做的一些跟图相关的优化。还有一大类是传统的普通的编译器本身就要做做的一些优化。同时在我们 AI 编译器当中也依然非常有效,只不过这时候我们的表达变成了一些图的中间的表达。这一类的也有,比如说常量的折叠,以及公共子表达式的消除,也就是我们公共子图的一些利用,以及死代码的消除等等。这些都是传统的编译器会涉及到的一些内容。好的,那对于神经网络自己的一些新的那些内容,主要是一些神经网络自己的算子,可能要做一些算子的融合,然后我们还希望能够做一些更好的内存分配等等。
我们刚才已经提到了,就是所有的这些 AI 的优化,编译器的优化都可以分 pass 去进行,那到底以怎样的一个顺序,就是他们每个都是一个单独的 pass,然后你以怎样的顺序去调度这些不同的 pass?其实就是很难提到,很难提出一个最优的一个策略,这有可能跟我们自己具体的代码都非常有关系。但是一个比较常见的 pass 就是他调度的顺序是我们图下面的这样的一个顺序。它也是借鉴了我们传统编译器的一个顺序。
可能最开始要做的事情是一些常量的折叠。然后接着就可能就是这样的话,你可能消除一大部分计算。然后接着我们做一些算子融合,让我们的计算都更加均匀一些。然后在这些融合了的算子之上,你有可能要做一些布局的转换,使得它的计算更加融洽。然后后面还有一些计算的一些简化,比如说公共子表达式的消除等等的这样的一些简化。
然后最后我们还会有这个就是等到最后一个 pass,我们一般会做死代码的消除。也就是说在我们所有优化完了之后,有可能这些死代码它未必是用户写出来的死代码,很有可能用户写的话代码并不会有这么多的冗余。但是就是经过了一遍一遍的优化之后,很可能会存在一些额外的并没有去并没有参与到最终结果计算的一些子图上。然后就可以去采用我们的这个死代码消除的这样的一些策略。好的,然后我们提到了这一大部分是跟这个神经网络相关的,一大部分是跟我们的传统的编译器相关的。
跟传统编译器最相关的我们也可以举一个例子,比如我们的这个常量折叠的这样的一个例子。那在传统的编译器当中的时候,我们很可能就会有一些常量,然后这些常量会参与后续的一些运算。所谓常量的折叠,其实就是说把这些常量按照它的数值的形式给它去做一个传播。然后替换掉这些它在代码当中的一个约束。甚至你可以直接把这个都是常量的运算直接去运行,这就是常量折叠。然后整体来说我们就会能够更加高效一些,这是传统编译器就有的一个举例。然后在我们 AI 当中,它就有了它自己的一些特例。
我们这里想举的一个例子也是一个 Batch Normalization。具体来说它主要是瞄准我们在比如说神经网络训练结束之后,然后可能这个时候我们的 Batch Normalization 这个层,它已经经过训练获得了它自己特定的一些均值和偏差等等这样的一些参数。这样一来就是在真正使用在训练结束之后的使用过程当中,这些值就都变成了一些常量,那么在这些常量上,其实就可以做一些相应的一些一些常量的一些转换。可以展示一下。那比如说我们认为就是比如说传统的你可能经常会有一个 Convolution 的计算,可以写成 的形式。然后接着你可能在这个 上面去做了一个 Batch Normalization,那就是 减去一个中间平均值,然后再除以方差的一些计算,然后再重新地去有一些成本来的一些参数在这里,那我就会发现可能整体有,这就是原始的做一遍 Convolution,做一遍 Specialization 的一些计算的操作。
而我们如果说想要用一些常量的转换的话,你就会发现这里面很多 Normalization 的一些相关的参数,在训练完了之后,我们都已经确定了它就是一些已经固定的这样的一些常数。而如果说你已经训练完了,那么你的卷积很可能也就是一个普通的线性的一个变换。那这样的话我们就可以把这两个算子去进行一定的融合,比如说我们就能把一些之前的线性的变换 直接带到我们的这个式子里。然后相应地你就可以去进行一些整理。最后我们就可以以一个比较简单的把一个比较简单的一个 Convolution 操作,统一地去完成我们之前说的既有 Convolution 又有 Normalization 的这样的一系列操作。
好的,举这个例子我们想说的是就是神经网络尤其是在训练完成之后,它的很大一部分,它的比如说权重等等都变成了一个常量。所以尤其是在 inference mode 的时候,这种常量的折叠传播都是非常能够去进行优化,有非常大优化空间的这样的一系列操作。好的,那就是在 inference mode 下,我们认为很多都是常量。
如果不是 inference mode,即使是在 training mode 下面,对于我们神经网络来说,比如说神经网络的形状、大小,然后 rank 这一类的,也有非常多的一些节点。它有可能是一些常量的节点。那么对于这些常量的节点,比如说你在训练的状态下,很有可能是可以把它替换成一个一个常量的。然后完成常量的替换以后,我们就可以做非常多的基于常量的大小相关的一些折叠的一些操作。然后可以实现一些操作的一些加速。这里面我们想提的是说,可能有一些它的这个形状会作为都是常量,但是它的 shape 有所不同。但是我们这个增量的相互计算当中,是有一定的 broadcasting 的一些技术的。所以对于一些各种各样的一些计算,我们也可以通过 broadcasting 然后来去进行我们常量的折叠的一些兼容的一些知识。
算子融合与课程展望(01:52:20)#
好的,除了常量的折叠,后续一些比较重要的事情是这个算子的融合。算子的融合这个部分就属于是我们神经网络带来的一些很重要的一些优化的策略了,我们后面想要举的一些例子主要是 XLA,就是我们前面提到的 TensorFlow 的一个 AI 编译器,他主要是提出了非常多的算子融合的规则。然后除此之外我们还想介绍的是一个 TVM 的一个基于基于这个数目的一些优化的一些策略。然后可能这个就放在我们明天的课程再去进行讲解。好的,那我们就提前一点点下课。然后接下来以后明天再去讲和算子的融合,谢谢大家。
附录#
关键点和注意事项#
torch.compile在使用前需要reset()来完成初始化。- 调试
torch.compile时,建议先使用backend="eager"模式,成功后再尝试backend="aot_eager",最后再使用默认的backend="inductor"。 - 大作业的实现更像 PyTorch 1.0 的动态建图和自动微分,依赖底层的算子库,而 AI 编译器是与之并列的另一条路径,关注图的构建、优化和硬件部署。
提示词和模型#
请严格按照以下要求处理并输出我的请求:
根据以上转录稿以及PPT的内容,帮我整理一份文字清晰、术语准确、逐字逐句与原文对应的转录稿。使用散文格式,禁止使用列表。原文本中有非常多识别错误,需要你结合上下文进行猜测。首先确定文本主题,然后想想都有哪些关键词,最后想一想原文中错误的字和哪个关键词相似,然后替换。整理过程中,你还需要去除重复的话和语气词。仅用三级标题 (###) 来创建清晰的章节结构,与主题内容相对应,并附带时间戳。如果转录稿中讲师特别提醒要注意某些内容,将这些内容汇总在输出的最后一个三级标题 ###关键点和注意事项 中。LaTeX 使用: 对正式的或复杂的数学公式、形状或符号(例如 k×k、[N,Cin,H,W])使用 inline LaTeX 格式进行标记。将转录文本和经过 MinerU 处理的课件放入提示词中。
模型:Gemini 2.5 Flash (via API).