深度学习模型训练和推理对 GPU 算力需求日益增长，精准的 GPU 性能预测 变得至关重要。ASPLOS 2025 会议上提出的 NeuSight 框架，正是为了解决这一难题而生。它无需实际执行，就能预测各种深度学习模型在未见过的 GPU 上的性能，无论是训练还是推理。NeuSight 通过分解问题，在硬件性能边界的约束下，进行更小粒度的 Tile-level 预测，并结合 GPU 硬件特征和软件库行为，显著提高了预测精度。例如，在 H100 GPU 上，GPT-3 延迟预测误差从先前工作的 121.4% 和 30.8% 降至仅 2.3%，且训练过程中无需 GPT-3 或 H100 的数据。本文将深入探讨 NeuSight 的核心机制和工作流程，揭示其如何实现精准的深度学习算力预测。
GPU 内核执行机制与 Roofline 分析
现代 GPU 库在执行 GEMM (General Matrix Multiply) 操作时，通常会将输出矩阵分割成多个均匀的 Tile。每个 Tile 代表输出矩阵的一个片段，负责加载相应的输入操作数，并计算关联的输出元素。这些 Tile 会被分配到 GPU 上的各个流式多处理器 (Streaming Multiprocessor, SM) 并行执行。
通常，可以并发执行的 Tile 数量受到 GPU 上 SM 数量的限制。因此，整个 GPU 内核以多个 Tile 波 的形式执行，如图所示。这种 Tile 策略通过将操作分解为多个较小的工作负载来实现 GEMM 的可扩展执行，每个工作负载在不同的输入和输出元素上执行相同的计算。
Roofline 分析 提供了一种简单而有效的方法，通过考虑其算术强度来估算 GPU 内核的大致性能。算术强度定义为浮点运算次数 (𝑓𝑙𝑜𝑝𝑠ₖ) 与内存事务大小 (𝑚𝑒𝑚ₖ) 之比。可实现的性能，也称为 Roofline 带宽，可以计算如下：
*Roofline Bandwidth = min(𝑓𝑙𝑜𝑝𝑠ₚ, 𝑚𝑒𝑚𝐵𝑊ₚ * Arithmetic Intensity)*
其中，𝑓𝑙𝑜𝑝𝑠ₚ 表示 GPU 的峰值每秒浮点运算次数 (FLOPs)，𝑚𝑒𝑚𝐵𝑊ₚ 表示其峰值内存带宽。Roofline 带宽指示了内核在 GPU 上可以实现的最大吞吐量。
例如，如果一个内核的算术强度较低，意味着它主要受内存带宽限制，那么即使 GPU 的浮点运算能力很强，也无法充分发挥。反之，如果一个内核的算术强度很高，意味着它主要受计算能力限制，那么即使 GPU 的内存带宽很高，也无法显著提升性能。
波的数量和 GPU 吞吐量之间存在着重要的关系。随着波的数量增加，每个 SM 处理的波的数量也增加，从而导致每个 SM 的线程计数更高。这种增加的线程级并行性有助于有效地隐藏 stall 延迟，从而提高整体性能。
基于 Tile-Level 的内核延迟预测
NeuSight 的核心思想是将内核延迟预测分解为 Tile-Level 预测。为了根据每个 Tile 的预测来估计内核的总延迟，NeuSight 使用以下公式：
*Total Kernel Latency ≈ num_waves * PerTileLatency*
其中，num_waves 表示 Tile 波的数量，PerTileLatency 表示每个 Tile 的平均延迟。
为了计算 num_waves，NeuSight 使用以下公式：
numwaves = ceil(prod(xᵢ / tᵢ) / numsm)
其中，xᵢ 和 tᵢ 分别表示第 i 个维度中输出和 Tile 的大小，N 是输出维度数。num_sm 表示 GPU 上可用的 SM 数量。Tile 维度可以从 PyTorch Profiler 等工具提供的元数据中获取。
在上述公式中，NeuSight 假设每个 SM 一次处理一个 Tile。内核总延迟通常与 Tile 波的数量呈线性关系，可以近似为这些波的顺序执行。
PerTileLatency 取决于 SM 的计算利用率和内核实现的内存带宽 (achievedBW)。它的下限受到 Roofline 带宽 (rooflineBW) 的物理约束。这种关系可以用以下公式表示：
PerTileLatency = max(tilecomputetime, tilememoryaccess_time)
其中，tile_compute_time 是 Tile 的计算时间，tile_memory_access_time 是 Tile 的内存访问时间。这些时间可以通过 Roofline 模型进行估计，考虑到计算强度和内存带宽的限制。
接下来，NeuSight 的目标是预测 GPU 利用率。它利用了这样一个观察结果：随着工作负载中可用线程数量的增加，GPU 吞吐量接近其最大值。为了模拟波的数量和利用率之间的关系，NeuSight 采用了一种机器学习 (ML) 方法。
基于机器学习的 GPU 利用率预测
NeuSight 使用一个包含五个专门的多层感知器 (MLP) 来预测不同类型 DNN 内核的 GPU 利用率（alpha 和 beta）：
- 批量矩阵乘法 (batched matrix multiplication)
- 全连接层 (fully-connected layers)
- 元素操作 (element-wise operations)
- Softmax
- 层归一化 (layer normalization)
每个 MLP 具有 8 个全连接层，带有 512 个隐藏单元，并使用 ReLU 激活函数。如下图所示，输入特征包括 GPU 硬件规格，例如内存大小、带宽、峰值 FLOPs 和 L2 缓存大小，这些规格经过每个 SM 的归一化处理，并表示为资源利用率。
这些特征之所以被选择，是因为它们在各个 GPU 代系中都可用，并且对算子性能有影响。对于不支持或未知的算子，NeuSight 假设它们是内存受限的，并相应地估计延迟。
MLP 模型的目标是预测以下等式的系数：
*Utilization = alpha * (1 – exp(-beta * num_waves))*
其中，alpha 代表 GPU 可以达到的最大利用率，beta 代表利用率上升的速度。num_waves 是前面计算的 Tile 波的数量。
通过训练这些 MLP 模型，NeuSight 能够学习内核延迟、GPU 特性和内核属性（如算子类型、Tile 维度和算术强度）之间的复杂非线性关系，从而实现更准确的 GPU 利用率预测。
算子融合支持
算子融合 将多个算子或内核组合在一起，以减少来自中间结果的内存流量并提高操作强度。通过融合算子，可以避免将中间结果写入内存，从而节省大量的内存带宽和延迟。
NeuSight 通过排除中间数据的内存占用量来计算融合内核的 𝑓𝑙𝑜𝑝𝑠ₖ 和 𝑚𝑒𝑚ₖ。这使得 NeuSight 能够更准确地评估融合内核的性能，因为它考虑到了融合带来的性能提升。
例如，考虑一个由两个算子组成的融合内核：一个卷积算子和一个 ReLU 激活函数。如果没有算子融合，卷积算子的输出将被写入内存，然后 ReLU 激活函数将从内存中读取该输出。通过算子融合，卷积算子的输出可以直接传递给 ReLU 激活函数，而无需写入内存。这可以显著提高性能，尤其是在内存带宽有限的情况下。
NeuSight 工作流程
NeuSight 估计深度学习模型在单 GPU 或多 GPU 系统上的性能，分为三个步骤：
- 预测每个内核的执行延迟： 使用前面描述的 Tile-Level 预测方法和机器学习模型来预测每个内核的执行延迟。
- 根据模型的数据流图聚合这些预测，以计算每个 GPU 的延迟： NeuSight 从 PyTorch 中使用 Torch.fx 提取算子图，收集内核元数据，并应用其内核级预测器来注释每个算子。通过对预测的延迟求和来获得最终的每个 GPU 延迟，反映了内核在 GPU 上的顺序执行。
- 合并集合通信和网络开销以估计分布式环境中的整体性能： NeuSight 通过基于并行策略（管道、数据或张量模型并行）使用适当的通信算子来增强 DNN 图，从而模拟多 GPU 服务器（例如，DGX 系统）上的分布式执行。它通过测量链路利用率并基于目标系统的峰值带宽进行外推来估计网络操作（例如，基于环的 all-reduce 和点对点发送/接收）的延迟。对于管道并行，NeuSight 使用用户提供的计划（例如，GPipe）来模拟执行并插入管道气泡，并根据微批大小和通信开销来估计其延迟。对于张量模型和数据并行，它插入 all-reduce 算子以同步激活或梯度，并结合通信和计算延迟来预测端到端性能。
通过这三个步骤，NeuSight 能够准确地预测深度学习模型在各种 GPU 和分布式环境中的性能，从而帮助用户优化模型部署和资源分配。
例如，考虑一个使用数据并行训练的 ResNet-50 模型。NeuSight 首先预测 ResNet-50 模型中每个卷积层、ReLU 激活函数和全连接层的执行延迟。然后，它根据数据流图聚合这些预测，以计算每个 GPU 的延迟。最后，它合并 all-reduce 算子的延迟，以估计整个分布式训练过程的性能。
NeuSight 的优势与局限性
NeuSight 相比于传统 GPU 性能预测方法，具有以下显著优势：
- 更高的预测精度： 通过采用 Tile-Level 预测和机器学习模型，NeuSight 能够更准确地捕捉内核延迟、GPU 特性和内核属性之间的复杂关系，从而实现更高的预测精度。
- 无需实际执行： NeuSight 无需实际执行即可预测 GPU 性能，这大大节省了时间和资源。
- 支持各种深度学习模型： NeuSight 支持各种深度学习模型，包括训练和推理。
- 支持分布式环境： NeuSight 能够模拟多 GPU 服务器上的分布式执行，并预测分布式环境中的整体性能。
然而，NeuSight 也存在一些局限性：
- 依赖于硬件规格和元数据： NeuSight 的预测精度依赖于准确的硬件规格和内核元数据。
- 需要训练机器学习模型： NeuSight 需要训练机器学习模型来预测 GPU 利用率。
- 可能无法准确预测某些特殊内核的性能： 对于某些特殊内核，例如自定义内核或高度优化的内核，NeuSight 可能无法准确预测其性能。
未来展望
NeuSight 作为一种新颖的 GPU 性能预测框架，具有广阔的应用前景。未来，NeuSight 可以进一步扩展和改进，以支持更多类型的 GPU 和深度学习模型，并提高预测精度。
以下是一些可能的未来发展方向：
- 支持更多类型的 GPU： 将 NeuSight 扩展到支持更多类型的 GPU，包括来自不同厂商的 GPU 和新型 GPU 架构。
- 支持更多类型的深度学习模型： 将 NeuSight 扩展到支持更多类型的深度学习模型，包括 Transformer 模型、图神经网络等。
- 提高预测精度： 通过改进 Tile-Level 预测方法和机器学习模型，进一步提高预测精度。
- 自动化参数调优： 将 NeuSight 与自动化参数调优工具集成，以帮助用户自动优化模型部署和资源分配。
- 在线性能预测： 将 NeuSight 部署到在线环境中，以实时预测 GPU 性能，并根据预测结果动态调整资源分配。
结论
NeuSight 是一种创新的 GPU 性能预测 框架，它通过 Tile-level 预测 和机器学习技术，实现了对深度学习模型在不同 GPU 上的精准性能预测，无论是训练还是推理。通过结合 GPU 硬件特征和软件库行为，NeuSight 显著降低了预测误差，为深度学习模型的部署和优化提供了强大的工具。尽管存在一些局限性，但 NeuSight 代表了深度学习算力预测领域的重要进展，并为未来的研究和发展方向指明了道路。 随着大模型对 GPU 算力 需求的持续增长，NeuSight 这类预测框架的重要性将日益凸显，助力我们更有效地管理和利用有限的计算资源，推动人工智能技术的持续发展。