作者/机构: Xizheng Wang (阿里巴巴云 & 清华大学), Qingxu Li (阿里巴巴云), Yichi Xu (阿里巴巴云), Gang Lu (阿里巴巴云), Dan Li (清华大学), Li Chen (中关村实验室), Heyang Zhou (阿里巴巴云), Linkang Zheng (阿里巴巴云 & 华南理工大学), Sen Zhang (阿里巴巴云), Yikai Zhu (阿里巴巴云), Yang Liu (阿里巴巴云), Pengcheng Zhang (阿里巴巴云), Kun Qian (阿里巴巴云), Kunling He (阿里巴巴云), Jiaqi Gao (阿里巴巴云), Ennan Zhai (阿里巴巴云), Dennis Cai (阿里巴巴云), Binzhang Fu (阿里巴巴云)

A1 主要贡献

本文介绍了一个名为 SimAI 的统一模拟器,旨在精确高效地模拟大规模语言模型(LLM)的训练过程,以解决因 LLM 训练所需 GPU 数量庞大而导致新设计、调优和优化验证困难的问题。现有模拟器通常只针对整个训练过程的特定粒度,导致精度不足。

核心问题与研究目标:
现有模拟器存在粒度不一、精度低、性能差等问题,导致成本性能分析不准确、大规模部署下优化能力受限,且开发测试复杂。为了克服这些挑战,本文旨在开发一个统一的模拟器,能够在单一框架内同时处理容量规划和性能调优,并满足以下目标:
1. 高精度工作负载生成: 灵活精确地生成能反映各种模型、参数和规模下真实 AI 训练行为的工作负载。
2. 高保真通信模拟: 精确模拟分布式 LLM 训练中使用的集体通信,并考虑关键优化。
3. 高保真计算模拟: 高效且精确地模拟不同 GPU 架构上的计算任务。
4. 快速可扩展的模拟速度: 能够高效地运行大规模 LLM 模拟。

核心创新与贡献:
SimAI 通过以下创新解决了上述挑战:
* 高保真全栈模型: SimAI 通过对训练框架、核心计算和集体通信库进行选择性和高保真的集成,实现了高精度的模拟。
* 精确的工作负载生成: SimAI 通过“劫持”主流训练框架(如 Megatron 和 DeepSpeed),使其在单台主机上运行,从而为任意规模的 LLM 训练生成精确的工作负载。
* 精确的计算模拟: 将工作负载分解为细粒度的核心(kernels),测量其在现有 GPU 上的执行时间,并将其映射到其他 GPU 类型,以实现精确的计算模拟。
* 精确的通信模拟: 通过“劫持”NVIDIA 集体通信库(NCCL),精确模拟集体通信的包级行为。
* 高效的模拟执行: 实施多线程加速和无锁全局上下文共享,以提高模拟器的执行速度。
* 实践应用与指导: SimAI 的模拟结果为新的主机设计和参数设置提供了有意义的指导,并已被工程团队采纳并应用于生产部署。

实验结果验证了 SimAI 的有效性,其在各种测试场景下与真实世界结果的平均对齐率达到98.1%,证明了其从小型实验室到大规模工业环境的鲁棒性和适应性。SimAI 已被开源。

A3 背景知识与设计动机

2.1 AI 训练基础设施

大规模语言模型(LLM)的训练依赖于专门的基础设施,通常涉及数十到数千个 GPU 协同工作。例如,训练一个1750亿参数的GPT-3模型需要1024个高端GPU连续运行34天。为了优化GPU利用率,主流训练框架如 Megatron【52】和 DeepSpeed【46】提供了多种并行化技术,包括数据并行(DP)、流水线并行(PP)和张量并行(TP)。这些技术通过集体通信库(CCL)来管理GPU之间的数据交换,例如使用 AllReduce 进行梯度同步或使用 AllGather 进行参数共享。CCL 将每个集体通信任务分解为一系列点对点的发送(Send)和接收(Receive)操作来完成数据传输。

在阿里巴巴的集群中,服务器内多个GPU通过高带宽的内部网络(如NVLink)连接,并通过网卡(NIC)连接到跨主机的RDMA网络。例如,A100服务器配备了8个NVIDIA A100 GPU和4个提供2x100Gbps带宽的NVIDIA CX6Dx NIC。每个GPU通过600GB/s的NVLink与其他GPU连接,并通过100Gbps的RDMA网络与其他服务器的GPU连接。在生产环境中,Megatron 和 DeepSpeed 是主要的训练框架,NCCL 是主要的CCL。

2.2 对统一模拟器的需求

