深度揭秘：Meta工程师如何构建超大规模AI训练网络？

最近，Meta的多个工程团队联合发表了一篇论文，描述了在引入基于GPU的分布式训练时，他们如何为其「量身定制」专用的数据中心网络。

前段时间发布的Llama 3.1 405B可谓是LLM界的良心开源。不仅公开了模型权重，而且在发布的论文中详细介绍了所用的算法和工程方法，比如模型架构、指令微调等等。

论文地址：https://ai.meta.com/blog/meta-llama-3-1/

此外，论文还难得地披露了训练基础设施的各方面细节，比如4D并行、集群通信、故障率和可靠性等等。

其中，关于集群意外中断及其归因统计更是让我们了解到，即使能用上最先进的H100 GPU，也要面对如此频繁的硬件故障。

但毕竟是1.6万块GPU组成的超大集群，工程量可想而知，即使Llama 3.1论文的篇幅有洋洋洒洒92页，也很难深入、详细地描述其构建过程。

因此，Meta的工程师们最近又发表了一篇论文，专门介绍如何大规模设计、实施和运营这个庞大的AI训练网络。

论文地址：https://dl.acm.org/doi/10.1145/3651890.3672233

这篇论文集结了Meta多个团队的智慧结晶，包括AI生产工程（AI Production Engineering）、AI与系统协同设计、AI硬件系统等。

选择RoCE的原因

分布式网络中的GPU间通信主要包括两个阶段。首先是在单个训练节点上的4～8个GPU之间进行「节点内通信」，使用NVLink高速互联方式。

如果训练工作需要额外GPU参与，就需要引入「节点间通信」，对此，业界通常有两种设计方式。

- 标准的TCP/IP网络或对其中的插件进行修饰，比如fastsocket

- 专有的互连技术，比如InfiniBand、NVSwitch、Elastic Fabric Adaptor、Inter-rack等

由于CPU开销和延迟的增加，前者容易造成集群性能的下降；后者尽管能提供较好的性能，但由于是专有技术，很难灵活部署。

因此，当Meta引入基于GPU的分布式训练时，工程师们决定为其量身定制数据中心网络，最后选择RoCEv2（RDMA over Converged EtherNet v2）作为主要的节点间通信机制。

RDMA全称Remote Direct Memory Access，可以在无需CPU参与的情况下实现互连GPU的内存共享。

基于TCP/IP的通信机制中，数据包必须先发送到内核才能拷贝至内存，而RDMA则绕过内核，信息可以直接到达或发送至应用内存。

RoCEv2则是实施RDMA的一种具体协议，规定使用以太网传输，数据包采用UDP格式，读写信息的封装和解封都由RDMA NIC硬件处理。

之所以选择RoCE，出于以下三方面的动机：

- RoCE与训练工作负载常用的RDMA一脉相承，确保已有设施的无缝衔接

- 使用以太网可以保留原数据中心相当比例的组件和工具，并能继续使用基于Clos的设计

- 整个技术栈都以开放标准为基础，确保网络基础设施的兼容和灵活

扩展后的RoCE网络中，每个集群可容纳数千甚至数万个GPU，而且可用于支持生产场景下的各种GPU工作任务，比如排名、内容推荐、内容理解、NLP和GenAI模型训练等。

拓扑结构

前后端分离

训练集群主要依靠两个互相独立的网络：前端网络（Frontend Network, FE）负责数据摄取、检查点和日志记录，后端网络（Backend Network, BE）用于训练。AI训练机架分别连接到FE和BE。

FE的网络层级中包含机架交换机 (RSW)、结构交换机 (FSW) 以及更高层的存储仓库，为GPU提供训练所需的输入数据。

BE是一种专用结构，以非阻塞架构连接所有RDMA NIC，在集群中的任意两个GPU之间，无论物理距离如何，都能提供高带宽、低延迟的无损传输。后端结构使用的协议即为RoCEv2。

FE和BE的分离是部署RoCE的早期做出的一项主要决策，主要是希望两个网络能够独立进化、互不干扰。此外，将用于AI训练的网络分隔出来，也能简化并加速路由和通信机制的迭代。

BE网络曾被多次修改。最初的GPU集群使用简单的星形拓扑，让一些AI机架连接到中央以太交换机（运行不可路由的RoCEv1协议）。

这种设计在GPU规模和交换机冗余方面有明显的限制。因此，后来迅速过渡到基于结构（fabric-based）的架构，以实现更好的扩展性和可用性。

AI Zone

AI机架的设计含有两层Clos拓扑，被称为AI Zone。

RTSW：Rack Training Switch，机架训练交换机 CTSW：Cluster Training Switch，集群训练交换机 ATSW：Aggregator Training Switch，聚合训练交换机

RTSW作为叶交换机，为机架内的 GPU 提供纵向扩展连接。主干层由模块化的CTSW组成，在集群中的所有机架之间提供横向扩展连接。

AI Zone旨在以非阻塞方式支持大量GPU互连，然而，Llama这种大型模型需要的GPU规模大于单个AI Zone。为了适应这一点，又加入了额外的ATSW层，用于连接数据中心内的CTSW，将RoCE域扩展到单个AI Zone之外。

跨AI Zone的连接在设计上做到了oversubscription，并使用ECMP平衡网络流量。

为了缓解跨AI Zone流量的性能瓶颈，调度程序也进行了改进，能学习到GPU服务器在网络拓扑结构中的位置。将训练节点划分到不同的AI Zone时，调度程序可以找到「最小切割」（minimum cut），从而减少流量。

