发表时间: 2025-01
加速LLM训练系统设计空间探索的多实验并行仿真
作者/机构: Fei Gui (清华大学, BNRist, and 清华大学深圳国际研究生院), Kaihui Gao (之江实验室), Li Chen (之江实验室), Dan Li (清华大学), Vincent Liu (宾夕法尼亚大学), Ran Zhang (之江实验室), Hongbing Yang (之江实验室), Dian Xiong (清华大学)
一句话结论 本文提出了首个基于 GPU 的大语言模型训练模拟器 Multiverse,通过单进程多实验并行策略与数据导向设计,在保证端到端仿真误差小于 3.0% 的前提下,将数万卡集群设计空间探索的仿真速度比现有基于 CPU 的方案提升了 43.1 至 73.2 倍。
要解决什么问题 随着大语言模型训练集群规模扩展至数万甚至数十万张 GPU,训练系统的设计空间(涵盖并行策略、网络拓扑和集合通信参数)急剧膨胀。为寻找最优配置(如搜索能让训练提速 3.4 倍的最优拓扑或微调并行组规模),工程师通常需运行上万次仿真实验。现有仿真方案卡在极低的执行效率上:采用多进程策略会因上下文切换和数据共享带来高昂同步开销;采用单进程单实验并行策略,同时启动大量实例会导致频繁的 CPU 上下文切换和严重缓存未命中。即使将多个实验合并到一个进程中执行,基于 CPU 的模拟器在处理海量独立实验时依然耗时过长,算力瓶颈和内存访问开销使得探索广阔设计空间的时间成本变得不可接受。
怎么做的 核心思路是将单进程多实验并行策略与数据导向设计结合,把整个仿真逻辑卸载到 GPU 上,利用单指令多数据特性进行大规模跨实验并行。为绕开 CPU 架构下的调度与缓存瓶颈,Multiverse 采用实体-组件-系统架构,将仿真对象拆分为纯粹的数据(组件)和处理逻辑(系统)。例如,训练任务和网络流被定义为实体,其状态按列式存储在全局内存表中,并带有隐式实验编号。这种跨实验的连续内存布局极大提升了相邻 GPU 线程访问数据的缓存命中率。在具体设计上,Multiverse 包含三个关键部件。首先是基于拉取的数据同步机制,用于解决网络仿真中常见的多对一写入冲突(如多个数据包同时发往同一端口)。传统做法需加锁或全局缓冲区,该机制将写入拆分为两步:发送方先在待办列表中注册计划,随后接收方主动从发送方拉取数据,实现完全无锁的并行处理。其次是高保真度的服务器内通信模型。由于张量并行等操作会在服务器内部产生大量通信,传统包级仿真极慢且现有分析模型误差大,Multiverse 引入校准后的分析模型 $y = \alpha + comm\_size / \beta$。其中 $\alpha$ 和 $\beta$ 是根据真实硬件(如 A100 或 H100)和具体集合通信操作类型离线拟合的经验参数,该模型不仅极快,还能结合重叠率精确模拟通信与计算争抢显存带宽导致的耗时延长。最后是内核融合技术。若为每个系统单独启动 GPU 内核,会产生巨大 CPU 与 GPU 同步开销。Multiverse 将所有仿真函数编译成统一的巨核,CPU 在每次批量仿真时仅需启动一次,GPU 内部所有线程便协同遍历整个系统执行图直至当前仿真步结束,从而彻底吃满 GPU 核心并消除频繁内核启动开销。
效果如何 实验在一台配备单块 H100 GPU 和 80 核 CPU 的服务器上进行,并以包含高达 1024 张 H100 GPU 的真实集群为基准测试平台。仿真的模型规模覆盖 13B 到 175B 参数。对比基线包括代表现有主流 CPU 路线的 ASTRA-sim 结合 UNISON(单进程单实验与单进程多实验)以及 ASTRA-sim 结合 DONS(单进程单实验与单进程多实验)。量化结果显示,在拓扑搜索、集合通信增强等真实用例中,Multiverse 仿真速度比所有基于 CPU 的基线快 43.1 至 73.2 倍。在模拟 54000 张 GPU 的超大规模集群时,单卡 Multiverse 依然保持 28.6 至 43.1 倍加速比。精度方面,得益于校准后的分析模型,Multiverse 服务器内集合通信耗时误差被控制在 1.2% 以下(基线误差高达 72.1%),端到端训练迭代时间与真实物理集群差异小于 3.0%。消融实验证明,基于拉取的同步和巨核技术分别带来最高 5.4 倍加速和 18.6% 耗时缩减。作者也指出其局限性:巨核设计要求每个执行节点有足够工作量,若并发实验数过少或集群规模太小会导致 GPU 利用率低下;此外,当单卡显存不足以容纳海量实验时,系统需跨多卡重新分配负载,此时跨卡执行是完全独立的。
核心问题: 随着大型语言模型(LLM)的快速发展,GPU集群的规模日益扩大,达到了数万甚至数十万级别。这极大地扩展了LLM训练系统的设计空间,涵盖了并行策略、通信参数、拥塞控制、网络拓扑等多个维度。为了寻找最优配置,研究人员需要进行大量的(例如,高达10k次)仿真实验。目前的仿真方法速度缓慢,无法满足高效探索设计空间的需求,而探索不足会导致训练性能显著下降(例如,次优拓扑可能导致训练迭代时间增加3.4倍)。
研究目标与创新点: 本文旨在解决设计空间探索中并行仿真实验效率低下的问题。
LLM训练基础设施的复杂性: LLM的训练依赖于包含数十到数十万个GPU的专用AI基础设施。构建高效的大规模训练系统需要在广阔的设计空间中进行决策,如图2所示。这包括:
仿真在设计空间探索中的作用: 由于物理测试平台成本高昂,而分析模型精度不足,仿真器成为验证系统设计的关键工具。两个实际案例说明了仿真的重要性:
多实验并行策略的分类: 论文将独立的仿真实验并行化策略分为四种类型:
四种策略的性能对比: 论文在一个配备80核CPU和一块H100 GPU的机器上,针对一个为128个GPU服务器训练GPT-3 13B模型寻找最优集合通信算法的场景,对这四种策略进行了实现和测试。实验结果如图3所示,得出以下结论:
Multiverse的核心思想: 基于上述观察,论文认为SPME+DOD是实现高性能探索工具的可行路径。DOD通过分离仿真逻辑与数据来提高SPME的执行效率。为此,论文引入了Multiverse,一个为AI训练系统实现SPME+DOD的模拟器,其核心思想包括:
Multiverse的组件: 如图4所示,Multiverse包含以下关键组件。用户只需编写应用代码,设置AI训练系统(包括工作负载、CCL参数、拓扑等),并指定要运行的多个实验,Multiverse会自动在GPU上执行多实验仿真。
实体与组件: 在AI训练的背景下,Multiverse保留了DONS中的网络实体(如Sender, IngressPort),并引入了新的关键实体,如训练任务(Task)和流(Flow)。实体的状态由其组件的值来表征。例如,如图5所示,定义Task实体的组件包括类型(计算或通信)、负载(计算时间或通信量)、前驱节点和后继节点。Flow实体则由状态、拥塞控制变量和发送缓冲区等组件来区分。具有相同组件的同类实体共享一个原型(archetype)。
系统与查询: 系统代表在实体集合上执行的数据并行计算(即仿真逻辑)。一个系统由一个查询和一个函数定义。查询用于指定输入中包含的组件数据,函数则对这些数据执行操作。例如,图5中的Schedule系统在每个Task上执行,它检查前驱节点的完成状态并激活后继任务节点。对于计算任务,它增加仿真时钟;对于服务器间通信任务,它在包级网络模拟器中输入新的流;对于服务器内通信任务,AnalyticalSys系统执行分析模型来估计通信时间。
系统执行图: 该图定义了在一个仿真步骤中需要执行的整套ECS系统,如图6所示。在一个仿真步骤中,Multiverse依次执行八个系统:Schedule、AnalyticalSys、SendSys、NICSndSys、ForwardSys、TransmitSys、NICRcvSys和ACKSys。Task实体执行Schedule系统后,服务器内通信由AnalyticalSys模拟,然后新流被注入网络模拟器。数据包经过发送、NIC、交换机转发路径,最终到达接收端,并根据需要返回确认(ACK)。此轮结束后,仿真进入下一个步骤,直至结束。只要仿真步长设置合理,这种执行模式可以保证仿真的正确性。
跨实验的组件存储管理: 在多实验仿真中,Multiverse集中管理所有组件数据的存储。它为所有实验中共享相同原型的实体构建一个单一的内存表,如图5所示。这些表以列式存储组织,以方便对连续组件的高效访问。为了处理每个实体的实验特定状态,Multiverse在每个表中包含一个隐式的ExpID组件。这种跨实验的存储方法提高了缓存命中率,因为相邻的GPU线程可以一致地访问组件数据,即使这些实体属于不同的实验。
在GPU中执行ECS系统: Multiverse使用NVIDIA CUDA C++编译器将系统编译到GPU上执行。每个系统调用直接映射到一个GPU线程,利用单指令多线程(SIMT)编程模型。当一个ECS系统被集成到执行图中时,Multiverse使用其查询来识别所有匹配的原型及其相应组件的列索引。这些信息使得Multiverse能够生成管理组件访问的代码,并在多个GPU线程上为每个匹配的实体调用ECS系统函数。例如,对于图5中的Send系统,如果查询在Flow表中识别出N行,Multiverse将调用入口函数N次,每次调用都会导致一个GPU线程执行ECS系统函数。
现有方法的局限性: 在LLM训练系统中,服务器内GPU间的TP流量占总网络负载的很大一部分,因此精确模拟TP流量至关重要。TP流量主要通过NVLink和PCIe等高带宽通道传输。现有的包级模拟器(如NS-3,DONS)由于通信协议的根本不同,难以准确模拟服务器内通信。而现有的AI训练模拟器(如ASTRA-sim)采用分析模型$y = \alpha + comm\_size / \beta$来估计通信时间,虽然速度快,但模型参数不准确,未能捕捉NCCL软件栈引入的开销,导致误差高达20%-72%。
Multiverse的校准分析模型: 论文通过大量实证测试发现,分析模型中的参数并非静态,而是需要根据不同场景变化。为此,论文测试了不同集合通信操作的通信大小与完成时间之间的关系,并通过拟合线性模型得出了修正后的参数。如图7所示,Multiverse提出的分析模型为每个集合通信操作精心调整了特定的参数($\alpha$和$\beta$),这些参数取决于服务器内GPU的数量和类型(如A100/H100)。校准后的模型使得服务器内通信的仿真误差仅为0.7%-1%,在不牺牲精度的前提下显著提升了速度。
模拟通信与计算的重叠: Multiverse还考虑了通信与计算重叠时因资源竞争(如SM和HBM带宽)对计算时间的影响。它根据计算与通信操作的重叠率以及由此导致的计算时间延长程度来调整计算时间。重叠率通过比较通信和计算操作的持续时间来计算,而计算时间的延长程度则通过一个考虑模型类型、计算操作和GPU类型的经验模型来评估。
多对一写入冲突问题: 在Multiverse中,nic_forward和forward等系统会频繁执行多对一的写操作。现有的ECS框架(如Unity ECS)由于缺乏原子操作,在这种场景下存在显著的局限性,通常需要将所有发送事件累积到全局缓冲区中,然后在主线程中顺序处理,这严重损害了并行效率。
基于拉取的同步策略: 为了解决这一挑战,Multiverse采用了一种基于拉取的数据同步策略。该方法将系统内的多对一写入过程分为两个阶段:set_write_plan和write。在set_write_plan阶段,待分发的数据包被初步列入一个待办队列。然后,在write阶段,每个目的地主动从源端拉取为其准备的数据包。例如,forward系统被分为set_forward_plan和forward两个子系统。前者为每个入口端口维护一个待办列表(通过位图实现),标记即将发生的转发事件。随后,在forward系统中,每个出口端口根据其待办列表独立地、逐一地从入口端口拉取数据包(如图8b所示)。这种方法使得所有端口能够以无锁方式执行forward系统,从而提高了并行效率。
传统内核启动的开销问题: 执行整个系统执行图的一个直接方法是为每个系统运行一个独立的CUDA内核。然而,这种策略存在一个显著问题:在GPU中,启动一个内核是一个耗时的过程,因为它需要由CPU初始化。如图9a所示,由于Multiverse包含众多ECS系统,这会导致大量的CPU-GPU同步开销,显著降低GPU效率。
巨核设计: 如图9b所示,Multiverse采用了一种由GPU驱动的方法和大规模内核设计。它将执行图中的所有系统编译成一个单一的CUDA内核,即巨核(megakernel)。该内核由CPU在每次批量仿真时启动一次,并在返回CPU之前完成整个执行图。为了管理系统执行和引擎级操作(如评估ECS查询),Multiverse创建了一个任务图。在执行过程中,所有GPU线程协同工作于任务图的同一个节点,完成后再移至下一个节点。虽然如果每个节点的工作量不足,这种方法可能导致GPU利用率低,但多实验的大批量仿真通常能为每个节点提供足够的工作量,从而轻松地充分利用所有GPU核心。
实现细节: Multiverse基于Madrona框架【52】实现,核心代码库约有13,000行C++代码,并已开源。目前,Multiverse支持TP和DP并行策略,以及包括Ring Allreduce、Allgather和Reducescatter在内的集合通信算法。此外,它还集成了DCQCN、HPCC、DCTCP等多种拥塞控制算法,以及ECMP和包喷洒等负载均衡算法。
min{link_delay},该设定借鉴自DONS【21】。
仿真速度对比: 如图10所示,在三个主要用例中,Multiverse始终显著优于所有其他方法。
最大规模与可扩展性:
服务器内通信仿真精度: 如图14所示,对于不同大小的集合通信操作,ASTRA-sim的误差可高达72.1%,而Multiverse的误差始终保持在1.2%以下,随着通信规模增大误差甚至低于0.8%。
端到端仿真精度: 如图15所示,在不同工作负载和集群规模下,Multiverse模拟的迭代时间与真实LLM训练任务的结果非常接近,即使在1,024 GPU规模下,差异也小于3.0%。随着模型增大(如从LLaMA 65B到GPT-3 175B),由于集合通信的消息尺寸更大,Multiverse的精度也相应提高。
各优化特性的贡献: 如图16所示,通过消融研究评估了Multiverse关键特性的影响:
本文论证了单进程多实验(SPME)并行策略通过减少调度开销和优化资源利用率实现了卓越的性能,但对于当前AI集群的规模仍显不足。为了提升SPME的效能,论文引入了Multiverse,一个新颖的基于GPU的AI训练模拟器。Multiverse利用DOD/ECS建模,并通过基于拉取的同步方法、高保真度的服务器内通信模型和巨核技术等优化措施,高效地利用了GPU的计算吞吐量。大量实验验证了Multiverse的准确性和效率,其速度比当前最先进的基于CPU的AI训练模拟器快43.1至73.2倍。