让AI简单且强大：深度学习引擎OneFlow技术实践

​

 

本文内容节选*msup主办的第七届TOP100summit，北京一流科技有限公司首席科学家袁进辉（老师木）分享的《让AI简单且强大：深度学习引擎OneFlow背后的技术实践》实录。

 

北京一流科技有限公司将自动编排并行模式、静态调度、流式执行等创新性技术相融合，构建成一套自动支持数据并行、模型并行及流水并行等多种模式的分布式深度学习框架，降低了分布式训练门槛、极大的提高了硬件使用率。该框架已经成功帮助众多头部互联网公司及人工智能企业提升了大模型训练效率，节约了硬件运营和使用成本，达到了降本增效的效果。一流科技是一家为企业客户提供面向大规模大计算大模型等深度学习框架的人工智能领域科技创新公司。

 

分享者袁进辉是北京一流科技有限公司创始人，任首席科学家。2008年7月在清华大学计算机系获得工学博士学位，获得清华大学优秀博士学位论文奖。2013年加入微软亚洲研究院从事大规模机器学习平台的研发工作。2014年发明了当时世界上最快的主题模型训练算法和系统LightLDA，只用数十台服务器即可完成以前数千台服务器才能实现的大规模主题模型，该技术成功应用于微软在线广告系统，被当时主管研究的全球副总裁周以真称为“年度最好成果”。2015年至2016年底，专注于搭建基于异构集群的深度学习平台，项目荣获微软亚洲研究院院长特别奖 (top 1%)。2017年创立北京一流科技有限公司，致力于打造分布式深度学习平台的事实工业标准。

 

 

编者按：2018年11月30日-12月3日，第七届全球软件案例研究峰会在北京国家会议中心盛大开幕，现场解读2018年「壹佰案例榜单」。本文为北京一流科技有限公司首席科学家袁进辉（老师木）分享的《让AI简单且强大：深度学习引擎OneFlow背后的技术实践》案例实录。

 

 

提纲：

 

• 研发OneFlow的动机

• OneFlow技术突破

• 总结

 

 

01研发OneFlow的动机

 

软件OneFlow简介

 

业界有人工智能浪潮的三驾马车之说，即数据、算法、算力。具体到算力，业界更多关注的是硬件，譬如GPU，甚至是TPU之类的AI专用芯片。但是，人们发现，有了更快的加速器之后，制约大规模分布式训练算力的瓶颈是软件。怎么帮助数据科学家和研究员们更轻松的把各种算法在底层硬件上跑起来，而且充分释放底层硬件的潜力，正是软件框架需要解决的问题。目前，已有的开源深度学习框架对数据并行场景解决的比较好，但遇到模型越来越大的场景就没有好办法。用户或者束手无策，或者只能付出极大成本基于开源框架做深度定制开发来满足需求。OneFlow团队的目标是研发一个通用框架自动解决这些问题，让那些没有框架研发能力的团队也能够享受分布式GPU集群带来的效率，这是我们历时两年多研发一套全新深度学习框架的初衷。

 

背后的动机：计算力是深度学习发展的最重要的推动力。

 

案例：2015 Microsoft Resnet

          2016 Baidu Deep Speech 2

          2017 Google NMT

 

2015年微软研究院发明的ResNet需要的计算量是7乘以10的18次方次计算（ExaFlops）。当然，可以推算一下用一颗24核的CPU来计算，大概需要多久能完成这些计算，也可以推算用几千个核心的GPU来算需要多长时间。可能是需要几个月或几个星期的时间。除了计算量在增长，模型大小也在增长，ResNet 这种CNN模型通常是几千万参数，需要几百兆字节的存储空间，百度研发的Deep Speech模型到了三亿参数的规模，然后Google的机器翻译模型NMT，已经到了几十亿参数，整个模型在一块GPU上已经放不下了。这种情况，数据并行无济于事，需要模型并行或流水并行来解决大模型的分布式训练问题。很不幸，目前还没有开源框架支持这些需求，一些大公司通过内部定制的系统来支持这种需求。

 

今年上半年Facebook发布了一个研究结果，收集35亿张弱标注图片，使用几百块GPU，经过接近一个月的时间，训练了一个用于图片分类的卷积神经网络模型，它能做到什么效果呢？能提高6个百分点的准确率。这是非常了不得的成绩，算法基本上没什么变化，仅仅是通过采用更多的数据和计算就能把top-1的准确率提高了这么多。要知道，对于一个商业价值很高的场景，提高0.5个百分点可能是一个团队一年的KPI。

 