LLM的快速发展对AI训练基础设施和优化方法提出了更高的要求,模拟在其中扮演着关键角色,主要服务于三个目标:

  • 对AI基础设施的综合评估:在新硬件部署前,需要从多个角度进行评估。
    • GPU选择:云服务提供商(CSP)需要评估新GPU模型在AI工作负载下的规模化性能,模拟可以在大规模部署前提供详细的性能预览,指导采购和扩展策略。
    • 网络架构设计:确定GPU后,需要优化网络架构以实现可扩展性。诸如每个GPU应配置多少网络带宽等关键问题,需要高精度模拟,因为大规模实验成本过高。
    • 主机架构设计:评估不同主机配置(如每台主机的GPU数量、内部互连方式)至关重要。模拟可以在无需构建昂贵硬件原型的情况下回答这些问题。
  • 对优化进行成本效益高的验证
    • 参数调优:模拟允许快速、低成本地测试各种模型参数和训练框架设置,以达到最佳性能。
    • 评估新机制:对于新的训练框架、集体通信方法或网络拥塞控制算法,模拟提供了一种在现实环境中低成本评估其有效性的方法。
  • 开发统一的模拟框架
    • 由于模拟需求遍及AI基础设施的不同组件和层次,一个统一的模拟框架至关重要。构建多个专用的模拟器会导致不一致或不完整的见解。例如,一个只关注GPU的模拟器可能会忽略关键的网络性能,导致不准确的结论。
    • 因此,我们的目标是开发一个统一的模拟器,在单一平台上解决所有这些需求,确保团队能够准确高效地验证新设计和优化。

2.3 我们的目标

  • 目标1:生成反映真实训练的工作负载。为获得准确的模拟结果,需要能捕捉训练框架详细行为的真实输入源。简单地根据浮点运算量估算工作负载过于粗糙。像Chakra【53】这样的追踪驱动方法有所改进,但仅适用于参数和规模相同的LLM,限制了模拟新模型或配置的能力。我们需要一个灵活且精确的工作负载生成器,能处理各种模型、参数和规模。
  • 目标2:高保真通信模拟。经典的包级网络模拟器(如NS-3【47】)不处理分布式LLM训练中使用的集体通信。集体通信库(如NCCL)为最大化性能会应用各种优化,这会影响流量模式。从头模拟这些优化可能导致保真度低。我们需要一个高精度的集体通信模拟器,它能整合关键的优化和增强功能。
  • 目标3:高保真计算模拟。现有的GPU核心计算模拟方案,如GPGPU-Sim【2】,在细节层面模拟,但对于大规模LLM模拟来说过于耗时。其他方法,如ASTRA-sim【45】,要么不支持不同的GPU,要么精度不足。我们需要一个高效的计算模拟器,既能提供精确度,又能支持大规模模拟的可扩展性。
  • 目标4:快速的模拟速度。使用现有方法的组合(例如PyTorch追踪生成器与ASTRA-sim),模拟一次128个GPU的GPT-3训练迭代需要一整天,而真实硬件上只需两秒。效率对于将模拟扩展到实际应用至关重要。模拟器不仅要满足前三个目标,还必须是可扩展的,并且能够高效地运行大规模LLM模拟。

A2 方法细节

3.1 SimAI 概述

SimAI 的关键组件如图1所示。每个模拟请求都包含详细的训练过程信息,例如模型本身及参数、训练框架配置、CCL参数以及内部/跨主机网络拓扑。

  • 工作负载生成器 (SimAI-WG):为每个模拟请求生成真实的工作负载(§3.2)。其输出是一个工作负载描述文件,该文件概述了算法模块、集体通信操作及其依赖关系。
  • 执行引擎: 接着,执行引擎处理该工作负载文件,将计算和通信操作的执行模拟为离散事件。我们利用计算模拟器 (SimAI-CP)通信模拟器 (SimAI-CM) 分别模拟计算和通信任务。
  • SimAI-CP: 将子模块转换为详细的核心(kernels),并使用一个自建的细粒度操作库提供精确的计算模拟(§3.3)。
  • SimAI-CM: 集成了 NCCL 的部分功能,将每个集体通信分解为点对点操作,以提供准确的通信模拟结果(§3.4)。
  • 加速: 此外,我们实现了多线程加速和无锁全局上下文共享,以进一步提升模拟速度(§3.5)。

Figure 1: SimAI 架构
Figure 1: SimAI 架构

3.2 工作负载生成器

3.2.1 生成工作负载文件

通过劫持现有框架生成工作负载。生成精确的工作负载文件至关重要,因为即使是微小的差异也可能导致模拟保真度降低。SimAI-WG 采用了一种“劫持”现有训练框架(如我们生产环境中流行的 Megatron 和 DeepSpeed)的方法,以生成与真实任务完全相同的工作负载。

