Optimizing All-to-All Collective Communication with Fault Tolerance on Torus Networks

作者/机构: Le Qin, Junwei Cui, Weilin Cai (The Hong Kong University of Science and Technology (Guangzhou)), Meng Niu, Yan Yang (Huawei Technologies Co., Ltd.), Jiayi Huang (The Hong Kong University of Science and Technology (Guangzhou))

A1 主要贡献

核心问题: 大规模分布式处理被广泛用于大型模型的推理和训练，如深度学习推荐模型（DLRM）和专家混合（MoE）模型。然而，All-to-All集体通信因其复杂的点对点通信模式和阻塞性，已成为分布式DLRM和MoE加速的主要性能瓶颈。此外，长时间的分布式处理经常遇到链路故障，这严重影响了系统的效率、可靠性和成本。与支持任意连接（如Clos网络）的交换式拓扑不同，环面（torus）网络上的All-to-All通信会因共享路由路径而相互干扰，造成严重的性能限制。

研究目标: 本文旨在解决上述挑战，提出了用于环面网络上具有容错功能的All-to-All优化的单维算法和多维调度策略。

创新点:
本文的贡献如下：
1. 针对无故障网络:
* 提出 HalfRing 算法：这是一种在带宽和延迟上均为最优的单环All-to-All算法。它利用双向链路构建最短通信路径，通过逐跳存储转发方式实现数据传输，并确保了带宽的充分利用。
* 提出 DimRotation 调度策略：该策略通过为不同数据块在多个维度上分配不同的通信序列来平衡流量，从而实现无通信气泡（bubble-free）和完全重叠的多维通信，极大地提升了All-to-All的效率。

  1. 针对有故障网络:
    • 提出 FoldedRing 算法：这是一种处理环内单链路故障的容错算法，它通过利用剩余链路构建一个逻辑上的“折叠环”来维持通信。
    • 提出 MATE 调度策略：该策略利用其他维度的链路带宽来加速故障环上的通信，确保即使在链路故障下也能实现可靠且高效的All-to-All通信。

核心成果:

  • 实验结果表明，与基准的环形（Ring）算法加流水线（pipeline）调度相比，HalfRing、DimRotation及其组合在无故障情况下分别取得了1.56倍、1.45倍和2.28倍的平均性能加速。
  • 对于存在单链路故障的All-to-All，MATE调度相比无故障基准可以实现1.37倍的加速。
  • 与TPUv4集群中现有的先进路由方法相比，本文提出的方法在无故障和容错场景下分别实现了1.57倍和1.61倍的加速。

(a) DLRM中嵌入层的All-to-All通信，嵌入表分散在2个设备上。训练在表查找时采用模型并行，在拼接操作后采用数据并行。

(b) 具有4个专家分散在2个设备上的top-2 MoE模型中的专家层。训练在门控前和聚合后采用数据并行，中间的专家计算采用专家并行。

(c) DLRM和Mixtral模型在8×8 TPUv3和8×8×8 TPUv4集群上训练(Train)和推理(Inf)的归一化时间分解。All-to-All部分占比从41.5%到95.7%不等。

图 1: 分布式DLRM (a) 和 MoE 模型 (b) 的All-to-All集体通信，以及模型训练和推理的归一化时间分解 (c)。评估方法在第5.1节中描述。

A3 背景知识

2.1 All-to-All实现概述

两种All-to-All实现方式。All-to-All的实现可分为两类：使用硬件路由的粗粒度编排和通过算法设计的细粒度编排。前者将All-to-All分解为每个节点的独立点对点通信（即多跳传输），编排它们的传输顺序，并依赖硬件路由执行每次多跳传输。这种方法在基于MPI的HPC集群中广泛使用，虽然简单且适用性广，但常常导致网络拥塞【27, Optimizing the bruck algorithm for non-uniform all-to-all communication, 2022, HPDC; 40, Adaptive and hierarchical large message all-to-all communication algorithms for large-scale dense gpu systems, 2021, CCGrid; 43, Optimization of all-to-all communication on the blue gene/l supercomputer, 2008, ICPP; 44, Optimal algorithms for all-to-all personalized communication on rings and two dimensional tori, 1997, JPDC; 54, Optimizing All-to-All Collective Communication on Tianhe Supercomputer, 2022, ISPA/BDCloud/SocialCom/SustainCom; 56, Bandwidth-optimal all-to-all exchanges in fat tree networks, 2013, ICS; 66, All-to-all personalized communication in multidimensional torus and mesh networks, 2001, TPDS; 70, Automatically tuned collective communications, 2000, SC; 78, Resiliency at Scale: Managing {Google’s} {TPUv4} Machine Learning Supercomputer, 2024, NSDI】。

