异构计算视角:深度学习框架的历史演进与展望

前言

人类对计算能力的需求永无止境。从历史上最早期手工连线用于计算弹道的大型机,到打孔纸带编码的科学计算程序,再到第一个高级语言fortran的出现,直至如今种类繁多的库、框架以及编程范式。硬件的运算性能越来越强,同时在软件上也逐步发展出更高效、更灵活的技术来充分利用硬件的能力。

近年来,随着深度学习的重焕新生,以及互联网发展所积累下的海量数据,使得传统的通用处理器已经无法满足对算力需求的指数级别增长。与此同时,真实业务中多样化的模型部署场景,例如对延时、功耗、成本的综合限制,也共同为软硬件技术提出了新的挑战。遍历市面上主流的深度学习框架,可以发现他们都包含了大量的抽象设计和代码,用于同时支持CPU和GPU甚至是其他更多种类的硬件架构。“以古为鉴,可以知兴替”。因此,站在异构计算的视角来对这些框架的设计思路和演进历史来进行梳理是很有意义的,籍此我们也能够大致预测深度学习相关基础技术的未来发展方向。

第一代

在Caffe出现之前,训练较大规模的卷积神经网络需要研究人员自行针对特定架构设计异构计算程序(例如早期的convnet),这需要一定的编程功底以及对相应体系结构的理解,事实上不仅开发效率不高,也不利于成果复用。随后,以Caffe为代表的第一代深度学习框架,用通用的protobuf文件作为模型描述方式,以及结构上相对分离的简洁、高质量的C++/CUDA代码实现,解决了如何快速搭建模型,并使之高效运行在GPU这种异构硬件上的问题,从而满足了计算机视觉领域巨大的算力需求。

然而,深度学习领域种类繁多的模型拓扑结构并不能全都用Caffe中简单的“pipeline of layers”的模型加以抽象。另一方面,Caffe也不能很好地满足多样、广泛的算子定制需求,新增一个算子需要分别开发其ForwardBackward过程,并且想要加速的话还得自己编写CUDA kernel,这依然对研究人员的异构编程功底提出了一定的要求。因此,各个大厂甚至是各个算法团队都会有fork出来自己维护的一套caffe框架。

第二代

以MxNet为例的第二代深度学习框架从两个方面对上一代技术的不足进行了改进。从用户API的层面,主要是使用静态数据流图,而不是简单的pipeline模式,来建立深度学习模型。这样的抽象使得用户可以用更灵活的方式来搭建模型,并且不用关心如何反向求导等细节;而从框架内部实现的角度,则能够从更底层、细粒度的层面来对整个数据流图进行变换优化,例如内存分配、数据重排、设备规划等等。

从异构硬件的使用方式的角度,这一代框架多数采用了Expression Template(例如mshadow)的设计模式,其思想主要在于将算子的数值算法实现与调度执行方案相解耦,从而使用模板语言封装掉异构硬件与CPU的平台差异。这样就可以在开发新算子时,用已有的模板封装来高效地写出异构代码,而不用将同样的算法分别实现两遍。然而,该技术的缺陷其实也很明显,注意到我们上面的用词是“高效写出”,而不是写出高效的异构代码。这是因为Expression Template毕竟是C++语言的编程模型,它是编译期静态的——在代码上下文中既不能得到运行时信息,也不能得到模型结构信息,导致这样实现的算子完全无法被优化。这种缺陷的根源在于,我们试图实现一个支持任意运行时参数的通用数值算法,但却只能使用编译期间静态确定,极其受限制的C++模板语言来描述它,所以除了实现功能以外,做不了更高级的事情。

整体而言,第二代框架提供了更加灵活友好的用户接口,并且也加入了分布式的支持,在性能和内存占用方面亦有所提升,可以说基本满足了常规模型训练用户的需求。

第三代

普遍观点认为Caffe和Tensorflow分别代表了第一代和第二代深度学习框架,但关于如何界定第三代技术,则存在一定的意见分歧。从用户接口层面,很多观点认为以Chainer和Pytorch为代表的,对“imperative program paradigm”,即所谓“动态图”的支持代表着下一代框架的主要特征。事实上,各个成熟框架也纷纷在上一代技术的基础上扩展出了类似的接口,例如Tensorflow Eager、MxNet Gluon、Paddle Fluid等等。但是,这些技术本质上是以部分程度牺牲性能来换取用户使用上的灵活性,并不能算是真正意义上的技术革新。