路由

上面讨论的计算能力和网络拓扑的扩展引发了大量训练流量的路由和负载均衡问题。具体来说，AI训练的工作负载体现出以下3个具有挑战性的特征：

- 低熵：与传统数据中心相比，AI工作负载的流的数量和多样性要小得多，流的模式通常是重复的和可预测的

- 突发性：在时间维度上，流量通常在毫秒级的时间粒度上出现或结束

- 「大象流」：每次流量爆发时，强度可以与NIC的线路速率相当

ECMP和路径固定

关于路由机制，最初考虑的是广泛采用的ECMP机制（Equal-Cost Multi-Path，等成本多路径），根据五元组的哈希值随机路由，其中五元组包括源和目标IP、源和目标UDP端口以及协议。

然而，正如预期的那样，由于低熵特征，ECMP的表现不佳。

另一种方案是在最初部署时设计的「路径固定」（path-pinning），根据目的地「切片」（即RTSW下行链路的索引）将数据包路由到特定路径。

如果每个机架都分配完全相同的作业且网络中没有故障，这种方法就会有不错的效果，但事实并非如此。

如图5所示，机架中的两台主机中只有一台使用上行链路带宽，这种不均衡的碎片化工作安排导致了特定RTSW上行链路的流量分布不均，使训练性能下降超过30%。

此外，上行链路或CTSW上可能发生故障，导致受影响的流通过ECMP被重新分配到其他CTSW，造成流量不均甚至与现有流的冲突，拉低了整体的训练效率。

要采用短期的缓解措施，可以将RTSW上行链路的带宽升级2倍，可以减轻性能影响但过于昂贵。

队列对扩展

接下来，团队对ECMP机制进行了进一步审视，利用NCCL等库中的队列对扩展（Queue Pair Scaling）功能，以实现路由机制的演进。

此外，将交换机配置为E-ECMP (Enhanced ECMP)，使用其UDF功能对RoCE数据包的目标QP字段进行额外哈希处理。

对于QP扩展，论文评估了两种策略。一是将每条消息拆分为通过单个QP发布而非多个QP（split），这会同时降低消息容量并产生多个ACK；另一种方法则是以循环方式将每条消息发布到不同的队列（round-robin）。

如图7所示，在AllReduce基准上，E-ECMP和QP扩展搭配使用时，性能相比基线可提升40%。在两种QP扩展策略中，round-robin的表现更好。

但值得注意的是，哈希操作的潜在概率性质是该方案的本质缺陷。此外，需要根据具体的工作负载定制QP扩展策略和参数，虽然短期内可行，但长期来看会提升操作复杂性。

拥塞控制

将集群过渡到400G网络部署时，团队尝试调整原有的DCQCN算法以适应新的网络速度和拓扑结构，然而却遇到了无法解决的问题，相比200G网络出现了性能下降。

因此，团队选择在没有DCQCN的情况下继续进行400G部署，在一年多的时间中仅使用PFC，没有任何其他传输级的拥塞控制机制。结果发现，训练集群表现稳定，并没有出现持续拥堵的情况。

接收方驱动的流量准入

为了缓解400G及以上网路的拥塞，论文对集群库和RoCE传输进行协同设计，强制执行接收方驱动的流量准入。

图14展示了这种GPU到GPU的通信架构，包含两个阶段的拷贝，以及接收方发起的通过NCCL集群库的通信。

每个GPU的HBM内存维护多个通道，用于并行传输消息块。

首先，发送方的GPU线程将数据从计算缓冲区（compute buffer）拷贝至可用的通道缓冲区（channel buffer）。

对于发送方的CPU代理线程，只有在收到来自接收方的CTS数据包（clear-to-send，包括大小和内存信息）后，才能发出RDMA写入请求。

之后，接收方的GPU线程将通道缓冲区中的内容复制到目标的计算缓冲区。

最后，双方的CPU代理线程回收通道缓冲区。通道缓冲区回收完毕并准备就绪后，接收方的CPU代理就会发送另一个CTS数据包，开启下一轮通信。

拥塞开始时，这种机制可以有效限制网络中的in-flight流量。然而，要实现正确的参数配置可能具有挑战性，比如：

- 由于大量的并发计算争夺GPU线程资源，通道数量受限

- 由于RoCE更粗粒度的流量控制，以及终端主机可能比较迟缓，设置通道缓冲区大小需要比Infiniband更仔细，在拥塞扩散和带宽利用率不足之间进行权衡。

拥塞控制一直是RDMA网络研究的焦点，其中DCQCN一直是以存储为中心的网络的黄金标准。然而，Meta团队所述的经验为定制拥塞控制算法提供了不同的视角。

总结

除了介绍工程设计方面的考量，论文也描述了使用的集群观测工具，以及一些故障排除案例。

通过分离FE和BE网络、采用不同的路由方案并优化集群流量模式，Meta团队构建了高性能且可靠的网络基础设施，强调了深入理解训练工作负载的重要性，并对相关的网络组件进行了「量身定制」。

参考资料：

https://engineering.fb.com/2024/08/05/data-center-engineering/roce-network-distributed-ai-training-at-scale/

https://web.stanford.edu/class/ee384y/Handouts/clos_networks.pdf

https://community.fs.com/article/roce-vs-infiniband-vs-tcp-ip.html