 
分享嘉宾：蓝昶博士 Google
编辑整理：何文婷 字节跳动
出品平台：DataFunTalk
导读：随着设备学习模型和数据规模的增长，大规模分布式设备学习训练的性能越来越作为公有云用户关注的问题。本文将介绍谷歌云 Vertex AI 平台在分布式设备学习训练性能优化方面做的一系列工作。
详细将围绕以下几点展开：
训练优化的背景
Fast Socket: NCCL的高性能网络栈
用Reduction Server加速梯度聚合 01训练优化的背景
1. Google Vertex AI平台简介
Vertex AI是Google的一站式托管云服务，是一个集成为了AutoML和AI Platform的AI设备学习以及服务平台。Vertex AI覆盖了从数据到特征工程、模型训练，超参调节和模型预测以及可预测性支持在内的一系列的需要。接下来的分享重点是聚焦在模型加速的技术部分。
2. 训练优化的背景
咱们先回顾一下咱们为何需要做分布式的训练，以及咱们做为一个平台为何需要关注这个问题。近期几年咱们能够看到深度模型最起始是在学术界的各个任务上取得了有些SOTA的突破，咱们近期几年亦看到这些应用起始渗入到企业级的用户，并且带来实实在在的业的提效。
云平台上的负载亦从最初的视觉模型起始，此刻变得非常多样，语言模型跟多模态模型亦逐步增加。随着大规模训练模型，语言模型在各个任务上的效果亦实现了突破的话。咱们看到用户的模型负载规模亦是在指数的增长。支持这些海量的大规模型的训练，亦给咱们带来有些新的挑战。简单来讲，单卡训练再也不是主流，训练大模型所需要的算力在指数增长，因此单卡不足以在足够的时间里面输出足够的算力。另外，因为GPU显存的限制，除非经过有些跟主存进行数据交换的有些trick，简单的单卡训练亦再也不可行。在这种状况下，分布式训练做为一种水平扩展的处理方式，是日前主流的选取。从日前扩展的方式来看，咱们最早遇到的是算力的限制，因此咱们最早是在主流的框架里面看到非常多数据并行的有些API支持。此刻随着模型规模的继续增大，模型并行以及混合并行亦起始进入主流的AI框架的研发路线。
3. 水平扩展的挑战：内存墙
图2水平扩展面临的一个基本的挑战是内存墙的问题。这个问题有两个方面，首要咱们能够看到，内存带宽是在长时间来看是慢于明显慢于算力的增长，在过去的20年里面计算的算力提高了九万倍，但内存的带宽只加强了30倍，这是一个基本的差距，并且差距会越来越大。其次，广义的内存带宽既包含片上的内存带宽，亦包含片间的互联带宽，还有网络带宽。这几类带宽的长时间趋势亦都面临内存墙的问题。咱们从这个图里面能够看到，内存带宽，高速互联带宽跟网络带宽之间亦是始终存在着数量级上的差距，因此随着训练规模的持续增大，训练的性能瓶颈最后会落到网络带宽上。因为有这两个方面的影响，倘若只是依靠硬件本身的发展，模型规模火速就会撞到内存墙，这亦便是为何咱们需要在框架跟算法方面做更加多工作的大背景。那样从一个云平台的方向来看，咱们做性能优化工作的首要目的当然是提高性能，为用户降低TCO，其次，咱们期盼做到非侵入式的优化，做到框架无关，让用户保持选取框架的自由度。
在实质的场景里面，咱们能够看到内存墙针对训练性能的影响。上图表示了一个例子，在单机跟多机的场景下，针对同一个模型的单步训练时间的比较，因为在多机跟单机之间计算跟参数更新的时间大体是恒定的，真正拖慢训练时间的是梯度聚合这一步，占用了大概2/3的时间，因此all-reduce常常是分布式训练的性能瓶颈。
4. 训练优化的技术路径
关于深度的分布式训练，重点工作主从技术栈上呈现从浅层到深层的一个过程。前三类的优化基本上是处在框架层，需要平台为用户供给基本的框架支持。例如说在计算图的并行化策略方面，咱们经过GSPMD和GPipe供给了包含数据并行、模型并行和流水线并行的通用化的并行化策略的抽象层。另外，咱们经过DeepSpeed来做支持，用ZeRO (Zero Redundancy Optimizer)来优化Optimizer的显存运用，以及咱们能够用低精度的压缩来加速参数的同步。在这次的分享里面，我会重点分享最后一类的优化，亦便是在集合通信层的有些优化。在AI框架的设计里面，这是讨论的比较少，然则针对一个平台来讲是非常重要的一类优化。
首要针对基本设备的优化常常是需要有一个全局的提效，其次这类优化针对用户跟上层的框架完全透明，不需要改动上层的代码就能够真正落地。咱们接下来经过有些详细的例子来看怎样以集合通信层为入手点来做这一类的优化。
5. NCCL简介
在GPU的训练场景里面，集合通信常常跟NCCL是同义词。NCCL是NVIDIA实现的一个集合通信库，它供给了像allreduce, allgather, broadcast等集合通信原语以及高性能的实现，用来支持多机多卡的训练。针对节点内通信，支持NVLink，PCIE和device P2P等方式；针对节点间通信支持像Socket和Infiniband等网络协议，并且网络协议能够支持经过用插件的形式来扩展。从软件栈的方向来看，NCCL对重点的训练框架都供给了支持，并且是主流框架默认运用的GPU的通信库。拿网络的协议站做一个类比的话，NCCL基本上跟IP协议同样，是全部协议栈的narrow waist的位置。
咱们认为NCCL尤其适合用来做为框架无关的性能提高的着力点。详细到几个通信的优化，咱们能够分为两类，一类是针对底层通信网络栈的优化，而另一类是针对几个通信算法的优化。那样接下来咱们将会用两个例子来分享咱们在这两个方面做的工作。 02Fast Socket：NCCL的高性能网络栈
首要咱们要介绍的工作是咱们对NCCL实现了一个高性能的底层网络栈。为何需要做这方面的工作呢？咱们前面说到NCCL用了非常多巧妙的设计和底层的优化来实现高性能的集合通信，然则它详细的实践还有非常多需要加强的地区，对此咱们对大信息和小信息做分别的讨论。
1. 提高挑信息的吞吐率
针对大信息的传输，重点是需要加强吞吐，但咱们实测发掘，在高带宽非RDMA的环境里面，NCCL的网络吞吐性能常常不尽如人意，因此在100G的以太网环境里面，实质用到的带宽远远达不到line rate，因此呢有巨大的提高空间。NCCL的默认实现其实亦思虑了怎样运用高带宽网络的问题。首要它的默认实现是在多个节点之间的每一个环，NCCL都会创立多个TCP的链接来并行传输信息，另外它本身亦会运用多个Ring来处理集合通信的请求（Ring是NCCL里面自己定义的一个概念，能够理解为是一个独立的通道，每一个通信通道亦需要独立的占用CPU跟GPU的kernel资源）。然则这种简单的运用海量连接的方式针对大信息的吞吐率效果并不睬想，原由有两个：
运用海量连接和Ring会占用CPU跟GPU的资源，反而会影响到计算本身的速度，并且会增多性能的抖动。
在实质的网络环境里面，区别的TCP流，它的带宽占用并不一致，因此会引起straggler effect，因此其他已然完成的流会等需要等待最后最慢完成的流，作为性能的瓶颈。
咱们能够仔细看一看NCCL的传输层的实现，简单把它抽象成右边这个图里面表示的这么一个结构。当NCCL得到最上层的传输请求之后，他会用Proxy的线程将准备传输的信息切片，而后用Round Robin的方式把切片后的数据交给Helper线程，经过区别的Socket发送到对面的节点。理想的状况下，Socket之间的带宽相会是相等的，因此这种方式看起来应该是无什么问题。但实质的数据中心网络里面影响带宽的方式非常多，像倘若是有multi-path的话，区别的connection会被路由到区别的path上，路由器的buffer占用亦不同样，因此就会引起区别的连接的带宽并不均等，因此大信息的吞吐率常常会被这般的straggler连接给拖慢。
为认识决这个问题，咱们采用了动态负载平衡的办法。首要咱们在NCCL原有架构的基本上，继续做细粒度的数据切分，流式地处理数据切片，其次，咱们改变了原来Proxy线程经过Round Robin分配负载的做法，经过感知每一个Socket当前的负载量和进度，把负载分配到最快的Socket之上，这儿的难点就应该怎样感知每一个Socket的负载。这儿一个小背景是Socket线程和Helper线程之间运用的是个一个无锁队列进行数据交互，一个很自然的想法便是说咱们期盼瞧瞧每一个线程对应的队列长度，缓冲区的长度越长，说明对面对应的Socket越慢，反之就越快。这亦是咱们最初的版本采用的一个做法，然则实质上咱们发掘效果加强有限。原由在于除了用户态的无锁队列之外，数据还能够在内核的Socket的缓冲区被缓冲，因此你倘若只看队列的长度，并不可准确地反映出Socket的当前的负载。修复这个问题亦很简单，咱们经过调节内核Socket的参数能够关闭发送的缓冲区，这般就能够让信号的反馈更加准确。另外，咱们经过压缩缓冲区的体积亦能够掌控在海量运用Socket的状况下。针对内核内存的占用，总体上针对性能亦是有帮忙的。
除了动态负载平衡，咱们还改进了NCCL针对请求的处理方式。NCCL本来的实现是不可并行处理传输的信息的，它只能在完成当前的请求之后才可处理下一个请求。咱们改进了这个方式，经过实现对请求队列的look ahead处理，能够流式地处理这个变行请求。另一，咱们开启了发送端的zero copy来降低用户态到内核态的拷贝开销，针对大于十KB的信息，咱们实测会有有显著的提高效果。
2. 降低小信息的延迟
针对小信息，咱们重点关注的是降低延迟。前面说到信息的发送是需要经过Proxy线程到helper线程的数据交换，每一次的发送都会导致多次的线程唤醒，跟上下文切换，导致比很强的开销。咱们处理这个问题的办法是经过Proxy线程直接掌控Socket的发送，避免线程切换的开销。
为了实现这个优化，咱们亦重新设计了NCCL掌控信息的结构以及传输掌控信息的办法。在数据信息足够小的时候，咱们能够相当于是把信息内inline在掌控信息里面，经过proxy直接发送。
最后咱们亦引入了内核的Busy polling用于掌控socket，让内核来动态的来poll socket能够明显地降低小信息的延迟跟抖动。
3. End to end 测试
之前介绍了咱们详细做的有些优化的办法，咱们亦测试了NCCL在64M到1G的信息体积上，针对all-reduce的网络吞吐率，能够看到经过Fast Socket的优化之后，NCCL在各个体积的带宽测试里面都取得了60%以上的加速比。在实质的100G以太网络里面，实质带宽亦能跑到将近line rate的数字。
针对end to end的性能测试，Fast Socket亦能取得一个比较显著的提速。图中表示的是在Fine-tune BERT-Large这种模型的时候，Fast Socket可在每秒训练的步数上会有大概30%以上的提速。这种提速是面向全平台的，因此咱们不需要用户侧做任何的改动，就能让用户实质的落地加速的效果。
4. 小结
咱们这里对Fast Socket做一个简单的小结。Fast Socket是咱们为NCCL在高带宽的网络环境里面实现的一个优化的网络栈，由于这些优化都位置于NCCL的通信层，因此支持所有的主流分布式的框架，并且能够做到全平台的加速。日前Fast Socket以插件的方式类似于Google Cloud的Deep Learning VM和Vertex AI等设备学习环境里，详细的实现代码，亦以开源的形式开放给社区，欢迎试用和交流。 03用Reduction Server加速梯度聚合
咱们前面的说到的Fast Socket重点是用工程的方式对集合通讯层做了非常多细致的底层优化。在这个工作里面，咱们换一个视角引出下一个工作，瞧瞧咱们怎样从改进集合通信算法的方向来加速梯度的聚合。
1. All-reduce简介
咱们先给一个all-reduce的详细的例子来回顾all-reduce的语义，这儿面每一个worker的节点有一个等长的数组，在训练里面一般是对应于某个参数的梯度，由于在数据并行的训练里面，每一个worker的训练输入的数据的批次不同样，梯度的数值亦区别，因此呢咱们需要对区别的批次的梯度求和，亦便是对应于all-reduce的操作。操作完成之后，每一个worker都会得到相同的结果。
 fill=%23FFFFFF%3E%3Crect x=249 y=126 width=1 height=1%3E%3C/rect%3E%3C/g%3E%3C/g%3E%3C/svg%3E)
