神奇：内存池化和分布式AI集群优化

作者： 扎波特的网线钳 2021-10-17 21:04:34

人工智能

分布式 分布式机器学习产生的原因很简单， 一方面是可供训练的数据越来越多，另一方面是模型自身的规模越来越大，所以必须要多个机器来搞。

[[429309]]

分布式机器学习产生的原因很简单， 一方面是可供训练的数据越来越多，另一方面是模型自身的规模越来越大，所以必须要多个机器来搞。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