在单台主机上生成工作负载。该“劫持”方法的关键挑战是在单台主机上运行,同时考虑到多个对等主机的交互。为此,我们对框架和 NCCL 进行了两项修改:
1. 让框架误以为它在一个拥有目标数量 GPU 的集群中运行,并设置相应的跨主机拓扑。
2. 跳过 NCCL 中所有真实的通信。对于涉及流水线并行的工作负载,SimAI-WG 必须配置正确的 rank 编号。
通过这种方式,训练框架能生成计算和通信操作的序列,包括算法子模块和集体或点对点通信。表1展示了LLM中典型的操作示例。然而,这些操作缺乏依赖关系说明,需要进一步明确。

Table 1: 计算和通信操作列表。⋆ 表示内存带宽密集型操作,△ 表示计算密集型操作。

Table 1: 计算和通信操作列表
Table 1: 计算和通信操作列表

定义操作依赖关系。由于计算和通信操作在执行期间会重叠,我们在工作负载文件中嵌入了依赖信息以反映这一点。这些依赖关系根据所使用的并行化框架确定,并在一个1024-GPU集群上使用 Nvidia Nsight Systems 进行了验证。图2展示了使用TP、DP和PP并行时的依赖关系示例。实线箭头表示严格的“先发生”关系,而虚线箭头表示重叠操作,即通信在计算开始后启动。该图显示了 Megatron 的 attention 和 MLP 反向传播阶段如何从重叠优化中受益。例如,MLPbackward 子模块仅在 Attentionbackward 子模块和 AllReduce 操作完成后才开始。

Figure 2: 算法子模块和通信操作之间的依赖关系演示
Figure 2: 算法子模块和通信操作之间的依赖关系演示

解耦工作负载与框架。我们选择将工作负载与训练框架解耦,以实现独立的模拟。与 Owl【9,Jason Flinn et al. OSDI 2022】等将模拟代码直接嵌入框架的做法不同,我们使用 SimAI-WG 提取模型工作负载,并在 SimAI 执行引擎中进行模拟。这种方法避免了修改训练框架带来的巨大开销,后者需要对监控工具、数据管道等关键组件进行大量更改。

大语言模型中的通信。在LLM训练中,元数据交换和屏障操作(小于1KB)的开销可以忽略不计,不是我们模拟的重点。通信时间主要由各种并行技术引入,如TP、PP、EP和DP。TP、PP和EP通常具有固定的通信模式和数据量,不受集群大小影响,涉及数十到数百MB的中等大小消息。相比之下,DP的通信范围随集群规模扩大而增加,通常涉及跨越数百到数千个节点的GB级集体通信(如AllGather/ReduceScatter)。

3.2.2 确定模型级基准套件

建立基准套件。为了对真实世界 LLM 训练场景的端到端性能进行全面评估,建立一个模型级的基准套件至关重要。

Table 2: 不同参数大小和规模的任务比例。

Table 2: 不同参数大小和规模的任务比例
Table 2: 不同参数大小和规模的任务比例

模型和框架选择。表2显示了过去六个月中模型参数和任务规模的分布。模型分为小、中、大三类,每类占比均超过20%。超过94%的LLM是GPT-3或LLaMA的变体,Megatron和DeepSpeed是主导框架。因此,我们的基准套件包含了一系列不同大小的开源LLM。

参数选择。我们关注影响工作负载模式的参数:
* 模型参数(如 hidden_size, num_layers, seq_len 等)
* 框架参数(如 world size, zero level, reduce_bucket_size, allgather_bucket_size, 以及并行策略如 TP, PP, DP, SP)
基于表2的统计数据,我们选择了一组最小的基准,覆盖了典型设置,如表3所示。用户也可以使用 SimAI-WG 生成自定义基准。我们相信这个基准套件能准确代表客户常用的真实LLM训练工作负载。

Table 3: SimAI 基准套件 v1.0

Table 3: SimAI 基准套件 v1.0
Table 3: SimAI 基准套件 v1.0

3.3 精确的计算模拟

精确模拟现有 GPU。通过在单台主机上运行模拟框架,我们获得了详细的子模块工作流。对于针对现有GPU的模拟,SimAI-WG 会输出所有子模块在相应GPU主机上的执行时间。由于在实际LLM训练中每个GPU专用于单个任务,我们可以通过遵循工作负载文件并填入每个子模块的执行时间来准确模拟整个计算过程。如 §4.3 所示,模拟精度在 96.9% 到 99.5% 之间。我们在 SimAI-CP 中维护一个操作数据库,记录基准套件中所有子模块的执行时间(见表4)。

Table 4: SimAI-CP 数据库格式

Table 4: SimAI-CP 数据库格式
Table 4: SimAI-CP 数据库格式