九月份Google发表了BigGAN模型，研究人员通过提高图片的分辨率来训练更大的GAN模型，CNN中间的activation和反向gradient会非常多，当然计算量也会大的非常多，基于TPU集群来完成训练。这个手段同样获得了比以前的GAN模型好的多的效果。

 

上个月，Google又发表了BERT模型，相当于一种大的多的transformer模型，在16个TPU上训练了4天，然后基于这个语言模型作为主干网络去解决各种常见的自然语言处理任务，发现在各任务上全面超越了以前的方法。很不幸，目前还没有出现在GPU集群上从零开始训练BERT-Large模型的办法，如果想在自己的业务里应用BERT，只能去下载Google预训练好的模型，然后做少量微调来使用。这倒不是资源不足的问题，即使是已经搭建了大规模的GPU集群的客户也无能为力，有钱也解决不了。

 

深度学习经过这几年的爆发式发展，特别引人注目的算法层面的创新越来越少了，今年比较吸引眼球的进步都来自于计算力，也就是人们常说的“大力出奇迹”的方式。怎么才能让更多的企业用户能享受到算力提升的红利，帮助算法科学家完成更多的KPI, 这是OneFlow非常关心的问题。常言道，工欲善其事必先利其器，框架在深度学习研究和落地的过程中就扮演了“工具”的角色，好的工具能大大加速人工智能研发的效率，甚至可能成为行业竞争的决胜法宝。从BigGAN和BERT等例子也可以看出来，当一家公司掌握了其他人不掌握的工具时，就可以引领算法研究的潮流，反过来，当一家公司的基础设施跟不上的时候，也就没办法做前沿探索，即使是做研究也只能跟在Google后面，因此称深度学习框架是人工智能制高点的战略武器一点不为过。

 

基于纯硬件的解决思路

 

 

案例：Nvidia DGX-2

          IBM Power9 Server

 

英伟达通过销售GPU成为这一波AI计算力红利的最大受益者，英伟达除了把单个设备做的越来越快，也做了服务器架构方面的诸多创新，出品了一系列超级计算盒子，每个盒子里面可以集成8个或者是16个计算力非常强的GPU（譬如DGX-1是P100，今年推出的DGX-2是V100），更特别的是，这些GPU之间使用了非常高速的互联，能够实现GPU之间点对点150GB以上的传输带宽，比常见的PCIe带宽要高一个数量级。这种设计使得DGX服务器能够使得16块GPU一起工作时几乎像一个单体芯片那样输出超强算力。

 

当然还有比DGX更特别的服务器，比如说IBM出的Power9 Server，它的独特之处在于他的CPU使用了不同于Intel x86 CPU的架构，而且支持CPU和GPU之间NV Link互连，意味着CPU和GPU之间的数据传输也能够做到150GB以上的带宽。目前世界排名第一的超级计算机Summit就使用了类似Power9 Server的架构。

 

基于这么强的硬件就能解决计算力的问题吗？

 

IBM和Nvidia一起搭建了世界上最强的超级计算机Summit，一共用了2万多块V100 GPU，还使用了最先进的互联技术(NVLink, Infiniband)，要说最强的硬件，除了TPU Cluster，应该没有更好的了，这是不是就够了呢？IBM首席科学家在今年的ASPLOS(计算机体系结构*会议)上做了一个特邀报告，主题是“只有很强的硬件，没有很好的软件还是不能解决扩展问题”。现在国内拥有几千块GPU乃至上万块GPU的头部公司不在少数，但基于开源框架能训练BERT-Large模型吗？不行，这也印证了软件框架瓶颈的问题：购买了很多的硬件，但用不起来，或者说不能很好的用起来。

 

理念：纵向扩展与横向扩展

 

 

纵向扩展

 

