Accelerating Design Space Exploration for LLM Training Systems with Multi-experiment Parallel Simulation
加速LLM训练系统设计空间探索的多实验并行仿真
作者/机构: Fei Gui (清华大学, BNRist, and 清华大学深圳国际研究生院), Kaihui Gao (之江实验室), Li Chen (之江实验室), Dan Li (清华大学), Vincent Liu (宾夕法尼亚大学), Ran Zhang (之江实验室), Hongbing Yang (之江实验室), Dian Xiong (清华大学)
A1 主要贡献
核心问题: 随着大型语言模型（LLM）的快速发展，GPU集群的规模日益扩大，达到了数万甚至数十万级别。这极大地扩展了LLM训练系统的设计空间，涵盖了并行策略、通信参数、拥塞控制、网络拓扑等多个维度。为了寻找最优配置，研究人员需要进行大量的（例如，高达10k次）仿真实验。目前的仿真方法速度缓慢，无法满足高效探索设计空间的需求，而探索不足会导致训练性能显著下降（例如，次优拓扑可能导致训练迭代时间增加3.4倍）。
研究目标与创新点: 本文旨在解决设计空间探索中并行仿真实验效率低下的问题。
1. 识别并验证了高效的并行仿真策略: 论文通过深入分析和实验，对比了四种多实验并行策略（如图1所示），发现单进程多实验（Single-process Multi-experiment, SPME）策略因其能减少调度开销和优化资源利用（如缓存友好）而性能最优。
2. 提出首个基于GPU的AI训练模拟器Multiverse: 为了充分发挥SPME的潜力，论文提出了Multiverse。该模拟器基于数据导向设计（Data-Oriented Design, DOD）原则，并利用GPU的单指令多数据（SIMD）特性进行大规模并行加速。
3. 设计了三种关键优化技术以提升GPU效率:
* 基于拉取（pull-based）的同步方法: 用于实现无锁数据传输，解决多对一写入冲突问题。
* 高保真度的服务器内通信模型: 通过对真实硬件进行剖析，建立了一个精确的分析模型来模拟GPU之间的通信，显著提高了仿真精度。
* 内核融合（kernel-fusion）技术: 将多个仿真函数编译成一个统一的GPU巨核（megakernel），大幅减少CPU-GPU同步开销，提升GPU效率。
4. 实验验证: 实验表明，Multiverse在多种用例和集群规模下，比当前最先进的基于CPU的模拟器快43.1至73.2倍。在高达54,000个GPU的集群仿真中，其结果与真实LLM训练的差异小于3.0%。
A3 背景知识与关键观察
2.1 探索LLM训练的设计空间
LLM训练基础设施的复杂性: LLM的训练依赖于包含数十到数十万个GPU的专用AI基础设施。构建高效的大规模训练系统需要在广阔的设计空间中进行决策，如图2所示。这包括：
* 并行策略: 主流框架如Megatron【54, Megatron-lm: Training multi-billion parameter language models using model parallelism, 2019, arXiv】和DeepSpeed【48, Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters, 2020, KDD】提供了数据并行（DP）、流水线并行（PP）和张量并行（TP）等多种策略。
* 集合通信: All-reduce和All-gather等操作在GPU间同步结果中扮演关键角色。
* 底层支撑: 传输协议和物理网络拓扑对训练性能有根本性影响。
仿真在设计空间探索中的作用: 由于物理测试平台成本高昂，而分析模型精度不足，仿真器成为验证系统设计的关键工具。两个实际案例说明了仿真的重要性：
* 拓扑搜索: 经典拓扑（如ROFT【43】、Torus【29】、Fattree【5】）可能无法满足新型模型和训练技术的需求。研究人员通常需要调整拓扑参数（如层数、端口数、连接关系）并进行超过1万次仿真实验来寻找最优拓扑。一个优越的拓扑可将训练速度提升3.4倍【58, 59】。
* 并行组规模优化: 将GPU分配到TP/PP/DP组是一个关键决策，因为TP需要高带宽，而PP需避免流水线气泡。根据模型大小和硬件拓扑，比较不同组规模通常需要近百次仿真实验【36】。尽管可以将优化分为长期（拓扑搜索）和短期（其他优化）阶段，但短期内仍需大量实验，因此提高仿真速度至关重要。
2.2 多实验并行化
多实验并行策略的分类: 论文将独立的仿真实验并行化策略分为四种类型：
1. 单进程单实验（SPSE）: 最常用的方法，启动多个进程，每个进程执行一个实验。例如，NS-3【49】和OMNeT++【57】默认使用单进程单线程，而DONS【21】和UNISON【8】使用多线程加速单个实验。
2. 多进程单实验（MPSE）: 使用多个进程运行一个仿真实验，如带有MPI的NS-3或OMNeT++。
3. 单进程多实验（SPME）: 在一个程序内执行多个实验，以节省进程固有的开销。
4. 多进程多实验（MPME）: 使用多个进程执行一个SPME程序。
四种策略的性能对比: 论文在一个配备80核CPU和一块H100 GPU的机器上，针对一个为128个GPU服务器训练GPT-3 13B模型寻找最优集合通信算法的场景，对这四种策略进行了实现和测试。实验结果如图3所示，得出以下结论：
* MP*E策略的同步开销高: MPSE和MPME性能最差，完成500个实验需要2000小时。基于进程的并行化因上下文切换导致数据共享开销巨大，高昂的同步开销降低了仿真性能。
* SPSE策略的缓存未命中率高: 尽管UNISON【8】和DONS【21】在单实验中通过多核并行获得了超线性加速，但在多实验场景下，简单地并行运行多个模拟器实例会导致频繁的CPU上下文切换和严重的缓存未命中，从而延长了每个实验的耗时。
* SPME策略性能更优: 通过在UNISON和DONS中配置多个拓扑来实现SPME，性能得到提升，这得益于跨实验的缓存未命中率和进程调度开销的降低。然而，即使采用SPME，DONS完成该案例仍需约370小时，速度仍然过慢。
* SPME的潜力需要DOD和GPU来释放: SPME对基于DOD的DONS性能提升显著。DOD将逻辑与数据分离，使得SPME+DOD能够发掘跨实验的批处理和并行化机会。这种模式与GPU的SIMD特性完美契合，GPU的大量核心可以进一步加速跨多实验的并行处理。如图3所示，SPME+DOD+GPU的性能比CPU上的最佳方案快70倍以上。
2.3 为AI训练系统实现并加速SPME+DOD
Multiverse的核心思想: 基于上述观察，论文认为SPME+DOD是实现高性能探索工具的可行路径。DOD通过分离仿真逻辑与数据来提高SPME的执行效率。为此，论文引入了Multiverse，一个为AI训练系统实现SPME+DOD的模拟器，其核心思想包括：
1. ECS建模: Multiverse采用实体-组件-系统（Entity-Component-System, ECS）架构【3, 17】来实现DOD原则。ECS将对象分解为独立的逻辑（系统）和数据（组件），这已被证明是一种实现DOD的有效方式。
2. 在GPU中实现SPME+DOD仿真: SPME+DOD允许模拟器利用跨实验的并行性和一致性，在GPU这样的高吞吐量并行处理器上实现高效率。Multiverse将所有仿真方面都委托给GPU，集中管理GPU内的存储，并将所有ECS逻辑编译到GPU上，自适应地利用所有核心执行跨所有实验的相同仿真逻辑。
3. 三项优化技术:
* 基于拉取的同步方法: 减少锁开销。
* 高保真度服务器内通信模型: 通过分析模型精确模拟GPU间通信。
* 内核融合技术: 将仿真函数编译成统一的GPU巨核（megakernel），显著提高GPU效率。
A2 方法细节
3.1 架构
Multiverse的组件: 如图4所示，Multiverse包含以下关键组件。用户只需编写应用代码，设置AI训练系统（包括工作负载、CCL参数、拓扑等），并指定要运行的多个实验，Multiverse会自动在GPU上执行多实验仿真。
* 系统模拟器 (System Simulator): 控制和调度AI训练仿真过程。它接收类似于ASTRAsim【47】的模型工作负载【2】作为输入，该工作负载描述了每个GPU的计算图。图中的节点是计算操作（如Embedding）或集合通信操作（如Allreduce）。计算操作的时间已被标注。Multiverse支持TP、PP、DP等典型并行策略。对于服务器间的集合通信，系统模拟器通过实现NCCL【23】的集合通信算法（CCA）生成一系列点对点收发流。为了提高精度，它通过劫持NCCL API来剖析每个通信的开始和结束时间，从而校准服务器间通信仿真的开销。
* 服务器内通信模拟器 (Intra-server Communication Simulator): 对于GPU间的服务器内集合通信（如TP通信），Multiverse直接使用提出的分析模型进行模拟。该模型根据操作类型和GPU类型具有不同的经验参数，实现了快速而准确的模拟。
* 服务器间网络模拟器 (Inter-server Network Simulator): 系统模拟器在该网络模拟器中注册跨网络的点对点通信需求。网络模拟器会进行离散事件仿真（DES）来严格执行包级事件，确保正确性，类似于NS-3【49】。当流完成后，它会通知系统模拟器。
* GPU内存模拟器 (GPU Memory Simulator): Multiverse支持模拟GPU内存使用情况，以检查某些实验设置（如并行组大小）是否会超出GPU内存限制。在实际仿真前，它会进行检查，如果超出限制则生成OOM错误。
* GPU运行时 (GPU Runtime): 运行时利用ECS抽象将多实验仿真任务高效地分配给GPU。它首先将ECS系统函数与组件管理的包装代码一起编译成一个统一的GPU巨核。然后，ECS系统的执行图被传输到GPU，并在每个仿真步骤中由巨核处理。GPU会自适应地利用所有核心来处理跨所有实验的相应实体。
3.2 AI训练系统的ECS建模
实体与组件: 在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）。此轮结束后，仿真进入下一个步骤，直至结束。只要仿真步长设置合理，这种执行模式可以保证仿真的正确性。
3.3 GPU中的多实验仿真
跨实验的组件存储管理: 在多实验仿真中，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系统函数。
4.1 服务器内通信仿真
现有方法的局限性: 在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类型的经验模型来评估。
4.2 基于拉取的数据同步
多对一写入冲突问题: 在Multiverse中，nic_forward和forward等系统会频繁执行多对一的写操作。现有的ECS框架（如Unity ECS）由于缺乏原子操作，在这种场景下存在显著的局限性，通常需要将所有发送事件累积到全局缓冲区中，然后在主线程中顺序处理，这严重损害了并行效率。
基于拉取的同步策略: 为了解决这一挑战，Multiverse采用了一种基于拉取的数据同步策略。该方法将系统内的多对一写入过程分为两个阶段：set_write_plan和write。在set_write_plan阶段，待分发的数据包被初步列入一个待办队列。然后，在write阶段，每个目的地主动从源端拉取为其准备的数据包。例如，forward系统被分为set_forward_plan和forward两个子系统。前者为每个入口端口维护一个待办列表（通过位图实现），标记即将发生的转发事件。随后，在forward系统中，每个出口端口根据其待办列表独立地、逐一地从入口端口拉取数据包（如图8b所示）。这种方法使得所有端口能够以无锁方式执行forward系统，从而提高了并行效率。
4.3 使用巨核（Megakernel）执行
传统内核启动的开销问题: 执行整个系统执行图的一个直接方法是为每个系统运行一个独立的CUDA内核。然而，这种策略存在一个显著问题：在GPU中，启动一个内核是一个耗时的过程，因为它需要由CPU初始化。如图9a所示，由于Multiverse包含众多ECS系统，这会导致大量的CPU-GPU同步开销，显著降低GPU效率。
巨核设计: 如图9b所示，Multiverse采用了一种由GPU驱动的方法和大规模内核设计。它将执行图中的所有系统编译成一个单一的CUDA内核，即巨核（megakernel）。该内核由CPU在每次批量仿真时启动一次，并在返回CPU之前完成整个执行图。为了管理系统执行和引擎级操作（如评估ECS查询），Multiverse创建了一个任务图。在执行过程中，所有GPU线程协同工作于任务图的同一个节点，完成后再移至下一个节点。虽然如果每个节点的工作量不足，这种方法可能导致GPU利用率低，但多实验的大批量仿真通常能为每个节点提供足够的工作量，从而轻松地充分利用所有GPU核心。
5. 实现
实现细节: Multiverse基于Madrona框架【52】实现，核心代码库约有13,000行C++代码，并已开源。目前，Multiverse支持TP和DP并行策略，以及包括Ring Allreduce、Allgather和Reducescatter在内的集合通信算法。此外，它还集成了DCQCN、HPCC、DCTCP等多种拥塞控制算法，以及ECMP和包喷洒等负载均衡算法。
* 仿真步长: Multiverse在一个仿真步骤中遍历执行图，并增量推进仿真时钟。为避免时间误差，步长（即并行DES中的lookahead概念）被设置为min{link_delay}，该设定借鉴自DONS【21】。
* 多GPU使用: 当单个GPU资源不足时，Multiverse可以利用多个GPU进行并行探索实验。它会评估是否需要将独立的实验重新分配到不同的GPU上，并在必要时自动平衡多GPU间的负载以最大化仿真速度。跨GPU的执行是相互独立的。
A4 实验环境
- 硬件配置:
- 仿真服务器: 一台Linux服务器，配备1块NVIDIA H100 GPU，80核Intel [email protected]，256GB内存。
- 真实测试平台: 一个基于Fattree拓扑和RoCEv2网络的集群，包含128台服务器。每台服务器配备8块NVIDIA H100 GPU和8个Mellanox ConnectX-7网卡（2x200Gbps）。服务器内GPU通过NVLink连接，提供900GBps带宽。集群规模包括128 GPU和1,024 GPU两种。
- 模型与数据集:
- 在真实测试平台上运行LLaMA 65B和GPT-3 175B模型进行训练，以验证仿真精度。
- 仿真实验涵盖了从13B到175B参数规模的模型。
- 软件配置:
- 对比方案: ASTRA-sim+UNISON (SPSE/SPME), ASTRA-sim+DONS (SPSE/SPME), Multiverse (SPSE)。
- Multiverse基于Madrona框架，用C++实现。
- 实验用例: 设计了四个真实的用例进行性能测试，涵盖了拓扑搜索、集合通信增强、并行组大小微调和拥塞控制算法评估，涉及的集群规模从128到54,000个GPU，模型规模从13B到175B参数。具体设置见表1。
A4 实验结果
6.1 Multiverse的速度优势
仿真速度对比: 如图10所示，在三个主要用例中，Multiverse始终显著优于所有其他方法。
* 在模拟128、1,024和8,192 GPU的集群时，Multiverse分别实现了高达73.2倍、67.6倍和57.4倍的加速。
* 与采用SPSE策略的Multiverse (SPSE)相比，SPME策略也带来了高达7.3倍的性能提升。这归因于DOD范式和跨多实验的批处理与并行仿真降低了缓存未命中率（如图11所示）和内存开销。
* 基于DONS的模拟器性能优于基于UNISON的，而Multiverse又远超基于DONS的，因为它利用了GPU的大量核心以及pull-based同步、快速分析模型和巨核技术等优化。
最大规模与可扩展性:
* 如图12所示，在模拟一个包含54,000个GPU、4,500个交换机和16.2万条链路的大规模集群时，单块GPU上的Multiverse仍比其他方法快28.6倍至43.1倍。
* 如图13所示，单块H100 GPU可以同时运行一个54k GPU集群的仿真实验，或者并发运行520个128 GPU集群的实验、70个1,024 GPU集群的实验。在足够大的集群规模或足够多的并发实验数下，Multiverse能够充分利用GPU的所有核心。
6.2 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的精度也相应提高。
6.3 Multiverse的消融研究
各优化特性的贡献: 如图16所示，通过消融研究评估了Multiverse关键特性的影响：
* 分析模型: 与使用包级DES模拟服务器内通信相比，校准的分析模型带来了1.7-1.8倍的加速。
* 基于拉取的同步: 与使用全局锁同步相比，pull-based同步带来了3.2-5.4倍的加速，因为它实现了无锁并行。
* 巨核: 与为每个ECS系统独立启动内核相比，megakernel将仿真时间减少了16.6%-18.6%，因为它降低了CPU-GPU的同步开销。
A5 结论
本文论证了单进程多实验（SPME）并行策略通过减少调度开销和优化资源利用率实现了卓越的性能，但对于当前AI集群的规模仍显不足。为了提升SPME的效能，论文引入了Multiverse，一个新颖的基于GPU的AI训练模拟器。Multiverse利用DOD/ECS建模，并通过基于拉取的同步方法、高保真度的服务器内通信模型和巨核技术等优化措施，高效地利用了GPU的计算吞吐量。大量实验验证了Multiverse的准确性和效率，其速度比当前最先进的基于CPU的AI训练模拟器快43.1至73.2倍。