细粒度核心模拟。在某些场景下,例如新的并行策略或优化对核心进行了重组,子模块级别的模拟可能不够用。在这种情况下,需要进行细粒度的模拟。我们设计了一个模块-核心转换器,将每个子模块分解为更小的核心,然后在不同的GPU上进行测试。表1的第三列展示了不同子模块中使用的核心示例。这进一步丰富了操作数据库,使得能够精确模拟高级优化和新功能。

支持未发布的 GPU。作为为云服务提供商(CSP)提供的模拟服务,模拟尚未发布的GPU有很强的需求。决策者需要评估购买新GPU是否值得投资。虽然无法物理访问未发布的GPU,但我们可能可以获得其核心规格。初步方法可能是通过从现有GPU的已知值进行缩放来估算计算时间,但这通常会导致高达25.1%的显著误差。

基于 Roofline 模型的估算。我们的分析表明,不同的核心有不同的性能瓶颈,通常分为计算密集型或内存带宽密集型。例如,Gemm 核心是内存带宽密集型,而 flash attention 是计算密集型(分类见表1)。为了提高准确性,我们基于操作数据库的数据和 Roofline 性能模型【58,Samuel Williams et al. Communications of the ACM 2009】,针对这两类核心提出了两个不同的估算公式。这些公式基于在现有环境中测量的计算密集型和内存密集型核心的执行时间(下标分别为 Comp_KnownMem_Known),用于计算它们在新GPU或配置上的执行时间(下标分别为 Comp_NewMem_New):
* 对于计算密集型核心:
$T_{Comp\_New} = T_{Comp\_Known} \times \frac{FLOPS_{Known}}{FLOPS_{New}}$
* 对于内存带宽密集型核心:
$T_{Mem\_New} = T_{Mem\_Known} \times \frac{Bandwidth_{Known}}{Bandwidth_{New}}$

为确保结果准确,我们建议使用架构相似的GPU作为基准。在我们的实验中,我们使用Nvidia A100作为计算其他Nvidia GPU核心执行时间的基准。我们将此方法称为 SimAI-CP-Model。

3.4 精确的通信模拟

集成 NCCL 关键流程。训练框架依赖于集体通信库,其中 NCCL【36】在生产中应用最广。NCCL 将集体通信操作(如 AllReduce)转换为网络级操作(如 Send、Receive),这个过程非常复杂,它会根据节点数、消息大小等因素选择最优算法。模拟与实际执行之间的微小不匹配都可能导致结果出现巨大偏差。为了准确复现这些算法选择,SimAI 直接集成了 NCCL 的关键程序。

Figure 3: SimAI 如何将集体通信操作转换为点对点通信操作列表的图示
Figure 3: SimAI 如何将集体通信操作转换为点对点通信操作列表的图示

复现 NCCL 的关键程序。SimAI-CM 使用一个名为 SimCCL 的 NCCL 修改版本来拦截关键操作。SimCCL 通过“劫持”技术,在单台主机上模拟详细的点对点操作,类似于 §3.2.1 中讨论的方法,但作用于集体通信层面。SimCCL 修改 NCCL 行为的步骤如下:
1. NCCL 初始化:SimCCL 使用 libhacknccl.so 拦截 NcclCommInitAll 函数,为每个GPU创建虚拟通信器,使系统表现得像在真实的多GPU集群中运行。
2. 拓扑发现:SimCCL 不搜索实际的PCIe设备,而是读取用户指定的拓扑文件。每个虚拟通信器独立处理拓扑,无需同步。
3. 主机内通信信道创建:SimCCL 在主机内的虚拟通信器之间建立信道,并将详细信息存储在隔离的信息表中。
4. 主机间通信信道创建:SimCCL 跳过使用 AllGather 从其他GPU收集信息的过程,因为它已经拥有所有主机的信息,直接创建主机间信道。
5. 集体通信转换:SimCCL 拦截集体通信调用,重建操作以追踪更低级别的通信。它跳过实际的数据传输,捕获GPU间的通信事件(数据大小、发送/接收方rank、路由等)来模拟RDMA层的行为。

支持所有 NCCL 参数。SimCCL 能够反映绝大多数 NCCL 参数【37】的通信行为。SimAI-CM 已被增强以支持特定功能,例如 PCI × NVLink (PXN)。PXN 允许 GPU 通过 NVLINK 连接使用同一节点内的非本地 NIC 进行数据传输,这在 rail-optimized 网络拓扑中很常见,可以实现消息聚合和网络流量优化。通过设置 NCCL_P2P_PXN_LEVEL,SimCCL 可以识别这些 PXN 流量模式并将其反映在输出的 FlowModel 中。

SimCCL 的设计与实现。SimCCL 模块旨在模拟 NCCL 或其他 CCL 中的通信操作处理流程,专注于将集体通信操作转换为易于解释的点对点通信,而不包含实际的数据验证。这种方法通常不影响端到端训练过程。在专家并行(EP)中,尽管门控模块的令牌分配受数据值影响,我们在模拟中假设其为均衡分布,这对模拟结果影响极小。

