谷歌分布式机器学习优化实践
分享嘉宾：蓝昶博士 Google
编辑整理：何文婷 字节跳动
出品平台：DataFunTalk
导读：随着机器学习模型和数据规模的增长，大规模分布式机器学习训练的性能越来越成为公有云用户关注的问题。本文将介绍谷歌云 Vertex AI 平台在分布式机器学习训练性能优化方面做的一系列工作。
具体将围绕以下几点展开：
训练优化的背景
Fast Socket: NCCL的高性能网络栈
用Reduction Server加速梯度聚合
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特别适合用来作为框架无关的性能提升的着力点。具体到几个通信的优化，我们可以分为两类，一类是对于底层通信网络栈的优化，而另一类是对于几个通信算法的优化。那么接下来我们将会用两个例子来分享我们在这两个方面做的工作。
首先我们要介绍的工作是我们对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等机器学习环境里，具体的实现代码，也以开源的形式开放给社区，欢迎试用和交流。
我们前面的提到的Fast Socket主要是用工程的方式对集合通讯层做了非常多细致的底层优化。在这个工作里面，我们换一个视角引出下一个工作，看看我们如何从改进集合通信算法的角度来加速梯度的聚合。
1. All-reduce简介
我们先给一个all-reduce的具体的例子来回顾all-reduce的语义，这里面每个worker的节点有一个等长的数组，在训练里面通常是对应于某个参数的梯度，因为在数据并行的训练里面，每个worker的训练输入的数据的批次不一样，梯度的数值也不同，因此我们需要对不同的批次的梯度求和，也就是对应于all-reduce的操作。操作完成之后，每个worker都会得到相同的结果。
所以总结来说，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/秒。
我们可以把算法带宽拆分为两项，第一项是输入数据的大小除以实际在网络里面传输的数据，那么我们称为算法效率，这一项对于一个特定的算法是一个常数，对，比如说对于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的样例，有兴趣可以继续参考。
总结一下今天分享的内容，内存墙问题是目前大规模模型训练很难避开的一个问题，并且很有可能是一个长期的问题，只靠硬件的自然演进应该是不够的，所以这也是我们在框架的基础上，技术栈的各个层面做性能工作的一个大背景。未来随着业务模型的变化，我们也看到将会围绕着更多其他并情化策略的优化工作。今天我主要是从云平台的角度看，我们需要怎么样的性能优化工作，侧重分享的是跟框框架无关的底层的一些优化。但是我们也可以看到很多跟框架或者甚至模型强强耦合的一些性能工作，比如说Deep Speed或者Horovod。这实际上代表的是两种范式，也可以引出很多讨论，比如说我们在做一些上性能方面的设计的时候，是应该尽量做到平台无关，还是对于某个优化应该推出一个新的框架，或者更进一步，更好的性能是不是一个AI框架的核心竞争力，我想这些都是一些非常有意思的问题。今天我的分享就到这里，如果有问题的话也欢迎提出，谢谢大家。
今天的分享就到这里，谢谢大家。
在文末分享、点赞、在看，给个3连击呗~
分享嘉宾：
专知便捷查看
便捷下载，请关注专知公众号（点击上方蓝色专知关注）
后台回复“分布式机器学习” 就可以获取《分布式机器学习专知资料合集》专知下载链接
请扫码加入专知人工智能群（长按二维码），或者加专知小助手微信（zhuanzhi02），加入专知主题群（请备注主题类型：AI、NLP、CV、 KG、论文等）交流~
相关内容
微信扫码咨询专知VIP会员