人工智能网络架构中光线路交换技术的应用优化

深度神经网络训练依赖于几种并行化策略，每种策略都表现出独特的通信行为。在数据并行中，每个GPU维护模型的完整副本，并处理分区的数据批次，通过前向和后向传递，并在下一次迭代开始之前通过全规约操作同步梯度参数。在张量并行中，每个GPU持有模型层的分区，处理完整的数据批次，并使用全收集、规约分散或全规约操作同步中间结果。在流水线并行中，模型层被分割成顺序阶段，每个阶段分配给一个GPU。GPU使用点对点通信与流水线中的相邻阶段交换数据。输入数据通常被划分为微批次，以便并发执行模型的不同部分并减少流水线气泡。在专家并行中，混合专家层中的专家分布在GPU之间。同步发生在数据并行节点之间，在每个混合专家层之前和之后需要全对全通信。

工作负载通常使用多种形式的并行，来自不同并行策略的流量保持分离。因此，每个GPU在三个不同的通信域中进行数据传输，即数据并行、张量并行和流水线并行域。

图1：(a) GPT3-175B模型映射到512个GPU的流量热图。(b) Giant OCS抽象，所有GPU直接连接到大型交换机，该交换机将GPU链路分配给三个逻辑连接的通信域。

研究团队通过分析在512个处理单元上训练GPT-3模型的流量热图，可视化了这种通信结构。热图清楚地显示了三个不同的通信组，流量完全保持在指定域内。这种结构化模式使得可以提取通信子图，该子图成为优化网络重构策略的基础。当所有GPU采用相同的并行化策略时，除了流水线并行域内的GPU外，所有GPU的计算图都相同。流水线并行涉及非对称的点对点通信（与另一个GPU来回），不同于张量并行和数据并行中使用的对称集体通信。因此，只需要基于具有不同流水线并行索引的GPU的计算图来优化重构过程，结果可以推广到其他GPU。

为了简化网络建模并专注于开发重构策略，研究团队首先提出使用Giant OCS抽象。在这种抽象中，每个具有可控光链路的GPU使用所有链路直接连接到单个大型非阻塞光线路交换机。Giant OCS然后可以在任何一对GPU之间分配任意数量的链路，前提是两个GPU都有足够的可用链路，这使得可以对任何任意拓扑进行建模。使用Giant OCS抽象，可以通过应用简单的带宽延迟流体模型来估计通信子图中每个操作的执行时间，前提是通信端点通过至少一个链路连接。然后可以对这一系列操作时间应用优化，以确定基于任何给定重构延迟重构网络的最佳权衡。

使用深度神经网络通信子图和Giant OCS抽象，可以将GPU组织成具有三个逻辑连接域的拓扑。如图1b所示，Giant OCS被配置为使每个GPU沿着三个维度与其他GPU直接连接，对应于三个通信域，定义为集合D = {TP, DP, PP}。根据并行化策略和使用的集体算法类型，Giant OCS可以在不同的通信域之间任意分配每个GPU上的链路。这有效地为每个通信域配置了GPU带宽，表示为总注入带宽的一部分，其中所有域的带宽总和等于1。这种逻辑拓扑表述简化了优化，通过消除将流量矩阵与目标拓扑匹配的复杂性，这个问题通常计算密集。相反，只需要根据选择的重构策略和集体通信算法将带宽分配给每个域。

Giant OCS抽象为拓扑评估提供了性能上限，因为在没有交换机端口数限制的情况下运行。在使用Giant OCS抽象和逻辑拓扑设计重构策略后，可以引入光线路交换机的端口数和可重构性约束。由于光线路交换机端口数有限，需要开发分层或多维拓扑来有效地在网络中物理互连GPU。这将需要对基于Giant OCS抽象开发的重构策略进行调整，并引入额外的功耗和成本考虑。

第一种策略仅基于通信子图中当前和先前通信节点的信息做出决策，允许以O(1)时间复杂度实时做出决策。研究团队首先考虑重构开销为零的理想场景。在这种情况下，应在每次域转换时将GPU带宽的100%分配给主动通信域，从而实现最小可能的通信时间。

