在深度学习领域,准确预测 GPU性能 对于优化模型训练和推理至关重要。然而,传统的 GPU性能预测 方法,尤其是针对新型模型和硬件,往往面临着准确率低下的挑战。近日,ASPLOS 2025 发表了一篇名为 “NeuSight” 的论文,提出了一种全新的 GPU性能预测框架,该框架通过精细化的瓦片级预测,并结合硬件性能边界的约束,大幅提高了预测准确性。特别值得一提的是,即使在未经过 GPT-3 或 H100 训练的情况下,NeuSight 也能将 GPT-3 在 H100 GPU 上的延迟预测误差从 121.4% 和 30.8%(此前工作)降低到仅 2.3%,展现出强大的泛化能力和实用价值。本文将深入探讨 NeuSight 的核心技术和工作流程,揭示其在大模型时代 GPU 性能预测方面的革命性突破。

传统GPU性能预测的局限性

传统的 GPU性能预测 方法主要依赖于机器学习模型直接预测内核延迟。这种方法的固有问题在于,对于未曾训练过的模型或新型硬件,模型的泛化能力往往不足,导致预测误差大幅增加。举个例子,假设我们使用一个在 V100 GPU 上训练的模型来预测在 A100 GPU 上的性能,由于两者的硬件架构存在差异,模型可能无法准确捕捉到 A100 的特性,从而导致预测结果偏差较大。此外,直接预测内核延迟通常忽略了 GPU 内部复杂的执行机制,如线程调度、内存访问模式等,进一步降低了预测的准确性。

NeuSight的核心:瓦片级预测与硬件边界约束

NeuSight 的核心创新在于将 GPU性能预测 问题分解为更小的、更易于管理的瓦片级预测。现代 GPU 库通常将矩阵乘法(GEMM)等操作划分为多个均匀的瓦片(tile),每个瓦片对应于输出矩阵的一部分,并由 GPU 上的流式多处理器(SM)并行执行。NeuSight 针对每个瓦片进行独立的延迟预测,然后将这些预测结果汇总,从而得到最终的内核延迟估计。

这种瓦片级预测方法具有以下优势:

  1. 精细化建模: 通过关注更小的计算单元,NeuSight 能够更准确地捕捉到 GPU 内部的执行细节和依赖关系。
  2. 硬件感知: NeuSight 在瓦片级预测中考虑了 GPU 的硬件性能边界,例如 SM 的数量、内存带宽等,从而确保预测结果的合理性。
  3. 降低复杂度: 将整体预测分解为多个瓦片预测,可以有效降低每个预测任务的复杂度,提高预测的准确性和效率。

例如,考虑一个矩阵乘法操作 C = A * B,其中 A 的维度是 1024×1024,B 的维度是 1024×1024,C 的维度是 1024×1024。传统的预测方法可能会直接预测整个矩阵乘法的延迟,而 NeuSight 则会将 C 矩阵划分为多个瓦片,例如 64×64 的瓦片,然后分别预测每个瓦片的延迟。通过汇总这些瓦片的延迟,NeuSight 可以更准确地估计整个矩阵乘法的延迟。

Roofline 模型:硬件性能的理论上限

为了约束瓦片级预测的合理性,NeuSight 借鉴了 Roofline 模型。Roofline 模型是一种简单而有效的性能分析方法,它通过考虑计算强度(arithmetic intensity)来估计 GPU 内核的近似性能。计算强度定义为浮点运算次数(𝑓𝑙𝑜𝑝𝑠ₖ)与内存事务大小(𝑚𝑒𝑚ₖ)之比。根据 Roofline 模型,内核的可达性能(roofline bandwidth)可以计算如下:

rooflineBW = min(𝑓𝑙𝑜𝑝𝑠ₚ, 𝑚𝑒𝑚𝐵𝑊ₚ * 计算强度)

其中,𝑓𝑙𝑜𝑝𝑠ₚ 表示 GPU 的峰值浮点运算速度(FLOPs),𝑚𝑒𝑚𝐵𝑊ₚ 表示 GPU 的峰值内存带宽。Roofline 带宽表示内核在 GPU 上可实现的最大吞吐量。NeuSight 利用 Roofline 模型为瓦片级预测设定了理论上限,确保预测结果不会超过硬件的物理限制。例如,如果一个瓦片的计算强度很高,但其预测延迟却非常低,那么 NeuSight 会根据 Roofline 模型对其进行调整,使其更符合实际情况。

利用率预测:机器学习的角色

NeuSight 还引入了机器学习模型来预测 GPU 利用率。作者观察到,随着工作负载中可用线程数量的增加,GPU 吞吐量会逐渐接近其最大值。为了模拟波的数量与利用率之间的关系,NeuSight 采用了一个多层感知器(MLP)来预测以下等式的系数:

Utilization = α + β * num_waves

其中,num_waves 表示每个 SM 可用的波的数量,αβ 是 MLP 预测的系数。MLP 的输入特征包括 GPU 硬件规格,如内存大小、带宽、峰值 FLOPs 和 L2 缓存大小,这些特征都经过归一化处理,并以资源利用率的形式表示。

NeuSight 针对不同类型的 DNN 内核使用了五种专门的 MLP:

  1. 批量矩阵乘法
  2. 全连接层
  3. 逐元素操作
  4. Softmax
  5. 层归一化