Table 5: 修改 NCCL 的工作量

Table 5: 修改 NCCL 的工作量
Table 5: 修改 NCCL 的工作量

移植工作量。构建图1中所示的模块花费了我们超过10,000行代码。其中大部分工作是为可重复使用而设计的,除了 SimCCL 模块。支持不同版本的 NCCL 或新的 CCL 需要重新适配 SimCCL,但我们的设计对原始 CCL 代码的修改是非侵入性的。如表5所示,基于原始 NCCL 代码库,仅需修改572行代码。

3.5 大规模模拟加速

性能瓶颈与多线程加速。在 SimAI 开发初期,模拟一次128个GPU的GPT-3训练迭代需要超过24小时,而物理集群上仅需2秒。为了能够模拟超过1000个GPU的AI基础设施,我们为 SimAI-CM 实现了多线程加速。我们选择了基于 NS-3 的开源方案 UNISON【3,Songyuan Bai et al. 2024】,因为它能自动划分网络拓扑并调度任务到合适的线程,且具有良好的可扩展性。

Figure 4: SimAI 中的多线程实现与优化
Figure 4: SimAI 中的多线程实现与优化

无锁共享全局变量。然而,集成 UNISON 带来了挑战。随着模拟规模的增加,大量全局配置和上下文需要在线程间共享。使用全局锁(甚至是原子锁)更新共享数据结构导致了性能瓶颈(如图4(a)所示)。这些全局变量跟踪节点间数据量、交换机队列长度等元数据,原子操作减慢了并发线程的执行。

无锁优化。为了解决这个问题,我们重构了 SimAI 以实现无全局锁运行。我们观察到大多数全局变量都由特定线程访问。通过以线程独立的方式管理这些变量,我们消除了对全局锁的需求。如图4(b)所示,SimAI 分离了这些变量。由于每个节点的模拟都在单个线程上运行,我们将全局变量隔离在一个以节点ID为索引的表中。这使得线程可以无需锁即可访问相关数据,减少了并发带来的性能问题。我们的无锁优化相比单线程版本实现了23倍的加速,相比多线程版本提升了15%。

A4 实验

4.1 实验环境

硬件平台。我们使用了两个测试平台,它们都采用了类似阿里巴巴HPN架构【44】的多轨、胖树、跨主机RoCEv2网络,如图5所示。
* 集群 A100:包含128个GPU主机,每台主机配备8个NVIDIA A100 GPU【1】和4个Mellanox ConnectX-6网卡(2×100 Gbps)。
* 集群 H100:包含128个GPU服务器,每台主机配备8个NVIDIA H100 GPU【13】和8个Mellanox ConnectX-7网卡(2×200 Gbps)。
主机内互连使用NVLink,A100的带宽为600 GBps,H100为900 GBps。模拟和物理集群共享相同的内部和跨主机网络拓扑。

Figure 5: H100 集群的多轨网络拓扑。对于 A100 集群,两个 GPU 共享一个 NIC
Figure 5: H100 集群的多轨网络拓扑。对于 A100 集群,两个 GPU 共享一个 NIC

基准测试。基准测试来自 SimAI 基准套件。我们将 ASTRA-sim 与一个能够解析 SimAI-WG 生成的工作负载文件的适配器进行了增强。物理GPU集群上的基准测试涉及真实的LLM任务,这些任务的模型、框架和NCCL参数与工作负载文件相匹配。我们的评估涵盖了128、512和1024个GPU的集群规模。我们展示了三个代表性模型的性能指标:GPT-3 13B、LLaMA 65B 和 GPT-3 175B。

4.2 实验结果

通信模拟精度

主机内通信。图6显示,SimAI比ASTRA-sim精确得多,与真实值非常接近。对于A100,SimAI和ASTRA-sim的平均偏差分别为3.9%和74.8%;对于H100,分别为2.3%和51.7%。ASTRA-sim在小消息场景下保真度很差,因为我们没有模拟运行时软件(如libibverbs)和NIC流水线,而ASTRA-sim对实际传输过程做了大量简化。

Figure 6: 主机内集体通信操作的总线带宽随不同消息大小和 GPU 类型而变化。请注意,性能与集群规模无关。
Figure 6: 主机内集体通信操作的总线带宽随不同消息大小和 GPU 类型而变化。请注意,性能与集群规模无关。

主机间通信。图7的结果再次证实,ASTRA-sim难以准确模拟小消息通信,而SimAI始终接近真实值。随着通信规模的扩大,ASTRA-sim的模拟与真实值之间的差距也越来越大。例如,在128个A100 GPU的集群上进行8MB消息的AllGather操作,ASTRA-sim的模拟时间误差为45.9%,当规模增加到512个GPU时,误差激增至530.2%。

