|
和深度学习框架打交道已有多年时间。从Google的TensorFlow, 到百度的PaddlePaddle,再到此刻腾讯的无量。很庆幸在AI技术爆发的这些年横跨中美几家公司,站在一个斗劲好的视角看着世界发生巨大的变化。在这些经历中,视角在不竭切换,从最早的算法研究,到后来的框架开发,到机器学习平台和更多基础架构,每一段都有分歧的感到感染和更深的领悟。
清明节这几天有些时间写了这篇文章,从我的视角,用几个深度学习框架串起来这些年历史上的一些有趣的插曲,和技术背后的一些故事,免得宝贵的记忆随着时间在脑中淡去。
TensorFlow
入门
故事开始在2015年底,我结束了在Google Core Storage和Knowledge Engine的工作,插手了Google Brain,在Samy Bengio下担任一名Research Software Engineer,简称RSWE。RSWE角色的发生主要是因为Google Brain和DeepMind发现Research Scientist很难在研究中解决复杂的工程问题,并最终技术落地。因此需要卷入一些工程能力斗劲强的Engineer和Scientist一起工作。而我斗劲“幸运”的成为Google Brain第一个RSWE。
插手新组的前一个周末,非常兴奋的提前探访了Google Brain的办公地址。想到能近距离在Jeff Dean旁边工作还是有些小感动,毕竟是读着MapReduce, BigTable, Spanner那些论文一路成长起来的。办公场所没有出格,Jeff和大师一样坐在一起,斗劲不测的是发现我工位旁几米的办公室门牌上写着谷歌创始人Larry Page&Sergey Brin,办公室被许多奖杯,证书,太空服之类的杂物包抄着。看来公司对于AI技术的重视程度真实非常的高。
言归正传,早期的TensorFlow斗劲缺模型示例,相关API文档还不太规范,于是先开始给TensorFlow搭建模型库。我花了一年时间把Speech Recognition, Language Model, Text Summarization, Image Classification, Object Detection, Segmentation, Differential Privacy, Frame Prediction等模型写了一遍,后来成为TensorFlow github上model zoo的雏形。那年还是个处处都是低垂果实的时候,没有GPT3这种极其烧钱的大模型,只要对模型做一些小的调整,扩大模型的规模,就能刷新State-Of-The-Art。Bengio大佬经常在世界各地云游,偶尔回来后的1v1还是能给我不少的指引。印象深刻的第一次面聊,这位写过几百篇论文的一字眉大神给方才入门的我在白板上手推了gradient descent的一些公式。此外一次1v1,他发给我一本Ian Goodfellow写的书(当时还是草稿pdf),然后我每天晚上就躺在茶水间的沙发上一边做尝试一边读书。
那年还发生了件有意思的插曲,AlphaGo大战人类棋手。DeepMind和Brain有非常紧密的合作关系,组里组织了一轮paper reading,仔细研读了相关的paper,然后大师带上啤酒和零食组织了不雅观战活动,感觉就有点像是在看球赛。那次学会了两件事,强化学习算法,还有围棋的英文是Go。
16年是TensorFlow高速成长的一年。Jeff的演讲里经常有TensorFlow代码被引用次数指数级暴涨的图。但是16年也是TensorFlow被喷的斗劲惨的一年。TF的Operator粒度长短常细的,据说这是从内部上一代框架DistBelif上吸取的教训。细粒度的Operator可以通过组合形成各种高层的Layer,具有更好的灵活性和扩展性。然而,对于性能和易用性来说倒是斗劲严重的问题,一个模型随便就有几千个甚至跟多的Operator。举几个例子:
当时我要实现第一个基于TensorFlow的ResNet,光为了写一个BatchNormalization(查了好几个内部版本竟然都有些问题),需要通过5~10个细粒度的算子通过加减乘除的方式组装起,1001层的ResNet有非常多的BatchNorm,整个ResNet有成千上万个Operator,可想而知,性能也不怎么样。不久后我伴侣Yao搞了个Fused BatchNormalization,据说能让整个ResNet提速好几成。
BatchNormalization只是初级难度,做Speech Recognition时为了在Python层用TensorFlow完成BeamSearch也花了不少功夫。当时写了个End-to-End的模型,用的是Seq2Seq with Attention,能够一个模型直接把声音转成文字。为了把搜索出产线上语音识此外数据训到最后收敛,用128个GPU整整花了2个月的时间。每天早上上班第一件事就是打开TensorBoard,放大后才能看到Loss又下降了那么一点点。
16年时TensorFlow训练模式主要是基于Jeff等几位的Paper,基于参数处事器的异步训练是主流。训练速度线性扩展性不错,但是今天基于ring的同步训练在NLP,CV这些范围的声音更响一些。记得第一次和Jeff单独交流是关于Speech Recognition分布式训练的尝试情况,加到128个GPU做异步训练基本能保证线性扩展,但是基于SyncOptimizer的同步训练速度会慢很多。当时Jeff问了下收敛效果有没有收到影响,我懵了一下,说没有仔细分析过,赶紧回去查一下。趁便八卦一下,Jeff真长短常瘦,握手的时候感觉几乎就剩皮包骨了。
开发过一些模型后发现算法研究员其实还有不少痛点。1. 不知道怎么Profile模型。2. 不知道怎么优化性能。为了解决这两个问题,我抽暇写了个tf.profiler。tf.profiler的道理斗劲简单,就是把Graph, RunMeta和一些其他的产物做一些分析,然后用户可以通过CLI,UI或者python API快速的去分析模型的布局,Parameter, FLOPs, Device Placement, Runtime等属性。此外还做了个内部数据的抓取任务,去抓算法研究员的训练任务的metrics,如果发现GPU操作率异常,网络通行量过大,数据IO慢时会自动发邮件提醒,并给出一些改削的建议。
让一个专心搞算法研究的人写一个白板的数学公式不难,但是让他去搞大白复杂的任务配置,分布式系统里的性能、资源、带宽问题确是件很困难的事。无论多么牛的研究员城市问为什么任务没能跑起来,是资源不够还是配置不合错误。记得有天薄暮,人不多,Geoffrey Hinton大神俄然走过来问到Can you do me a favor?My job cannot start...(正当我筹备承诺时,Quoc Lee已经抢先接单了,真是个精神的越南小哥。。。)
Moonshots
Google Brain每年会组织一次Moonshots提案,许多后来斗劲成功的项目都是这样孵化出来的,比如AutoML,Neural Machine Translation等等。团队成员会提出一些当时技术斗劲难达到的项目,大师组成类似兴趣小组的形式投入到这些项目中。
此刻火的一塌糊涂的AutoML有点因为商业化或者其他原因,感觉已经对原始的定义做了极大的拓展。当时Brain孵化这个项目的时候有两队人在做LearningToLearn的项目,一个小队但愿通过遗传算法来搜索更优的模型布局,另一个小队则决定使用强化学习算法搜索。遗传算法小队在使用资源时斗劲谨慎,凡是只使用几百个GPU。而另一个小队则使用了几千个GPU。最后强化学习小队更早的做出了成果,也就是Neural Architecuture Search。而另一个小队虽然后来使用更多的GPU也达到了类似的效果,但是要晚了不少。
一个斗劲有趣的插曲是Brain虽然很早就有几万张GPU,但是每当论文截稿的前一段时间总是不够用,此中搞NAS的同学常常在邮件中被暗示。为了解决资源的分配问题,带领们被卷入了一个非常长的email,后来概略解决方案是每个人会被分配少量的高优先级GPU和适量的竞争级GPU资源。而NAS的同学因为已经完成了成本的原始堆集成为了一个很火的项目,得到了特批的独立资源池。为了撑持这个策略,我又开发了个小东西,此刻回头想想还挺吃力不奉迎的。
动态图
快速成长的时间总是过得很快,Megan插手Brain后,我被放置向她陈述请示,当时的RSWE团队已经有十几人,而Google Brain也从几十人变成了几百人。
2017年初,经Megan介绍,TensorFlow团队一位资深专家Yuan Yu找到我,问有没有存眷Pytorch,约我调研后一块聊聊。于是我就去网上搜集了一下Pytorch的资料,又试用了一下。作为一个TensorFlow的深度用户,我的第一反映就是Pytorch解决了TensorFlow很大的痛点,用起来非常的“自然”。
和Yuan聊完后,我们快速的决定在TensorFlow上也测验考试撑持类似Pytorch的imperative programming用法。Demo的开发过程还算斗劲顺利,我概略花了一个多月的时间。记得当时我把项目定名为iTensorFlow, short for imperative TensorFlow。(后来被改名成eager,感觉好奇怪)。
Demo的设计思路其实也不复杂:1. TensorFlow graph可以被切分成任意粒度的Subgraph,可以通过函数调用的语法直接执行,2. TensorFlow对用户透明的记录执行过程以用于反向梯度计算。给用户的感觉就就类似Python native的运行。
进而发生几个推导:1. 当Subgraph的粒度是operator时,基本等价于Pytorch。2. 当Subgraph粒度由多个operator组成时,保留了graph-level optimization的能力,可以编译优化。
最后再埋个伏笔:1. tf.Estimator可以自动的去融合Subgraph,形成更大的Subgraph。用户在开发阶段基于imperative operator-level Subgraph可以简单的调试。用户在部署阶段,可以自动融合大的Subgraph,形成更大的optimization space。
做完之后,我非常兴奋的和Yuan演示成果。Yuan也说要帮我在TensorFlow里面推这个方案。当时Pytorch的成长速度非常的快,TensorFlow的Director也召集了多名专家级的工程师同时进行方案的探索。 当时我还没能进入TensorFlow的决策层,最终得到的结论是1. 让我们成立一个虚拟组专门做这个项目。2. 之前的Demo全部推倒从头做,TensorFlow 2.0作为最重要Feature 发布,默认使用Imperative Mode (后改名叫Eager Mode,中文常常叫动态图)。我则作为团队的一员在项目中贡献来一些代码。
后来Brain来了位新的大神,Chris Lattner,在编程语言和编译范围研究的同学估计很多认识他。他提出来但愿用Swift来实现Deep Learning Model的Progamming,也就是后来的Swift for TensorFlow。理由概略是Python是个动态的语言,很难静态编译优化。后来我和他深入讨论来几次,从技术上非常附和他的不雅概念,但是也明确的暗示Swift for TensorFlow是一条很难走的路。用Python并不是因为Python语言多么好,而是因为很多人用Python。和Chris的一些交流中我对编译过程中的IR和Pass有了更深的了解,对后来在PaddlePaddle中的一些工作发生了不少的影响。
一个插曲是某位TensorFlow团队的资深专家有次暗暗和我说:Python is such a bad language。这句话我品味了好久,不外和他一样没有勇气大声说出来。。。
当时动态图的项目还延展出两个斗劲有趣的项目。有两个其他团队的哥们想对Python做语法分析,进而编译control flow。我很委婉的暗示这个方案做成通用解决方案的可能性不太大,但是这个项目依然被很执着的做了很长一段时间,而且进行了开源,但是这个项目也就慢慢寿终正寝了。另一个很酷的项目是完全用numpy来构造一个deep learning model。通过隐式的tape来完成自动的求导,后来项目仿佛逐渐演化成来JAX项目。
API
后面我逐渐转到了TensorFlow做开发。记得2017年还发生了一件印象深刻的事情,当TensorFlow收获海量用户时,网上一篇“TensorFlow Sucks"火了。虽然那篇文章很多不雅概念我不能苟同,许多想法斗劲肤浅。但是,有一点不能否认,TensorFlow API是斗劲让人蛋疼的。1. 同一个功能往往几套反复的API撑持。2. API经常变换,而且经常发生不向后兼容的问题。3. API的易用性不高。
为什么会发生这个问题呢?可能要从Google这个公司的工程师文化说起。Google长短常鼓励自由创新和跨团队贡献的。经常会有人给另一个团队贡献代码,并以此作为有影响力的论据参与晋升。所以在早期TensorFlow还不是出格完善的时候,经常有外部的团队给TensorFlow贡献代码,此中就包含了API。此外,在Google内部的统一代码仓库下,放出去的API是可以很容易的升级改削的,很多时候只需要grep和replace一下就行。但是github上放出去的API完全纷歧样,Google的员工不能去改削百度,阿里,腾讯内部的TensorFlow使用代码。对此TensorFlow团队早期的确没有非常有效的方案,后来才呈现了API Committee对public API做统一的把关和规划。
在我做视觉的时候,和Google内部一个视觉团队有过很多合作,此中一个是slim API。这个视觉团队非常的强,当年还拿了CoCo的冠军。随着他们模型的推广风行,他们的tf.slim API也被广为传布。slim API的arg_scope使用了python context manager的特性。熟悉早期TensorFlow的人知道还有tf.variable_scope, tf.name_scope, tf.op_name_scope等等。with xxx_scope一层套一层,复杂的时候代码几乎没有什么可读性。此外就是各种global collection,什么global variable, trainable variable, local variable。这在传统的编程语言课里,全局变量这种东西可能是拿来当背面教材的。然而,算法人员的视角是纷歧样的,with xxx_scope和global collection能减少他们的代码量。虽然我们知道合理的法式设计方式也可以做到,但是算法专家估计需要把时间用来读paper,不太愿意研究这些法式设计的问题。
记得在早期内部还有两个流派的争论:面向对象和面向过程的API设计。
基于我教育历史的洗脑,感觉这个是不需要争论的问题。Keras的Layer class和Pytorch的Module class这些面向对象的接口设计无疑长短常优雅的。然而,其实当时的确发生了非常激烈的争论。一些functional API的作者认为functional的调用非常节省代码量:一个函数就可以解决的问题,为什么需要先构造一个对象,然后再call一下?
在TensorFlow动态图能力开发的早期,我们也反复讨论了2.0里面接口的设计方案。作为炮灰的我又接下了写Demo的工作。
闭关两周后,我给出了一个方案:1. 复用Keras的Layer接口。2. 但是不复用Keras的Network,Topology等其他更高层的复杂接口。
原因主要又两点:1. Layer长短常简洁优雅的,Layer可以套Layer,整个网络就是一个大Layer。Layer抽象成construction和execution两阶段也非常自然。2. Keras有很多历史上为了极简设计的高层接口。我个人经验感觉很难满足用户灵活的需求,并不需要官方提供。而且这样可能会导致TensorFlow API层过度复杂。
后来方案被采纳了一半,大佬们但愿能够更多的复用Keras接口。其实没有完美的API,只有最适合某类人群的API。有个小插曲,当时Keras的作者François也在Google Brain。为了在TensorFlow 2.0的动态图和静态图同时使用Keras的接口,不得不在Keras API内做很多改造。凡是François在Review代码时都非常的不情愿,但是最后又往往妥协。很多时候,出格是技术方面,底细可能在少数不被大大都人理解的人手上,需要时间来发现。
TPU
感觉互联网公司那几年,真正把AI芯片做得成熟且广泛可用的,只有Google一家。TPU一直都是Google Brain和TensorFlow团队存眷的重点。原因可能是Jeff老是提起这件事,甚至一度在TensorFlow搞GPU优化是件很没前途的事情。
TPU有个斗劲出格的处所,在于bfloat16的类型。如今bfloat16,还有英伟达最新GPU上的TF32都已经被广为了解了。在当时还是个不小的创新。
bfloat16的道理非常简单,就是把float32的后16bit全部截掉。和IEEE的float16对比,bfloat16的mantissa bits会少一些,但是exponential bits会多一些。保留更多的exponential bits有利于gradients很接近0时不会消掉,保证bfloat16训练时能够更好的保留模型的效果。而传统基于float16训练时,往往需要做loss scaling等调试才能达到类似的效果。因此bfloat16能让AI芯片更快的运算,同时又确保收敛效果凡是不会有损掉。
为了在TensorFlow上全面的撑持bfloat16,我当时花了不少的功夫。虽然之前有基于bfloat16通信的方案,但是要在所有处所都无缝打通bfloat16,还有非常多的工作要做。比如eigen和numpy都不撑持bfloat16这种特殊的东西。幸好他们都可以扩展数据类型(就是文档太少了)。然后还要修复成百上千个fail掉的unit tests来证明bfloat16可以在python层完备的使用。
TPU是一个非常高难度,跨团队,跨技术栈的复杂工程。据说Google有位非常优秀的工程师,为了在TPU上撑持depthwise convolution一个TPU kernel,花掉了半年的时间。
其实这一点也不夸张,除了底层的硬件设计,单是将tensorflow graph编译成硬件binary的XLA项目早期就至少卷入几十人。从HLO到底层的target-specific code generation,几乎又重写了一遍TensorFlow C++层,远比之前的解释型执行器复杂。
TPU的训练在底层跑通后,我基于底层接口的基础上完成Python层的支撑API,然后去实现几个模型。当时碰到了好几个难题,有些在几周时间内解决了,有些持续到我不再团队后好些年。这里举几个例子。
- 当时一个TPU Pod(仿佛是512个chips)算得太快了,即使是很复杂模型的计算也会卡在数据的IO和预措置上。后来搞了个分布式的data processing,通过多个CPU机器来同时去措置数据,才能喂饱TPU。
- 早器的TPU API易用性斗劲弱。凡是一个model需要在TPU上train几百步然后再返回python层,否则TPU的性能会飞快的退化。这对于算法人员是很不友好的,这意味着debug能力的缺掉,以及大量复杂模型无法实现。记得当年OKR被迫降低为撑持常见的CV模型。
- TPU如何撑持动态图。记得我当时迫于TPU的约束,做了个所谓的JIT的能力。就是Estimator先在CPU或者GPU上迭代N步,完成模型的初法式试,然后再自动的deploy到TPU上。从算法人员角度,既满足单法式试的能力,又能在主要training过程用上TPU。
团队
Google Brain是个很神奇的团队,斗劲不客气的说,在2015年后的几年间包揽了全世界在深度学习范围一半以上的关键技术打破,比如TPU,TensorFlow, Transformer, BERT, Neural Machine Translation, Inception, Neural Architecture Search, GAN,Adverserial Training, Bidrectional RNN等等。这里不只有深度学习范围的图灵奖获得者,还有编程语言、编译器、计算机体系布局、分布式系统的顶级专家,甚至还有生物,物理学专家。Jeff将这些人放在一起后,发生了神奇的化学反映,加快了技术改变世界的法式。
PaddlePaddle 飞桨
Paddle其实诞生时间斗劲早,据说是大约13~14年的时候徐伟老师的作品。后来据说Andrew Ng感觉Paddle叫一次不外瘾,就改名成了PaddlePaddle。Paddle和阿谁年代的框架Caffe有类似的问题,灵活性不够。很多处所用C++写成斗劲粗粒度的Layer,无法通过Python等简单的编程语言完成模型的快速构造。
后来17年下半年,团队开始完全从新写一个框架,但是担任了Paddle的名字。2017年底的时候,Paddle国内的团队找到了我,邀请我担任Paddle国内研发团队的负责人。抱着打造国产第一框架的抱负,我接受了邀请,一个月后就在北京入职了。
早期设计
插手团队的时候,新的Paddle还是一个斗劲早期的原型系统,里面有一些设计已经被开发了出来。我发现此中有些设计理念和TensorFlow有明显的差异,但是实现的时候却又仿照了TensorFlow。
仿编程语言
设计者但愿设计一种编程语言来完成深度学习模型的构建(有点类似Julia等把深度学习模型的特性嵌入到了编程语言中)。然而在实现上,我发现其实和TensorFlow斗劲类似。都是通过Python去声明一个静态模型布局,然后把模型布局交给执行器进行解释执行。并没有发现一种新的深度学习编程语言。
这块我基本没有对设计进行调整。本质上和TensorFlow早期静态图的没有区别。但是在细节上,TF基于Graph的模型可以通过feed/fetch选择性的执行任意一部门子图,更加灵活。Paddle中与Graph对应的是Program。Program就像正常法式一样,只能从头至尾完整的执行,无法选择性的执行。因此Paddle在这块相对简化了一些,但是可以通过在Python层构造多个Program的方式补全这部门灵活性的缺掉,总体来说表达能力是足够的。
Transpiler
Transpiler是对Program进行直接改写,进而可以让模型能够被分布式运行,或者进行优化。初衷是斗劲好的,可以降低算法人员的使用难度。然而在实现上,最开始是在Python层直接对Program布局进行改写。后来我从新设计了IR+Pass的Compiler体系,通过一种更系统性的方式做了实现。
LoDTensor
可能是因为团队的NLP和搜索布景斗劲强,对于变长序列的重视程度很高。Paddle的底层数据是LoDTensor,而不是类似其他框架Tensor。LoDTensor相当于把变长序列信息耦合进了Tensor里面。这可能导致斗劲多的问题,比如很多Operator是完全序列无关的,底子无法措置序列信息在输入Tensor和输出Tensor的关系,进而斗劲随机的措置,给框架的健壮性埋下隐患。虽然我一直想敦促序列信息和Tensor的解耦合,但是因为种种原因,没有彻底的完成这个重构的方针,但愿后面能改掉。
性能
18年初的时候,Paddle还是个原型系统。由于OKR方针,团队已经开始初步接入一些业务场景。其实一个斗劲大的痛点就是性能太差。单机单卡速度非常慢,单机4卡加速比只有1.x。但是性能问题的定位却非常困难。我花了些时间写了些profile的东西,比如timeline。一些明显的性能问题可以被快速的定位出来并修复。
但是单机多卡的速度还长短常慢,timeline分析后发现此中有个ParallelOp,存在大量的Barrier。最后改写成了ParallelExecutor,把Program复制了N份部署在多张卡上,在此中插入AllReduce通信算子,然后这N倍的算子基于图依赖关系,不竭把ready的算子扔进线程池执行。即使这样,我们也发此刻多卡的性能上,分歧模型需要使用分歧的线程调剂策略来达到最优。很难有一种完美的one-fits-all的方案。后面我们再聊如何通过IR+Pass的方式插件化的撑持分歧的算子调剂策略。
分布式的训练也碰到不少的问题。一开始使用grpc,花了挺大的功夫做并行请求,然后又切成了brpc,在RDMA等方面做了不少的优化。分布式训练的性能逐步得到了提升。此外为了做到自动化分布式部署,前面提到的Transpiler随着场景的增加,Python代码也变得越来越复杂。
模型推理在公司内碰到来非常强劲的对手。Anakin的GPU推理速度的确很快,让我吃惊的是他们竟然是用SASS汇编完成大量基础算子的开发,针对Pascal架构做了异常极致的优化,甚至在某些场景远超TensorRT。我一直主张训练和推理要尽量用一样的框架,并不需要一个单独的推理框架来解决性能问题。使用分歧的框架做推理会造成很多不测的精度问题和人工开销。
因为推理性能的问题,我们和兄弟团队发生来旷日持久竞赛,作为狗头军事,我充实阐扬来团队在CPU这块的技术堆集、以及和Intel外援的良好关系,在CPU推理场景常常略胜一筹。在GPU方面苦于对手无底线使用汇编,和我方战线太多、人员不够,只能战略性放弃了部门头部模型,通过撑持子图扩展TensorRT引擎的方式,操作Nvidia的技术优势在许多个通用场景下展开进攻。此刻回想起来真实一段有趣的经历。
Imtermediate Representation&Pass
Imtermediate Representation+Pass的模式主要是从LLVM的架构上借鉴来的。在编译器上主要是用来解决把M个编程语言中任意一个编译到N个硬件设备中任意一个执行的问题。简单的解决方案是为每个编程语言和硬件单独写一个编译器。这需要M*N个编译器。显然这对于复杂的编译器开发来说,长短常高成本的。
Intermediate Representation是架构设计中抽象能力的典型浮现。分歧编程语言的层次纷歧样,或者仅仅是纯挚的撑持的功能有些差异。但是,这些编程语言终归需要在某种硬件指令集上执行。所以在编译的过程中,他们会在某个抽象层次上形成共性的表达。而IR+Pass的方式很好的操作了这一点。其基本思想是通过多层Pass (编译改写过程),逐渐的把分歧语言的表达方式在某个层次上改写成统一的IR的表达方式。在这个过程中,表达方式逐渐接近底层的硬件。而IR和Pass可以很好的被复用,极大的降低了研发的成本。
深度学习框架也有着非常类似的需求。
- 用户但愿通过高层语言描述模型的执行逻辑,甚至是仅仅声明模型的布局,而不去关心模型如安在硬件上完成训练或者推理。
- 深度学习框架需要解决模型在多种硬件上高效执行的问题,此中包罗协同多个CPU、GPU、甚至大规模分布式集群进行工作的问题。也包罗优化内存、显存开销、提高执行速度的问题。
更具体的。前文说到需要能够自动的将用户声明的模型Program自动的在多张显卡上并行计算、需要将Program拆分到多个机器长进行分布式计算、还需要改削执行图来进行算子融合和显存优化。
Paddle在一开始零散的开展了上面描述的工作,在分布式、多卡并行、推理加速、甚至是模型的压缩量化上各自进行模型的改写。这个过程非常容易发生反复性的工作,也很难统一设计模式,让团队分歧的研发快速理解这些代码。
意思到这些问题后,我写了一个Single Static Assignment(SSA)的Graph,然后把Program通过第一个基础Pass改写成了SSA Graph。然后又写了第二个Pass把SSA Graph改写成了可以多卡并行的SSA Graph。
后面的事情就应该可以以此类推了。比如推理加速可以在这个基础上实现OpFusionPass, InferenceMemoryOptimizationPass, PruningPass等等,进而达到执行时推理加速的目的。分布式训练时则可以有DistributedTransPass。量化压缩则可以有ConvertToInt8Pass等等。这一套东西基本解决了上层Program声明到底层执行器的Compiler问题。
这个过程中的确碰到了不少的阻力。比如分布式早期通过Python完成了这个逻辑,需要迁移到C++层。压缩量化的研发更喜欢写Python,而IR&Pass是基于C++的。分歧Pass间挨次依赖和Debug等。
全套深度学习框架东西
TensorFlow Everywhere原本是TensorFlow团队时的一个标语,意思是TensorFlow需要撑持深度学习模型在任意的场景下运行,进而达到AI Everywhere的方针。可以说深度学习框架但愿成为AI的“操作系统”,就像鱼离不开水、App离不开iOS/Android一样。
Paddle作为全面对标TensorFlow的国产深度学习框架,自然也但愿提供全套的解决方案。在早期的时候,Paddle和公司其他团队合作了PaddleMobile,提供了移动端的推理能力。后来又开展了Paddle.js,撑持在H5、Web等场景的推理能力。为了在toB,在Linux的基础上又新增了Windows的撑持。为了撑持无人车等设备、又撑持了在更多分歧设备上运行。
举个PaddleMobile的例子。深度学习框架想再移动设备上部署面临这斗劲多的挑战。手机的空间和算力都比处事器小很多,而模型最开始在处事器训练好后体积相对较大,需要从很多角度下手。1. 使用较小的模型布局。2. 通过量化,压缩等手段削减模型体积。
此外移动段深度学习框架是凡是基于ARM CPU,GPU则有Mali GPU, adreno GPU等等。为了最求斗劲极致性能,常常需要使用汇编语言。有个同学写到后面几乎怀疑人生,感觉本身大学学的东西不太对。为了不显著增加APP的体积,框架编译后的体积需要在KB~几MB的级别,因此需要基于部署的模型布局本身用到的算子进行选择性编译。极端的时候甚至需要是通过C++ Code Gen的方式直接生成前向计算必需的代码,而不是通过一个通用的解释器。
回顾
随着项目的复杂化,很多棘手的问题逐渐从深度学习的范围技术问题改变成了软件工程开发和团队打点分工的问题。随着团队的不竭变化,本身有时候是作为一个leader的角色在措置问题,有的时候又是以一个independent contributor的角色在参与讨论。很庆幸本身经历过这么一段,有些问题在亲身经历后才能想得大白,想得开。时代有时候会把你推向风口浪尖,让你带船队扬帆起航,在更多的时候是在不竭的妥协与摸索中寻找前进的标的目的。
无量
无量是腾讯PCG扶植的一个深度学习框架,主要但愿解决大规模保举场景下的训练和推理问题。深度学习在保举场景的应用和CV、NLP、语音有些纷歧样。
- 业务会持续的发生用户的行为数据。当用户规模达到数千万或者上亿时就会发生海量的训练数据,比如用户的画像,用户的点击,点赞,转发行为,还有Context等等。
- 这些数据是高度稀疏的,凡是会编码成ID类的特征进而通过Embedding的方式进入模型训练。随着业务规模的提升和特征工程日渐复杂,比如累计用户数,商品,内容增加,以及特征交叉的使用,Embedding参数的体积可以达到GB,甚至TB级。
- 保举场景是实时动态变化的,新用户,内容,热点不竭发生。用户的兴趣,意图逐渐变化,因此模型需要持续不竭的适应这些变化,时刻保持最好的状态。
调整
19年中这个项目时概略有2~3人。团队但愿开发一个新的版本,基于TensorFlow进行扩展加强,使得无量可以复用TensorFlow已有的能力,而且能够撑持保举场景下的特殊需求。无量一开始采用的是基于参数处事器的架构。TensorFlow被复用来提供Python API以及完成基础算子的执行。而参数处事器,分布式通信等方面则是本身开发,没有复用TensorFlow。
这个选择在团队当时的情况下是斗劲合理的。如果选择另一种标的目的,基于TensorFlow底层进行改造,研发难度会斗劲大,而且很可能与社区版TensorFlow走向分歧的标的目的,进而导致TensorFlow版本难以升级。而把TensorFlow作为一个当地执行的lib则可以在外围开发,不需要了解TensorFlow内部的复杂逻辑,也可以复用一些其他开源组件,比如pslib。
早期在软件开发的流程上相对斗劲欠缺。为了保障工程的推进,我先辅佐做了些基础工作,比如加上了第一个自动化测试和持续集成,对一些过度封装和奇怪的代码做了重构和简化。
此外,在接口层也做了一些调整。本来框架开始执行后就进入C++执行器,无法从python层提供或者返回任何执行成果,也无法在python层执行逻辑进行插件化的扩展。为了满足预期用户将来需要进行调试的需求,我模拟tf.Session和tf.Estimator对执行层的接口做了重构。这样用户可以通过feed/fetch的方式单法式试执行的过程。也可以通过Hook的方式在执行前后扩展任意的逻辑,提高框架的适用场景。
此外一个问题是python层基本完全是全局变量,很难进行多模型的封装。像TensorFlow有Graph实例或者Paddle有Program实例。因为python层需要重构的量斗劲大,我暂时先插手了Context的封装,勉强将各种状态和配置封装在了Context下。考虑到短期可能不会有更复杂的需求,暂时没有把这件事做完。
reader那块也做了一些重构。最开始那块的线程模型异常复杂,一部门是因为分布式文件系统等基础设施无法提供斗劲好的SDK,导致许多逻辑不得不在深度学习框架里面,比如文件的当地缓存。考虑到特征加工的逻辑斗劲复杂,以及一些老的TensorFlow用户可能习惯于tf.Example和tf.feature_column等基础算子库,我在reader层引入了基于TensorFlow的tf.dataset。不外后来发现用户似乎更关心性能问题,喜欢自定义C++ lib的方式来解决特征措置的问题。
API设计是个老大难的问题。TensorFlow,Paddle,无量都没能幸免。在一个多人协同的团队里,每个研发更多还是存眷每个独立功能是否完成开发,而功能的接口往往需要考虑到整体的API设计风格,易用性,兼容性等许多因素,常常在高速迭代的过程中被忽略掉。不幸的是API常常不能像内部实现一样后期优化。当API被放给用户使用后,后续的改削往往会粉碎用户代码的正确性。很多时候只能本身评审一下。
升级
无量颠末一年基础能力的打磨,逐渐的成为来整个事业群统一的大规模保举模型训练和推理框架,支撑数十个业务场景,每天都能出产数千个增量和全量的模型。简单的完成功能已经不能满足业务和团队成长的需求,需要在技术上更加前沿。
数据措置
数据格式上要从本来的明文转到更高效的二进制。此外基于CSR编码的稀疏数据可以进一步的减少数据措置时的拷贝等额外开销。
流水线
尽量挖掘训练中可以并行的处所,通过流水线的方式提高并发度,进而提高训练的速度。比如在数据读取的过程中,就可以提前按照参数处事器的规模对数据进行预切分,并奉告参数处事器需要提前筹备哪些参数。这样当pull/push的时候能够更快的完成计算,进而提高每个minibatch的速度。
同样,当使用GPU训练时,也可以在数据IO的并行过程中,预计算未来需要用的的Embedding参数。这样GPU训练下一轮的数据时,需要用到的Embedding已经提前被计算好,可以直接开始训练,减少来等待的时间。
定制化参数处事器
由于无量解决的一个关键问题是保举模型的海量参数问题,因此参数处事器必需是高度优化过的。而且应该合理的将保举模型的范围常识引入到设计中,通过特殊的策略进一步发生差异化的优势。
定制化的线程模型,内存打点和HashMap。由于参数是被切分归属到分歧线程上,所以可以通过无锁化的把每次pull/push的参数措置好。此外由于海量参数消耗较大硬件成本,内存空间都需要通过定制化的内存池来打点。否则很可能有大量的空间碎片在默认内存库中无法及时偿还给操作系统。此外也有无法精细化控制内存清理机制,导致内存OOM或者浪费。定制化的内存打点可以解决这些问题,甚至通过特殊的内存裁减策略,在不损害模型效果的基础长进一步降低内存的开销。高性能HashMap则是需要解决Embedding快速的增删改查的问题。
Embedding向量的打点也是有非常多可以改良的处所。1. 动态的改变Embeding向量的长度来撑持模型的压缩,提高模型效果。2. 扩展Embedding的元数据来记录热度,点击展现等统计值,有助于提高训练推理时高级分布式架构的Cache命中率,已经模型的训练效果。3. 模型的恢复和导出机制在大规模Embedding场景对于Serving时能够实时加载模型更新重要。此外还需要考虑到任务掉败重启后资源伸缩等问题。
GPU训练
传统PS架构的训练模式下,由于单台机器的计算能力有限,需要几十甚至上百个实例进行分布式训练。但是这样会导致大量的计算被用在来无效的开销上。比如稀疏特征在网络通信两边的措置。这种额外开销甚至经常超过有效计算。
GPU和相应的高速网络链接可以解决这一问题。单台8卡机器通过NVLink连接起来,速度甚至可以超过几十台物理机,有更高的性价比。但是由于几百GB,甚至TB级的参数问题,还有Embedding的GPU计算问题,导致GPU一直都没有被广泛的用起来。
然而尝试发现其实稀疏特征存在显著的Power-law分布,少部门Hot特征使用远多于其他大量不Hot的特征。因此,通过在数据措置时统计特征,然后批量将将来新需要的Embedding换入GPU,就可以让GPU长时间的进行持续训练,而不需要频繁的和CPU内存交换参数。
GPU预测
随着保举模型复杂度提高,引入传统CV,NLP的一些布局需要消耗更多的计算。CPU往往很难在有效的时间延迟下(几十毫秒)完成大量候(几百上千)选集在复杂保举模型的推理。而GPU则成为了一个潜在的解决方案。
同样,GPU推理也需要解决显存远小于Embedding参数的问题。通过在训练时预先计算Hot Embedding,然后加载如推理GPU,可以必然程度的缓解这个问题。在推理时仅有少部门的Embedding没有在GPU显存中缓存,需要通过CPU内存拷贝进入GPU。
而通过模型的量化和压缩能进一步减少Embedding参数的规模。尝试表白当大部门Embedding参数的值控制为0时,模型依然能够表示出本来的效果,甚至略优。
总结
深度学习算法的成长和深度学习框架的成长是相辅相成,互相促进的。从2002年时Torch论文发表后,框架的技术成长相对迟缓,性能无法显著提升导致无法探索更加复杂的算法模型,或者操作更加大规模的数据集。
在2010年后逐渐呈现了Caffe, Theano等框架,通过将更高性能的GPU引入,可以训练更加复杂的CNN和RNN模型,深度学习算法的成长呈现来显著的加速。
到了2014~2017年几年间,TensorFlow的呈现让用户可以通过简单的Python语言将细粒度的算子组装各种模型布局。而且模型可以简单的被分布式训练,然后自动化部署在处事器,手机,摄像头等各种设备上。而Pytorch的动态图用法满足了研究人员对易用性和灵活性更高的要求,进一步推进算法研究。
国内的深度学习框架技术在这股浪潮中也紧跟这世界的法式。Paddle在14年摆布发生,在国内堆集了必然的用户,在当时基本能比肩其他的框架。虽然在TensorFlow和Pytorch等更先进的框架呈现后,国内错过了宝贵的几年技术升级的窗口以及社区生态培育时机,但是我们看到从18年到20年间,新版的PaddlePaddle,OneFlow, MindSpore等深度学习框架陆续开源,技术上逐渐赶了上来。
保举场景在电商,视频,资讯等众多头部互联网公司的火爆导致保举系统对AI硬件的消耗在互联网公司超过了传统NLP,CV,语音等应用的总和。许多公司开始针对保举场景(以及广告,搜索)的特殊需求对深度学习框架进行定制优化。百度的abacus是斗劲早期的框架,和其他早期框架一样,在易用性和灵活性上较弱。无量,XDL等框架则进行了改良,兼顾了社区兼容性,算法易用性和系统的性能等纬度。
深度学习的框架的触角其实远不止我们常见到的。随着AI技术的推广,Web、H5、嵌入式设备、手机等场景下都有许多优秀的深度学习框架发生,如PaddleMobile, TFLite,tensorflow.js等等。
深度学习框架的技术也逐渐从更多纬度开始拓展。ONNX被提出来作为统一的模型格式,虽然离方针有很长的距离和问题需要解决。但是从它的风行我们能看到社区对于框架间互通的巴望。随着摩尔定律难以维持,框架开始更多的从新的硬件和异构计算范围寻求打破。为了撑持海量的算子在CPU、FPGA、GPU、TPU、NPU、Cerebras等众多AI芯片上运行,TVM、XLA等借鉴编译技术几十年来的堆集,在更加艰巨的道路长进行来持续的探索,经常能传来新进展的好动静。深度学习框架也不再仅应用于深度学习,还在科学计算,物理化学等范围发光发烧。 |
|