谢志强作者

OSDI 2020 | RAMMER如何进一步“压榨”加速器性能

编者按:传统的深度学习框架为了模块化设计,通常使用互不感知的两层调度模型,导致无法充分发挥硬件的计算性能。针对现有深度学习框架的局限,微软亚洲研究院和北京大学、上海科技大学合作提出了 RAMMER:一种可以成倍甚至几十倍地提升深度学习计算速度的编译框架。这套编译框架的背后是微软亚洲研究院打造的深度神经网络编译器 NNFusion,目前已在 GitHub 上开源。

传统深度学习框架由于其自身的局限性,如今还远没有充分发挥出硬件的计算性能。微软亚洲研究院的研究员们在一些测试集上发现,现有的深度学习模型只能用到 GPU 2%到40%的性能。传统的深度学习框架通常把神经网络计算抽象为由算子(operator)与依赖关系构建而成的数据流图(data flow graph, DFG),并按照拓扑序将算子逐个调度给硬件,或像 MXNet 中的依赖引擎(dependency engine)可以将多个算子同时调度到 GPU(算子间并行性,inter operator parallelism)不同的流(stream)执行。而在此之下又存在另一层调度器负责充分发掘算子内并行性(intra operator parallelism)并将计算任务映射给更小粒度的处理单元。

这样的两层调度法尽管为系统设计带来了一些简洁,但在实际的部署中,两层调度器互相不感知会导致几个问题:运行时的调度开销很大,算子间并行性没有被有效利用,以及忽视了算子内和算子间两种并行性的相互影响。

图1:(a) 传统低效率的调度方案,(b) 优化后的调度方案

为了能够打破这种僵局,微软亚洲研究院北京大学、上海科技大学合作提出了一种可以成倍甚至几十倍地提升深度学习计算速度的编译框架 RAMMER。研究员们将原数据流图中的算子解析为 rOperator 并将其分解为更小的调度单元 rTask,将底层的硬件抽象为由多个虚拟执行单元(virtualized execution units, vEU)组成的 vDevice。在这套新的抽象下,用户可以通过更细的 rTask 粒度将数据流图调度到多个 vEU 之上,兼顾了计算任务中的两种并行性与底层计算资源的协调。整个调度方案在编译期生成并“静态”映射到硬件计算单元上,因此可以天然地消除掉许多原本存在的调度开销。

图2:(a) 传统深度学习框架,(b)RAMMER 深度学习编译框架

尽管上述介绍中用了不少 CUDA 的概念,但不难发现 RAMMER 整个设计是硬件通用的。它可以很好地适配到诸如 GPU、IPU 等主流深度学习加速器上。研究员们在 NVIDIA GPU、AMD GPU 和 GraphCore 上都评估了这套编译技术所能取得的性能收益。与 TensorRT 相比,RAMMER 在部分模型上实现了最高3.1倍的性能加速。

图3:在 NVIDIA V100 GPU 上批次大小设置为1的端到端的模型推理时间对比

RAMMER的设计与实现