当重构开销非零时（如实际场景），重构的性能优势会被相关开销抵消。表征这种抵消的一个直观策略是评估抵消是否超过指定的重构阈值。该阈值可以在各种情况下定义，例如占空比或节省的时间。在本研究中，重点关注节省时间的情况，其中如果不重构的通信时间超过重构的通信时间与重构延迟之和达到给定阈值百分比，则触发重构。在这种情况下，当重构阈值为0时，只要通信时间有任何减少，就会发生重构。

图2：(a) 在H100 GPU上训练GPT-3 175B模型的通信间隔累积分布函数。(b) 作为变化重构延迟函数的暴露重构百分比。

当决定在节点之后不重构时，必须为与下一个节点相关的域保留一些非零带宽，否则在到达下一个节点之前需要重构。为了解决这个问题，引入了参数最小域带宽，定义每个域在任何时候的最小带宽量。这确保即使不执行重构，下一个域始终有一些非零带宽进行通信。由于这种保留带宽，可分配给给定通信域的最大带宽不再是100%。例如，5%的最小域带宽确保每个GPU始终至少维持5%的带宽到三个通信域中的每一个，并且每个域的最大可控带宽变为1 - 2 × 5% = 90%。

分析显示，较低的最小域带宽值在低重构延迟下表现良好，因为在频繁重构期间将更多可控带宽分配给主动通信域。相反，在较高的重构延迟下，较大的最小域带宽值表现更好，因为当重构不频繁时，将更多带宽分配给非通信域。将最小域带宽设置为0会恢复到基线情况，其中在每次域转换时重构100%的带宽，导致在每次域转换时重构，并导致通信时间呈指数增长。为了获得最佳性能，提出通过扫描不同的最小域带宽值并选择具有最小通信延迟的值来近似帕累托最优解。

图3：(a) 实时策略中作为重构延迟函数的通信时间，针对各种最小域带宽值。(b) 单次策略中最小延迟和比例带宽分配的GPU注入带宽函数的通信时间。

实时策略具有低时间复杂度，但只产生局部最优解。鉴于整个通信图在工作负载执行开始之前是已知的，可以通过在开始时预先计划所有重构来做出全局最优决策，并将重构事件插入到深度神经网络的执行调度中。

这种策略的基线情况涉及在工作负载执行开始时进行单次重构，其中根据每个域的总流量分配带宽，通过遍历通信图中的所有通信节点并对每个节点中的流量进行求和获得。直观地，分配给每个通信域的带宽应该与该域中的流量成比例。然而，这种比例带宽分配方案实际上并不能最小化整体通信时间。研究团队通过将带宽分配问题制定为网络效用最大化问题来展示这一点。

假设每个通信域在分配带宽时产生效用。目标是确定最大化效用的带宽分配，受容量约束的限制。为了最小化总通信时间，效用函数应该定义为每个通信域中流量大小与分配带宽之比的负数。这个效用函数是负数，因为目标是最大化效用，这对应于最小化由流量与带宽比表示的通信时间。可以使用拉格朗日函数求解这个优化问题。通过取拉格朗日函数的偏导数并将其设置为零，可以推导出拉格朗日乘数的表达式，然后求解带宽分配。

这个解决方案确保了对通信图上的完全或部分域流量需求的最优带宽分配，并将在需要带宽分配时在整个工作中应用。比较显示，前者始终设定性能上限，而比例分配仅在较高注入带宽时接近该上限。

图4：具有3个节点的通信子图中最优重构的动态规划算法示例，重构延迟为1 ms，每个GPU的可控带宽为1 GBps。

为了实现考虑工作负载执行期间重构的全局最优设计，蛮力方法涉及穷举搜索所有决策组合并选择最小化总通信时间的组合。在每个通信节点，重构决策是二元的：可以重构或不重构。因此，这个搜索算法具有指数算法复杂度O(2^N)，其中N是通信图中通信节点的总数。当N扩展到数千时，搜索空间的大小使得该算法难以处理。