从异构计算的角度,可以认为Tensorflow XLA(及其衍生品JAX)、TVM、Tensor Comprehension(下简称TC)等深度学习编译器代表了深度学习框架的第三代技术,他们在保留了灵活的上层用户接口的同时,尝试通过更高级的编译技术来彻底解决异构算法程序的开发与优化问题,充分挖掘硬件的计算性能。

Tensorflow XLA将深度学习领域一些常用的算子抽象为HLO IR(High Level Optimizer),其本质上就是一系列特别细粒度的算子,用它们可以组合出任意的实际算子,例如,常见的batch_norm可以由一系列的broadcastreduceelement_wise等操作所组成。在XLA中,将HLO各个算子映射到LLVM IR的思路可以概括为:在信息更丰富的上下文中,用更抽象的方式来写C++模板代码。XLA的主要不足也在于这里,它本质上是采用LLVM IR来描述每一个XLA OP的具体算法的,虽然能直接复用LLVM中已有的一系列优化,但部分程度上限制了其扩展到其他硬件架构的能力。事实上,LLVM IR本身的设计目标是用于抽象不同的处理器架构,这导致从低层次的IR中,我们已经完全丢失了算法本身的结构信息。

TVM采取了更加激进的技术路线,它不仅包含类似HLO的细粒度计算图表示NNVM(后续会迭代到Relay IR),同时也加入了Halide IR作为一种比HLO更加底层的IR,从而能够直接用定义清晰的领域特定语言(DSL, Domain Specific Language)来描述所有算法。后端可以针对已有的算法描述来创建调度方案(Schedule),并采用一系列调度原语来进行优化,从而得到较好的性能。优化过程既可以是由了解硬件架构的开发者手动编写,也可以通过AutoTVM针对硬件架构抽象出通用搜索规则,从而实现半自动化的代码优化。采用抽象DSL的最大优势在于容易扩展到不同类型的硬件架构。

TC的技术则是其中最为激进的,它同样是基于Halide IR,但是在代码生成阶段直接采用编译器理论中的多面体模型优化技术,通过仿射变换将多重循环映射到高维向量空间,并基于硬件架构特性来进行全自动化的循环变换与优化,整个过程可以完全无需人工干预。

上述这些工作试图利用不同类型的编译器技术,来尝试解决异构程序的优化问题,并且在很多场景下能够取得高于手工优化的效果。无论是从技术深度还是实际应用场景上面,比起“动态图”这样用户接口层面的改进,它们更能够代表这一代框架技术。

展望未来

在深度学习框架百花齐放百家争鸣(总感觉下一句是“引蛇出洞”),并且整体上趋于稳定成熟的今天,也有观点认为这只是各个大厂在AI时代自己造轮子推标准,试图抢占开发者市场先机的行为。而且大家多数无非都是cuDNN、cuBlas的高层次封装而已,其区别仅仅在于上层API是否好用灵活,却难以掩盖整体上同质化严重,缺乏新技术与新思路。但如果我们回顾以上深度学习框架的发展史,站在软件与硬件结合的角度去思考,就可以清晰地理出一个脉络,看到软件框架是如何与异构硬件共同迭代进步的,以及接下来的创新点何在。

第一代深度学习框架,解决了如何用异构硬件来满足重焕新生的深度学习领域巨量的算力需求。相比之下,一些算法特性上面并不特别适合异构加速的传统科学计算领域,例如物理海洋学中的大规模流体动力学数值模拟,仍然不得不借助于MPI集群等传统并行计算技术——整个集群的峰值浮点性能可能还比不上一块Kepler计算卡。第二代框架采用更高级的抽象机制与设计模式,极大地降低了开发异构算法的复杂度,但受限于C++本身的特性,这样的实现在性能上很难达到最优。第三代框架引入了自动化的代码调优方式以及更加广泛的异构硬件平台,极大地降低了深度学习模型的移植、部署成本。

