RailX: A Flexible, Scalable, and Low-Cost Network Architecture for Hyper-Scale LLM Training Systems

日益增长的大型人工智能工作负载对超大规模基础设施提出了迫切需求；然而，传统的互连网络架构在可扩展性和成本效益方面都存在明显不足。\textit{铁路优化}（Rail-optimized）等基于树形结构的网络代价极其高昂，而像\textit{Torus}这样的直接互连拓扑又缺乏足够的二分带宽和灵活性。本文中，我们提出了\textit{RailX}，一种基于节点内直接连接和节点间电路交换的可重构网络架构。该架构将节点和光交换机以二维物理结构组织，实现了比现有集中式电路交换网络更佳的可扩展性。我们提出了一种基于\textit{哈密尔顿分解}（Hamiltonian Decomposition）理论的新型互连方法，将多个基于铁路的独立环形结构组织成\textit{全连接}（all-to-all）拓扑，同时优化了环形集体通信和全连接通信。通过一个扁平的交换层，\textit{RailX}可以互连超过10万片芯片并实现高达TB级别的带宽，而网络直径仅为2到4跳。在每单位注入/All-Reduce带宽的成本方面，\textit{RailX}不到Fat-Tree的10%；在每单位二分/全连接带宽的成本方面，也仅为Fat-Tree的一半以下。具体而言，只需约13亿美元即可将20万片芯片以1.8TB带宽互连。\textit{RailX}同样适用于机器学习即服务（MLaaS）场景，能够灵活映射不同形状、规模和并行策略的单个或多个训练任务，并具备绕过故障的能力。

论文试图解决传统互连网络架构在面对日益增长的大规模AI工作负载时，扩展性不足和成本过高的问题。随着AI模型规模的不断增长，现有的拓扑结构如Fat-Tree、Torus等无法有效满足高带宽和灵活互连的需求。这是一个新兴且重要的问题，因为超大规模AI训练对基础设施提出了前所未有的挑战。

论文提出了一种名为RailX的可重构网络架构，结合了节点内直接连接和节点间电路交换的优势。其关键思路是基于Hamiltonian分解理论，将多个基于Rail的环形结构组织成全连接（all-to-all）拓扑，从而在保证高带宽的同时提升网络的扩展性和灵活性。相比传统树状或直接拓扑结构，RailX在成本和性能之间取得了更好的平衡。

{支持超过100K芯片的互连，仅需一个扁平交换层，网络直径为2~4跳，显著降低延迟。,相比Fat-Tree，RailX的每单位注入/All-Reduce带宽成本低于其10%，每单位二分/全连接带宽成本低于其50%。,实验证明，仅需约13亿美元即可互联20万芯片并提供1.8TB带宽，展示了其出色的性价比。,支持ML-as-a-Service场景，能够灵活映射不同形状、规模和并行策略的训练任务，并具备容错能力。,该架构具备可扩展性和灵活性，为未来超大规模AI训练基础设施设计提供了新方向。}

{"Jiang et al., \"Slim Fly: A Cost Effective Low-Diameter Network Topology,\" in IEEE/ACM Supercomputing, 2017.","Al-Fares et al., \"A Scalable, Commodity Data Center Network Architecture,\" in ACM SIGCOMM, 2008.","Li et al., \"Bandana: A Disaggregated Memory-aware Network Architecture for Accelerating Distributed DNN Training,\" in MLSys, 2023.","Nan et al., \"Circuit Switching for Scalable Data Center Networks,\" in IEEE INFOCOM, 2021.","Zhang et al., \"Optical Switching for High-Performance AI Clusters,\" in IEEE MICRO, 2022."}