为了降低时间复杂度，研究团队提出了一种从最优序列分割算法改编的动态规划方法。问题可以重新表述为找到最小化总通信时间的通信节点的最优分区，每个分区以网络重构开始。考虑通信图中的最后一个通信节点。这个最后的节点属于最优分区中的单个段，该段从某个较早的节点开始，在该节点发生重构。这种表述建议了一个递归解决方案：如果最后一个段的起点已知，可以将其消除并递归求解使用较早节点的最优分区。

从这里，可以定义递归关系OPT，其中OPT[j]表示段中所有节点的最小通信时间。段的总延迟表示为单三角矩阵L，其中L[i][j]是从节点i到节点j的总通信时间，重构发生在节点i。然后，完整的递归关系OPT可以表示为OPT[j]等于L[i][j]加上重构延迟加上OPT[j-1]的最小值，其中i从1到j。基本情况OPT[0]为0。注意这里每个节点的通信延迟是通过使用最小延迟带宽分配方案计算的。延迟矩阵L可以在O(N^3)时间内预先计算。

通过一个简单的例子来说明这个动态规划算法的操作。通信图有三个通信节点，位于三个不同的域中，每个节点的通信大小为1 MB。重构延迟设置为1 ms，每个GPU的可控带宽为1 GBps。在第一步中，计算并缓存带宽分配表，其中每个(i,j)项表示为从i开始到j结束的段分配给每个域的带宽比例。值得注意的是，当所有流量大小相等时（如本例中），比例带宽分配等同于最小延迟带宽分配。接下来，计算并缓存段延迟表L，其中每个(i,j)项表示使用带宽表中定义的带宽分配，从i开始到j结束的段内所有节点的通信延迟。

按照递归关系填充OPT数组。基于这些OPT值，发现最优重构策略需要在每个节点开始之前的每个域转换处进行重构，总通信时间为6 ms。如果只在开始时重构一次（即单次重构），将总带宽的三分之一分配给三个域中的每一个，则总通信延迟将为1 ms + 3 × 3 ms = 10 ms。如果重构延迟大于或等于3 ms，则单次重构将与在每个节点重构同样有效或更有效。

使用Giant OCS抽象来说明所提出重构策略的结果。图5显示了在4096个GPU上训练GPT3-175B模型的单次训练迭代期间的GPU域带宽分布，在1微秒和1毫秒重构延迟下比较实时和预先计划的策略。沿底部的红色垂直线表示由每种重构策略确定的各个重构事件。

图5：在单次训练迭代期间的域带宽分布，比较实时重构策略（a，c）与预先计划策略（b，d），针对不同的重构延迟（1微秒和1毫秒）。

对于实时重构，扫描一系列最小域带宽值并为每个重构延迟选择帕累托最优值。在重构延迟为1微秒时，帕累托最优值被选择为最小域带宽等于0，这使得在每次域转换时都能重构。由于重构开销低，预先计划的策略同样在每次域转换时重构，导致实时和预先计划策略的性能相同。在重构延迟为1毫秒时，与1微秒情况相比，实时方法重构频率较低，因为开销增加，但由于决策是局部最优的，因此仍然比预先计划的方法重构更频繁。在这种情况下，帕累托最优最小域带宽等于0.2，为三个域中的每一个保留20%的带宽，只允许分配总带宽的40%。相比之下，预先计划的策略最佳地平衡重构开销和域带宽分配，只需要少量重构即可最小化整体训练迭代时间。观察到数据并行域在整个迭代过程中保持带宽分配，因为专家并行全对全流量也映射到数据并行域。

然后定量比较不同重构策略的性能。图6a展示了实时和预先计划策略的迭代通信时间作为重构延迟的函数。对于实时策略，遍历一系列最小域带宽值，并为每个重构延迟选择最佳值以近似帕累托解。单次策略作为基线比较包括在内，代表当前最先进的策略，带宽分配根据优化公式给出。由于不在工作负载执行期间重构，因此在不同的重构延迟下保持恒定。

图6：实时、预先计划和单次重构的通信时间，作为(a)在450 GBps注入带宽下的重构延迟函数，以及(b)在1微秒重构延迟下的GPU注入带宽函数，其中单次不涉及工作负载内重构。

