GPU利用率背后的性能真相

2024-09-05 AIGC 0

一般而言，机器学习团队理解GPU使用情况的常见度量标准是GPU利用率，通常通过在终端中运行nvidia-smi来查看。许多集成的可观测性工具也将GPU利用率作为其主要性能指标进行跟踪。

然而，AI Infra团队Trainy在实操中发现，GPU利用率并不总是理解GPU性能的最佳指标。实际上，在不做任何计算的情况下读取/写入内存，就可达到100%的GPU利用率！

本文作者Roanak Baviskar在本文中讲述了他们是如何发现这一点的，以及在这一过程中的其他发现。

（本文由OneFlow编译发布，转载请联系授权。来源：https://trainy.ai/blog/gpu-utilization-misleading）

作者｜Roanak Baviskar

OneFlow编译
题图由SiliconCloud平台生成

在Trainy，我们致力于管理GPU集群的基础设施，因此花费了大量时间思考这些问题。去年，我们与一家基础模型公司合作，以扩展和提高他们LLM训练的效率。我们遵循了几乎所有关于PyTorch性能调优指南中提到的基本步骤，即：

通过更改数据加载器默认值（num_workers，batch_size，pin_memory，预取因子等）来充分利用GPU。
通过使用混合精度（fp16，bf16）最大化使用张量核心
使用来自apex/deepspeed（例如 FusedAdam、FusedAdamW 等）的融合优化器
使用专为训练（H100SXM、A100SXM）设计的实例/网络。同时，可能的话，使用更新的实例H100 > A100 > V100

这些简单的更改使我们达到了100%的GPU利用率，并具有出色功耗，这真是太好了！为了检查是否还有更多提升空间，我们计算了训练工作负载的MFU（模型算力利用率）。

MFU，即模型FLOPS利用率，是理解GPU性能的最佳指标之一，这是在Google的PaLM论文中介绍的。它是“观察到的吞吐量（每秒Token数）与系统在峰值FLOPS运行的理论最大吞吐量的比例”。用更简单的话来说，它表示的是，相比GPU最大能力，你的工作负载每秒执行的浮点运算次数。MFU唯一的真正缺点是，与GPU利用率等指标相比，计算MFU可能有些困难，因为它依赖于参数和框架。

很遗憾，上述模型训练只达到了约20%的MFU。作为参考，目前大多数LLM训练达到了大约35%-45%的MFU。所以问题变成了：我们如何只用了GPU计算能力理论最大值的20%，同时GPU利用率却达到了100%？

为了回答这个问题，我们需要更好地理解GPU利用率实际上在跟踪什么。

1 GPU利用率到底是什么？

GPU利用率在Nvidia文档中宽松地定义为“当前报告的利用率针对的是GPU的计算资源和内存接口。”这种表述非常模糊。

更好的定义（出人意料地）可以在Datadog的NVML文档（https://docs.datadoghq.com/integrations/nvml/#metrics）中找到，"在过去采样周期内，GPU上一个或多个kernel执行的时间百分比。”为了确定这个定义为什么具有误导性，我们需要快速了解一下 GPU的工作原理。

GPU具有核心（core）和多处理管理器（https://cvw.cac.cornell.edu/gpu-architecture/gpu-characteristics/kernel_sm）。在Nvidia GPU中，这些多处理管理器被称为流多处理器（SM），而在AMD硬件中，这些被称为计算单元（CU）。以下是GH100 GPU的示意图，共有144个SM。

这些多处理管理器可被视为一组工人的工头，在这种情况下，就是核心（core）。当你启动一个CUDA kernel时，工作将由一个或多个SM在CUDA core上执行。如下面所示，GH100芯片上的单个SM包含许多CUDA core。

这意味着度量指标GPU利用率仅衡量在特定时刻是否有kernel在执行。它无法表明你的kernel是否充分利用了所有可用core，或者是否将工作负载并行化到GPU的最大能力。在最极端的情况下，通过仅在执行0 FLOPS的情况下对内存进行读写，可获得100%的GPU利用率。

现在我们澄清一下：这只会误导那些没有系统背景的人（比如许多机器学习工程师）。正如这里（https://arthurchiao.art/blog/understanding-gpu-performance/#24-the-use-methodology）提到的，GPU利用率的定义在“USE（https://www.brendangregg.com/usemethod.html）”方法论下是有意义的。

