这是Ruta架构用于分布式机器学习集群的算法介绍，它也是Ruta设计之初的一部分，通过主机路由的方式和3D-Torus组网避免了使用交换机带来的incast，同时它包含了一个随路数据同步机制，这是实现一个大规模分布式多业务路由器时用到的流表同步算法，稍微一点改动构成了topology aware的allreduce算法用于分布式机器训练,它也成为智能网卡的另一个落地场景，关于Ruta架构可以参考如下： 《Ruta？智能插座？重新发明路由器！》 |
0.概述
随着模型参数的增加和训练数据的增加，单台机器算力无法满足，存储无法满足，所以要分布式机器学习，而分布式集群算力随着各家公司搞GPU/NPU性能已经提升了很多倍，然后把瓶颈推给了网络。本文主要讲述一种基于Ruta的主机路由协议的方式构建3D-Torus或者Tofu-D(6D)拓扑的方式来构建超大规模机器学习集群的网络优化方式。第一章，一些分布式机器学习基本介绍和已有通信架构的回顾及相关缺陷。第二章，基于Ruta架构的3D-Torus拓扑结构的allreduce算法。第三章，基于Ruta架构的分布式机器学习通信架构实现建议。注：第一章是一些State-of-Art的介绍，主要是给一些不太懂分布式机器学习通信流程的网络架构师做通信业务模式的背景介绍，可能很多熟悉分布式机器学习架构的同学们可以跳过第一章。
1.分布式机器学习架构综述
由于数据规模和模型规模的扩大必须利用计算机集群来构建分布式机器学习框架，并行处理的方式主要是在两个维度。数据规模比较好理解随着IOT和数据采集，可供于训练的数据集通常可以达到数十TB，而模型规模主要是业界竞争的重点，例如GPT-3参数已经达到1700亿，而常见的各厂家模型也到了数十亿级(nVidia MegatronLM 83亿，微软Turing-NLG 170亿）
通常多个工作节点没有共享的公共内存，容量受限于本地内存，训练数据的规模非常大，无法本地存储，因此需要对数据集进行划分，分配到各个工作节点上，然后工作节点依据各自分配的局部数据对模型进行训练。数据的划分可能还在整个训练过程中进行多次随机采样或者shuffle的算法在工作节点之间改变训练样本。这些操作对存储和通信都带来了一些压力，业界通常采用的是SSD来降低数据读取延迟，而最近随着Intel Persistent Memory的出现，有部分公司已经将数据集放置于PMem中，而通信上因为是一个标准的分布式存储的应用，有些企业直接采用了RoCE的方式处理，这也是各家做智能网卡的必争之地。
还有训练样本例如图片等本身可能需要解码等操作，有一些特殊的优化方式。
如果机器学习模型的规模较大，无法本地存储，则需要对模型进行划分。通常的划分方式，例如部分决策树类的模型可以基于维度训练子模型，然后采用集成学习的方式处理。而神经网络类模型则可以横向划分或者纵向划分的方式分解工作量：
由于大规模机器学习需要对数据和模型分块，那么通信主要就集中在两部分：一部分是模型的参数同步，另一部分是数据样本的混洗。
数据分块的混洗通常就是简单的多个节点之间的随机交换数据集，通信模式相对简单，而参数的更新同步则主要是采用AllReduce的方式处理，例如常见的Model Average(MA)需要将每个子模型的参数平均，那么需要在多个节点中将参数求和的处理方式：
最简单的方法就是大家都集中到某一个节点，加好(Sum Reduce)了再广播出去，Reduce集合的地方就是被称作Parameter Server的节点。
1.2.2 Recursive halving and doubling聚合到一点可能会存在一些性能瓶颈，那么抽像成一个树状模型，分层聚合呢？也就是halving-doubling的算法，当然再这里还有一种变种的Butterlfy算法就详细展开了，在阿里2020 HPCA投稿的<EFLOPS:Algorithm and System Co-design for a High Performance Distributed Training Platform>中便采用了这样的算法，只是在通信上增加了Rank-Mapping，构成了HDRM算法。
为了进一步利用带宽，还出现了Ring Allreduce方式，具体如下：
Ring可能在大规模的情况下带来延迟，我们也可以采用Segmented Ring Allreduce的方式，分段的汇总。这是腾讯2018年提出的分层结构，即利用GPU的算力和nvlink的带宽，在组内进行Reduce操作，然后跨机采用组间Allreduce，同步完成后再组内进行broadcast的方式，这也是一种非常取巧的利用组内低延迟高带宽来降低组间通信量的做法。
当然还有一些特殊的通信拓扑，例如Sony 2D-Torus，Google 2D-Mesh以及IBM 3D-Torus等。例如SONY在两年前提出的2D-Torus拓扑结构，在组内做Scatter-Reduce，然后Ring Allreduce。
除了使用那么降低通信量的另一种方法就是滤波和压缩了。一种是基于时域的滤波，也就是说每个子模型迭代几轮以后再更新一次参数：
但是在处理神经网络等非凸模型时缺乏理论保证， 往往需要通过调节一些超参数来取得好的训练效果， 比如本地迭代的轮数、 各个模型使用的学习步长等。另一种方式是采用空域滤波的方式，即压缩通信量,一种比较直观的方法是对模型参数进行过滤。其基本思想是，如果一次迭代中某些参数没有明显变化，便可以将其过滤掉，从而减少通信量。实践中，在训练的后期，众多的参数会趋于收敛，只需要保留少量的参数更新信息，整个模型学习的结果就可以有效地保留下来。以上很多内容都是参考的如下这本书，以及文末附上的参考资料：以往HDD慢的时候，压力都在模型传输上，后期随着SSD逐渐改善并随着RoCE以及NVMe-over-fabric等技术消除了存储的瓶颈，而计算的瓶颈随着异构的各种NPU出现也逐渐消除，网络的瓶颈则是随着近几年数据中心的建设也逐渐地释放开，那么还有什么可以改进的地方呢？一方面是主机内部多GPU构建时的通信瓶颈，阿里在EFLOPS论文描述过的三种主机内的拥塞：
解决办法自然很简单咯，给每个GPU配一个网卡，反正网卡也不贵，买螺丝的Connect-X 4我淘宝买的也就3000 RMB一块还带2个100Gbps的接口
然后阿里的做法是继续使用基于CLOS架构的数据中心，然后对端口进行编号的方式构建一个层次化的拓扑结构：
2.基于Ruta架构的3D-Torus Allreduce算法
2.1 数据中心Incast
有很多反直觉的事情又发生了，数据中心CLOS架构本身是无阻塞的而且用了很多年了，拥塞出现在哪呢？在一个多对一的通信机制下，最容易出现的就是incast了：
为了解决Incast的问题大家也是费心了，例如阿里去年SIGCOMM的论文HPCC便是使用Inband Telemetry的方式实时监控buffer，而Google今年SIGCOMM的Swift则是解耦交换机buffer和接收方buffer。另一方面，阿里好像最近还在招人做Multipath TCP(MPTCP)和MPQUIC，并使用MultiStream等方式想通过CLOS交换机集群ECMP的方式降低Incast的影响。而Fungible则是说可以通过智能网卡调度避免：
IBM提出的名为BlueConnect的算法，其主要思想是考虑了节点间不同switch的带宽不同，从而做出不同的分解利用3D-Torus结构根据带宽权重分解流量，降低CLOS架构交换机的影响，尽量的避免拥塞。
还有一些工作是基于P4的可编程交换机实现交换机上的KV Store(NetCache)，或微软实现的《Scaling Distributed Machine Learning with In-Network Aggregation》通过交换机的可编程能力来缓解。
反直觉的事情又一次发生了，所有的人都在CLOS拓扑上瞎折腾，花那么大功夫折腾测量延迟精确管控Buffer，又是Multipath ECMP，还搞可编程交换机，搞这么复杂干嘛，倒推到问题的本源，不用交换机不就没有Incast了么？ |
其实这个想法是在一个很偶然的情况下，桌上放了一本Visual Group Theory的书复习一下抽象代数，没想到的是突然被几张图弄来了灵感，既然Ring Allreduce是一个环形拓扑，为什么还需要交换机呢？
2.2 分布式集群的异构网络--其实你并不需要交换机
我第一次看到TPU的板子和集群的时候就非常惊奇，它的板子上明显有很多IO：
三种颜色的线和整个集群的部署印证了一点，TPU在内部构建一个集群互相通信，可以注意到CPU节点都放置在两侧，而TPU节点放置在中间,大家可以明显的观测到上下层互联的线缆：
其实就是一个3D-Torus的结构，很多超算其实都是采用这样的拓扑互联的，而IBM还有5D-Torus， Fujitsu的Post-K超算还用了6D-Torus。
其实这样的拓扑结构在2000年路由器构建集群的时候Avici（也就是后来的华为NE5000系列路由器）就尝试过，后来Omni-path也在弄，但是似乎没啥消息了，和CLOS架构相比如下：
而Google也提出了2D-Mesh算法也证明了这点，水平和垂直两个方向直接做Ring Allreduce完成。
| 3D-Torus的拓扑结构使得工作节点之间可以互相直接连接，因为Allreduce本来就要逐跳同步参数，用交换机反而显得有点多余，并且通过这样的方式一方面摆脱了交换机带来的incast影响，另一方面与传统的CLOS架构相比，另一方面从结构在扩展性上与系统规模呈线性关系，同时多冗余的数据通路也可以实施容错等操作。唯一的缺点是主机路由协议设计上。虽然RDMA(Infiniband) 上已经有一些3D-Torus的实现，或者Intel OmniPath也有一些类似的实现，但是容易被Vendor Locking。因此我们还是期望于以太网上通过标准软件的方式实现，并且有后期通过智能网卡硬件加速的能力。Ruta路由协议便可以用于解决这样的问题，它建立在标准的以太网传输介质之上，采用邻居自动发现机制自动识别Torus拓扑，同时可以实施Segment Routing源端选择路径。并且伴随着网络规模的扩大还可以升级支持6D-Torus的拓扑结构降低通信延迟。 |
2.3 Ruta构建3D-Torus
控制面：draft-zartbot-srou-signalling.txt转发面：draft-zartbot-sr-udp.txt为了实现3D-Torus拓扑，每台机器需要3张2x100GE网卡，分别连接X+，X-，Y+，Y-，Z+，Z-方向.这些链路可以采用服务器之间网卡直连的方式，当然也可以为了灵活部署全部接在交换机上划分VLAN/VXLAN的方式连接。
每条链路上的地址随便取/30的网段划分就好，最终Ruta会通过链路发现的方式将linkstate和互联拓扑同步到ETCD中，建议采用一个板载的10G接口或者千兆接口实现ETCD和管理网络的连接。节点的SystemIP编号采用可以采用 10.X.Y.Z，172.X.Y.Z的方式，这样的目的是可以根据XYZ的值做基于目的地址的路由，当然聪明的人看到后刚好X.Y.Z为24bit，还可以构造成一个MPLS Label的方式（这个就不细说了，毕竟技术扶贫还是要有度的）。Ruta可以根据目的归属的Label或者SystemIP自由选择向dX/dY/dZ方向的方式完成路由。数据包的封装用SRoU实现即可。2.4 Ruta 3D-Torus Ring Allreduce
Ruta协议的最大优点是两个，一个是可以使用linkstate获悉全网拓扑和拥塞程度以及链路失效情况，另一个是完全自主的路径决策，并通过Segment Routing的方式可以构建指定路径的转发。这一个功能针对Ring Allreduce算法也非常有用，正好可以在随路转发的时候进行Ring Allreduce。先从简单的一维Ring谈起，如下图所示，节点寻找经过>N/2个节点的某个节点作为Segment-List中的一段，保证顺时针的数据流向，而最后一个Segment指向自己。
由于目的节点为4，报文沿着环路发送到节点2，2将参数保存并更新，然后转发给3，以此类推，节点4完成参数同步操作后，弹出4标签，灌入1标签继续往1发。最终1收到该报文证明整个环上已经完成同步, 因此也不需要防环了，正好借助于环构建可靠传递机制。由于操作仅是同步参数，并不需要随路进行reduce scatter操作，因此操作是幂等的。如果1没有收到回来的包，则证明路径上已经被丢弃，1可以选择超时重发。这样的Ring Allreduce操作可以使用UDP报文实现，因此也不需要使用TCP重复的进行ACK了。还有一种做法可以节省一半的时间和大量的节点计算代价，即环上可以双向发送，一种方法是针对同一种参数，让报文带有Revision,例如顺时针采用[4,1]-->2, 逆时针采用[3,1]->5的方式，当某个节点发现Revision和本地一致后，则丢弃报文，或者改用一个特殊的ACK报文继续传递表明Ring Allreduce已经完成. 另一种做法则是逆向环传递不同的参数。
针对3D-Torus环境，则可以采用固定一个维度进行Ring Allreduce的方式，例如先全部节点对X维度进行Ring AllReduce，完成后对Y维进行Ring AllReduce，这样整个2D平面都完成了AllReduce，最后对Z维进行Ring AllReduce完成对整个立方体的AllReduce。数据的混淆也可以在整个3D-Torus集群上采用随机选择节点，选择disjoint path的方式进行传递，也可以采用和3D-Torus Ring AllReduce的算法，对各个节点进行随路混淆。看上去似乎可以用SRv6来做，一方面是控制平面上的快速收敛，另一方面是Linkstate的主机实现，还有就是SRv6很大程度上需要Kernel的支持，对于应用程序而言，并不是非常的友好。而Ruta是纯UDP的实现，也就是userspace可以很容易的处理。3.基于Ruta架构的3D-Torus Allreduce实现
3.1 纯软件解决方案
通过ETCD获取拓扑和Linkstate以及节点路由等算法的代码，用Golang实现大概7000行左右，而转发平面可以使用BPF进行处理，因为Ruta中继节点的报文处理非常简单。同时建议通信和计算线程解耦的方式：
3.2 智能网卡解决方案
土豪们可能对于通信延迟要求更为极致，例如nVidia有了卖螺丝的网卡后，说不定哪天直接就NVLink带外连接网卡了，也不需要像阿里EFLOPS那样为每一组GPU和网卡配置PCIE-Switch了，其实已经有一些端倪了，如下图所示：
而Ruta协议的报文封装和路由和VXLAN或者SRv6实现基本上一致，硬件加速上用P4实现或者Verilog写FPGA也非常容易。因此可以建议有自己FPGA网卡的某几个云自己玩，没有的喜欢P4的买点Pensando来自己开发：
而都没有的但是精通C/C++的某厂去买Fungible吧。注意可以借助他们卡上的HBM实现和GPU的参数同步和本地缓存等操作。具体细节就不多说了，技术扶贫还是不光要扶，还是要靠各位自己上路的~4.未来之路
分布式机器学习本质上对系统架构师的考验是全方面的。一部分是算法本身，它需要大量的计算数学矩阵代数的知识，很多直觉上看似正确的近似都会带来算法的精度下降或者严重到不收敛的地步，而有时候超参数调参更是碰运气。另一方面是计算机体系架构上的理解，指令集的并行上以及如何解耦I/O密集型和计算密集型事务并采用很好的异构措施。最后就是对通信架构的理解，如何进行随路计算和随路数据同步，如何高效的通过编码规约实现特定计算任务。4.1 6D-Torus
3D-Torus的结构随着GPU规模的扩大也可能存在单个Ring的瓶颈，同时节点内可能存在多块GPU和nvlink，则可以做类似于Fujstu Post-K Computer那样的Tofu-D 6D-Torus架构的尝试：
4.2 Learning on Learn
第二章中，Ruta实现了一种最简单的基于维度独立的Ring Allreduce算法，而实际上Ruta协议是完全的Topology free的协议，并且在userspace给程序员提供了非常灵活的编程手段，少量的几个header编码便可完成灵活可控的路由选择。甚至可以在训练中采集数据利用AIOps的思路Learning on Learn的处理方式自动优化Ruta的路径选择降低拥塞。4.3片上网络和主机网络的深度融合
随着光互连等技术的出现以及芯片封装工艺的变革及Chiplet概念的出现，片上网络也是一个非常值得研究的话题，同样以Post-K的A64FX为例，如何高效的利用HBM2片上网络和片间通信，它们采用的TofuD有很多值得学习的地方。
前路有太多的未知，我们也有太多的差距，无论是芯片技术上，还是数学算法上，我只是比一般人多读了几本书，多看了一点世界，仅做一点微小的贡献，希望能够帮助到一些人，也算是完成我技术扶贫的一点小心愿吧。 世界很大，眼界更大一些，少一点比友商快几秒，毕竟生如蝼蚁当有鸿鹄之志，命如纸薄应有不屈之心，大家一起打怪兽改变世界不是更好么
|
[1]分布式机器学习：算法、理论与实践, 作者: 刘铁岩 / 陈薇 / 王太峰 / 高飞[2]腾讯机智团队分享--AllReduce算法的前世今生[4]EFLOPS: Algorithm and System Co-design for a High Performance Distributed Training Platform[5]Avici System, Scalable Switching Fabrics for Internet Routers