对于小于100微秒的重构延迟，通信时间几乎保持恒定，因为大多数重构事件有效地隐藏在通信间隔后面，允许工作负载内重构比单次方法实现高达2.3倍的改进。实时策略在较低的重构延迟（小于1毫秒）下密切跟踪预先计划的策略，因为在这种情况下的最佳策略是在每次域转换时重构。相比之下，单次策略在较高的重构延迟（大于或等于1毫秒）下接近预先计划策略的性能，因为高开销阻止了更频繁的重构。这表明可以在较低的重构延迟下使用实时近似最优性能（即预先计划），而在较高的重构延迟下使用单次。

图6b显示了当重构延迟为1微秒时，改变注入带宽时不同策略的性能。观察到随着每个GPU的总注入带宽增加，重构策略之间的性能差距缩小。这是因为较高的带宽减少了每个通信节点的通信时间，使得重构延迟在整个通信周期中相对更显著。当每个GPU配备足够高的带宽（即1600 GBps = 12.8 Tbps）时，如新兴的光电共封装链路技术所实现的，所有三种策略都收敛，表明在这种情况下不需要快速切换和工作负载内重构。

在使用Giant OCS抽象探索了各种重构策略之后，下一步是构建一个实用的拓扑，该拓扑考虑了光线路交换机端口数的约束和系统功耗。所提出的重构策略适用于任何集成了光线路交换机的拓扑。在本研究中，提出了一种多维直接连接拓扑，在网络规模和能效方面都提供强大的可扩展性。

常见的集成了光线路交换机的拓扑采用Fat-Tree或基于Torus的架构变体。在基于Fat-Tree的光网络中，光线路交换机插入在分层电子分组交换器之间或替换一些骨干电子分组交换器，带宽重构仅限于较高的聚合级别，导致低重构粒度。在深度神经网络工作负载中，每个GPU仅与其他GPU的有限子集通信，完全配置的Fat-Tree拓扑导致带宽利用不足和由于其广泛的聚合而增加的功耗。相比之下，以TPUv4 Pod为代表的基于Torus的光线路交换机拓扑只能重构外部环绕链路。这限制了只能基于预定义的并行维度重构连接，从而阻止在工作负载执行期间跨域进行更细粒度的带宽重构。

图7：(a) 具有n个维度和每个光线路交换机k个端口的OCSBCube拓扑。(b) n = 2，k = 2，l = 2的OCSBCube示例。(c) 2维OCSBCube中不同边缘可重构性程度的最大域带宽。

为了解决这些限制，提出将光线路交换机集成到BCube拓扑中，以创建具有所有维度全对全直接连接的可重构多维网络，与混合并行的流量模式对齐以提高带宽利用率。BCube(n, k)是一个递归定义的网络拓扑，其中k是交换机端口数，n是拓扑中的维度数。基本单元BCube(n = 1)通过将k个服务器连接到一个k端口电子分组交换器来构建。对于OCSBCube，用光线路交换机替换每个电子分组交换器，用GPU替换每个服务器，GPU配备可控光链路，这些链路通过嵌入式光收发器均匀分布在n个维度上。在每个维度内，每个链路连接到不同的光线路交换机。

一般的OCSBCube(n, k, l)由k端口光线路交换机构成，连接k个OCSBCube(n - 1, k, l)，总共有p = k^n个GPU和n级光线路交换机。OCSBCube(n, k)的直径为n，通常很小（例如，n = 2或3），因为端点数量随n呈指数增长。当每个维度的链路数大于或等于k - 1时，OCSBCube在每个维度中连接的GPU之间提供直接的全对全光连接。这有效地实现了具有n个维度的广义扁平蝶形拓扑，允许每个GPU直接与n(k-1)个对等体通信。