本文采用的细粒度方法。后者将每个点对点通信进一步分解为多个单跳传输并调度其顺序。通过存储转发机制，中间节点在关联的端点临时存储数据，然后转发给下一个节点。这消除了多跳路由硬件复杂的拥塞控制开销，并通过每一步的精确链路分配避免了网络竞争。这种细粒度调度在机器学习系统的集体通信中被广泛研究，能有效防止网络拥塞【16, Synthesizing optimal collective algorithms, 2021, PPoPP; 32, Communication algorithm-architecture co-design for distributed deep learning, 2021, ISCA; 45, Enhancing Collective Communication in MCM Accelerators for Deep Learning Training, 2024, HPCA; 61, Themis: A network bandwidth-aware collective scheduling policy for distributed training of dl models, 2022, ISCA; 64, {TACCL}: Guiding Collective Algorithm Synthesis using Communication Sketches, 2023, NSDI; 75, TACOS: Topology-aware collective algorithm synthesizer for distributed training, 2023, arXiv】。

多维环面网络上的优化策略。本文工作采用后一种方法来优化分布式机器学习计算中的All-to-All。对于多维环面网络，如TPUv4【38, Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings, 2023, ISCA】集群中使用的3D环面，我们将All-to-All分解为跨每个维度的顺序阶段。我们的优化集中在两个关键方面：维度内算法和跨维度调度。第2.2和2.3节分别介绍了算法和调度的基线设计。

2.2 单维算法

基线Ring算法。在All-to-All中，每个进程向所有其他进程发送唯一的数据【68, Optimization of collective communication operations in MPICH, 2005, IJHPCA】。用于All-to-All通信的常见算法包括Ring【34, Optimal bucket algorithms for large MPI collectives on torus interconnects, 2010, ICS】、Direct【68, Optimization of collective communication operations in MPICH, 2005, IJHPCA】、Halving-Doubling【25, Eflops: Algorithm and system co-design for a high performance distributed training platform, 2020, HPCA】和Bruck【13, Efficient algorithms for all-to-all communications in multi-port message-passing systems, 1994, SPAA】。这些算法的性能因网络拓扑而异。Ring算法常用于网格和环面等直接拓扑，因为它提供了良好的可扩展性和零竞争【61, Themis: A network bandwidth-aware collective scheduling policy for distributed training of dl models, 2022, ISCA】。

Ring算法的执行过程。图2展示了在四节点环形网络中使用Ring算法的All-to-All过程，可分为三个阶段。如图2a所示，每个节点的数据被分为八个部分，每个方向传输四个部分。在All-to-All中，每个节点的目标是接收所有索引与其自身节点索引匹配的数据部分。这些索引在图2中用颜色表示。为实现此目标，通信被组织为三个All-to-All阶段，每个阶段具有不同的跳数距离。如图2a-f所示，阶段1-3分别对应1到3的跳数距离。采用存储转发方法，多跳传输通过多个子阶段完成以避免竞争。图2d-2f展示了阶段3的详细过程，其中3跳传输涉及三个转发子阶段以到达其目的地。例如，在阶段3-1，节点1必须通过节点2和3转发其紫色数据块以到达节点4；类似地，其蓝色数据块在相反方向被转发。同时，其他节点也并发执行类似的转发操作。

图 2: 在环形（1-D环面）拓扑中使用Ring算法的All-to-All。在具有双向带宽的网络中，Ring算法将数据在两个方向上平均分配以充分利用带宽。图中，右侧块表示顺时针传输的数据，而左侧块（带白点）表示逆时针传输的数据。四种颜色表示每个节点接收的最终All-to-All数据，而颜色深浅表示每个节点最初持有的数据。每条链路的颜色反映了该时刻正在传输的数据块。Ring算法采用存储转发方法进行多跳传输以避免拥塞，在每个子阶段标记为Fwd。

2.3 多维调度

多维网络上的All-to-All分解。多维网络拓扑在大型机器学习系统中普遍使用【2, An overview of the BlueGene/L supercomputer, 2002, SC; 5, The tofu interconnect d, 2018, CLUSTER; 6, Tofu: A 6D mesh/torus interconnect for exascale computers, 2009, Computer; 11, The Blue Waters super-system for super-science, 2017, Contemporary high performance computing; 38, Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings, 2023, ISCA; 39, A domain-specific supercomputer for training deep neural networks, 2020, CACM; 78, Resiliency at Scale: Managing {Google’s} {TPUv4} Machine Learning Supercomputer, 2024, NSDI】，这使得在多维上高效调度通信成为一个挑战。在N维网络上的All-to-All可以分解为N个顺序的单维All-to-All阶段。每个阶段的通信使用第2.2节介绍的算法实现。通过将数据划分为多个数据块进行流水线调度【61, Themis: A network bandwidth-aware collective scheduling policy for distributed training of dl models, 2022, ISCA】，可以提高整体网络带宽利用率。

基线X-Y调度。如图3所示，2D环面网络上的All-to-All涉及两个阶段，按顺序跨X和Y维度传输。以图3a中的节点1为例，其数据被分为九个部分，分别发送给节点1-9。图3b显示了X维度通信阶段的结束状态，其中每个部分被发送到其目的地对应的列（例如，节点1从节点1-3接收目标为第1列的数据）。每个阶段的通信可以使用单维算法（如Ring）来执行。如图3c所示，在到达相应列后，每个数据部分通过Y维度阶段被发送到其最终目的地。最后，如图3d所示，节点1在阶段2结束时从所有九个节点收集到第一部分数据，从而完成All-to-All通信。

图 3: 在2D环面网络上使用X-Y调度的All-to-All。(b)-(c)显示了每个阶段的结束状态，其中文本表示数据来源。

固定基数网络中的数据量恒定。对于各维度节点数相同的固定基数网络，每个阶段的通信成本因在各维度传输的数据量固定而保持不变。例如，在阶段1，节点1发送三个需要到达第二和第三列的数据部分给节点2和3。类似地，节点1也从节点2和3接收三个目标为第1列的数据部分。因此，在阶段1结束时，节点1仍然持有九个数据部分。在阶段2，节点1传输源自节点1-3的数据，这些数据需要发送到节点4和7，每次传输仍然包含三个数据部分。

多维调度优化的重要性。优化多维网络中的通信调度对于随着网络规模扩大而最大化带宽利用率至关重要。虽然先前的工作【32, Communication algorithm-architecture co-design for distributed deep learning, 2021, ISCA; 61, Themis: A network bandwidth-aware collective scheduling policy for distributed training of dl models, 2022, ISCA; 76, Image classification at supercomputer scale, 2018, arXiv】已经探索了此类网络中高效的All-Reduce调度，但跨多网络维度的All-to-All调度仍未得到充分探索。

2.4 容错集体通信

网络中的常见故障类型。互连网络中的常见故障包括节点和链路故障【23, Principles and practices of interconnection networks, 2004, Elsevier; 26, Interconnection networks, 2003, Morgan Kaufmann】。通常，机器学习系统通过用健康组件替换故障组件来恢复训练【15, MoC-System: Efficient Fault Tolerance for Sparse Mixture-of-Experts Model Training, 2025, ASPLOS; 49, Understanding and improving failure tolerant training for deep learning recommendation with partial recovery, 2021, MLSys; 73, Gemini: Fast failure recovery in distributed training with inmemory checkpoints, 2023, SOSP】。然而，这会产生巨大的数据备份和检查点成本，并中断训练过程。容错路由通过创建新的通信路径绕过故障区域，提供了一种更高效的解决方案，实现了无需更换设备或中断训练的容错【26, Interconnection networks, 2003, Morgan Kaufmann; 78, Resiliency at Scale: Managing {Google’s} {TPUv4} Machine Learning Supercomputer, 2024, NSDI】。图4a展示了在节点和链路故障期间的自适应路由。端点的故障可能导致输出端口的路由功能丧失，从而引起与这些连接相关的链路故障【26, Interconnection networks, 2003, Morgan Kaufmann】。

OCS故障导致的链路故障。此外，光路交换机（OCS）的故障可能导致所有连接的链路失效，产生周期性的链路故障模式【78, Resiliency at Scale: Managing {Google’s} {TPUv4} Machine Learning Supercomputer, 2024, NSDI】。如图4b所示，4×4×4的TPUv4 pod需要48个OCS用于可重构互连，每个pod在X-Y、Y-Z和X-Z平面上的表面节点连接到16个不同的OCS。例如，负责X-Z平面互连的OCS 34故障会导致四个相应的链路失效，包括两个环绕链路：(2,3,0)-(2,3,3)和(6,3,0)-(6,3,3)。图4中的两种故障类型都导致了链路故障，凸显了为解决链路故障而进行容错优化的迫切需求。

针对集体通信的容错设计。传统的容错路由是为通用的点对点通信设计的，而不是为模型训练中常见的集体通信设计的。集体通信涉及跨多个设备的数据传输，具有固定的模式和沉重的开销。虽然简单的容错路由可以维持单个节点的数据传输，但路由的改变会影响链路分配，从而降低其他节点的效率。为集体通信量身定制的容错方案【42, Highly Available Data Parallel ML training on Mesh Networks, 2020, arXiv】可以最小化对通信效率的影响，从而为更可靠和高性能的模型训练提供基础。

(a) 节点和链路故障可能影响原始数据传输，需要新的路由来绕过故障。 (b) 在一个由两个TPUv4 pod组成的8×4×4环面网络中，一个OCS故障导致的两条链路故障。
图 4: 网络中常见的故障类型。

2.5 故障模型与假设

故障检测与容忍机制。我们遵循谷歌TPUv4上的故障检测和容忍机制【78, Resiliency at Scale: Managing {Google’s} {TPUv4} Machine Learning Supercomputer, 2024, NSDI】。机器故障在任务分配前通过飞行前检查被识别和解决。在任务执行期间，一个名为healthd的软件守护进程通过持续检查所有TPU间的链路连接性来监控硬件健康状况。如果发生链路或OCS故障，healthd会通知集群管理服务Borg，触发从上一个检查点重启作业。受影响的TPU会将其芯片间互连（ICI）路由表从无故障模式切换到容错模式。从这时起，包括All-to-All在内的所有通信都在网络故障的情况下继续进行。

本文的故障假设。我们优化了无故障和有故障环面网络上的All-to-All。在有故障的情况下，我们假设在随机位置发生单链路故障。由于环面网络的几何对称性，故障位置对我们的方法呈现一致的视角。我们的方法也可以推广到更复杂的故障场景，例如OCS故障和多链路故障，这在第4.2、5.4和5.7节中讨论。

A2 方法细节

3 无故障All-to-All

本节中，我们从两个方面优化无故障All-to-All。第3.1节介绍用于单维环面（环）中All-to-All的HalfRing算法；而第3.2节涵盖跨多维度的通信调度DimRotation。

3.1 单维HalfRing算法

Ring算法的次优性。在环形拓扑上的Ring算法中，所有节点在每个阶段同时以相等的距离传输数据。该算法始终能实现全链路带宽利用。然而，由于不必要的长传输路径导致额外的链路带宽消耗，其性能并非最优。如图2d-2f所示，按照顺时针顺序，从节点1到节点4的数据传输需要3个子阶段，尽管它们之间仅相隔一跳。在典型的双向网络中，Ring算法独立使用顺时针和逆时针链路，并同时在两个方向上通信。在阶段3中，顺时针传输需要三跳，而逆时针仅需一跳。这种不对称性导致某些阶段的性能次优，因为非最短路径的转发消耗了额外的链路带宽。

HalfRing算法设计。基于这一观察，我们提出了HalfRing算法，该算法利用双向链路进行最短通信距离的路径分配。在每个阶段，HalfRing根据发送方和接收方之间的实际距离确定传输方向。由于所有节点都通过最短路径通信，每个阶段只在一个方向上消耗带宽，从而使另一方向的带宽可用于具有相同通信跳数的另一个阶段。以图5a中的节点1为例，HalfRing分别选择逆时针和顺时针方向传输两个紫色和两个绿色数据块（均为一跳）。由于没有链路冲突，两个方向的通信同时进行。在一个有奇数个节点（$N=2k+1$）的环中，有$2k$个阶段，All-to-All可以在k对中完成。对于一个有偶数个节点（$N=2k$）的环，有$2k-1$个阶段，导致一个不成对的阶段，如图5b中的阶段2。HalfRing将不成对阶段的数据平均分配，并双向发送，从而充分利用带宽。

HalfRing的性能与正确性。通过在每个多跳阶段全面利用双向带宽，HalfRing确保了全带宽利用和最短通信路径，在环形拓扑上实现了带宽和延迟方面的最优性能。算法的性能分析如表1所示，算法在算法1的HalfRing_Generator中有详细说明。以表1中的Ring算法为例，共有$N(N-1)$次单跳数据传输（即子阶段），每次数据大小为$S/N$。给定双向带宽$2B$，传输时间的计算方法是传输次数乘以每次传输的数据大小再除以带宽。与Ring算法类似，HalfRing将所有数据传输分解为相邻节点间的单跳传输，并明确编排每一跳，从而消除了死锁（无多跳传输）、活锁（无绕路）和网络竞争（无链路共享）的可能性。

图 5: 使用HalfRing算法的All-to-All。阶段1传输的数据量是Ring算法阶段1的两倍。在这种情况下，HalfRing将All-to-All时间缩短为四个子阶段（图5a-5c），而Ring算法需要六个子阶段（图2a-2f）。

表 1: 不同All-to-All算法在环（1-D环面）上基于线性成本模型【68, Optimization of collective communication operations in MPICH, 2005, IJHPCA】的性能分析。总通信时间由两部分组成：启动时间和数据传输时间。启动时间是指由算法和系统延迟引起的固定延迟，与消息大小无关。传输时间反映了数据传输的持续时间。在深度学习训练的All-to-All集体通信中，传输时间占主导地位。参数定义：S：每个节点的数据大小，N：环中节点数，B：单向带宽，α：每跳传播延迟。Ratio表示相对于基准Ring算法的性能。

† FoldedRing的性能是在具有单个随机链路故障的网络中分析的，而其他算法则在健康网络中分析。

3.2 多维DimRotation调度

DimRotation调度机制。N维拓扑上的All-to-All可以分为N个阶段，按顺序跨N个维度进行，每个维度内的通信独立采用其All-to-All算法。尽管存在这种顺序依赖性，数据可以被划分为多个数据块，每个数据块采用不同的通信维度顺序。DimRotation表示每个数据块的通信维度顺序，如图6中的示例所示。在2D环面中，数据被分为两个数据块，数据块1遵循X-Y顺序，数据块2遵循Y-X顺序，从而确保了全带宽利用。对于数据块1，目标为特定列y的数据在阶段1通过HalfRing发送到该列y与X维度的交点节点。然后，在阶段2，数据在Y维度上使用HalfRing中继到目的地。

图 6: 在2D环面网络上使用DimRotation的一个例子。

DimRotation与流水线调度的对比。我们通过随时间变化的比较进一步说明DimRotation的优势。3D环面上的All-to-All需要跨X、Y和Z三个维度进行三个阶段，每个阶段都充分利用相应维度中所有环的带宽。每个维度的环在每个阶段同时传输数据。流水线调度将数据分成多个数据块【61, Themis: A network bandwidth-aware collective scheduling policy for distributed training of dl models, 2022, ISCA】，然后按相同的X-Y-Z维度顺序依次通信。如图7a所示，流水线通过在不同维度上同时运行多个数据块来增强带宽利用率。然而，流水线调度不可避免地会引入“气泡”，阻碍网络带宽的充分利用。此外，在流水线调度中选择合适的块大小也具有挑战性【16, Synthesizing optimal collective algorithms, 2021, PPoPP】。块太大，流水线无法充分重叠不同块在各维度上的时间，导致性能不佳。相反，块太小，大量的数据块会引入显著的调度成本和通信初始化开销，增加总体延迟。

图 7: 在3D网络上比较3个数据块的流水线和DimRotation调度。X-1表示数据块1的X维度通信。

DimRotation的优势与适用性。为解决这些挑战，我们提出了DimRotation调度，以确保无气泡且完全重叠的多维All-to-All。对于N维环面，数据被平均分为N个数据块，第i个数据块按维度i, i+1, ...的顺序通信。图7b展示了3D环面上的DimRotation时间线，其中DimRotation允许三个数据块执行无冲突、全覆盖的多维通信，实现了完全的带宽利用。同时，数据块的数量被设置为完全重叠通信所需的最小值，显著减少了调度开销。对于具有异构带宽或混合基数（维度间节点数不同）特征的网络，DimRotation也能提供最优的调度性能。当异构性导致各维度通信开销不同时，总时间消耗总是受限于性能最差的维度。与All-Reduce【61, Themis: A network bandwidth-aware collective scheduling policy for distributed training of dl models, 2022, ISCA】不同，改变执行顺序或块大小无法优化性能，因为每个维度的数据量是固定的。DimRotation确保总All-to-All时间不超过在性能最差维度上传输完整数据的通信时间。根据Pod分区，All-to-All拓扑在某些维度可能缺少环绕链路，导致类似于混合基数环面的带宽异构性，可以类似地处理。该算法在算法1的Scheduler中给出。

4 容错All-to-All

本节中，我们从两个方面介绍我们的容错All-to-All方案。第4.1节介绍容错的FoldedRing算法。第4.2节涵盖MATE和MATEe（MATE增强版），它们加速了故障环上的通信。

4.1 单维FoldedRing算法

单链路故障下Ring算法的失效。Ring算法要求每两个节点之间有直接链路进行数据传输。如果一条链路发生故障，两端的节点无法传输数据，导致Ring算法失败。

FoldedRing算法设计。对于单链路故障，我们发现故障链路的端点节点并未被隔离，它们仍然可以通过绕行建立连接。基于此，我们提出了FoldedRing算法，该算法使用所有可用的逆时针链路为故障链路构建一个补偿连接。连同剩余的顺时针链路，一个“折叠环”被创建出来，以Ring的方式继续通信。

FoldedRing的执行过程。如图8b所示，当节点1和4之间的链路发生故障时，从节点4到节点1的逻辑连接可以通过三个蓝色的物理逆时针链路构建。通过利用与蓝色链路方向相反的其他顺时针链路，这四个节点的Ring算法得以恢复。FoldedRing算法利用所有可用的链路带宽进行通信。值得注意的是，FoldedRing可以进一步适应其他集体通信中的容错环算法，如All-Reduce、Reduce-Scatter和All-Gather。FoldedRing算法的性能分析见表1。

(a) 无故障环上的Ring算法 (b) 故障环上的FoldedRing算法
图 8: 比较图2a中All-to-All阶段1的Ring和FoldedRing算法。当节点1和4之间的链路发生故障时，所有逆时针链路带宽被重新用于补偿故障链路。通信以折叠环的方式继续进行。FoldedRing完成阶段1所需的时间是Ring的两倍。

4.2 多维MATE调度

FoldedRing面临的挑战。虽然FoldedRing能够在故障环上实现容错的All-to-All，但它仍面临几个挑战。首先，如表1所示，FoldedRing的性能仅为基准Ring算法的一半，远低于我们提出的HalfRing算法。此外，FoldedRing需要更长的启动时间来建立故障链路端点之间的连接，进一步扩大了性能差距。如图10a所示，故障环上的慢速传输导致DimRotation调度出现不匹配。X维度的传输因故障环而减慢，这反过来又影响了其他维度的传输，导致整体All-to-All性能下降。其次，FoldedRing算法只能处理环内的单链路故障。对于两个或更多故障，FoldedRing无法在故障环上建立连接。因此，尽管故障链路两端点之间的通信仍可通过其他方向的路由继续，但All-to-All仍然被迫中断。

MATE调度机制。这些挑战催生了MATE，一种针对容错All-to-All的调度优化。MATE利用其他维度的链路来加速故障环上的数据传输，从而提供更高效、更鲁棒的All-to-All容错解决方案。

MATE加速阶段详解。图9展示了处理节点(0,1)和(1,1)之间链路故障的加速阶段。最初，部分数据通过故障环使用FoldedRing传输。随后，MATE利用其他X维度环的链路为故障环上的每个节点构建双向连接，从而能够使用HalfRing高效地通信剩余数据。以节点(0,1)和(1,1)之间的传输为例，(0,1)-(0,2)-(1,2)-(1,1)中的三条红色链路用于连接(0,1)和(1,1)，下方的三条红色链路建立从(1,1)到(0,1)的连接。这六条红色链路构成了端点之间的双向连接。类似地，蓝色和绿色链路在故障环内其他相邻节点之间建立双向链路以供HalfRing使用。因此，通过将数据分发到这些额外的X维度链路上使用HalfRing同时传输，并在故障环内使用FoldedRing，All-to-All得到了增强。

(a) DimRotation调度，故障环中使用FoldedRing算法，这减慢了整个X维度的通信。

(b) MATE调度：故障环中的通信中断，X维度其他环中的通信保持。

(c) MATEe调度：故障环利用FoldedRing在X维度维持分配数据量的通信，减少了加速阶段的数据量$M_e$。

图 10: X维度存在链路故障时的调度。故障的X维度环使用FoldedRing，其他X维度环使用HalfRing。由于加速需要其他维度的链路，它作为一个单独的加速阶段执行，表示为$M$或$M_e$。

图 9: MATE在2D环面网络上加速阶段的链路利用情况。在(0,1)和(1,1)之间发生链路故障时，MATE利用其他维度的链路来加速故障环上的慢速通信。在2D环面网络上，除故障环内使用FoldedRing的原始通信外，可以利用一组额外的双向链路来加速每次数据传输。

MATE的可靠性与通用性。MATE能够在N维环面网络上可靠地加速容错All-to-All，主要原因有二。首先，包含故障环的每个平面都可以为该环构建无冲突的双向连接。本质上，故障环中的数据传输被卸载到同一维度内的其他环。我们只需确保这些无故障的环被识别，并且数据可以从故障环传输到它们。例如，在图9中，这是通过Y维度的链路将故障环上的节点与两个无故障环的节点连接起来实现的。这两个环是(0,2)-(1,2)-(2,2)和(0,0)-(1,0)-(2,0)。其次，每个加速平面之间的链路应该是无冲突的。这一特性由环面拓扑的正交性天然保证。因此，理论上我们可以在N维环面网络中利用$N-1$个平面来加速单链路故障的通信。MATE充分利用故障环上每个点的带宽来提供加速。

MATEe增强调度。MATEe还在正常阶段将一部分数据分配给故障环，使得部分传输在加速阶段之前完成，从而提高了带宽利用率。

MATE处理复杂故障的能力。MATE可以应用于更复杂的故障场景以增强网络弹性。首先，当单个环内发生多个故障时，FoldedRing无法维持容错传输。MATE通过加速阶段，利用其他环的链路重新路由通信来解决这个问题。其次，MATE也可以处理跨多个环的故障。当每个故障的加速阶段所需链路冲突或故障发生在不同维度时，MATE为每个故障分配单独的加速阶段。如果同一维度中的故障不共享相同的链路，多个故障环可以同时执行加速阶段。以图4b中的OCS故障为例，MATE可以直接应用，时间开销与单链路故障时保持一致。此外，MATE也可以应用于其他集体通信中的多维调度。

A4 实验环境

我们使用ASTRA-SIM模拟器【60, Astra-sim: Enabling sw/hw co-design exploration for distributed dl training platforms, 2020, ISPASS】来实现我们提出的优化。我们同时使用了ASTRA-SIM内置的分析网络后端（用于无竞争场景）和集成的GARNET模拟器【3, GARNET: A detailed on-chip network model inside a full-system simulator, 2009, ISPASS】（用于周期精确的竞争场景模拟）。

数据集与模型:
* 工作负载: 评估了四种代表性的DLRM和MoE模型：DLRM【51, Deep learning recommendation model for personalization and recommendation systems, 2019, arXiv】、Wide & Deep【18, Wide & deep learning for recommender systems, 2016, RecSys】、Deepspeed-1.3B+MoE-128【59, Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation ai scale, 2022, ICML】和Mixtral 7Bx8【4, Mixtral of experts, n.d., http://mistral.ai】。
* 非均匀All-to-All: 使用AI2 Reasoning Challenge (ARC)数据集【19, Think you have solved question answering? try arc, the ai2 reasoning challenge, 2018, arXiv】训练Mixtral 7B×8模型，提取专家选择分布以评估非均匀通信。

硬件配置:
* 合成与可扩展性实验: 2D/3D/4D环面拓扑，每链路带宽32 GB/s，网络延迟100 ns。
* 与Google TPUv4路由比较: 4×4×4和8×4×4环面拓扑（模拟单/双TPUv4 pod），每链路带宽56 GB/s，包大小512字节，flit宽度256比特。
* 真实工作负载实验: 模拟TPUv3（8×8，82 GB/s带宽）和TPUv4（8×8×8，56 GB/s带宽）的网络配置。
* 真实机器实验: 使用两台NPU节点，每台配备8个设备，共16个NPU。通过200Gb/NPU的全连接RoCE ToR交换机进行节点间通信，并配置为模拟4×4环面拓扑。

软件配置:
* 模拟器: ASTRA-SIM，集成了分析后端和GARNET后端。
* 基线方法:
* 无竞争：Ring算法 + 流水线调度。
* 竞争：XYZ维度顺序路由（DOR）和wild-first路由（WFR），基于【78, Resiliency at Scale: Managing {Google’s} {TPUv4} Machine Learning Supercomputer, 2024, NSDI】。

  • 真实机器实现: 使用PyTorch Distributed模块【8, Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation, 2024, ASPLOS】实现。

表 2: 系统配置

A4 实验结果

5.2 合成实验

  • 性能加速比 (图11): 在2D、3D和4D环面网络上，与基线（Ring+Pipeline）相比，无故障场景下，HalfRing算法、DimRotation调度以及两者结合（HalfRing+DimRotation）分别实现了1.56倍、1.45倍和2.28倍的平均性能加速。在单链路故障的容错场景下，MATE和MATEe甚至超越了无故障基线，分别实现了1.36倍和1.37倍的平均加速，而FoldedRing+Pipeline的性能则下降至基线的0.55倍。
  • All-to-All带宽 (图12): HalfRing算法因其最短路径特性实现了显著更高的带宽。MATE和MATEe的带宽随着网络维度的增加而提高，因为有更多链路可用于加速阶段。
  • 维度利用率 (图13): 在5×5×5环面网络上，DimRotation消除了流水线调度中出现的通信“气泡”，实现了完美的链路利用。MATEe通过在正常阶段分配部分通信任务给故障维度，改善了MATE中故障维度链路利用率低的问题。

图 11: 在2D、3D和4D环面网络上，不同数据大小下的All-to-All性能加速比。Ring_Pipeline、Ring_DimRotation、HalfRing_Pipeline和HalfRing_DimRotation用于无故障网络，而FoldedRing_Pipeline、FoldedRing_DimRotation、MATE和MATEe用于第一维度存在一个链路故障的网络。

图 12: 在2D、3D和4D环面网络上，不同数据大小下的All-to-All带宽，其中FoldedRing、MATE和MATEe用于第一维度存在一个链路故障的网络。我们将All-to-All带宽定义为：每个节点的通信大小除以All-to-All时间，这与先前工作[25, 32, 75]一致。

图 13: 在5×5×5环面网络上，8MB All-to-All通信下不同算法和调度设置的维度利用率。

5.3 可扩展性研究

  • 节点数量扩展 (图14左): HalfRing和FoldedRing算法在每个维度节点数增加时表现出良好的线性可扩展性。MATE和MATEe的性能随着维度增加而提高，显示了良好的跨维度可扩展性。
  • 网络维度扩展 (图14右): DimRotation的性能几乎不受网络维度增加的影响。MATE和MATEe的通信时间随着维度增加而变得更加稳定，因为加速阶段更快。

图 14: 可扩展性研究：左图显示在2D环面网络上，每个节点4MB All-to-All大小的可扩展性。由于MATE和MATEe在不同维度下性能不同，图中也展示了它们在2D、3D和4D网络上的结果。右图显示了每个节点4MB大小和每个环4个节点时，随维度增加的可扩展性。

5.4 与TPUv4路由设计的比较

  • 与DOR/WFR比较 (图15): 在模拟的单个TPUv4 pod（4×4×4）和双pod（8×4×4）上，与Google的维度顺序路由（DOR）相比，HalfRing+DimRotation在无故障情况下平均提速1.57倍。在容错场景下，与wild-first路由（WFR）相比，MATE和MATEe在单pod上平均提速1.26倍和1.24倍，饱和时最高提速1.46倍和1.61倍。这是因为DOR/WFR存在流量不均衡和拥塞问题，而我们的方法能有效平衡流量。

(a) 在4×4×4 TPUv4 pod上的All-to-All。(b) 在8×4×4 TPUv4 pod上的All-to-All。WFR-x/y/z表示在维度1/2/3应用WFR处理链路故障。F1/F2指由OCS故障引起的1/2个链路故障，均在维度2。
图 15: 在单个和两个TPUv4 pod上，无故障和容错All-to-All的性能比较。DOR和HalfR+DR（HalfRing + DimRotation）用于无故障场景。WFR、MATE和MATEe用于容错。

5.5 真实模型性能

  • 端到端应用加速 (图16): 在DLRM和MoE模型的训练和推理中，与基线（Ring+Pipeline）相比，HalfRing+DimRotation、FoldedRing+Pipeline、MATE和MATEe分别实现了1.97倍、0.54倍、1.24倍和1.38倍的All-to-All平均加速比，并带来了1.64倍、0.63倍、1.20倍和1.29倍的端到端总时间平均加速比。