纵向扩展是通过把单个设备或者是单个机器做的越来越强，或通过编译器优化的手段让作业在一个设备上或者是一个机器内部把硬件性能发挥到极致来满足现在日益增长的计算需求。硬件从多核架构CPU发展到众核架构GPU，GPU从P100到V100, 为了追求更高的效率，甚至研发FPGA或ASIC芯片来获得更高算力。当前最知名的AI芯片是Google的TPU，寒武纪，华为，阿里，百度等本土公司也在研发AI芯片。AI芯片的主要问题是有物理限制（譬如制程，功耗，同步时钟等等约束），人们不能生产出计算力任意大的芯片。也有人把这个现象称为硅基扩展瓶颈（Silicon Scaling）。除了提高单个芯片的吞吐率，英伟达的DGX也是纵向扩展的例子，DGX通过在一个机器内部高速互联手段实现芯片之间点对点极高的传输带宽，从而使得多芯片间协作起来更加高效。

 

 

横向扩展

 

如果纵向扩展仍不能满足需求，人们继续把多台服务器通过高速以太网或Infiniband连接起来组成集群来实现更高算力。如果能投入多少硬件资源，就得到多少计算力，那计算力瓶颈就迎刃而解了。理想很丰满，现实很骨感。一方面，芯片间互联带宽要比片内数据访问带宽低一到两个数量级，在芯片间搬运数据成为瓶颈，另一方面，编写在多芯片上高效运行的软件非常挑战，以深度学习为例，神经网络的结构不同，效率最高的并行方式（逻辑任务向物理计算单元的映射）也不同。在集群上实现线性加速比纵向扩展更有想象空间，但实现难度更大。一个理想的横向扩展方案，不管底层实际使用了多少松散耦合在一起的芯片，在上层用户眼里就像在一个专门为当前任务打造的巨大单体芯片一样，编程简单而且任务运行时能把底层每一个独立的芯片都利用充分。要实现这个目的，必须依靠软件框架。

 

逻辑任务到物理拓扑之间的最优映射复杂多变

 

给定一个特定的神经网络模型和一批计算资源，从任务到设备之间的映射有多种方式，但不同的映射方案运行效率不同。哪种方案最优既取决于作业本身的特性，也取决于底层硬件的拓扑。神经网络由很多局部计算（通常称为kernel）搭建组成，每一个局部计算是采用数据并行，还是模型并行取决于这个局部任务的计算传输比。现在业界讨论比较多的卷积运算参数量很小，但中间结果量大，所以最划算的方法是对数据进行切分，不同的设备处理不同的数据，在设备之间偶尔进行参数同步，这种数据并行方法基本上是一个已经被解决的问题。还有一些运算，中间计算结果相对于参数量更少，就适合模型并行。还有一些网络参数量很大或中间计算结果都很大，可能采用流水并行（也就是接力的形式）是最优的。模型并行和流水并行中通信的数据路由要比数据并行复杂，同时，怎么重叠计算和传输从而提高设备利用率也非常挑战，现有开源框架对这些更复杂的并行模式的支持还比较初级。

 

通信密集，延迟敏感

 

左图展示了一个常见的大数据处理引擎的架构，集群中的计算资源一般分成用于中心调度的Master节点和用于处理数据的Worker节点。Master节点以有向无环图（DAG）的方式管理整个作业的进度，同时监控所有Worker的资源使用状况，在合适的时机把一个子任务（Task）分配给某个Worker去做，某个Worker在完成一个子任务之后，会向Master节点汇报，等待Master分配新的任务。在传统大数据处理中，Worker执行一个子任务的时间量级一般在几十秒钟或数分钟。其它开销，诸如发生在Master节点那里的排队开销，Master和Worker之间对话的时间开销，以及数据传输开销都是数十毫秒，相对于Worker的工作时间可以被忽略。但是深度学习训练的负载与此不同，深度学习训练更接近流式计算，一方面是因为随机梯度下降算法采用的小批次训练，计算粒度小，另一方面是因为底层硬件吞吐率可能是CPU的数十倍，计算太快。直接后果就是，数据处理时间越来越短，每个子任务可能几百毫秒就完成了，在这种情况下，之前可忽略的那种几十乃至几百毫秒的开销就非常显著了，如果不能通过技术手段把这些开销消除或掩盖掉，整个系统的性能就非常低。

 

 

02OneFlow技术突破

 

 

基于静态调度的流式计算引擎

 