默认情况下，GPU注入带宽（即光纤）在n个维度上均匀分布。然而，正常BCube中的带宽不能跨维度重构，当流量模式需要跨维度的不均匀带宽时，会导致带宽利用不足。因此，提出使用波长可重构多端口收发器来替换GPU上嵌入的多个收发器。这种收发器设计利用密集波分复用梳状激光器技术提供的大量波长并行性，能够跨不同光纤重构波长，从而在网络边缘实现动态带宽分配。通过参数ε ∈ [0, 1]来表征边缘可重构性的程度，该参数反映了跨链路波长分配的灵活性。较高的ε表示更大的波长分配灵活性但成本增加，反之亦然。

使用这个参数，可以定义可分配给每个域的最大带宽为每个维度默认带宽的(1 + ε)倍。当没有边缘可重构性（ε = 0）时，OCSBCube(n = 2)中可分配给维度的最大带宽是总注入带宽的一半。相比之下，当所有带宽都可以从一个维度重构到另一个维度（ε = 1）时，最大可控带宽是全带宽。当每个维度的一半带宽可以重定向时，每个维度可以接收全带宽的最大75%。

除了边缘重构之外，还提出了一种多维作业映射策略，其中每个通信域在OCSBCube拓扑中尽可能多地映射到多个维度。具体来说，如果一个域需要pd个GPU，并且OCSBCube拓扑有n个维度，其中pd是n的整数倍且大于n，则该域映射到所有n个维度中的GPU，每个维度分配pd/n个GPU。与Google的TPUv4 Pod映射策略相比，后者将每个通信域分配给3D Torus内的单个维度，并将TPU注入带宽划分到这三个维度，所提出的映射策略能够同时在所有可用维度上传输单个域中的流量。这有效地使用全注入带宽并消除了域间重构的需要。

接下来将OCSBCube与Giant OCS基线进行比较，将GPT3-175B工作负载映射为需要边缘重构和多维作业映射以获得最佳性能。图8说明了OCSBCube和Giant OCS之间作为重构延迟函数的通信时间差异，在不同的边缘可重构性级别下进行评估。

图8：在不同OCSBCube边缘可重构性下，实时（a）和预先计划重构策略（b）的重构延迟函数通信时间。

对于两种重构策略，随着边缘可重构性达到最大值（即ε = 1），OCSBCube的性能收敛于Giant OCS。随着边缘可重构性降低，OCSBCube的性能下降，因为每个通信域中可重构带宽的数量减少了。当边缘可重构性达到ε = 0时，OCSBCube的性能大约减半，因为只有一半的注入带宽可以重构到通信域。注意到随着重构延迟增加，低和高边缘可重构性之间的性能差距缩小，表明随着重构频率降低，系统对边缘重构程度的敏感性降低。此外，发现边缘可重构性的变化对预先计划策略的影响小于对实时策略的影响，因为预先计划方法有效地降低了重构频率，以在较低的边缘可重构性下维持最佳性能。

数据中心网络成功的最终衡量标准涉及平衡性能与功耗。研究团队对OCSBCube、基于两级光线路交换机的Fat-Tree拓扑和TPUv4 Pod 3D-Torus设计进行了全面比较。这些评估考虑了完整的系统功耗，包括处理单元、收发器、交换机和网络组件，在不同的带宽配置级别上。

假设每个GPU的功耗为700 W。为了获得链路的功耗，将功率密度（焦耳/比特）乘以工作链路带宽。TPUPod和基于Fat-Tree的光网络中的电链路由电串行器/解串器供电，假设消耗5 pJ/bit。另一方面，假设光链路使用光电共封装收发器，根据最近在光I/O chiplet方面的工作消耗5 pJ/bit。假设72端口NVSwitch类型电子分组交换器的功耗为275 W，并应用线性缩放来考虑不同的端口数和每端口带宽。通过改变每端口功耗来建模光线路交换机功耗，遵循表1中显示的范围。通过将组件数量乘以相应的组件功率来计算每个网络的总功耗。

图9：在不同每端口光线路交换机功率下，具有4096个GPU的不同拓扑的功耗分解。