每个 MLP 都有 8 个全连接层,包含 512 个隐藏单元,并使用 ReLU 激活函数。选择这些特征的原因在于,它们在各个 GPU 代系中都可用,并且对算子性能有显著影响。对于不受支持或未知的算子,NeuSight 假设它们受内存限制,并相应地估计延迟。

例如,如果 MLP 预测到一个瓦片的利用率非常低,那么 NeuSight 会相应地提高该瓦片的延迟,以反映 GPU 的空闲时间。反之,如果 MLP 预测到一个瓦片的利用率很高,那么 NeuSight 会相应地降低该瓦片的延迟,以反映 GPU 的高效利用。

算子融合的支持

算子融合是一种常见的优化技术,它通过将多个算子或内核组合在一起,减少中间结果的内存流量,并提高运算强度。NeuSight 通过排除中间数据的内存占用来计算融合内核的 𝑓𝑙𝑜𝑝𝑠ₖ 和 𝑚𝑒𝑚ₖ。这种处理方式能够更准确地反映融合算子的性能,避免高估其延迟。例如,考虑一个由卷积层、ReLU 激活函数和池化层组成的融合算子。NeuSight 会将这三个算子视为一个整体,计算其总的浮点运算次数和内存事务大小,从而更准确地预测其性能。

NeuSight 的完整工作流程

NeuSight 估计深度学习模型在单 GPU 或多 GPU 系统上的性能,其工作流程包括以下三个步骤:

  1. 预测每个内核的执行延迟: NeuSight 利用其瓦片级预测方法和硬件边界约束来预测每个内核的执行延迟。
  2. 根据模型的数据流图汇总这些预测: NeuSight 使用 Torch.fx 从 PyTorch 中提取算子图,收集内核元数据,并应用其内核级预测器来注释每个算子。最终的每个 GPU 延迟通过对预测延迟求和获得,反映了内核在 GPU 上的顺序执行。
  3. 结合集体通信和网络开销来估计分布式环境中的总体性能: NeuSight 通过根据并行策略(流水线并行、数据并行或张量模型并行)用适当的通信算子增强 DNN 图来模拟多 GPU 服务器(例如,DGX 系统)上的分布式执行。它通过测量链路利用率并根据目标系统的峰值带宽进行外推来估计基于环的 all-reduce 和点对点发送/接收等网络操作的延迟。对于流水线并行,NeuSight 使用用户提供的调度(例如,GPipe)来模拟执行并插入流水线气泡,并根据微批大小和通信开销来估计它们的延迟。对于张量模型和数据并行,它插入 all-reduce 算子来同步激活或梯度,并将通信和计算延迟结合起来以预测端到端性能。

例如,考虑一个使用数据并行策略在多个 GPU 上训练的模型。NeuSight 会在模型的算子图中插入 all-reduce 算子,用于同步不同 GPU 上的梯度。通过测量 all-reduce 操作的延迟,并将其与计算延迟相结合,NeuSight 可以更准确地预测整个训练过程的性能。

实际案例与数据验证

论文中展示了 NeuSight 在多个深度学习模型和 GPU 上的实验结果,包括 GPT-3 在 H100 GPU 上的延迟预测。结果表明,NeuSight 在预测准确性方面显著优于以往的方法。例如,NeuSight 将 GPT-3 在 H100 GPU 上的延迟预测误差从 121.4% 和 30.8%(此前工作)降低到仅 2.3%,这表明 NeuSight 能够准确捕捉到 GPT-3 模型和 H100 GPU 的特性。

此外,NeuSight 还被用于预测 ResNet-50、BERT 等模型的性能,并在不同 GPU 上进行了验证,包括 V100、A100 和 H100。实验结果一致表明,NeuSight 能够提供准确可靠的 GPU性能预测,为深度学习模型的优化和部署提供了有力的支持。

NeuSight的未来展望

NeuSight 作为一种新颖的 GPU性能预测框架,在大模型时代展现出巨大的潜力。未来的研究方向包括:

  1. 支持更多类型的 GPU 和模型: NeuSight 可以扩展到支持更多类型的 GPU,例如 AMD GPU、Intel GPU 等,以及更多类型的深度学习模型,例如 Transformer 模型、GAN 模型等。
  2. 集成到深度学习框架中: NeuSight 可以集成到常用的深度学习框架中,例如 PyTorch、TensorFlow 等,为用户提供更便捷的 GPU性能预测 功能。
  3. 应用于自动调优: NeuSight 可以应用于深度学习模型的自动调优,根据预测的性能自动选择最佳的配置参数,例如批大小、学习率等。
  4. 云原生环境下的性能预测: 云原生环境下,GPU资源的管理和调度更加复杂,需要更精确的 GPU性能预测 来优化资源利用率。NeuSight 可以针对云原生环境进行优化,提供更准确的性能预测。

总之,NeuSight 的出现为 GPU性能预测 领域带来了革命性的突破。通过精细化的瓦片级预测、硬件边界约束和机器学习模型的结合,NeuSight 能够提供准确可靠的 GPU性能预测,为深度学习模型的优化和部署提供了有力的支持,推动了人工智能技术的进一步发展。在AI算力日益重要的今天,NeuSight 让我们能够更清晰地了解算力的真实性能,助力AI更好地发展。