分布式机器学习产生的原因很简单, 一方面是可供训练的数据越来越多,另一方面是模型自身的规模越来越大,所以必须要多个机器来搞。RoCE一类的通信协议自然被用到了,这其实也是nVidia要买Mellanox的根本原因,而并行的方法主要有如下两种:
数据并行很容易解释,主要是如何存储训练样本,并且在多机器之间传递混淆样本,基本上大家大同小异的都在采用SSD、分布式存储解决这些问题,当然还有内存池化的需求.
另一个问题便是模型并行,当单个工作节点无法存储时,就需要对模型本身进行分割。当分布式训练每轮迭代完成都需要将参数进行同步,通常是将每个模型对应的参数加总求和再获得平均值,这种通信被称为AllReduce
最开始的时候,是采用一个集中式的参数服务器(Parameter Server)构建,但是很快就发现它成了整个集群的瓶颈,然后又有了一些环形拓扑的All-Reduce
而对于nVidia而言,它们极力的扩大NVLink的带宽,同时也快速的迭代NCCL,都是为了解决这个AllReduce的问题,但是这些只在单机或者一个极度紧耦合的集群内部。另一方面主机间的通信,自然就选择了超算中非常常见的RDMA ROCE了。
但是即便如此,AllReduce的延迟还是极大的影响了整个训练集群的规模:
Allreduce算法简介可以参考鹅厂总结的:
腾讯机智团队分享--AllReduce算法的前世今生[1]
另一个工作:EFLOPS
阿里在HPCA2020上发布了一篇论文
阿里其实也看清楚了这个问题,PCIe的拥塞,内部调度的拥塞,网卡的拥塞:
然后解决方案很简单,反正钱多,一个GPU配一个网卡就好,然后网口多了,交换网也改成两套Fat-Tree
阿里的文章中有一个结论
也就是说即便是用了HDRM,也就40Gbps的带宽了,那我先告诉你们NetDAM的一个结论100Gps轻松跑满,单个Alveo U55N可以跑满200Gbps,赠送一句话: In me tiger sniffs the rose.
NetDAM实现AllReduce
首先不谈AllReduce的算法和相应的拓扑,在带宽一定的条件下的约束是通信延迟和计算延迟. 如果采用RoCE,从一台机器读和写都要经过一次PCIe,所以从根源上要解决这个问题就是内存前置,延迟不就下来了么?
通信延迟降下来了,我们再来看计算延迟,传统的方式要怎么加:
而在计算域内,CPU嘛,AVX512加咯还能怎么样,带上Cache延迟抖动都不好控制,丢GPU上还要多一次Memory Copy,即便是直接使用GPU-Direct不也要过一次PCIe么?所以你跑不到线速100Gbps很正常
直接在网卡上放置大量ALU,收到包的时候,包还在SRAM buffer中,这个时候ALU根据包头的NetDAM Instruction,可以多个ALU同时去load本地DRAM,然后add到相应的SRAM里。加完以后,整个包改个IP头直接就转发,这样一个9000B的报文可以承载2048个float32,等同于AVX(32*2048)的SIMD-Add,所以我当然比你CPU快咯,而且加的时候没有DRAM的Store,只有最后一跳才会Store,又省了多少?
除此之外,针对AI训练的场景,还有很多可以直接通过NetDAM ALU过滤的方法,例如当一个SIMD内部的2048个float32有一半以上的0时,我可以很简单的使用
而在AllGather阶段,也就是说算好数据需要再次分发的时候,RoCE的组播似乎只是一个概念上的东西,而NetDAM则可以在这个阶段充分发挥以太网组播或者广播的能力,当然具体的丢包重传,这些都在NetDAM之间就可以完成,FPGA检测到Seq丢失直接产生一个READ报文给源就行了,压根不需要CPU参与,具体内容明天讲拥塞控制的时候详细说。
关键还不止这一点,它还内带了一个Segment Routing头,可以做链式反应,就像原子弹那样~嘣~~~
链式反应另一个特点就是,打开了通向3D-Torus拓扑的新空间,毕竟连交换机延迟都省了,而且用RingAllreduce跑满带宽还不需要考虑incast,漂亮不?香不香?
而Google TPU集群为什么要用Torus-Ring,甚至一些超算用6D-Torus,想明白了么?其实就是在扩展性上,Non-Blocking成本很高,而且临时扩大或者缩小集群规模需要添加额外的设备构成FatTree,Incast也不好控制,而Torus虽然是有阻塞的但是可以通过通信模式来避免阻塞。
即便是用Fat-Tree的数据中心,我们也给你们准备了Ruta的方案来做流量工程,比起那啥搞什么PortRank,更加简单直观的是哪儿不堵走哪儿~ 拥塞控制,明天给你们安排~
NetDAM实现内存池
NetDAM是一个标准的UDP协议,NetDAM可以独立于主机单独部署, 因此可以构成一个非常大规模的内存池:
因此普通主机 用户态不需要任何特殊的开发套件,直接一个UDP Socket就可以控制整个内存集群,爽不爽?
而当你主机自己有了NetDAM卡了以后,可以玩的更High, 分区全局地址空间(partitioned global address space:PGAS)了解一下, 在这种场景下,我们可以把一个交换机芯片改造成MMU,对外提供一个虚拟的IP地址和UDP端口,然后构成一个大的虚拟化池隐藏内部拓扑。而每个netDAM报文访问的内存地址由交换芯片查表做地址转换到最终的NetDAM。这种情况下,交换机MMU还可以采用Interleave编址来解决内存局部使用过热的问题...
继续从分布式AI训练集群来看,对内存池的需求主要是一个是训练数据集的分发和混淆,另一个是参数和梯度的更新。所以这次HotChip中Cerebras提供了一个Memory-X套件:
计算任务上,MemoryX还添加了Optimizer
结论 NetDAM也可以同样的实现这个功能:)
预告....EFLOPS谈完了,我们来谈谈HPCC?当延迟为确定性时,只需要考虑Buffer深度了,那么算法就更简单了:
Reference
[1]腾讯机智团队分享--AllReduce算法的前世今生:
https://zhuanlan.zhihu.com/p/79030485