为了对任意作业和资源都达到类似巨大单体专用芯片的效果，OneFlow首创了静态调度（左图）和流式执行（右图）架构的深度学习框架。静态调度是什么思路呢？它来自于计算机体系结构。我们熟知的CPU芯片中真正做算术运算的器件只占很小比例的面积，大部分面积在做乱序执行，流水线和缓冲区的管理。学界和工业界很久以前就开始探索怎么让芯片的有效面积尽可能多的做算术运算，静态调度的思路应运而生，基本上是把流水管理，指令排布之类的工作从硬件转移至编译器。这样硬件复杂度就可以大幅降低，当然编译器复杂度肯定会提高很多。有一个叫VLIW（超长指令集架构）的指令集就采用了这种思路。OneFlow的静态调度体现在两方面，首先，编译器自动解决从逻辑任务到硬件资源的映射，包括数据并行，模型并行，流水并行的设备分配以及数据路由方案，大大降低了分布式编程的复杂度，用户只需要关心任务的逻辑结构以及本次任务可使用的硬件资源，而不用去编程实现数据在硬件资源中的流动机制；其次，静态调度把所有能在正式运行之前得到的调度策略，资源管理策略等问题都在编译阶段解决，运行时就不需要在线求解最优的调度方案，从而大大降低运行时开销。

 

经过静态编译，每个设备负责运行的子任务是预先可知的，每个子任务的上下游依赖也预先可知，在运行任务时，就不再需要中心调度器，只需要支持上下游任务之间局部的握手信号即可，即生产者向消费者发送的请求以及消费者向生产者发送的确认，整个系统以全链路异步的方式运行。这个思路也来自于芯片设计领域一种叫异步电路的技术。OneFlow另一个区别于其它深度学习框架的特色是把数据搬运看成一等公民，在静态分析阶段就把磁盘IO，主存和设备之间数据搬运，节点间数据搬运看作和计算同等重要的任务，在代价分析和调度策略里作为一等公民进行显式建模，从而得到重叠传输和计算的最优方案。与此相对，已有开源框架把数据搬运当成二等公民处理，编译期的注意力主要集中在计算的优化上。熟悉软件定义网络（SDN）技术的朋友可以发现，OneFlow编译器相当于网络的控制平面，用于获取数据计算和转发策略，运行时相当于网络的数据平面，执行体依照控制层面的策略去转发和处理数据。

 

产品对比

 

OneFlow历经两年的研发，2018年10月份才推出1.0版本，还是一个很年轻的系统，目前正在客户的生产环境里面试用和迭代。实事求是的讲，我们在模型的丰富程度，易用性，多语言支持等方面还有比较大的提升空间，目前是落后于其它框架的。但是，OneFlow在企业级大规模应用上是称得上遥遥领先的：（1）分布式最容易使用，用户在写程序的时候是感受不到多机和单机的区别的；（2）OneFlow支持数据并行，模型并行和流水并行，而其它框架只支持最容易支持的数据并行；（3）OneFlow在分布式训练时的扩展能力，加速比是最优秀的。这些特点也正是OneFlow作为企业级深度学习框架，比已有开源深度学习框架优秀之处。

 

人有我优，用数据并行加速CNN训练

 

卷积神经网络（CNN）作为最容易解决的一个问题，是大家最喜欢拿来做基准测试的应用。在过去一年，很多公司用数据并行方法，已经可以用数千块GPU做到几分钟就在ImageNet数据集上训练好ResNet模型。如果发现TensorFlow的参数服务器不给力，上层使用Horovod，底层使用Nvidia NCCL已经可以做到很漂亮的结果。需要注意的是，以前社区有一个认识是TensorFlow并行做的不好，速度比其它框架慢，实际上今天已经不是这样了，TensorFlow团队的benchmark项目（https://github.com/tensorflow/benchmarks）针对CNN做了很多优化，做数据并行已经是开源框架里最优秀之一了。我们使用完全一样的算法和硬件(V100 GPU, 100Gbps RDMA网络)，和TensorFlow benchmark对比会发现，无论是基于单机多卡，还是多机多卡都是比TensorFlow快。上图左边是OneFlow，右边是TensorFlow，除了AlexNet遇到硬件瓶颈，OneFlow都能做到线性加速，TensorFlow在单机多卡和多机多卡上与OneFlow还是有一定的差距。

 

 

阿姆达尔定律

 