Figure 7: 主机间集体通信操作的总线带宽随不同消息大小和集群规模而变化。
Figure 7: 主机间集体通信操作的总线带宽随不同消息大小和集群规模而变化。

计算模拟精度

不同GPU上的精度。我们在Nvidia A100、H100和H20上测试了SimAI的计算模拟精度。如图8所示,SimAI-CP对所有计算核心的估算都非常准确,总执行时间的差距在0.5%-3.1%之间。然而,SimAI-CP-Model的偏差较大,在13%-15%之间,特别是在Attention-QKV和MLP-linear核心上。因此,我们建议用户仅在无法获得GPU时使用SimAI-CP-Model。相比之下,我们无法在ASTRA-sim使用的SCALE-sim【50】中运行所有计算核心,仅矩阵乘法核心成功运行。ASTRA-sim在模拟其他GPU时误差巨大,H100的总误差为49.8%,A100为117.9%,H20为224%。这说明由于GPU计算流水线的封闭性,这些方法难以精确估算计算时间。

Figure 8: GPT-3 175B 的计算核心在不同 GPU 上的执行时间。
Figure 8: GPT-3 175B 的计算核心在不同 GPU 上的执行时间。

端到端模拟精度

不同规模下的精度。我们验证了SimAI在不同集群规模下的端到端性能精度。图9显示,使用SimAI的所有工作负载的迭代时间都非常接近真实LLM任务在真实集群中的时间。即使在1024个GPU的规模下,差距也小于4%。然而,ASTRA-sim在A100和H100 GPU的迭代时间上都不精确。SimAI与真实值之间的偏差小于3.9%,比ASTRA-sim好36.1倍。此外,随着模型尺寸的增加(从GPT-3 13B到GPT-3 175B),SimAI的精度也相应提高,这是因为GPT-3 175B的集体通信操作的消息尺寸更大,其模拟更为精确。

Figure 9: 不同工作负载在 ASTRA-sim、SimAI 和真实集群上运行的迭代时间。
Figure 9: 不同工作负载在 ASTRA-sim、SimAI 和真实集群上运行的迭代时间。

A7 补充细节

5.1 在生产中的应用:指导新GPU的主机设计

H100 主机设计。2023年,当H100 GPU发布时,我们面临一个关键问题:为充分利用其增强的性能,需要配置多大的网络带宽,同时最小化总资本支出?我们使用 SimAI 模拟了一个包含1024个H100 GPU的训练集群在不同网络容量(100Gbps、200Gbps、400Gbps)下的性能。如图10所示,从200Gbps升级到400Gbps带来了19%的性能提升,相对于资本支出而言是显著的。这些模拟结果促使我们最终决定为H100主机设计采用每GPU 400Gbps的配置(即8个H100 GPU和8个2x200Gbps的Mellanox-CX7网卡)。

Figure 10: GPT-3 175B 在不同网络带宽下的迭代时间。
Figure 10: GPT-3 175B 在不同网络带宽下的迭代时间。

H20 主机设计。对于2023年11月发布的H20,我们再次使用SimAI进行主机设计。如图10所示,由于H20的计算性能低于H100,所需的网络带宽也随之降低。将网络带宽从100Gbps增加到200Gbps,端到端训练性能提升了11%;但从200Gbps增加到400Gbps,性能仅提升6%。在综合评估网络设备投资与性能提升后,我们决定为每个H20配置200Gbps的网络带宽(即每台主机配备8个H20 GPU和8个CX6网卡)。这再次证明了即使对于无法获取的GPU设备,SimAI-CP-Model也能提供有意义的参考。

5.2 在生产中的应用:量化扩展效益

TP 参数调优。在生产中,端到端性能受模型和框架参数(特别是并行策略)的显著影响。使用 SimAI,我们可以自动识别最佳参数设置。一个令人惊讶的发现与张量并行(TP)配置有关。传统观点认为,增加单个TP组中的GPU数量会增加通信开销,通常将TP组大小设置为与主机内通过NVSwitch连接的GPU数量相匹配。

实验与发现。为了量化分析此参数设置的权衡,我们在SimAI中对8、16、32和64个GPU(均通过NVSwitch互连)的配置进行了实验。图11展示了关键发现:
1. 在8-GPU主机中,GPT-3 13B、LLaMA 65B和GPT-3 175B的最佳TP大小分别为4、8和8。
2. 对于任何给定的TP大小,主机中GPU越多,性能越好。
3. 即使GPU数量更多,如果TP大小不合适,性能也可能下降。

