RFold:通过作业折叠与拓扑重构优化环面集群AI训练调度
1. 项目概述在构建大规模AI训练集群时我们常常面临一个核心矛盾如何让形态各异的机器学习作业高效地“住进”一个结构固定的硬件“宿舍”里。这里的“宿舍”指的是环面Torus拓扑集群它由成千上万个计算单元如GPU、TPU等统称XPU通过规则网格连接而成因其高带宽、低延迟的邻居通信特性非常适合分布式训练中频繁的AllReduce等集体通信操作。而“作业形态”则是指一个分布式训练任务对计算资源的排布要求比如一个需要做4路数据并行和6路模型并行的作业其理想的资源形状就是一个4x6的二维矩形。问题在于用户提交的作业形状千差万别而集群的物理拓扑是固定的。传统的调度策略比如“首次适应”First-Fit算法就像在停车场里找连续的空车位很容易因为形状不匹配而产生大量“碎片化”的空闲资源——明明总空闲车位很多但你的大巴车就是停不进去。另一种“尽力而为”Best-Effort的策略虽然能塞进更多车但可能把你的作业拆得七零八落导致内部通信要穿越其他作业的“领地”引发严重的网络拥堵拖慢整个训练过程。今天要聊的RFold就是卡内基梅隆大学和哈佛大学的研究团队提出的一种新思路。它不再把作业形状和集群拓扑看作一成不变的铁板而是让两者“双向奔赴”。简单说RFold干了两件聪明事一是“折叠”作业即在不影响通信效率的前提下灵活变换作业的资源排布形状二是“重构”集群利用光路交换技术动态调整部分网络连接。通过这种协同适配RFold能在多租户环面集群中同时实现高资源利用率和高作业性能。根据他们在4096节点模拟器上的初步评估相比传统方法RFold能将集群绝对利用率提升57%并将作业完成时间最多缩短11倍。这对于动辄拥有数万张加速卡、电费惊人的大型AI集群来说意味着巨大的成本节约和效率提升。2. 核心问题与现有方案的局限2.1 环面拓扑与ML作业的“形状约束”要理解RFold的价值得先看清它要解决的核心难题。现代大规模机器学习训练普遍采用数据并行、张量并行、流水线并行等策略。这些并行策略直接决定了作业对计算资源的拓扑形态需求。例如一个采用4路数据并行和2路张量并行的作业其最优的资源分配就是一个4x2的网格。在这个网格里同一行的4个XPU进行数据并行的梯度同步同一列的2个XPU进行模型并行的张量通信。这种对特定“形状”的需求我们称之为形状约束。环面拓扑天然适合这种网格状的通信模式。在一个3D环面中每个XPU在X、Y、Z三个维度上都与邻居直接相连并且首尾相接形成环。这种结构为邻居通信提供了高带宽、低延迟的通道。然而它的优势也带来了调度上的挑战为了获得最佳性能作业希望被分配到一个连续的、形状匹配的XPU子集上。这里的“连续”指的是在环面网格中物理位置相邻“形状匹配”则是指分配的子集在三维空间上的长宽高恰好符合作业的并行需求。2.2 传统调度策略的两难困境面对形状约束现有的调度器主要采取两种策略但各有明显短板。策略一首次适应First-Fit及其变种。这种策略严格遵循形状约束只在集群中寻找一块连续的、形状完全匹配的空闲区域来放置作业。它的优点是能保证作业获得独占的、无竞争的通信链路性能最优。但缺点也极其突出资源碎片化。想象一下集群就像一个三维的俄罗斯方块游戏盘不断有各种形状的方块作业落下和消失。First-Fit策略会导致盘面上出现大量零散的小空洞即使总的空闲格子数足够放下一个新方块也可能因为找不到一块连续的、形状合适的区域而无法放置。结果就是集群利用率低下作业排队等待时间变长。策略二尽力而为Best-Effort放置。这种策略为了追求高利用率放宽了“连续性”要求。它允许将一个作业的XPU分散在集群的不同位置。这样做确实能填满更多碎片但代价是通信性能的严重下降。在环面中非相邻的XPU通信需要经过中间节点转发。当作业的XPU分散各处时其内部通信流将不可避免地与其它作业的通信流共享并竞争网络链路带宽。对于通信密集型的ML训练作业这种网络争用会显著增加每次迭代的通信时间从而拖慢整个训练过程。这就形成了一个经典的“鱼与熊掌不可兼得”的局面要保证性能无竞争就得牺牲利用率产生碎片要提高利用率填满碎片就得牺牲性能引入竞争。现有的调度器通常只能优化其中一个目标。2.3 可重构环面引入新的可能性转机来自于硬件层面的创新以谷歌TPU v4集群为代表的设计引入了光路交换技术。在这种设计中一个大规模的静态环面如16x16x16被分解为许多个硬连线的、更小的立方体如4x4x4的“Cube”。这些立方体之间的连接不再是固定的而是通过光路交换阵列动态建立。这就赋予了集群一种新的能力拓扑重构。调度器可以为了适配某个特定形状的作业临时将几个小立方体“拼接”成一个更大、形状不同的逻辑环面。例如一个需要4x4x32形状的作业在静态16x16x16环面中无法放置因为Z维度32超过了最大值16但在可重构环面中可以通过将8个4x4x4立方体在Z维度上首尾相连来实现。可重构性为解决形状不匹配问题打开了一扇门但它并非万能。首先重构的粒度受限于硬连线立方体的大小。其次重构本身也有成本并且过于频繁或细粒度的重构可能不切实际。最关键的是即使有了重构如果作业形状不是立方体维度的整数倍仍然会面临内部无法形成完整通信环即缺少“回绕链路”的问题导致性能损失。这就需要软件调度层面更智能的配合而RFold正是将硬件可重构性与软件智能调度相结合的产物。3. RFold的设计哲学与核心技术RFold的核心思想非常直观既然作业形状和集群拓扑都可以在一定程度上改变那么何不让它们相互适应共同寻找一个“最佳契合点”这个“最佳契合点”需要同时满足两个条件一是让作业的XPU能形成高效的通信环避免性能损失二是尽可能紧凑地利用集群空间减少碎片。RFold通过两大技术支柱来实现这一目标折叠与重构。3.1 技术支柱一作业形状折叠“折叠”是一种改变作业资源排布形状但同时保持其通信语义等价的技术。它的理论基础是图同态。我们可以把一个作业所需的通信模式抽象成一个图其中节点代表XPU边代表需要高带宽通信的XPU对。同样集群中可用XPU的物理连接关系也可以构成一个图。折叠的本质就是在集群图中寻找一个子图使其与作业的通信图同构或同态然后将作业映射到这个子图上。这听起来有些抽象我们来看几个具体例子1D作业折叠一个纯数据并行的作业其形状为A x 1 x 1意味着A个XPU需要形成一个一维环进行AllReduce。在静态环面中如果A不是环面某维度大小的整数倍就无法获得首尾相接的回绕链路通信性能会打折扣。通过折叠我们不再强求这A个XPU必须在一条直线上。只要能在集群的可用XPU中找到任意A个节点它们能通过物理链路连接成一个环那么这个作业就能被高效放置。这相当于把一个一维的线“折叠”成了二维或三维空间中的一个环。2D作业折叠考虑一个1 x 6 x 4形状的作业它在Y和Z两个维度上进行通信。直接放置可能面临两个问题Y维度的6个XPU无法在4x4x4的立方体边界形成回绕环同时可能产生资源碎片。通过折叠我们可以将这个2D作业变换为一个4 x 2 x 3的3D形状。原来在Z维度的通信被映射到新形状的X维度原来在Y维度的通信则被巧妙地映射到新形状的Y和Z维度上形成的一个虚拟环中。这样作业就能被完整地塞进一个4x4x4的立方体既解决了回绕问题又消除了碎片。3D作业折叠对于三个维度都大于1的“真3D”作业折叠空间较小但依然存在可能。例如一个4x8x2的作业通过折叠可以变形为4x4x4。其核心在于对通信路径进行重映射将原有作业中不同层的通信分别映射到新形状中通过回绕链路或内部链路形成的环上。折叠的威力在于它极大地扩展了调度器的搜索空间。对于一个形状为S的作业调度器不再只寻找与S完全匹配的空闲区域而是可以寻找所有与S通信等价的形状变体所匹配的区域。这就像玩拼图时不仅考虑原图块还能考虑其旋转、翻转后的形态大大提高了找到合适位置的概率。3.2 技术支柱二拓扑动态重构折叠是从软件层面让作业去适应集群。而重构则是从硬件层面让集群去适应作业。如前所述基于光路交换的可重构环面允许调度器在运行时改变部分网络连接。RFold利用这种能力来服务于一个核心启发式原则最优的放置方案应该消耗最少数量的可重构立方体和OCS链路。OCS链路是宝贵的、可灵活配置的资源而立方体内部的硬连线链路是固定的。因此调度算法会尝试多种将大作业形状分解为小立方体组合的方案。例如放置一个4x4x32的作业算法会生成诸如“8个4x4x4立方体排成一条线”或“4个4x4x8的逻辑块”等多种分解选项。然后它为每一种分解方案在集群中寻找匹配的、空闲的立方体集合。最后选择那个所需立方体总数最少、且所需OCS连接最少的方案并触发光路交换完成拓扑重构。重构与折叠是相辅相成的。重构为折叠提供了更灵活的“画布”——通过拼接立方体可以创造出更多样化的逻辑拓扑来容纳折叠后的作业形状。而折叠则能弥补重构的不足——当作业形状无法被立方体整数倍分解时通过折叠将其变形可能使其恰好能放入单个立方体或更少的立方体中从而减少对OCS链路的依赖并可能获得更好的内部连接性。3.3 RFold调度算法的工作流程结合了折叠与重构两大技术后RFold的调度流程可以概括为以下几步形状变体生成当一个新的作业到达时调度器首先基于图同态理论枚举出该作业原始形状所有可能的、通信等价的折叠变体。例如对于一个18x1x1的作业可能的变体包括能形成18个节点环的所有2D或3D排布方式。虚拟重构与方案搜索对于每一个形状变体调度器会在当前集群的空闲资源图上虚拟地执行重构操作。即尝试寻找一种分解该变体形状为若干立方体组合的方式并检查是否存在对应的空闲立方体集合能满足这种组合。这个过程会生成多个候选的分配计划。方案评估与排序对每一个候选分配计划计算其“成本”核心指标就是消耗的可重构立方体数量和OCS链路数量。优先选择成本最低的计划。决策与执行选择最优计划并执行两件事一是通过OCS进行实际的物理拓扑重构如果需要二是在重构后的逻辑拓扑上将作业的各个进程映射到具体的XPU上。这个流程的核心是协同优化。它不是在固定的拓扑上硬塞作业也不是盲目地重构拓扑去迎合原始形状而是让“折叠”和“重构”两个自由度协同搜索找到全局最优的匹配点。4. 实现考量与性能分析4.1 系统实现的关键挑战将RFold从理论转化为实际可用的调度器需要解决几个工程上的挑战1. 同态形状的高效枚举穷举一个复杂通信图的所有同态子图可能是指数级复杂的。RFold需要利用ML作业通信模式的规律性通常是网格或环来设计高效的算法。对于标准的DP、TP、PP组合形成的网格状作业其形状变体可以通过整数因子分解和维度重映射来系统性地生成而非通用的图同构搜索。2. 资源搜索与碎片评估在拥有数千个节点的大规模集群中快速搜索可行的放置方案是关键。这可以借鉴高性能计算中处理器分配算法的思想如基于空间填充曲线如希尔伯特曲线将3D资源映射到1D进行快速线段查找再结合折叠和重构的约束进行剪枝。评估“碎片化”也需要一个合适的度量RFold采用“消耗立方体数”作为核心启发式因为它直接关联到OCS资源的使用和未来作业的放置难度。3. 与现有调度框架的集成RFold主要是一个放置策略。在实际系统中它需要与上层的作业接纳策略协同工作。例如可以与基于排队论的接纳控制器结合当RFold判断当前集群状态无法高效放置某个作业时可以反馈给接纳层决策是让其排队等待还是考虑采用“尽力而为”模式启动。RFold也可以与支持弹性伸缩的调度器如Optimus、Gandiva结合在作业运行期间根据资源碎片情况动态调整其资源形状和位置。4. 拓扑重构的开销光路交换的重构时间通常在毫秒到秒级。虽然相比长达数小时甚至数天的ML训练作业来说很短但频繁重构仍需考虑。调度器需要权衡重构带来的放置收益与重构本身的开销及对其它作业可能造成的短暂扰动。一种策略是进行批处理将一段时间内到达的多个作业的放置需求一并考虑规划一次重构来同时满足多个作业从而摊薄开销。4.2 模拟实验与结果解读研究团队通过一个自定义的作业级离散事件模拟器来评估RFold。他们构建了一个4096节点的可重构3D环面集群由64个4x4x4立方体组成并扩展了公开的ML作业追踪如Microsoft Philly trace来生成包含不同形状需求的作业流。他们对比了四种策略FirstFit在静态环面中使用首次适应算法。Folding在静态环面中仅使用折叠技术。Reconfig在可重构环面中仅使用重构技术。RFold在可重构环面中结合使用折叠与重构。评估的核心指标有三个作业完成率成功被调度执行的作业比例。作业完成时间从作业提交到完成的时间。集群利用率集群中XPU处于忙碌状态的时间比例。模拟结果清晰地展示了RFold的优势在作业完率上FirstFit由于严格的形状约束只能调度约10.4%的作业。仅使用折叠Folding能将完成率提升到44%因为它让更多形状不匹配的作业得以“变形”安置。仅使用重构Reconfig立方体大小为4x4x4则能达到100%的完成率因为理论上任何形状都可以通过拼接立方体来实现。RFold同样能达到100%的完成率。这里的关键差异体现在后两个指标上。在作业完成时间上RFold展现了巨大优势。在4x4x4立方体的配置下RFold的中位作业完成时间比仅使用重构的策略快了11倍尾部p99作业完成时间也快了2倍。这是因为RFold通过折叠使得作业更有可能被放置到单个或少数几个立方体内从而获得了完整的、无竞争的环内通信链路。而仅靠重构的作业其XPU可能分布在多个立方体跨立方体的通信需要通过OCS链路带宽和延迟可能不如立方体内硬连线链路且可能受其他作业影响。在集群利用率上FirstFit和Folding由于调度成功率低平均利用率很难超过40%。Reconfig凭借高调度成功率能将利用率提升到较高水平。而RFold在Reconfig的基础上通过折叠进一步优化了资源打包密度将平均集群利用率额外提升了20%。这意味着RFold能用同样的硬件同时运行更多的作业或者让作业更快完成从而释放资源。这些结果强有力地证明折叠与重构的协同效应不是简单的叠加而是产生了“112”的效果。它同时攻克了资源碎片化和网络争用这两个长期困扰环面集群调度的问题。4.3 实际部署的注意事项如果你正在管理或设计一个环面拓扑的AI集群并考虑引入RFold或类似思想以下几点值得注意硬件依赖性RFold的拓扑重构能力依赖于OCS或类似的可重构网络技术。如果你的集群是传统的静态互连那么只能应用折叠技术收益会打折扣但依然有价值。作业特征分析RFold的收益大小与作业负载的特征密切相关。如果集群中运行的都是形状规整如2的幂次方且尺寸较小的作业传统FirstFit的碎片化问题可能不严重。反之如果作业形状多样、大小不一RFold的优势将非常明显。因此在部署前需要仔细分析历史作业轨迹。调度器开销RFold的搜索空间比传统策略大得多可能增加调度决策的延迟。需要设计高效的启发式算法和剪枝策略确保调度决策能在毫秒级完成避免成为系统瓶颈。与弹性训练的结合新兴的ML框架支持动态调整并行度。未来调度器或许可以与训练框架联动不仅调整作业的放置还能在合理范围内建议或指导作业调整其并行策略即形状实现更深层次的协同优化。5. 未来展望与扩展思考RFold为环面集群的调度打开了新的思路但这项研究远未结束它引出了更多值得探索的方向。超越严格连续放置RFold目前追求的是无竞争的连续放置。但在实际中有时“立即以非连续方式启动作业承受一定的网络竞争减速”可能比“长时间排队等待一个连续位置”更划算。未来的调度器可能需要一个更精细的模型能够量化网络竞争带来的性能损失并与排队延迟进行权衡做出最优决策。这有点类似操作系统CPU调度中的“响应时间”与“吞吐量”的权衡。重构粒度的权衡本文主要基于4x4x4的立方体进行讨论。但立方体的大小是一个关键的设计选择。更大的立方体如8x8x8意味着更少的OCS端口需求集群可扩展性更好但重构灵活性下降。更小的立方体如2x2x2则相反能提供更细粒度的重构能力可能进一步提升资源利用率但对OCS的端口数和控制复杂度要求更高。如何为特定的工作负载和成本预算选择最优的立方体大小是一个需要深入研究的问题。超越3D环面当前工作和主流硬件都聚焦于3D环面拓扑和最多3维的作业形状。但随着模型规模和并行策略的演进是否会出现更高维度的通信模式RFold中“折叠”的思想本质上是图映射理论上可以扩展到更高维度的拓扑和作业形状。此外其他网络拓扑如胖树、Dragonfly与可重构技术的结合也是值得探索的领域。与异构资源调度的融合现代AI集群往往是异构的包含不同代际、不同内存容量的加速卡。RFold目前主要关注拓扑和形状未来可以与考虑计算异构性的调度器如Gavel、Sia结合在满足形状和拓扑约束的同时还能将作业匹配到最合适的硬件类型上实现多维度的资源优化。从我个人的实践经验来看集群调度问题的复杂性在于它处于系统软件、网络硬件和应用需求的交叉点。RFold的出色之处在于它没有孤立地看待任何一个层面而是通过“折叠”和“重构”这两个杠杆巧妙地撬动了应用需求与硬件供给之间的僵局。它告诉我们在软硬件协同设计的大趋势下调度器的创新不能只停留在算法层面更需要深刻理解底层硬件的能力和上层应用的特征。对于正在构建或运营大规模AI基础设施的团队来说关注这类拓扑感知的调度技术很可能是在不增加硬件资本支出的前提下获取显著性能提升和成本节约的下一个关键。