但回到我们面临的问题，这个定义确实解释了我们所看到的GPU利用率与MFU利用率之间的差距！确实还有更多的性能有待挖掘，我们只是需要找到它。

2 深入挖掘

下一步为了搜寻更多性能提升，当然就是对模型的训练循环进行性能分析。我们使用了PyTorch Profiler来审视训练循环，以便更好地了解情况。

如下图所示，Softmax kernel显示出很高的GPU利用率，但对于一个称为SM效率的指标来说却很低。这已经为我们敲响了警钟，因为朴素Softmax是LLM中一个臭名昭著的瓶颈，许多像FlashAttention这样的kernel融合的出现都是为了解决其内存受限特性。鉴于这些信息，SM效率统计指标可能指出我们模型执行中的低效问题。

3 但SM效率代表什么？

SM效率（也称为SM活动）是描述Nvidia GPU在特定时间间隔内有多少百分比的SM处于活跃状态的指标。如我们之前提到的，SM可以被视为一组CUDA核心的工头。例如，Nvidia H100 GPU有132个SM，每个SM有128个核心，总共提供了16896个核心。通过测算SM效率，我们可以确定CUDA核心是否在使用我们的流多处理器。例如，如果有一个CUDA kernel连续运行10秒，但只使用了1个SM，那么在H100上，这将记为100%的利用率，但SM效率将是1/132=0.7%。

太好了，这就是我们正在寻找的！我们可以逐层监控SM效率，以确定在优化方面的潜在收益中哪些是低垂果实。

4 优化

现在我们可以轻松识别哪些kernel在GPU上没有被充分利用，可以着手优化这些层。由于这是一个transformer堆栈，大多数收益将通过融合在transformer block定义中的层获得。下面的图表总结了我们优化的内容。

这里的融合意味着，不使用PyTorch原生定义的一组层，而是用一个GPU kernel替换它，该kernel使用CUDA或 Triton实现，将所有层合并为一个kernel。这种加速的结果是，与某些层（例如Softmax）的数学运算时间相比，每个kernel读/写GPU内存的时间减少了。Flash Attention就是一个这样的融合kernel的例子。需要融合的其他kernel包括MLP、dropout层归一化以及残差和操作。

我们自己编写了这些kernel吗？没有。大多数kernel已经在Flash Attention这样的库中实现，它们提供了nn.Modules层的实现，这样你就不用从头开始使用kernel实现torch.autograd.function。此外，这些实现通常已经进行了硬件优化，因此不仅更快，而且使用的内存更少。

最大的挑战在于，确定在你的代码中需要在何处替换合适的层。虽然torch.compile尝试自动完成这一任务，但截至目前，torch.compile与较新的分布式策略（如FSDP）不兼容（https://dev-discuss.pytorch.org/t/torch-compile-fsdp-dec-8th/1718），并且在实践中由于图断裂（graph break）并未提供很多承诺的加速。希望在未来，torch编译器能够为我们自动完成这一任务，但目前我们仍需手动添加融合实现。

在结果方面，我们实现了将训练时间提速 4 倍，该客户的MFU从最初20%的MFU提升到了38%。我们所做的大部分优化来自于这些融合kernel，以及在考虑到他们的模型大小和可用的3.2 Tbps Infiniband带宽的情况下，找到合适的“级别（level）”的模型并行度。

5 结论

我们强烈建议大多数AI团队在GPU集群以及GPU利用率上跟踪SM效率。这能提供更准确的性能指标，表明你从GPU中榨取了多少性能，而GPU利用率则可以衡量机器是否处于闲置状态。当然，计算MFU也很棒，但这并非是一个可以随时监控的指标。同时，Nvidia DCGM（数据中心GPU管理器）默认提供SM活动数据。

也有更精细的度量指标，例如SM占用率（或PyTorch Profiler中的已实现占用率Achieved Occupancy），这些指标告诉我们每个SM正在执行多少工作。然而，理解这些指标比尽可能提高SM效率要复杂得多。如果你对了解更多信息感兴趣，我建议你查看PyTorch Profiler博客、DCGM文档、Nsight的内核性能分析指南以及Nsight文档。

祝你好运，充分榨干你的GPU的性能！

其他人都在看