上面的评测结果中，在32卡时，OneFlow仍是线性加速，当卡数增加到一定程度，譬如几百或者是上千时迟早会遇到天花板。并行效率不同的系统，只是遇到天花板时间早晚的问题，这是阿姆达尔定律所揭示的规律。比如说上图绿色曲线表示一个并行度（parallel portion）为95%的任务，什么时候遇到天花板呢？可以计算出来，加速到20倍的时候就到了天花板了，后面投入再多的资源进去它也不可能再加速了。假设系统的并行度不随卡数变化，在卡数少时，大部分系统还是比较接近线性的，各个系统之间差别很小，但当卡数增多时，系统迟早会遇到天花板，即使增加再多的GPU也不会进一步提升吞吐率。这表明，在卡数比较少时实现线性加速比不一定能在卡很多时还能实现线性加速，但在卡数较少时就实现不了线性加速，在卡数更多时肯定距离线性加速更远。由此可见，把系统的运行时开销优化到极致，对于大规模集群训练效率是至关重要的。

 

人无我有，分布式训练BERT-Large模型

 

BERT-Large是谷歌最近推出的一个学习语言模型的大型神经网络，基本上在常见的自然语言处理任务上都显著超越了以前的方法。BERT-large有24层，整个模型大概1.3G，每一层中间结果都蛮大的，如果不做内存优化，对于32GB显存的V100，一次也就处理八九个句子。虽然BERT是个大杀器，但客户想基于自己语料重新训练一个BERT-Large模型，却不可能。谷歌在TPU Cluster上用16个TPU训练BERT-Large需要4天时间。没有TPU的用户，只能使用GPU，最主要的是，现在还没有开源的分布式解决方案，谷歌放出来TensorFlow代码只支持单GPU卡，如果用户做一些定制去支持分布式，很遗憾，加速比也很不理想。如左上角图所示，即使是在有NVLink互联的单机八卡服务器上，TensorFlow也只能实现四五倍的加速，按这种加速比去推算一下，即使是使用几十块V100也是需要一个月以上的时间。在Google BERT论文发表后不久，我们团队就基于OneFlow实现了和TensorFlow准确率一样的BERT，在单机八卡服务器上数据并行接近线性加速，在8机64卡的配置下，也能跑到50倍以上的加速比。这还不是线性加速比，我们正在做一些优化工作，不久以后对于BERT-Large在多机多卡也能实现线性加速比。OneFlow现在的实现在单精度条件下只需要8天就能训练出来BERT-Large模型，如果加上半精度支持，时间会再缩短一半，只需要三四天。需要指出的是，Google BERT的词典只有4万个单词，当词表达到几十万或上百万级别时，embedding层就无法用数据并行计算了，必须做模型并行，而后续的层次可以继续使用数据并行，也就是混合并行，OneFlow可以很方便的支持起来。最近，我们已经开始为几家头部互联网公司提供BERT训练服务，在客户自己的数据集上训练BERT-Large模型。

 

以训练安防领域的大规模人脸识别模型为例，当人脸类别达到百万级时，最后的全连接层必须使用模型并行，要解决这个问题，用户就不得不深度hack 已有开源框架，此时会面临易用性和高效性的难题。词嵌入和广告/推荐系统领域也存在许多大模型的问题，模型容量可达几十GB甚至几百GB乃至TB，也只有少数头部企业不计研发成本才能做一些定制开发来支持这些需求。OneFlow 可以很方便高效的支持这些需求，大大节省用户成本，帮助用户完成以前搞不定的事情。

 

 

05总结

 

一路走来，我们深切体会了do right things, do things right 如此重要。在诸多方向里，我们经过反复论证，认为这个领域最关键也最难的问题是横向扩展，从公司成立之初，就立下解决这个业界公认难题的目标。不同于其它框架的技术路线，OneFlow以软硬协同设计为指导思想，从芯片设计领域借鉴了大量有益的思路，在纯软件层面解决了横向扩展难题。我们深信现在OneFlow的技术路线是解决深度学习横向扩展难题的必由之路，在我们走通这条路径之后，很高兴看到技术社区其它团队已经开始沿着这个方向进发。创新和创造是OneFlow决胜的法宝，仅仅follow已有框架走过的路是不可能实现超越的，唯有创新才有机会。最后，我们深感真正有价值的事都是长跑，除了技术因素，情怀和坚持也必不可少，seeing is believing, believing is seeing。