功耗分析显示了消耗模式的显著差异。基于Fat-Tree的光拓扑随着注入带宽的增加表现出快速扩展的功耗，主要由电子分组交换器和串行器/解串器功耗需求驱动。相比之下，OCSBCube和基于Torus的设计都显示出更渐进的功耗扩展，因为主要依赖带宽独立的光线路交换机。在较低的交换机功率水平下，收发器功耗主导直接拓扑的总消耗，而随着每端口消耗增加，交换机功耗成为主要贡献者。

为了评估光线路交换机拓扑的性能和功耗权衡，研究团队检查了三个不同规模的基于Transformer的训练工作负载：Megatron-5B（64个GPU = 4TP×2DP×8PP）、Anthropic-52B（256个GPU = 4TP × 4DP × 16PP）和GPT3-175B（4096个GPU = 4TP×32DP×32PP）。假设重构延迟为1微秒，每端口光线路交换机功率为3.8 W。

图10：不同拓扑在各种训练工作负载（Megatron-5B、Anthropic-52B、GPT3-175B）下的迭代时间、能耗和每焦耳token比较。

图10显示了迭代时间（顶行）、能耗（中行）和每焦耳token（底行）作为GPU注入带宽的函数。观察到OCSBCube实现了与基于Fat-Tree的光网络相当的迭代时间，但能耗降低（高达1.72倍），因为OCSBCube避免使用电子分组交换器和电串行器/解串器组件，这些组件的功耗随着带宽增加而快速扩展。另一方面，虽然OCSBCube和TPUPod在较高带宽下表现出相似的能耗，但OCSBCube能够实现更快的迭代时间（高达1.84倍），这归功于其更细粒度的执行内重构能力。这使得所有通信域都能有效地使用所有可用带宽，与TPUPod的固定3D-Torus拓扑和基于环的集体操作形成对比，后者将每个通信域的带宽使用限制为仅两个链路。

此外，OCSBCube在每焦耳token比率方面始终优于其他拓扑，改进高达1.75倍。这种效率优势源于拓扑通过细粒度重构精确匹配通信模式的能力，同时避免了过多电交换基础设施的功耗开销。在功率受限的数据中心环境中，这些改进直接转化为在固定功率预算内部署更多处理单元的能力，或者在保持性能的同时降低运营成本。

本研究建立了一个全面的框架，用于评估和设计专门针对人工智能工作负载优化的光可重构网络。定量方法为系统架构师提供了明确的指导，帮助他们在重构延迟、带宽配置、交换机技术和拓扑设计之间的复杂权衡中导航。随着光线路交换技术继续进步，具有更快的切换速度、更高的端口数和更低的功耗，该框架使得能够就哪些技术和策略组合将为特定部署场景提供最佳性能做出明智的决策。

研究结果表明，人工智能网络的未来可能会根据可用技术和系统约束遵循不同的路径。对于采用传统光线路交换机且重构延迟为毫秒级的系统，单次重构策略结合高带宽光链路提供了最佳的性能复杂性权衡。然而，随着新兴硅基光电子技术实现微秒或亚微秒级切换，更激进的工作负载内重构策略变得可行，可能在不产生过多开销的情况下提供显著的性能改进。

研究还强调了网络拓扑、重构策略和物理链路技术协同设计的重要性。像OCSBCube这样提供具有细粒度重构能力的多维直接连接的拓扑可以更有效地适应混合并行策略的结构化通信模式。当与网络边缘的波长可重构收发器结合时，这种设计在保持出色能效的同时实现了带宽分配的灵活性。

展望未来，从这项基础性工作中涌现出几个研究方向。多租户训练集群在资源共享和网络隔离方面提出了新的挑战，需要扩展这些重构策略以处理多样化的并发工作负载。随着系统扩展到数万个处理单元，故障弹性变得越来越关键，需要研究可重构网络如何在保持性能的同时适应组件故障。最后，开发专门针对可重构拓扑优化的集体通信算法可能进一步提高分布式训练工作负载中的数据移动效率。

ACTINA框架为实现光可重构网络在人工智能系统中的全部潜力提供了有价值的参考。通过提供严格的分析工具并展示实现优越性能功耗权衡的实用拓扑，该研究为能够满足人工智能需求同时保持在可接受功耗约束内的下一代数据中心网络提供了技术路径。

NOTICE