因此总结来讲，all-reduce的语义是：
规约所有节点上的数组，并且把这个结果返回到所有节点。
All-reduce有非常多详细的实现，一般是能够由两步组合而成，一般分别是由于reduce-scatter和all-gather的两步组合而成。Reduce-scatter完成之后，每一个节点各自持有N分之一完整规约过后的数据。
在右图这个例子里面，每一个节点需要最少要发送(n-1)/n的数据，这一点非常容易证明。例如说在图中的例子，节点1，2，3分别要来自需要来自节点0的1/4的数据，因此这个节点0最少需要送发送3/4的数据，一样属于节点0规约的1/4的数据块需要来自节点1，2，3相同位置的数据块，因此它亦最少需要接受3/4的数据。
下一步的all-gather则是将各个节点上1/4的规约结果发送到所有的节点。效果上等价于四次的broadcast。咱们亦很容易证明all-gather的每一个节点亦需要收发(n-1)/n的数据，因此呢all-reduce里面每一个节点需要传输大概两倍于原始输入的数据。这个是个非常关键的结论，咱们待会会回到结论来瞧瞧怎样优化这一点。
2. All-reduce性能分析
咱们能够经过这么一个简单的模型来分析all-reduce的性能。咱们能够定义算法带宽为输入数据的体积除以all-reduce执行的时间。例如说每一个节点的输入数据是1G，而后all-reduce用了一秒，算法带宽便是1G/秒。
咱们能够把算法带宽拆分为两项，第1项是输入数据的体积除以实质在网络里面传输的数据，那样咱们叫作为算法效率，这一项针对一个特定的算法是一个常数，对，例如说针对ring all-reduce来讲，这一项便是n/(2(n-1))，当节点数N非常大的时候，这个数相当于1/2，第二项便是实质传输的数据除以执行时间，亦便是实质的网络或说总线的吞吐率，咱们叫作之为总线带宽。这一项亦受到实质的硬件和协议栈的限制。
为了加强全部算法带宽，有两种思路：
用更新更好的硬件加强总线开关，而后再加上更优化的网络栈实现，例如说像咱们刚才说到的fast socket，或说用Infiniband RDMA等等。然则从咱们刚才说到的内存墙的趋势，我们能够看到硬件带宽的增长始终是有限度的。
经过加强算法的效率，在这个场景下，亦便是降低数据的传输量。一个能够证明的结论是all-reduce的算法效率理论上界是ring all-reduce日前的水平，亦便是当N非常大，当工作节点的数量非常大的时候，大概便是1/2。咱们的reduction server工作调节了all-reduce的一个设定，经过一个稍微区别的思路，算法效率说到了加强到日前的两倍。
3. Reduction Server
Reduction Server启发于parameter server的通信方式。在parameter server的架构里面，worker的节点在每一个iteration只传输一次的参数数据。受到启发。便是说实质上咱们能够想在all-reduce的框架里面，咱们是不是亦能够经过这种方式来做集合运算。咱们的方式是引入了reduction server节点，它的通信拓扑跟parameter server是一致的，然则节点的实现更加简单，例如说节点，它不需要保留参数，亦不需要计算梯度，它只要在每一个action里面规约来自worker节点的梯度数据，并且返回给worker节点，经过这种方式，worker亦只需要发送跟接收一次完整的完整的梯度数据，并且咱们能够经过实现流式的规约，隐匿worker收发之间的延迟，充分的利用双向的带宽。
另一一个非常重要的点是，咱们虽然增多了额外的节点数量，然则这些节点都是轻量级的CPU节点，总开销，倘若用公有云的价格来看的话，总开销的节点开销是GPU节点的10%以内，并且还有一个优良是咱们能够把这些轻量级的CPU节点跟他其他网络利用率低的节点混合安排，因此实质上增多的成本能够忽略不计。这个表里总结了跟传统的all-reduce算法相比，reduction server能做到的优良，首要它能把数据的传输量减半，亦便是说算法带宽相当于是原来的两倍，另一一个优良是它能把延迟的量级从ring all-reduce的O(N)降到O(1)，这一点是针对小信息的性能亦是非常重要的。
咱们来瞧瞧咱们实现的方式，那样在worker的节点端，咱们基于NCCL以下的通信层实现了一个到reduction server的通信层，因此咱们能够在不改动框架的状况下实现从all-reduce到reduction server的无缝切换跟加速。在reduction server的节点端，咱们基于Fiber实现了高性能的网络通信层。在这之上是一个轻量级的规约引擎。规则引擎的重点工作是经过高性能SIMD优化过后的算子对输入的数据进行规约。咱们在规约引擎亦实现了完整的数据类型支持，并且能够支持混合精度的压缩、规约。
4. 训练性能&TCO
最后咱们瞧瞧reduction server针对训练性能的提高，咱们能够看到它针对各个体积的信息的all-reduce操作都会有显著的性能提高。
在end to end的测试里面，相针对传统的ring all-reduce, reduction server能够把训练速度提高75%上下。除了纯粹性能方面的优化，下面这个表的例子里面咱们能够看到，由于相当于说咱们的速度变快了，咱们花费的时间亦少，因此便是说它实质上总成本亦能降低。即使思虑了额外的CPU节点的开销在内，用户仍然能够大幅降低训练的TCO。
日前reduction server已然集成到Vertex AI的平台，用户无需改动代码就能够很方便地为日前自己已有的分布式训练任务开启reduction server的支持。咱们在Vertex AI的平台的网站上亦发发布了相应的博客文档以及Notebook的样例，有兴趣能够继续参考。 04总结与展望
总结一下今天分享的内容，内存墙问题是日前大规模模型训练很难避开的一个问题，并且特别有可能是一个长时间的问题，只靠硬件的自然演进应该是不足的，因此这亦是咱们在框架的基本上，技术栈的各个层面做性能工作的一个大背景。将来随着业务模型的变化，咱们亦看到将会围绕着更加多其他并情化策略的优化工作。今天我重点是从云平台的方向看，咱们需要怎么样的性能优化工作，侧重分享的是跟框框架无关的底层的有些优化。然则咱们亦能够看到非常多跟框架或乃至模型强强耦合的有些性能工作，例如说Deep Speed或Horovod。这实质上表率的是两种范式，亦能够引出非常多讨论，例如说咱们在做有些上性能方面的设计的时候，是应该尽可能做到平台无关，还是针对某个优化应该推出一个新的框架，或更进一步，更好的性能是不是一个AI框架的核心竞争力，我想这些都是有些非常有意思的问题。今天我的分享就到这儿，倘若有问题的话亦欢迎提出，谢谢大众。
今天的分享就到这儿，谢谢大家。
在文末分享、点赞、在看，给个3连击呗~
分享嘉宾：
关于咱们： DataFun：专注于大数据、人工智能技术应用的分享与交流。发起于2017年，在北京、上海、深圳、杭州等城市举办超过100+线下和100+线上沙龙、论坛及峰会，已邀请近1000位专家和学者参与分享。其公众号 DataFunTalk 累计生产原创文案500+，百万+阅读，13万+精细粉丝。
| 
发表于 2024-7-28 01:18:52