RAMMER 背后是微软亚洲研究院在过去一年多时间里打造的一套名为 NNFusion(https://github.com/microsoft/nnfusion)的深度神经网络编译器(DNN compiler)。NNFusion 能够将现有的模型编译为对应设备的可高效运行的源码,同时支持用户自行替换内核实现或自动从外部导入高性能的内核实现。为了方便地与现有的代码库和 GPU 程序设计模型兼容,RAMMER 采用的是源代码转换的方式,而不是像 TVM、Tensor Comprehension [1] 一样定义新的计算抽象需要用户提供算子的计算逻辑

我们先来了解一下什么是 rTask。rTask 是组成 rOperator 的互相独立的更小的任务单元,也是 RAMMER 抽象中最小的调度单元。在 NVIDIA GPU 上,对于用户提供的内核 CUDA 实现,通过一个小的解析器,可以将每个线程块(thread block)转化为一个 rTask。所以 rTask 可以利用原本内核实现中的语义,虽然这样 rTask 在实现上是与原本的程序设计模型耦合的,但是也大幅度降低了所需的工作负担。

那么要如何创建 vDevice 和 vEU 呢?目前,根据硬件的特性再配合简单的启发式搜索就可以创建 vDevice 和 vEU 了。如果 V100 中有80个 SM,每个 SM 最多能够运行32个线程块,那么就可以创建一个包含有2560个 vEU 的 vDevice,而后根据 rTask 所构成的数据流图与 vDevice 的情况,通过一些简单的策略(譬如直接将 rTask 平铺上 vEU)就能够生成足够高效的调度计划。

由于硬件和编程模型的限制,GPU 运行时的分发器(dispatcher)和调度器(scheduler)并不对用户开放可编程接口,所以研究员们采用了持续线程(persistent thread)[2]的方式,巧妙地以一个相对小的开销将 vEU 与 SM 绑定起来。这样就可以将调度计划“静态”映射给硬件设备了。

RAMMER背后的思考:探索问题的本质

一篇好的系统论文,不仅可以优化性能,更要阐明一个问题。起初,研究员们只是想改善一个具体的神经网络推理时 GPU 利用率偏低的问题(出于对延迟的保障,很多场景下设置小的批次大小其实是标准做法),而除了优化算子实现以外,朴素的想法就是将多个算子一同交给 GPU 设备同时执行,但这并不是一个新的问题。CUDA 很早就引入了流的概念对其提供支持,GPU 社区之前也有一些效果不错的工作(concurrent kernel execution [3]、elastic kernel [4]等),那么在深度神经网络的场景下,为什么大家对这个问题的认知不足?

像上文提到的 MXNet 中有依赖引擎一样,TensorFlow 开发早期也有支持多个流的尝试。但是到后来都接近弃置了,主要原因可能有以下几个方面:

  • 不同的 CUDA 流在运行时采用空间分片(spatial multiplexing)的方法来调度不同流队列(stream queue)上的算子,粒度更粗而彼此之间又极易产生相互干扰影响最终性能 [5]。

  • GPU 的 SM 在不断增加。现在 Ampere GA100 中有128个 SM,但几年前 Kepler GK180 中仅有15个 SM,所以在早期,无论是 GPU 社区还是 DNN 的框架开发,在现有的 GPU 编程模型下都已经形成了硬件对于算子间并行性并没有太多加速潜力的印象。

  • 早期的神经网络结构比较简单,如 AlexNet 等本身在算子间并行性上也没有更多发挥的空间。但随着 AutoML 的出现,网络结构趋于复杂,此外也有ResNext [6]、ResNeSt [7] 等工作引入了新的神经网络设计模式,这个问题正变得更重要。

只是将算子间并行性挖掘起来会是一个好的性能优化,但不足以成为一个好的系统工作。在之前的工作中,微软亚洲研究院的研究员们已经完成了初步的实现并且在一些模型上获得了较好的加速效果,但是当时还没有完全对问题进行清楚地定义,而且因为没有 NNFusion 代码库的支持,实验相对简陋,没有取得很好的反馈。

重新定义一个问题和定位一个工作并不是在用不同的写法来写“茴”字。之前的工作只是在做一个广义上的内核融合,并没有设立起 rTask 和 vEU 的抽象。而此次确定本质的问题在于原本系统中两层调度的差距以后,新的抽象很快探明了更大的优化空间:首先是将原本的通过成本模型来选择子图进行融合的问题,转变为了以更细粒度下的调度和资源分配问题。而得益于绝大部分情况下,神经网络计算的特征(DFG, 算子和张量)在编译时间是已知的,因此可以将调度的开销移交给编译器,这既提升了搜索的效率也简化了系统设计。

更重要的是,让算子间并行性与算子内并行性相互影响这个问题走进研究员们的视野。举个例子,如果对于同一个算子有两种内核实现,其中一个比另一个多消耗三倍的资源(如 CUDA Cores、Shared Memory 等),但是只取得两倍的加速,这在并行计算中是很常见的一个现象。而在此前单个算子独占整个硬件的情况下,毫无疑问会选择更快的实现。而研究员们的实验表明,在算子间和算子内两种并行性协同调度的情况下,选择资源“性价比”最高的实现而非“最快”往往是更优的选择。这其实挑战了之前许多生成高性能算子的工作如 AutoTVM [8] 等的一个基本假设,单个算子独占整个硬件表现出的计算性能是否真的是性能调优的金标准?显然,子图替换TASO [9] 加上高性能内核实现 TVM 两个“最优”相结合,并没有带来真的最优。

研究员们基于新的抽象,仅尝试了简单的策略就在一些场景下获得了超过现有 SOTA 的性能。在此,欢迎大家基于这个抽象进一步尝试更多调度策略,来探索对于一个数据流图(或者其子图)搜索算子间和算子内并行性相互影响下的更高性能的整体实现。

NNFusion 现已在 GitHub 开源:https://github.com/microsoft/nnfusion。目前已经发布了0.1 版本, 该版本支持 TensorFlow 和 ONNX 在内的主流模型格式以及 CUDA GPU 等设备,并提供了丰富的性能优化策略,支持端到端的模型到源代码的 AOT 编译来消除运行时的开销,同时还消除了对第三方库或框架的依赖。如果你有更深入的研发需求,可以直接修改 NNFusion 生成的代码来进行模型的定制化优化。

欢迎前往体验 NNFusion,也期待你可以在 NNFusion 中贡献真知灼见,一起“压榨”加速器的性能!

参考文献

[1]Tensor Comprehensions https://arxiv.org/abs/1802.04730
[2] Persistent Thread https://ieeexplore.ieee.org/document/6339596
[3] Concurrent Kernel Execution https://ieeexplore.ieee.org/document/5999803
[4] Elastic Kernel http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.649.8875&rep=rep1&type=pdf
[5] Dynamic Space-Time Scheduling for GPU Inferencehttps://arxiv.org/pdf/1901.00041.pdf
[6] ResNext https://arxiv.org/abs/.054311611
[7] ResNeSt https://arxiv.org/abs/2004.08955
[8] AutoTVM https://arxiv.org/pdf/1805.08166.pdf
[9] TASO https://github.com/jiazhihao/TASO
微软研究院AI头条
微软研究院AI头条

专注科研19年,盛产黑科技

产业OSDI 2020
相关数据
微软亚洲研究院机构

微软亚洲研究院于1998年在北京成立,是微软公司在亚太地区设立的基础及应用研究机构,也是微软在美国本土以外规模最大的一个研究院。微软亚洲研究院从事自然用户界面、智能多媒体、大数据与知识挖掘、人工智能、云和边缘计算、计算机科学基础等领域的研究,致力于推动计算机科学前沿发展,着眼下一代革命性技术的创新,助力微软实现长远发展战略。通过与微软产品部门紧密合作,微软亚洲研究院将众多创新技术转移到了微软的核心产品中,如Office、Windows、Azure、Bing、Visual Studio、Xbox Kinect以及小冰、Cortana和Microsoft Translator等人工智能产品。

https://www.msra.cn/
深度学习技术

深度学习(deep learning)是机器学习的分支,是一种试图使用包含复杂结构或由多重非线性变换构成的多个处理层对数据进行高层抽象的算法。 深度学习是机器学习中一种基于对数据进行表征学习的算法,至今已有数种深度学习框架,如卷积神经网络和深度置信网络和递归神经网络等已被应用在计算机视觉、语音识别、自然语言处理、音频识别与生物信息学等领域并获取了极好的效果。

感知技术

知觉或感知是外界刺激作用于感官时,脑对外界的整体的看法和理解,为我们对外界的感官信息进行组织和解释。在认知科学中,也可看作一组程序,包括获取信息、理解信息、筛选信息、组织信息。与感觉不同,知觉反映的是由对象的各样属性及关系构成的整体。

调度技术

调度在计算机中是分配工作所需资源的方法。资源可以指虚拟的计算资源,如线程、进程或数据流;也可以指硬件资源,如处理器、网络连接或扩展卡。 进行调度工作的程序叫做调度器。调度器通常的实现使得所有计算资源都处于忙碌状态,允许多位用户有效地同时共享系统资源,或达到指定的服务质量。 see planning for more details

TensorFlow技术

TensorFlow是一个开源软件库,用于各种感知和语言理解任务的机器学习。目前被50个团队用于研究和生产许多Google商业产品,如语音识别、Gmail、Google 相册和搜索,其中许多产品曾使用过其前任软件DistBelief。

张量技术

张量是一个可用来表示在一些矢量、标量和其他张量之间的线性关系的多线性函数,这些线性关系的基本例子有内积、外积、线性映射以及笛卡儿积。其坐标在 维空间内,有 个分量的一种量,其中每个分量都是坐标的函数,而在坐标变换时,这些分量也依照某些规则作线性变换。称为该张量的秩或阶(与矩阵的秩和阶均无关系)。 在数学里,张量是一种几何实体,或者说广义上的“数量”。张量概念包括标量、矢量和线性算子。张量可以用坐标系统来表达,记作标量的数组,但它是定义为“不依赖于参照系的选择的”。张量在物理和工程学中很重要。例如在扩散张量成像中,表达器官对于水的在各个方向的微分透性的张量可以用来产生大脑的扫描图。工程上最重要的例子可能就是应力张量和应变张量了,它们都是二阶张量,对于一般线性材料他们之间的关系由一个四阶弹性张量来决定。

神经网络技术

(人工)神经网络是一种起源于 20 世纪 50 年代的监督式机器学习模型,那时候研究者构想了「感知器(perceptron)」的想法。这一领域的研究者通常被称为「联结主义者(Connectionist)」,因为这种模型模拟了人脑的功能。神经网络模型通常是通过反向传播算法应用梯度下降训练的。目前神经网络有两大主要类型,它们都是前馈神经网络:卷积神经网络(CNN)和循环神经网络(RNN),其中 RNN 又包含长短期记忆(LSTM)、门控循环单元(GRU)等等。深度学习是一种主要应用于神经网络帮助其取得更好结果的技术。尽管神经网络主要用于监督学习,但也有一些为无监督学习设计的变体,比如自动编码器和生成对抗网络(GAN)。

启发式搜索技术

计算机科学的两大基础目标,就是发现可证明其运行效率良好且可得最佳解或次佳解的算法。而启发式算法则试图一次提供一个或全部目标。例如它常能发现很不错的解,但也没办法证明它不会得到较坏的解;它通常可在合理时间解出答案,但也没办法知道它是否每次都可以这样的速度求解。

映射技术

映射指的是具有某种特殊结构的函数,或泛指类函数思想的范畴论中的态射。 逻辑和图论中也有一些不太常规的用法。其数学定义为:两个非空集合A与B间存在着对应关系f,而且对于A中的每一个元素x,B中总有有唯一的一个元素y与它对应,就这种对应为从A到B的映射,记作f:A→B。其中,y称为元素x在映射f下的象,记作:y=f(x)。x称为y关于映射f的原象*。*集合A中所有元素的象的集合称为映射f的值域,记作f(A)。同样的,在机器学习中,映射就是输入与输出之间的对应关系。

逻辑技术

人工智能领域用逻辑来理解智能推理问题;它可以提供用于分析编程语言的技术,也可用作分析、表征知识或编程的工具。目前人们常用的逻辑分支有命题逻辑(Propositional Logic )以及一阶逻辑(FOL)等谓词逻辑。

MXNet技术

MXNet是开源的,用来训练部署深层神经网络的深度学习框架。它是可扩展的,允许快速模型训练,并灵活支持多种语言(C ++,Python,Julia,Matlab,JavaScript, Go,R,Scala,Perl,Wolfram语言)

批次技术

模型训练的一次迭代(即一次梯度更新)中使用的样本集。

深度神经网络技术

深度神经网络(DNN)是深度学习的一种框架,它是一种具备至少一个隐层的神经网络。与浅层神经网络类似,深度神经网络也能够为复杂非线性系统提供建模,但多出的层次为模型提供了更高的抽象层次,因而提高了模型的能力。

AMD机构

超威半导体(中国)有限公司专门为计算机、通信和消费电子行业设计和制造各种创新的微处理器(CPU、GPU、主板芯片组、电视卡芯片等),以及提供闪存和低功率处理器解决方案,公司成立于1969年。AMD致力为技术用户——从企业、政府机构到个人消费者——提供基于标准的、以客户为中心的解决方案。

https://www.amd.com/zh-hans
北京大学机构

北京大学创办于1898年,初名京师大学堂,是中国第一所国立综合性大学,也是当时中国最高教育行政机关。辛亥革命后,于1912年改为现名。2000年4月3日,北京大学与原北京医科大学合并,组建了新的北京大学。原北京医科大学的前身是国立北京医学专门学校,创建于1912年10月26日。20世纪三、四十年代,学校一度名为北平大学医学院,并于1946年7月并入北京大学。1952年在全国高校院系调整中,北京大学医学院脱离北京大学,独立为北京医学院。1985年更名为北京医科大学,1996年成为国家首批“211工程”重点支持的医科大学。两校合并进一步拓宽了北京大学的学科结构,为促进医学与人文社会科学及理科的结合,改革医学教育奠定了基础。

官网,http://www.pku.edu.cn/
暂无评论
暂无评论~