GPU Utilization is a Misleading Metric
原文作者：Roanak Baviskar
译者：菜小鸟魔王
01 What is GPU Utilization, really?
GPU 利用率在 Nvidia 的官方文档中定义得较为宽泛：“GPU 计算资源（compute resources of the GPU）和内存接口（memory interface）的当前利用率。”该定义颇为笼统。
然而，在 Datadog 的 NVML 文档中，我们可以找到一个更为精确的定义：“在最近一次采样周期内，GPU 上至少有一个 kernels 在运转的时间所占的百分比。”要理解这一定义为何可能会引起误解，我们需要简要了解 GPU 的工作机制。
GPU 包含了处理核心（cores）和 multiprocessing managers （译者注：在 GPU 中，这些管理器负责调度和协调多个处理核心的工作，确保它们可以高效地执行并行计算任务。）。在 Nvidia 的 GPU 中，这些 multiprocessing managers 被称为 streaming multiprocessors（SMs），而在 AMD 的 GPU 上，则被称为 compute units（CUs）。下图是拥有 144 个 SM 的 GH100 GPU。
这些 multiprocessing managers 可以类比为管理一群工人（即处理核心（cores））的领班。启动 CUDA kernel 时，kernel 任务会被分配到一个或多个 SMs ，然后在这些 SMs 的 CUDA 处理核心（cores）上完成。如下图所示，GH100 芯片上的每个 SM 都配备了大量的 CUDA 处理核心（cores）。
这个指标“GPU 利用率”实际上只是在检测在特定时刻是否有 kernel 正在执行，并不能反映 kernel 是否充分利用了所有可用的处理核心（cores），或者是否将 GPU 的并行处理能力发挥到了极致。最极端的一种情况 —— 在浮点运算次数为 0 FLOPS 时，仅通过内存的读写操作，也能达到 100% 的 GPU 利用率。
现在我们需要明确一点：对于不太了解底层系统原理的人，特别是那些专注于机器学习但不一定具备深厚底层系统知识背景的工程师来说，GPU 利用率的定义可能会造成误解。如此文所述，GPU 利用率的定义在“USE”方法论下是有其合理性的。
回到当前的问题来，这个定义确实可以解释 GPU 利用率与 MFU 为什么存在差异！显然，还有更多潜在性能等待我们去充分利用，关键在于如何找到并利用它们。
02 Digging Deeper
为了寻求获得更高的性能，下一步是分析模型的训练过程。我们通过 Pytorch Profiler 工具对训练过程进行了深入观察，以期获得更多信息。
如下图所示，Softmax kernel 的 GPU 利用率较高，但其所谓的 SM 效率指标却较低。这已经给我们敲响了警钟，因为传统的 softmax kernel 在大语言模型（LLMs）中是一个知名的性能瓶颈，许多 kernel fusions 技术（如 FlashAttention 等）都是为了解决内存受限问题而开发的。基于这些信息，SM 的效率低下可能表明模型在执行过程中存在某些环节效率不足的问题。
03 But What does SM Efficiency represent?
SM efficiency（亦称为 SM activity）是 Nvidia GPU 上的一个性能指标，反映了在特定时间间隔内活跃的 SMs 所占的百分比。如前所述，SMs 可以类比为管理 CUDA kernels 的领班。以 Nvidia H100 GPU 为例，它配备有 132 个 SMs，每个 SM 包含 128 个处理核心（cores），总计有 16,896 个处理核心（cores）。通过监测 SM efficiency，我们可以得知 CUDA kernels 是否充分利用了这些 SM。假设 CUDA kernel 持续运行了 10 秒钟，但仅使用了一个 SM，那么在 H100 GPU 上，利用率为 100%，但 SM efficiency 却只有 1 / 132 = 0.7%。
太棒了，这正是我们想要找的性能指标！我们可以逐层监控 SM efficiency ，以此来识别哪些可优化点是最容易提升性能的 “low-hanging fruits” 。
04 Making Optimizations
现在，我们能够轻松确定哪些 kernels 在 GPU 上的运行速度不高，我们可以开始对这些模型层进行优化。由于这是一个 transformer stack ，因此大部分性能提升来自于合并 transformer block definition 中的各个模型层。下图总结了我们的优化工作。
所谓的融合（fusing），是指我们不再使用 PyTorch 内置的模型层定义，而是用一个 GPU kernel 来替代它，这个新的 GPU kernel 是使用 CUDA 或 Triton 编写的，其作用是将所有模型层合并为一个 kernel。每个 kernel 读/写 GPU 内存的时间比某些模型层（如 Softmax）进行数学运算的时间要少，因此速度得以提升。Flash Attention 就是一个这样的 fused kernel 实例。其他需要 fuse 的 kernels 包括 MLP 和 dropout 层的 norm residual add operations（译者注：“norm residual add operations” 指的是在一个神经网络的层中，首先对输入进行层归一化，然后将归一化后的结果与原始输入通过残差连接相加的过程。）操作。
这些 kernels 是我们自己编写的吗？并不是。这些 kernels 大多数已经在像 Flash Attention 这样的库中实现，它们提供了模型层的 nn.Modules 实现方式，让我们无需使用 kernels 从零开始实现 torch.autograd.function。而且，这些库中实现的模型层通常已经针对特定硬件进行了优化，不仅运行速度更快，而且占用的内存也更少。
最大的挑战在于识别代码中哪些地方需要替换适当的模型层。尽管 torch.compile 试图自动完成这项工作，但在撰写本篇文章时，torch.compile 无法与最新的分布式策略（例如 FSDP）兼容，并且在实际应用中也并未带来预期的加速效果，主要因为存在 graph breaks 的问题。期待未来 torch 编译器能够自动完成这项工作，但目前，我们只能手动添加融合实现的模型层,实现性能优化。
至于加速 🧪 成果，我们为一位客户实现了 4 倍的训练时间加速，并将 MFU 提高了 38% ，从最初的 20% 提升到了 58%。我们所做的大部分优化都来自这些 fused kernels，以及根据模型大小和可用的 3.2 Tbps Infiniband 网络找到合适的“模型并行级别（“level” of model parallelism）”。
05 Conclusion
我们强烈建议大多数 AI 团队监控 GPU 集群上的 SM Efficiency 和 GPU 利用率。SM Efficiency 能更真实地反映 GPU 的性能，而 GPU 利用率则可以作为衡量机器是否空闲的指标。当然，计算 MFU 也很重要，但这并不是一个可以时时、逐层监控的指标。与此同时，Nvidia 的 DCGM（Data Center GPU Manager）默认就提供了 SM 的活动数据。
还有更细致的指标，比如 SM occupancy（或 Pytorch Profiler 中的 Achieved Occupancy），能够告诉我们每个 SM 完成了多少工作。然而，理解这些指标并不像尽可能提高 SM Efficiency 那样直接。如果你对这些内容感兴趣，想了解更多内容，我建议你查看 Pytorch Profiler 的博客、DCGM 的文档、Nsight 的 kernel 分析指南和 Nsight 的文档。
Thanks for giving this a read, and good luck squeezing every bit of performance out of your GPUs!