Figure 11: 不同主机和不同TP大小下的训练性能。x轴表示单个主机中的GPU数量。由于模型较大且GPU内存有限,GPT-3 13B的TP应大于2,GPT-3 175B和LLaMA 65B的TP应大于4。我们不能为GPT-3 175B运行TP=64或TP=48,因为TP大小应为前馈神经网络(FNN)层数(96)和world size(1024)的因子。我们不让TP超过单个主机中的GPU数量,因为RDMA网络是瓶颈。
Figure 11: 不同主机和不同TP大小下的训练性能。x轴表示单个主机中的GPU数量。由于模型较大且GPU内存有限,GPT-3 13B的TP应大于2,GPT-3 175B和LLaMA 65B的TP应大于4。我们不能为GPT-3 175B运行TP=64或TP=48,因为TP大小应为前馈神经网络(FNN)层数(96)和world size(1024)的因子。我们不让TP超过单个主机中的GPU数量,因为RDMA网络是瓶颈。

实践指导。这些发现为实际LLM训练提供了两个宝贵的指导方针:
1. TP参数应适应整个模型层。但增加TP大小可能会降低端到端吞吐量;更好的策略是设置更多的数据并行(DP)组进行并行训练。
2. GPU主机设计必须考虑LLM的演进。如果已知一个层的最大尺寸,主机内的GPU数量应据此确定,优先考虑增强横向扩展性能。

6.1 模拟器的演进历程

  • 从一开始就追求精度。我们最初使用ASTRA-sim作为原型,但其内置工作负载导致模拟结果与实际存在92.1%至143.9%的巨大差异。为了确保精度,我们从头开始构建 SimAI-WG。起初通过手动分析框架代码来重建工作流,但这种方式耗费大量人力且难以维护。最终,我们转向了本文详述的“模拟框架”(mock-framework)方法,实现了模拟精度的显著提升。
  • 提升计算模拟精度。我们开发了 SimAI-CP,其操作库的构建分为两个阶段。第一阶段,我们提取LLM的关键子模块并测量其执行时间,达到了满意的精度。第二阶段,我们将子模块分解为按计算和内存访问强度分类的更细粒度的核心,这使我们能够根据性能规格估算不可用GPU上的操作时间。
  • 提升集体通信模拟精度。在使用 ASTRA-sim 时,我们发现其集体通信算法与 NCCL 不一致,导致模拟不准确。即使实现了 NCCL 的算法,结果仍不理想。因此,我们放弃优化 ASTRA-sim,转而从头开始模拟 NCCL 的关键行为。
  • 加速整个模拟器。在实现高精度后,我们面临模拟执行效率的挑战。SimAI 基于 NS-3 构建,我们集成了 UNISON【3】进行加速。在解决上下文访问冲突后,我们通过分析上下文依赖性并实现无锁访问,获得了高达23倍的性能提升,使模拟器具备了实用性。

6.2 模拟即服务(Simulation as a Service)

  • 无GPU运行。为了使 SimAI 成为一个可以在任何主机上运行的真正模拟服务,我们通过识别并挂钩所有CUDA相关函数,将 SimAI 与GPU软件栈解耦。
  • 任务管理与调度。作为一个模拟服务,SimAI 必须响应并发的模拟请求。我们维护一个主机集群,并使用 Kubernetes【29】将每个请求作为独立任务调度到负载最轻的主机上。
  • 参数调优。SimAI 允许对影响模型训练的各种参数进行微调,包括拓扑、模型和并行策略参数,以优化整体性能。
  • 细粒度监控。对于每次模拟,我们不仅生成包含总体训练时间和吞吐量的报告,还集成了广泛的监控统计数据,并通过一个前端系统展示模拟过程中不同指标的变化。

7 相关工作总结

本文在多个领域借鉴并改进了现有工作:
* AI基础设施模拟:SimAI 受 ASTRA-sim【45, 60】启发,但在工作负载生成、计算/通信模拟保真度和可扩展性方面做出了显著改进。
* GPU与计算模拟:与 GPGPU-Sim【2】等指令级模拟器相比,SimAI 在保证精度的同时,避免了大规模部署中过长的模拟时间。
* 网络模拟器与可扩展性:SimAI 的加速基于 UNISON【3】,并通过无锁全局变量共享进行了增强,解决了并行离散事件模拟(PDES)中的性能瓶颈。
* AI基准:与 MLPerf【31】等现有基准套件不同,SimAI 的基准套件专注于代表真实生产环境中的 LLM 训练负载。
* 基于学习的模拟器:与需要大量训练数据的端到端性能估计器(EPE)【20, 63】不同,SimAI 采用离散事件模拟来模拟模型训练行为,无需大量先验数据。
* 故障模拟:未来工作计划加入模型训练中的故障模拟功能,以评估容错机制。

A5 结论