站在用户需求的角度,这三代深度学习框架的发展史,就是从满足算力需求,到提升开发效率和灵活性,再到追求(各个架构上)最佳硬件利用率的发展过程。回顾先前的过程,其实一直是软件开发者在想方设法用(gui)好(tian)异构硬件,但这些已有的硬件架构并不一定在各个场景下面都能很好满足需求。那么自然而然地,下一代深度学习框架应该不再仅仅是软件框架,而是一整套用于通过软件与硬件的协同设计来将各种场景下的异构算力发挥到极致的完善解决方案。事实上,在上一次神经网络发展的兴盛时期,学术界就已经在这方面进行尝试和探索。但在当时,通用处理器的发展如日中天,如果开发者和用户只需要坐等一年就可以得到翻倍的模型性能,那么大家必然不会选择通用性和易用性都满是坑的专用加速器。时至如今,随着摩尔定律的失效,越来越多的科学计算领域已经不能再指望着这样天上掉馅饼的好事,引入领域特定架构(DSA, Domain Specific Architecture)必定是未来的发展趋势。Patterson近年来的报告和讲座在不断强调这一点,认为这种趋势代表着计算机架构的黄金时代再次到来。

但是,站在现实一些的角度,领域专用架构必须达到足够低的流片成本、足够高的硬件开发效率以及足够易用的软件工具链,才能够在科研以及实际应用中真正具备竞争力。在很多场景下,计算只是手段而不是目的,如果不能以可接受的开发成本来实现目标,那么用户宁可转向性能较差,但开发效率更高的通用方案。目前为止,在科学计算方向仅有深度学习领域出现了TPU这样具有足够竞争力的定制架构,这是因为这一领域巨大的经济效益能够支撑得起当前芯片和完整软件栈研发的巨量投资——事实上XLA的研发成本可能还要高于硬件。

DARPA所提出的电子复兴计划(Electronics Resurgence Initiative),正是在试图解决上述的前两点问题,希望在未来十年内大幅度降低流片成本,提高硬件研发流程的效率,并建立如同软件一般繁荣的开源IP生态。但是,即便我们能够降低设计新硬件架构的成本,依然存在的一个问题就是,我们不可能在流片前实现完所有的算法,确保它们都能够以高效率运行在这个架构上面,更不可能在短期内迭代出一套完整易用的软件工具链。既然存在这样的制约,又如何保证领域特定架构具有真正的竞争力呢?在这方面,Intel的VLIW(超宽指令字)架构就是一个典型的失败例子——硬件架构师认为可以将解决数据依赖的重担完全交给编译器,但事实上问题并没有这么简单,对于通用计算的场景,编译器没有能力横跨多个上下文,大幅度进行指令甚至代码逻辑重排来消去片外访存的随机性。因此,硬件架构的设计不能做无根据的简化,必须充分考虑现有编译技术的能力上限

站在软件的角度,给定数值算法,要求其高效运行在不同硬件架构上,这个过程本质上还是编译器中的调度规划问题。不同硬件架构具有不同的调度特性,而近年来兴起的一些领域特定语言的设计目标恰好就是抽象出这些特性,从而将算法实现与硬件调度相解耦。因此,领域特定架构应该针对领域特定语言的描述范围来进行设计。这样,已有的DSL编译工具链能够遍历所有支持的算法描述,并且引入一些机器学习以及启发式搜索算法,针对目标架构生成、优化可执行的代码,从而快速评估架构的效率和性能。最终,当确定一个基础的硬件架构以后,基于DSL的设计流程将能够形成这样一个全自动的闭环迭代:

  1. 通过深度学习模型调节硬件架构中的超参数
  2. 使用高级硬件描述语言(如Chisel)快速生成硬件RTL与模拟器
  3. 编译器基于算法的DSL描述,快速为硬件架构生成所有代码实现
  4. 评估算法性能,将相应结果反馈给第一步的模型,从而开始下一轮迭代

这样,我们就有可能解决软件易用性与优化效率的问题,使得领域特定架构具有真正的生命力。此外,需要注意的一点是,DSL这类技术本质上是一把双刃剑。并非在任何场景下都适用,例如这篇博文就对那些不适用的场景进行了批判。但是,站在异构程序开发的角度,如果我们能够用一套DSL来抽象出多数算法,并且也能大幅度提高其开发、优化效率,更能够辅助专用硬件架构的设计,那么就有理由认为这样的DSL是必要的、成功的。

站在更长的时间尺度上来预测,对于各个科学计算领域,未来有可能抽象出各自不同的DSL,并且能覆盖其中绝大部分的算法描述、调度和优化需求。例如,以矩阵乘累加运算为主的深度学习领域,代表其计算特性的DSL一定会与涉及大量偏微分方程迭代计算的流体力学领域大相径庭。未来,当专用硬件设计迭代的流程更加成熟,流片成本大幅降低,且与软件工具链深入整合之后,我们就能真正迎来用异构计算取代通用计算的新时代。

comments powered by Disqus
Published:
2018-12-31
分类:
Tag: