SuperMesh: Energy-Efficient Collective Communications for Accelerators

作者/机构: SABUJ LASKAR, PRANATI MAJHI, ABDULLAH A M MUZAHID, EUN-JUNG KIM, Texas A&M University, College Station, TX, United States


A1 主要贡献

核心问题

随着深度神经网络(DNN)模型复杂性和计算需求的不断增长,基于小芯片(chiplet)的加速器已成为一种流行的可扩展解决方案。这些加速器通常采用2D网格(mesh)拓扑结构。然而,现有的集体通信算法在2D网格拓扑中常常因边界节点连接性有限而遇到通信瓶颈,导致性能下降。边界节点的链路数量少于内部节点,但在集体通信中需要处理相同量的数据,这造成了链路利用率不均衡,内部节点常常因等待数据而空闲。

研究目标

本文旨在解决2D网格拓扑中边界节点的通信瓶颈问题,从而提升基于小芯片的加速器在集体通信(如AllReduce, ReduceScatter, AllGather)中的性能和能效。

创新点与主要贡献

本文的核心观察是,无需全局修改网格或引入昂贵的长距离链路,仅通过增强边界节点的局部连接性,即可有效解决通信瓶颈。基于此,本文提出了一种名为SuperMesh(SM)的新型拓扑及其两种变体,并通过协同设计的集体通信算法充分利用其优势。

  1. 提出SuperMesh拓扑:

    • SuperMeshBi (SMBi):在现有网格拓扑的所有外围链路旁,并行添加双向链路。
    • SuperMeshAlter (SMAlter):在外围链路上交替添加双向链路。
    • 这两种拓扑的设计原则是,仅在网格的外围区域增加短链路,以保持原始结构的延迟、可扩展性和能效,同时与现有的小芯片加速器设计兼容。
  2. 协同设计集体通信算法:

    • 为充分利用SuperMesh拓扑的额外链路,本文协同设计了针对AllReduce (AR)、ReduceScatter (RS)和AllGather (AG)的流水线集体通信算法。
    • 这些算法通过在SuperMesh上构建四个不相交的通信树(而传统网格上通常只能构建三个),并利用流水线机制,显著提升了通信效率和带宽利用率。
  3. 显著的性能和能效提升:

    • 与传统的2D网格拓扑相比,本文提出的算法和拓扑在AllReduce上实现了平均1.18-1.33倍的加速,在ReduceScatter和AllGather上实现了1.77-2.22倍的加速。
    • 通过缩短通信时间,SuperMesh在完成相同通信任务时的总能耗仅为传统网格的0.72-0.84倍。

如下图所示,SuperMesh(SM)在带宽和链路能效方面均显著优于其他SOTA(state-of-the-art)互连拓扑。

图1:SOTA中介层(interposer)拓扑的带宽(BW)和链路能效(每焦耳带宽)。DB = Double Butterfly, Kite(S) = Kite_Small, Kite(M) = Kite_Medium, FT = Folded Torus, SM = SuperMesh (Proposed)。
图1:SOTA中介层(interposer)拓扑的带宽(BW)和链路能效(每焦耳带宽)。DB = Double Butterfly, Kite(S) = Kite_Small, Kite(M) = Kite_Medium, FT = Folded Torus, SM = SuperMesh (Proposed)。

下图直观地展示了2D网格拓扑的瓶颈所在。左图显示,在AllGather操作中,边界附近的链路利用率远高于内部链路。右图的数据遍历模式解释了这一点:内部节点有四个链路可以同时传输数据,而角落节点只有两个,导致边界链路长时间被占用。通过在边界增加链路(如图2a右侧所示),可以显著改善链路利用率,缓解瓶颈。

(a)SOTA Mesh(左)与边界带宽加倍的Mesh(右)的链路利用率百分比。(b)数据遍历模式。图2:6×6网格在AG期间的链路利用率;橙色和较粗的链路表示使用率较高,而蓝色和较细的链路表示使用率较低。
(a)SOTA Mesh(左)与边界带宽加倍的Mesh(右)的链路利用率百分比。(b)数据遍历模式。图2:6×6网格在AG期间的链路利用率;橙色和较粗的链路表示使用率较高,而蓝色和较细的链路表示使用率较低。

A3 背景知识

2.1 基于小芯片的加速器

机器学习模型复杂性驱动硬件可扩展性。随着机器学习(ML)模型复杂性的迅速增加,开发可扩展的硬件加速器变得至关重要,而小芯片(chiplet)技术已成为一个有前景的解决方案。该技术通过将多个小芯片组装到基板上(如有机基板【45】或硅中介层【23, 75】),并通过高密度链路互连,使它们能够作为一个整体协同工作。小芯片技术的优势包括:相比大规模单片芯片有更高的良率,利用中介层等先进基板可获得更大的芯片面积【23】,以及可以灵活地在不同应用中复用小芯片。

小芯片技术的挑战。然而,小芯片技术也带来了一些挑战,主要源于Die-to-Die(D2D)接口。这些挑战包括:D2D链路比片上连接需要更多功耗,导致能耗增加【50, 60】;以及由于可用引脚(pin)有限,带宽受到限制【7】。凸点(bump)的最小间距限制了每平方毫米小芯片面积上的凸点数量,从而限制了D2D链路的带宽。由于D2D链路限制了片间互连的数据宽度,它们通常在高频率下运行。为在如此高的频率下运行而不引入不可接受的误码率,这些链路的长度必须最小化【1】。

小芯片加速器的架构与拓扑选择。此外,基于小芯片的加速器通常由两类小芯片组成:计算小芯片和I/O小芯片【7】。在片间通信中,一个芯片的D2D发射器对数据进行编码,并将其转发给另一个芯片上相应的D2D接收器。为了增强可扩展性,D2D接口被策略性地放置在小芯片的四周。虽然已经为中介层提出了各种拓扑结构,但为了确保D2D链路的信号完整性,互连距离和交叉连接存在限制,这因此限制了某些拓扑(如环面(torus)和蝴蝶(butterfly))的使用。因此,最好保持点对点的并行互连。考虑到这些因素,大多数现有的瓦片式加速器【24, 26, 37, 58, 69, 80】目前都采用网格(mesh)互连网络。

2.2 2D-mesh的集体通信算法

现有算法及其在Mesh拓扑上的挑战。集体通信算法如ReduceScatter(RS)、AllGather(AG)和AllReduce(AR)对于在分布式DNN训练和推理期间同步数据至关重要。为了提高通信效率,已经开发了各种算法,但它们的性能因底层网络拓扑而异。对于网格拓扑,这些算法面临着重大挑战。基于环(Ring)的方法,包括ring【15】、bidirectional ring【35】和hierarchical ring【82】,随着网格规模的增大延迟会增加,而为其他拓扑优化的HDRM【12】和C-Cube【10】在网格上表现不佳。诸如double binary tree【55】之类的算法因为与拓扑无关而受限,MultiTree【22】和TACOS【77】由于对网格不对称性的处理不佳而表现不佳。像SCCL【8】、TACCL【57】和TE-CCL【3】这样的解决方案也未能有效扩展。

Pipeline AR(TTO)算法的优势。在为2D网格设计的集体算法中,pipeline AR【35】(也称为Three Tree AllReduce或TTO)是一种非常有前途的方法。例如,在一个4×4的网格中(如图3所示),它构建了三个不相交的树,根节点分别为0、15和3。与传统的基于树的算法(为每个节点创建一棵树,此例中为16棵)不同,pipeline AR只构建了三棵树,每棵树需要处理大约五倍的数据。其主要优势在于其流水线机制:一旦一个数据块完成在某条链路上的传输(例如,在树0中从15到14),下一个数据块可以立即利用该链路,而不会受到其他树的干扰,因为这些树是不相交的。与SOTA AR算法中数据块需要因共享链路而顺序流动不同,pipeline AR利用不相交的树,使得在第一个数据块传输后,后续数据块能够更快地传播。这显著减少了总通信时间并提高了链路利用率,使其特别适用于处理大量数据的深度学习工作负载。

图3:使用pipeline AR在4x4网格中形成不相交树和数据块流水线。
图3:使用pipeline AR在4x4网格中形成不相交树和数据块流水线。

Pipeline AR的局限性。尽管pipeline AR提供了显著的加速,但它有两个关键限制。首先,它只使用N-1个节点进行训练,牺牲了一个小芯片的计算能力来优化通信(图中用红叉标记)。这个限制是因为在2D网格中无法使用所有N个节点形成三个不相交的树,因为角节点的链路只有两个。其次,pipeline AR不支持SOTA的RS或AG,这些操作要求所有节点都参与数据的收集和分发。它只在三个节点上执行归约操作,这限制了其适用性。为了改进pipeline AR,TidalMesh【38】被提出,它利用所有N个小芯片,并通过叠加和合并数据块来获得更好的性能。然而,它的贡献主要针对AR,对RS或AG等其他集体操作没有性能提升。在现代工作负载如LLMs中,张量并行【61】在每个transformer块中需要两次AR,而ZeRO【51】训练则使用RS和AG来访问分布式权重和梯度,因此在小芯片加速器中加速这些集体操作至关重要(更多例子见表1)。为了区分pipeline AR修改版的RS和AG操作与它们的SOTA对应版本,我们下文将称之为Semi-RS和Semi-AG。

表1:分布式DNN工作负载中的常见集体通信
表1:分布式DNN工作负载中的常见集体通信

2.3 现有的片上和中介层拓扑

片上拓扑及其对中介层系统的不适用性。虽然已经为片上和基于中介层的系统提出了许多拓扑,但2D-mesh因其可扩展性和与当前集成技术的兼容性而仍是使用最广泛的。然而,其由大量跳数引起的高通信延迟催生了替代设计。像MECS【17】和DIA-Torus【72】这样的方法通过长链路、旁路链路、快速通道或对角线连接来增强路由器间的连接性。类似地,ADAPT-NoC【84】引入了自适应链路以改善数据包延迟,而其他设计则将快速链路嵌入到网格拓扑中【47】或重新配置网络连接【43, 63】以增加对剖带宽并减少长途通信。尽管这些方法能有效减少片上通信延迟,但快速链路或旁路对于中介层系统来说不太适用,因为长导线延迟和多周期跳数限制了它们相对于片上网络的效果【5】。

中介层拓扑及其在集体通信中的低效性。对于基于中介层的系统,像DB、ButterDonut、Kite和FT这样的拓扑也使用长链路或对角线链路来减小网络直径。这些方法在点对点通信中能有效降低平均数据包延迟,但对于所有节点同时交换数据的集体通信而言效率较低。在集体通信中,大量数据流经网络,而更长的链路会增加每跳延迟和每比特的链路能量,这与导线长度成正比。因此,虽然减小的网络直径减少了总跳数,但其好处被链路延迟比2D-mesh增加2-4倍【5】以及能效显著降低所抵消。

集中式设计与2D-Mesh的优势。此外,许多这些拓扑采用集中式设计(例如,每个路由器连接四个加速器),这虽然能降低延迟,但对于需要高吞吐量的集体通信来说是次优的。虽然集中可以减少路由器功耗,但它也增加了路由器的度数。由于这些原因,2D-mesh仍然是基于小芯片的加速器最常用的拓扑,因为它在可扩展性、能效和与集成技术的契合度之间取得了平衡【37, 58, 69】。在芯片多处理器(CMPs)中,采用XY路由时,内部链路会面临拥塞。HeteroNOC【42】通过在中心区域引入异构路由器和额外带宽来解决这个问题,但这并不能改善受边界节点链路限制的集体通信性能。相比之下,SM在2D-mesh框架内以最小的修改增强了集体通信。与其他拓扑不同,它优先考虑集体性能而不减小网络直径,使其易于适应现有加速器设计,同时保持能效。

与ARIES和HexaMesh的比较。ARIES【81】沿所有水平和垂直连接的小芯片引入了旁路链路,使得链路数量相比SOTA网格网络增加了一倍。尽管它可以自适应地选择必要的链路,但对于各种集体通信算法,这些增加的链路大部分仍未被使用。由于ARIES提供了每行每列小芯片之间所有可能的链路,因此可以禁用所有内部链路,将其转换为我们提出的拓扑。然而,我们可以在一个8×8的网格中禁用多达84%的附加链路,并且仍然保持与我们提出的拓扑相同的性能,没有发现这些额外链路带来任何额外的好处。HexaMesh【25】将每个小芯片连接到六个邻居,呈六边形布局,但边界小芯片仍然只有3-4个链路,导致集体通信中的不平衡。此外,一些现有的加速器如AMD AIE【45】在计算节点的边界设有内存节点。一种潜在的策略是利用这些内存节点的连接在边界提供额外的路径,而不是引入专用链路。然而,这种方法会增加附加链路的延迟,因为它们需要通过内存节点经过多个跳数,而保持统一的链路延迟对于流水线数据块时的高效集体性能至关重要【35】。


A2 方法细节

3 SuperMesh拓扑

3.1 概述

设计动机与原则。设计SM(SuperMesh)的主要动机是解决2D-mesh网络的瓶颈问题,以最小的拓扑修改实现集体通信性能的显著提升,同时保持功耗和延迟效率。由于边界节点链路数量有限是导致集体性能下降的主要原因,SM在网格的外围区域引入了额外的链路,为数据移动提供更多路径。SM的设计专门针对集体通信算法,并基于以下基本原则:

  • 额外链路仅连接相邻节点:由于硅中介层中的D2D链路必须保持短距离以确保在高工作频率下的低误码率,并且长链路会增加延迟和功耗,因此SM增加的所有额外链路都只放置在相邻节点之间。这确保了新链路的延迟和能量特性与现有链路相似。
  • 对网格拓扑的最小改动:由于能效至关重要,且2D-mesh拓扑因其在小芯片加速器中的可扩展性和能效而被广泛使用,SM通过仅在外围区域添加链路,同时保留网格的内部网格结构,确保了最小的能量开销和高可扩展性。

架构假设。为了说明SM拓扑,我们采用了一个类似Simba的架构,其中只有计算小芯片使用网格互连。在互连中集成内存或I/O瓦片的情况将在3.4节中进一步讨论。

图4:提出的SuperMesh拓扑。SuperMeshBi在网格周界引入双向链路,SuperMeshAlter则交替引入双向链路。
图4:提出的SuperMesh拓扑。SuperMeshBi在网格周界引入双向链路,SuperMeshAlter则交替引入双向链路。
3.2 SuperMeshBi拓扑

拓扑结构。我们的第一个提议拓扑,SMBi(SuperMeshBi),通过在所有外围链路旁并行添加双向链路来扩展SOTA网格互连。如图4(左)所示,对于一个N×N的系统,SMBi保留了固有的网格拓扑。然而,它在所有水平和垂直边界小芯片的外围链路旁增加了额外的双向链路,而内部区域则未作修改。这种布局确保了每个内部小芯片与基线小芯片加速器中的完全相同。例如,在网格拓扑中,小芯片0(左上角)使用两个D2D接口连接小芯片1和4。在SMBi中,小芯片0剩余的两个D2D接口被用来增加额外的链路,分别连接到小芯片1和小芯片4。因此,所有角小芯片现在都有四个链路,并且增加的链路保持在核心区域之外,确保了短D2D链路的无缝放置。类似地,其他边界小芯片连接到两个相邻的水平或垂直小芯片。然而,这些边界小芯片的路由器——除了角小芯片——需要在路由器中增加一个端口以支持与五个相邻小芯片的连接,而无需对内部小芯片的路由器进行任何更改。

3.3 SuperMeshAlter拓扑

拓扑结构与优势。SuperMeshBi确保每个节点至少有四个链路,但它需要一些边界节点连接到五个其他邻居。然而,对于偶数尺寸的网格,可以在不改变任何路由器的情况下为所有节点确保四个链路。这是通过在外围节点之间交替添加链路来实现的,我们称之为SuperMeshAlter,如图4(右)所示。这种设置保证了所有角节点都与两个相邻邻居增加了双向链路,而其他边界节点只与一个外围邻居增加了双向链路。这种安排在增强集体性能的同时,仍然只连接四个其他邻居。然而,SMAlter的有效性主要体现在具有偶数维度的网格中。在具有奇数维度的网格中,角节点只有三个连接到相邻节点。

3.4 SM与Mesh的比较分析

拓扑特性比较。在表2中,我们对SM的两种变体和SOTA网格拓扑进行了比较分析。只有SMBi需要部分边界节点的路由器度为6,而网格和SMAlter都保持度为5。SM的直径和对剖带宽与网格相似;然而,当与集体通信算法协同设计时,SM中的额外链路显示出优势。此外,增加的链路比例对于较小的拓扑更高,并随着拓扑规模的增加而减小,从而最小化了面积开销。与网格相比,SMBi和SMAlter的Network-on-Interposer(NoI)面积开销分别为1.15-1.23倍和1.04-1.11倍。此外,通过减少集体通信时间,SMBi的总能耗仅为网格的0.74–0.84倍,SMAlter为0.72–0.80倍。

与其他中介层拓扑的比较。与DB、FT、Kite Small和Kite Medium等其他中介层拓扑相比,SMBi的NoI面积开销为1.1-1.79倍,能耗降低1.2-3.67倍。同样,SMAlter的NoI面积开销为1.05-1.6倍,能效提高1.24-3.79倍。由于DB和Kite等拓扑使用集中化设计,它们需要更少的路由器和NoI链路,从而面积更小。因此,SM和网格都显示出更高的NoI面积开销。我们只报告了NoI的面积和功耗,因为所有其他组件在各种拓扑中都是相同的。

表2:mesh和SuperMesh之间的比较。
表2:mesh和SuperMesh之间的比较。

路由与死锁避免。由于集体通信遵循固定的时间表,并且在SM中,源和目的地是相邻节点,因此可以避免死锁。像MultiTree【22】和TTO【35】这样的现有工作,通过在每个加速器中维护一个调度表来记录路由细节,该表存储集体操作、源、目的地和依赖关系;SM采用了类似的方法。对于点对点通信,当有额外链路可用时,维度顺序路由可以通过拥塞感知自适应路由【16, 41, 52】来增强。

奇数尺寸网格与内存小芯片的考量。虽然SMBi适用于奇数尺寸的网格,但它需要一个小芯片连接到五个邻居。交替的双向链路可以减少大多数节点的连接到四个,但至少有一个小芯片必须保持五个链路以在角落处保留四个连接。当内存小芯片连接到边界节点时,一些边界需要第六个连接。为了适应这种情况,我们采用了HexaMesh【25】架构,它支持每个小芯片六个链路,非常适合边界节点。这种六端口小芯片仅在SMBi和SMAlter中内存与计算小芯片共享同一互连时才需要。这种配置对于SM来说是足够的,即使在所有侧面都放置了内存或I/O小芯片。

4 Co-designed Collectives for SuperMesh

由于SM在标准网格拓扑的基础上引入了外围链路,这些新增链路的优势需要通过协同设计的集体通信算法来体现,这些算法能更好地利用新链路。本节将探讨我们提出的流水线AR、RS和AG算法如何显著提升性能,超越SOTA网格。现有的集体通信综合算法,如MultiTree、TACOS【77】、TACCL【57】和TE-CCL【3】,它们根据拓扑生成调度,但在网格环境中常常表现不佳(详见第6节)。相比之下,pipeline AR【35】在AR集体通信中显示出巨大改进,使其成为2D网格拓扑的一个有前景的算法。然而,它仅限于AR,不支持RS和AG。此外,它在计算过程中会有一个节点未使用,造成资源浪费。本节将展示我们为SM协同设计的流水线AR算法如何解决这些问题。

4.1 优化流水线AllReduce

网格拓扑中树数量的限制。在网格拓扑的pipeline AR算法中只形成三棵树的主要原因是可用链路数量有限。具体来说,在一个$N \times N$的网格中,有$4N^2 - 4N$条链路。然而,要建立四棵不相交的树,需要$4N^2 - 4$条链路。对于任何$N \geq 2$,$4N^2 - 4N$总是小于$4N^2 - 4$,这表明构建四棵不相交的树缺少链路。即使配置三棵不相交的树,也需要将一个节点排除在计算之外。这是因为在集体通信中,每个节点都必须能够与所有其他节点进行收发。因此,当树的根节点设在三个角节点时(如图3所示),有一个节点(节点6)缺乏支持所有不相交通信所需的三个输入和输出链路。

SuperMesh的解决方案。所提出的拓扑不仅确保了所有节点在计算中都被利用——通过为四个角节点提供至少三个输入和输出链路,并为所有其他节点提供四个链路,从而在不改变网格内部结构的情况下创建四棵不相交的树。这种方法提高了AR性能并最大化了计算效率。

在SuperMesh上构建四棵不相交的树。在图5中,我们展示了在4×4的SMBi和SMAlter拓扑中,以角节点0、3、12和15为根的四棵不相交的树,用于Semi-RS和Semi-AG。对于SMBi的Semi-RS,我们首先分别从南、西、东、北方向为每棵树连接节点,仅使用新增的链路。接下来,我们分别从东、南、北、西方向将节点连接到已连接的节点上,使用现有的链路。例如,在树0中,节点4、8和12首先从南边使用新增链路连接到节点0,然后使用现有链路向东连接。Semi-AG的过程类似,但起始方向不同。然而,如图5b所示,SMAlter拓扑由于链路是交替连接的,不允许如此直接地构建树,需要一种自动的树连接策略。虽然我们在这里展示了手动建树,但我们将在4.2节中介绍一个算法,用于为任何给定的拓扑自动形成四棵不相交的树。由于总共有四棵不相交的树,每棵树处理$D/4$的数据,其中$D$是总数据大小。

图5:使用4x4 SuperMesh的pipeline AR的Semi-RS和Semi-AG不相交树。蓝色表示新增链路,黑色表示默认链路。
图5:使用4x4 SuperMesh的pipeline AR的Semi-RS和Semi-AG不相交树。蓝色表示新增链路,黑色表示默认链路。
4.2 优化ReduceScatter和AllGather

网格拓扑对RS和AG的限制。在一个有N个节点的系统中,传统的RS和AG要求每个节点归约和广播总数据的$1/N$。然而,由于pipeline AR在四个角节点上归约所有数据,这种方法不直接适用于RS和AG。如果我们在网格拓扑中尝试为RS和AG构建根节点不在角落的不相交树,我们最多只能形成两棵不相交的树。这是因为角节点此时需要发送和接收数据,但它们只有两个输入和输出链路,这导致在集体通信过程中许多链路未被利用,从而严重影响性能。

SuperMesh对RS和AG的优化。由于边界节点链路数量有限是造成此问题的原因,而SM在边界增加了额外的链路,因此有机会将流水线算法的优势应用到RS和AG上。在图6中,我们展示了在4×4 SMAlter拓扑中形成四棵不相交的树,突出了边界节点和内部节点可能的树形成方式。上图显示了根在边界节点(非角节点)的树,而下图显示了根在中心节点的树。由于每个节点都有四个输入和输出链路,构建以内部和边界节点为根的四棵不相交的树是可行的。例如,当树的根在边界时,树的高度为五。然而,当树的根在内部节点时,树的高度会增加,因为每个节点必须只用一条链路连接到每棵树。优先从内部连接边界节点有助于形成对称且不相交的树,因为边界节点现在有了额外的链路。

流水线集体通信中树高度的影响。虽然增加树的高度通常会延长其他基于树的算法【22, 55】的集体通信时间,但这对于流水线集体通信来说问题不大。在流水线场景中,尽管初始数据块的传输可能需要更长时间,但后续的数据块会以流水线的方式快速流动,因此不会显著影响整体集体通信时间。这对于处理大模型尺寸和大量数据块的现代DNN系统尤其有利。在流水线RS和AG中,每个树集传输$D/N$的数据量,其中$D$是总数据大小。

图6:使用4x4 SuperMeshAlter拓opor的pipeline AG的不相交树。蓝色表示新增链路,黑色表示默认链路。上图显示了4棵根在边界节点的不相交树,下图显示了根在内部节点的生成树。
图6:使用4x4 SuperMeshAlter拓opor的pipeline AG的不相交树。蓝色表示新增链路,黑色表示默认链路。上图显示了4棵根在边界节点的不相交树,下图显示了根在内部节点的生成树。

奇数维度SMAlter的限制。尽管在SMBi和偶数维度节点的SMAlter拓扑中总是可以创建四棵不相交的树,但在奇数维度节点的SMAlter中情况有所不同。在这种情况下,由于角节点的可用链路数量有限,每个集合只能形成三棵树。

数据块调度。在图7中,我们展示了在AG集体通信期间,数据块如何在不同的树集之间进行调度。为了说明这个概念,考虑图6中根为节点1的树。在一个包含多个数据块的集体通信中,当第一个数据块通过链路(1 → 0)、(1 → 5)和(1 → 2)到达第1层的节点时,这些链路就可供数据块2使用,因为这些树是不相交的,意味着这些链路不被其他树使用。因此,当数据块1前进到第2层时,数据块2前进到第1层,其他数据块也相应地继续传输。对于一个有$N$个节点和数据大小为$D$的拓扑,每个不相交的树处理$D/N$的数据。当四棵不相交的树同时活动时,网络中一次传输$4 \times D/N$的数据。这确保了所有RS操作在$\lceil N/4 \rceil$个流水线集合中完成,每个节点最终剩下$D/N$的数据。当数据量很大时,流水线气泡(开始和结束时未使用的链路)被最小化,从而最大化了流水线AR的优势。

图7:使用不同树根的pipeline AG中多个数据块的调度。
图7:使用不同树根的pipeline AG中多个数据块的调度。

流水线RS的优势。为了更好地理解流水线RS的好处,考虑一个4×4的SMAlter拓扑,数据量为512MB,块大小为16KB,每个节点必须归约2048个块。使用像MultiTree这样的算法,总延迟将是$S \times 2048 \times T$,其中$S$是步数,$T$是每块的延迟,因为链路在不同步骤之间不是不相交的。相比之下,使用流水线RS,树高最多为七(图6),延迟变为$4 \times (2048+6) \times T$,其中因子4对应四个树集(图6)。因此,如果$S > 4$,流水线AR有明显优势。我们的实验表明,对于4×4的网格,MultiTree需要九步来完成RS,这证明了流水线RS实现了超过2倍的加速。

图8:在pipeline RS和AG期间不相交树形成的根选择策略。
图8:在pipeline RS和AG期间不相交树形成的根选择策略。

树根选择策略。尽管使用流水线RS和AG创建四棵不相交的树是可行的,但随机选择根节点可能导致一个集合内树的高度不均匀,影响通信在各树之间的均匀性。为了解决这个问题,我们按特定顺序策略性地选择树根,旨在使一个集合中所有根的树高均匀。这种方法确保所有树同步开始和结束通信,从而最大化效率。这个策略在图8中用6×6的SMAlter进行了说明,其中不同的集合用不同颜色高亮,并显示了每个集合的树高。我们发现,按系统顺序选择根——从每一侧选择一个,并从外边界向内进行——能使每个根有效地利用边界上的新增链路,并产生大小均匀、树高最小的链路。对于像奇数维度节点的SMAlter这样的配置,一次只能形成三棵树,我们采用从三个不同侧面选择三棵树的方法。

调度生成算法。算法1概述了使用给定的拓扑、根和维度$N$为SM生成Semi-RS和Semi-AG调度的步骤。在第3-6行,我们首先创建一个空的树列表,每个树用根进行初始化。如果根不在边界上,它们首先被连接到最近的边界节点(第7-9行)。例如,在图6(底部)中,节点1被连接到根5,从而高效地利用了新增的边界链路。第10-13行使用BFS将来自($N-1$)行和($N-1$)列的节点连接到每棵树,同时从拓扑中移除已使用的链路。例如,在图6中,节点9、0、2、4、6、8和10被连接到以5为根的树。这个$N-1$的限制确保每棵树都有足够的路径到达边界节点,并高效利用新增链路。接下来,第14-19行迭代地将与每棵树现有节点相邻的孤立节点加入,确保所有节点都被连接。最后,第20-21行通过对树进行自底向上遍历生成Semi-RS调度,通过自顶向下遍历生成Semi-AG调度。通过从拓扑中移除已连接的链路并优先使用新增链路,该算法确保了不相交树的形成和高效的链路利用。

1: 输入: 包含所有链路的拓扑,一组根节点,维度 N
2: 输出: Semi-RS 调度 (SemiRSSchedules) 和 Semi-AG 调度 (SemiAGSchedules)
3: 初始化空列表 Trees 来存储每个根的一棵树
4: for all root ∈ roots do
5:     Trees[root] ← [root]
6: end for
7: if 根节点不在边界 then
8:     将每个根连接到其最近的边界节点
9: end if
10: for all root ∈ roots do
11:     使用 BFS 将 (N-1) 行和 (N-1) 列的节点连接到 Trees[root]
12:     从拓扑中移除这些连接中使用的链路
13: end for
14: while 所有节点都未连接到树 do
15:     for all root ∈ roots do
16:         连接一个与 Trees[root] 当前节点相邻的剩余节点
17:         从拓扑中移除使用的链路
18:     end for
19: end while
20: 通过对 Trees 进行自底向上遍历生成 SemiRSSchedules
21: 通过对 Trees 进行自顶向下遍历生成 SemiAGSchedules
22: return SemiRSSchedules, SemiAGSchedules

A4 实验环境

硬件配置
* 计算单元 (PE): 每个小芯片(chiplet)包含16个PE,每个PE内含一个128×128的MAC阵列(用于主要实验)或一个32×32的MAC阵列(用于小型配置对比实验),工作频率为1 GHz。
* 互连网络:
* 链路带宽: 每个链路25 GBps。
* 跳延迟: 每跳20 ns。
* 流量控制: 虚通道(Virtual Channel, VC)直通流控,每个物理通道配置4个VC,VC缓冲深度为318。
* 拓扑: 实验涵盖了从3×3到16×16(即9到256个小芯片)的各种规模的Mesh和SuperMesh(SMBi, SMAlter)配置。

  • 功耗模型: 基于DSENT工具,采用32nm bulk/SOI low-Vt工艺节点,小芯片间链路的互连能耗为1.75 pJ/bit。

软件配置
* 模拟器:
* 网络模拟: BookSim,已修改以支持并行链路。
* 计算模拟: SCALE-Sim,采用其支持训练时间评估的扩展版本。
* 功耗模拟: DSENT。
* 集成: 使用Python接口连接SCALE-Sim和BookSim。

  • 数据流: 采用输出固定(output stationary)数据流。
  • 对比算法:

    • TTO (Three Tree AllReduce): 针对2D-mesh的流水线AllReduce算法。
    • TACOS: 拓扑感知的自动集体通信综合器。
    • TE-CCL: 基于多商品流的集体通信综合器。
    • MultiTree: 基于树的集体通信算法。
  • 调度机制: 在每个加速器中维护一个调度表,记录集体操作、源、目的地和依赖关系。该表在初始化时计算一次并加载到网络接口中复用。对于64节点拓扑,调度表开销为3.2KB。

数据集与模型
* 带宽测试: 数据量从1MB到512MB不等。
* 可扩展性测试: 总数据量为 375KB × N,其中N为小芯片数量。
* DNN模型:
* 模型: InceptionV4, Densenet201, ResNet152, VGG19。
* 数据集: 模拟ImageNet规模(120万样本)。
* 批大小: N节点拓扑下,mini-batch大小为16 × N。

  • LLM模型:
    • 模型: BLOOM 176B, Gemma 7B, Falcon 40B, LLaMA 405B, GPT-3, Qwen3-32B。
    • 配置: 模拟预填充(prefill)阶段,批大小为16,序列长度1024,精度为FP16。
    • 计算时间: 使用LLMCompass获取。

A4 实验结果

6.1 带宽利用率

  • AllReduce (AR): 如图9所示,SuperMesh(SM)通过构建4棵不相交的树,相比于只能构建3棵树的TTO(在mesh上),实现了1.18-1.33倍的带宽提升,相比TACOS提升1.25-2.13倍。SMBi和SMAlter性能相近。
  • ReduceScatter (RS) & AllGather (AG): 如图10所示,SM显著优于TACOS。SMBi性能最佳;偶数尺寸的SMAlter与其相当,但奇数尺寸的SMAlter因只能构建3棵树而性能稍低。对于小数据量,TACOS因没有流水线启动开销而表现更佳;但随着数据量增加,SM的优势愈发明显,最高可达1.97倍的加速。
图9:4x4, 5x5, 8x8和9x9 Mesh与SM在不同数据大小下的AR带宽使用情况。
图9:4x4, 5x5, 8x8和9x9 Mesh与SM在不同数据大小下的AR带宽使用情况。
图10:4x4, 5x5, 8x8和9x9 Mesh与SM在不同数据大小下的RS和AG带宽使用情况。上排为RS结果,下排为AG结果。
图10:4x4, 5x5, 8x8和9x9 Mesh与SM在不同数据大小下的RS和AG带宽使用情况。上排为RS结果,下排为AG结果。

6.2 可扩展性分析

  • 如图11所示,在从9节点扩展到256节点的过程中,SM的集体通信算法表现出线性的扩展性。
  • 对于AR,SM相比TTO平均有1.3-1.33倍的加速,相比TACOS最高可达1.99倍。
  • 对于AG,SM相比TACOS实现了1.42-1.85倍的加速。
图11:从9节点到256节点的可扩展性,数据大小为375 x N KB,结果归一化到3x3 mesh和SM。
图11:从9节点到256节点的可扩展性,数据大小为375 x N KB,结果归一化到3x3 mesh和SM。

6.3 模型性能结果

  • DNN训练: 如图12所示,在6×6拓扑上,SM相比TTO有两个优势:1) 使用全部36个小芯片(TTO使用35个),减少了迭代次数;2) AR通信本身快30%。这使得在128×128 MAC阵列配置下,端到端训练加速比达到1.27-1.33倍。
  • LLM推理: 如图13所示,在4×4拓扑上进行张量并行推理,SM将AR通信时间减少约30%,带来了1.08-1.12倍的端到端推理加速。
图12:在6x6 mesh和SM上使用pipeline AR训练流行DNN模型一个epoch的时间分解。AR、前向+反向传播和端到端训练加速均归一化到mesh。
图12:在6x6 mesh和SM上使用pipeline AR训练流行DNN模型一个epoch的时间分解。AR、前向+反向传播和端到端训练加速均归一化到mesh。
图13:在4x4 mesh和SM上使用pipeline AR对流行LLM模型单个transformer块的预填充阶段进行推理时间分解。AR、计算时间和推理加速均归一化到mesh拓扑。
图13:在4x4 mesh和SM上使用pipeline AR对流行LLM模型单个transformer块的预填充阶段进行推理时间分解。AR、计算时间和推理加速均归一化到mesh拓扑。

6.4 能效分析

  • 如图14所示,SM虽然因更高的链路利用率而动态链路功耗稍高,但由于通信时间显著缩短,总的静态路由器能耗大幅降低,最终总能耗更低。SMBi和SMAlter的总能耗分别仅为mesh的0.74-0.84倍和0.72-0.80倍。
  • 如图15所示,与其他中介层拓扑(如DB, FT, Kite)相比,SM的能效优势巨大。这些拓扑因采用长链路和集中式设计,导致链路和路由器能耗都更高。SM比其他拓扑节能1.2-3.79倍。
图14:8x8 mesh和SM进行AR通信的归一化功耗和能量分解。
图14:8x8 mesh和SM进行AR通信的归一化功耗和能量分解。
图15:在8x8系统上,各种中介层拓扑进行AR通信的归一化能量分解。缩写:DB=Double Butterfly, FT=Folded Torus, Kite(S)=Kite Small, Kite(M)=Kite Medium。
图15:在8x8系统上,各种中介层拓扑进行AR通信的归一化能量分解。缩写:DB=Double Butterfly, FT=Folded Torus, Kite(S)=Kite Small, Kite(M)=Kite Medium。

6.5-6.11 其他性能评估

  • 对其他算法的增益 (图16): SM拓扑同样能提升TE-CCL、MultiTree和TACOS等算法的性能,因为其解决了mesh的根本瓶颈。例如,TACOS在SMBi上的性能比在mesh上提升了1.88倍。
  • 链路利用率 (图18): 在6×6系统上,SMAlter的链路利用率高达97%,优于SMBi和TTO(均为87%),是最高效的配置。
  • 与其他拓扑的性能对比 (图19): SM的AR性能远超其他中介层拓扑,例如比DB快7.9倍,比mesh快1.33倍。
  • AllToAll性能 (图21): SM在AllToAll通信中同样优于mesh,SMBi和SMAlter分别实现了高达1.72倍和1.53倍的加速。
  • 平均包延迟 (图22): 在点对点通信负载下,SM也表现出更低的平均包延迟,在饱和负载下,SMBi可将延迟降低高达74%(均匀流量)。