图 16: 在典型TPU网络配置下，推荐模型和MoE模型的归一化时间分解。以Ring算法和流水线(PL)调度组合为基线，我们展示了HalfRing(HalfR) + DimRotation(DR)的无故障性能，以及在网络出现单链路故障时FoldedRing(FoldR) + pipeline、MATE和MATEe的性能。

5.6 非均匀All-to-All

  • 非均匀通信性能 (图17): 在模拟MoE模型训练中常见的非均匀All-to-All通信时，无故障情况下，HalfRing+DimRotation相比DOR路由平均提速1.27倍。在单链路故障情况下，MATEe相比Google WFR路由平均提速1.17倍。

图 17: 在模拟的TPUv4 pod上，使用AI2 Reasoning Challenge (ARC)数据集[19]训练Mixtral7B×8[4]的专家选择分布轨迹下的非均匀All-to-All性能。

5.7 对多重故障的弹性

  • 多故障场景性能 (图18): 在三种不同的多链路故障模式下，与Google WFR相比，我们提出的方法分别实现了1.43倍、1.14倍和1.55倍的平均性能加速。

图 18: 在单个TPUv4 Pod上，三种链路故障下的容错All-to-All性能。

5.8 真实机器性能

  • 真实硬件验证 (图19): 在一个由16个NPU模拟的4×4环面网络上，我们的无故障设计（HalfRing+DR）相比基线（Ring+Pipe）最高提速1.84倍，容错设计（MATE）则达到基线性能的0.77倍。运行时分解显示，MATE由于其多路径绕行的复杂性，引入了更多的中断和更高的通信时间。

图 19: 在16-NPU模拟的4×4环面网络上的All-to-All。

A5 结论

本文指出了大规模分布式深度学习计算中日益增长的All-to-All开销问题，并针对无故障和容错通信场景探索了算法和调度优化。对于无故障的环面网络，我们提出了HalfRing来提高单维传输效率，并使用DimRotation来增强整体带宽利用率。对于存在链路故障的环面网络，我们引入了FoldedRing以实现基本的容错功能，并进一步提出了MATE，该方法利用多维链路来加速故障环上的通信。评估结果显示，HalfRing和DimRotation分别实现了1.56倍和1.45倍的平均加速，两者结合时最高可达2.28倍。在单链路故障下，MATE相比无故障基线实现了1.37倍的加速。与Google的路由方法相比，我们的方法在无故障和容错场景下分别实现了1.57倍和1.61倍的加速。

A6 附录

A.1 摘要

工件内容。该工件包含了在环面网络上实现容错All-to-All的代码，以及其设置和运行说明。我们提供了复现论文中主要结果的指令和一键运行脚本，特别是复现图11至图19的结果。

A.2 工件清单（元信息）

  • 程序: Python = 3.7 (分析后端)，Python = 2.7 (GARNET后端)，Astra-SIM
  • 运行环境: Ubuntu = 22.04
  • 实验: 包含了运行模拟和真实机器测试的脚本。
  • 指标: 评估了性能加速比、通信带宽、带宽利用率、性能可扩展性、与Google路由的比较、端到端时间分解、非均匀All-to-All、多故障下的性能以及真实机器通信时间。
  • 输出: 工件的输出是PDF格式的图表，复现了论文的主要结果。
  • 所需磁盘空间 (大约): 约2 TB。
  • 准备工作流所需时间 (大约): 约40分钟。
  • 完成实验所需时间 (大约): 所有使用分析后端的模拟实验大约需要9小时。使用GARNET后端的模拟时间与通信大小成正比，短则几分钟，长则可达6天。所有实验可能需要长达两周时间，建议并行模拟。真实机器实验约需20分钟。
  • 是否公开可用: 是。

A.3 描述

访问方式。该工件已存档在Zenodo，也可从GitHub访问。

硬件依赖。模拟实验在24核Intel Xeon Gold CPU、512 GB DRAM和2 TB磁盘的服务器上进行。真实机器实验使用了16个Ascend 910B4 NPU。

软件依赖。实验在Ubuntu 22.04 LTS上运行。模拟实验需要Python环境。真实机器测试需要CANN 8.2.RC1和torch_npu 2.1.0.post12。完整的依赖说明在README.md文件中。

A.4 安装

安装步骤。为分析后端和GARNET后端分别提供了Anaconda环境创建、环境激活和ASTRA-SIM编译的脚本。

A.5 实验工作流

运行脚本。为分析后端和GARNET后端的模拟实验以及真实机器实验都提供了一键运行的脚本。最后，还提供了生成图表的脚本。

A.6 评估与预期结果

结果位置。结果和图表分别可以在./analytical_backend/examples/results/*、./garnet_backend/examples/results/*和./src/User/Chimera/experiment/Pictures目录下找到。

A.7 备注

附加信息。工件的README.md文件提供了关于代码组织和运行实验详细步骤的附加信息。