本文介绍了 SimAI,一个为大规模语言模型(LLM)训练设计的统一模拟器,旨在提供准确和高效的模拟。SimAI 通过集成主流训练框架和集体通信库的关键流程,在各种测试场景下实现了与真实世界结果平均98.1%的一致性。为了增强可扩展性和性能,SimAI 采用了无锁、多线程的优化。该模拟器已在阿里巴巴云生产环境中广泛使用,并为制定有价值的运营指南做出了贡献。

A6 附录

9.1 NCCL 环境变量

表6列出了官网【37】记录的所有NCCL环境变量,以及它们在SimAI模拟器中的支持情况。可以看出,除了我们认为对模拟不必要的变量外,目前只有四个环境变量尚未支持,它们与自适应路由和Infiniband SHArP相关功能有关。我们将这些作为SimAI通信模拟器未来工作的方向。

Table 6: NCCL 环境变量及其如何影响通信模式

Table 6 Part 1
Table 6 Part 1

(续下页)

Table 6: NCCL 环境变量及其如何影响通信模式 (续)

Table 6 Part 2
Table 6 Part 2

(续下页)

Table 6: NCCL 环境变量及其如何影响通信模式 (续)

Table 6 Part 3
Table 6 Part 3

(续下页)

Table 6: NCCL 环境变量及其如何影响通信模式 (续)

Table 6 Part 4
Table 6 Part 4

方法细节中的参考文献引用汇总

在《方法细节》章节中,作者引用了以下关键文献来支撑其设计选择和实现方法:

  • 段落 3.2.1 - 定义操作依赖关系

    • 引用: 【9】Owl: Scale and flexibility in distribution of hot content. (Jason Flinn et al. OSDI 2022)
    • 原文描述: "Instead of embedding simulation code directly into frameworks like Owl [9], we used SimAI..."
    • 引用说明: 此处引用 Owl 是为了对比 SimAI 的设计选择。SimAI 选择将工作负载与训练框架解耦,而不是像 Owl 那样将模拟代码直接嵌入框架,从而避免了修改框架的巨大开销。
  • 段落 3.3 - 支持未发布的 GPU

    • 引用: 【58】Roofline: an insightful visual performance model for multicore architectures. (Samuel Williams et al. Communications of the ACM 2009)
    • 原文描述: "To improve accuracy, we propose using two equations tailored to these kernel types, based on data from our operation database and the Roofline performance model [58]."
    • 引用说明: 引用 Roofline 模型是为了提供理论基础,说明 SimAI 如何根据计算密集型和内存带宽密集型这两种不同的性能瓶颈,为未发布的 GPU 设计不同的性能估算公式。
  • 段落 3.4 - 精确的通信模拟

    • 引用: 【36】Nvidia collective communications library (nccl). (NVIDIA)
    • 原文描述: "Training frameworks rely on collective communication libraries, with NCCL [36] being the most widely used in production."
    • 引用说明: 引用 NCCL 是为了指出它是生产环境中最主流的集体通信库,因此 SimAI 的通信模拟器选择精确地模拟和集成 NCCL 的行为是合理且必要的。
  • 段落 3.4 - 支持所有 NCCL 参数

    • 引用: 【37】Environment variables of nvidia collective communications library (nccl). (NVIDIA Docs)
    • 原文描述: "SimCCL reflects the communication behavior for the vast majority of NCCL parameters [37]."
    • 引用说明: 引用 NCCL 环境变量文档是为了表明 SimAI 对 NCCL 的支持是全面且详细的,能够反映绝大多数官方参数对通信模式的影响。
  • 段落 3.5 - 大规模模拟加速

    • 引用: 【3】Unison: A parallel-efficient and user-transparent network simulation kernel. (Songyuan Bai et al. 2024)
    • 原文描述: "We chose UNISON [3] for three key reasons: (1) It is open-source and builds on NS-3. (2) It automates the partitioning of network topologies and schedules each task to the appropriate thread. (3) It has superior scalability, as demonstrated in their evaluations."
    • 引用说明: 引用 UNISON 是为了解释 SimAI 加速方案的技术来源。SimAI 选择了 UNISON 作为其多线程加速的基础,并在此之上进行了无锁优化,以解决 UNISON 在大规模模拟中遇到的性能瓶颈。
    • 引用: 【10】Parallel and Distributed Simulation Systems. (R.M. Fujimoto. Wiley 2000), 【18】Synchronization methods in parallel and distributed discrete-event simulation. (Shafagh Jafer et al. Simulation Modelling Practice and Theory 2013)
    • 原文描述: "For example, parallel discrete-event simulation (PDES) [10, 18] and UNISON [3] distribute network topologies across multiple logical processes..."
    • 引用说明: 引用这些文献是为了介绍并行离散事件模拟(PDES)这一更广泛的技术背景,并将 UNISON 定位为其中的一种实现方法,以说明 SimAI 选择该技术路线的合理性。