图16:在4x4 mesh和SM上,使用TE-CCL、MultiTree和TACOS进行AG的带宽使用情况。
图16:在4x4 mesh和SM上,使用TE-CCL、MultiTree和TACOS进行AG的带宽使用情况。
图17:在不同4x4 SM拓扑上,各种数据大小的AR带宽使用情况。
图17:在不同4x4 SM拓扑上,各种数据大小的AR带宽使用情况。
图18:在6x6 mesh和SM上,使用512MB集体数据大小的所有基准测试的链路利用率百分比。X轴是周期数。
图18:在6x6 mesh和SM上,使用512MB集体数据大小的所有基准测试的链路利用率百分比。X轴是周期数。
图19:不同中介层拓扑的AR带宽使用情况。缩写:DB = Double Butterfly, FT = Folded Torus, Kite(S) = Kite Small, Kite(M) = Kite Medium。
图19:不同中介层拓扑的AR带宽使用情况。缩写:DB = Double Butterfly, FT = Folded Torus, Kite(S) = Kite Small, Kite(M) = Kite Medium。
图20:在4x4部分SMBi和SMAlter拓扑上,不同数据大小的AR带宽使用情况,其中三边添加了额外链路,一边未修改用于I/O。
图20:在4x4部分SMBi和SMAlter拓扑上,不同数据大小的AR带宽使用情况,其中三边添加了额外链路,一边未修改用于I/O。
图21:在3x3 mesh和SM上,使用TE-CCL的不同数据大小下的AllToAll带宽使用情况。
图21:在3x3 mesh和SM上,使用TE-CCL的不同数据大小下的AllToAll带宽使用情况。
图22:4x4 mesh和SM在均匀和龙卷风流量模式下的平均数据包延迟。
图22:4x4 mesh和SM在均匀和龙卷风流量模式下的平均数据包延迟。

A5 结论

本文针对网格(mesh)拓扑中因边界节点链路有限而导致的集体通信性能瓶颈问题,提出了两种新的拓扑结构:SMBi和SMAlter。SMBi在所有边界节点间添加双向链路,SMAlter则交替添加双向链路。这两种设计在保留网格结构的同时,通过在边界节点之间引入额外的短链路,保证了可扩展性并提升了能效。

为充分利用新拓扑,本文协同设计了针对AR、RS和AG的集体通信算法。实验结果表明,与传统的2D-mesh拓扑相比,这些算法在AR上实现了1.18-1.33倍的加速,在RS和AG上实现了1.77-2.2倍的加速。

未来的工作将集中于将这些设计扩展到其他集体通信操作,并为任意分区设计系统以加速混合并行计算。


A6 附录

A.1 摘要

工件评估指南。本节为论文的工件评估提供指南。大部分复现结果的说明都在代码仓库的README.md文件中提供。我们首先提供一个评估清单,然后描述所需的硬件和软件依赖。最后,安装和实验流程部分概述了如何设置环境和运行脚本以进行全面验证。通过遵循README.md中的步骤,可以完全复现论文中呈现的所有结果。

A.2 工件清单(元信息)

  • 程序: Python = 3.9.5, ScaleSim, BookSim2
  • 运行环境: Ubuntu = 22.04
  • 硬件: 1台装有Ubuntu = 22.04的机器
  • 输出: 包含模拟结果的Json文件
  • 实验: 带宽利用率、可扩展性结果、DNN基准测试性能
  • 所需磁盘空间(大约): 约20GB
  • 准备工作流所需时间(大约): 约15分钟
  • 完成实验所需时间(大约): 每个模拟通常需要10分钟到48小时。所有实验可能需要几天时间。建议并行模拟
  • 是否公开可用: 是
  • 是否已存档(提供DOI): https://doi.org/10.5281/zenodo.167321
  • GitHub链接: https://github.com/sabujlaskar/SuperMesh_AE

A.3 描述

A.3.1 如何访问。代码可以通过Zenodo和GitHub访问。Zenodo记录提供了代码的ZIP文件,解压后即可获得最终代码;而GitHub仓库则包含最新且持续维护的代码库。

A.3.2 硬件依赖。单台机器可以处理单个模拟。然而,为了同时执行多个模拟并高效地复现所有结果,建议使用多核机器或服务器。我们的模拟不需要GPU。

A.3.3 软件依赖。我们在Ubuntu 22.04 LTS操作系统上运行了实验,但其他版本的Ubuntu也应该可以工作。Python是模拟的先决条件。关于所有软件依赖的更多信息可以在README.md文件中找到。

A.4 安装

安装指南。请参阅代码仓库的README.md文件,其中包含了详细的“设置”分步指南。

A.5 评估与预期结果

实验复现。运行实验和生成本文结果的说明可以在README.md文件中找到。具体来说,我们提供了所有的模拟结果,以及用于创建论文中图表的脚本。结果存储在micro_2025文件夹中,生成图表的Python脚本位于utils/supermesh_python_scripts文件夹。此外,我们在utils/supermesh_runscripts文件夹中提供了复现所有模拟输出文件的脚本。由于每个shell脚本都会在后台启动多个并发模拟,建议使用核心数较多的机器。新模拟的结果将保存在micro_2025_new文件夹中。