如何把 PyTorch 的 GPU 利用率提升到 100% ?

时至今日，GPU 的珍贵程度无需多言，再加上当前特殊的大环境，手头拥有的 GPU 更显得是一种稀缺资源。在过去的几年里，我们在许多 PyTorch 案例中发现，用户手中的 GPU 并未得到充分的发挥，存在着大量的浪费现象。这不仅是对昂贵硬件资源的一种浪费，也限制了计算任务的效率和速度。因此，我们要尽可能把手中的 GPU 充分利用起来。

GPU 利用率 Link to heading

GPU 利用率 (Utilization) 有多种衡量方式，其中最常见的一种是 GPU 上有计算和图形活动的时间占总运行时间的比例。如果 GPU 利用率不足100%，则说明 GPU 在程序运行的某些时间处在空闲状态，没有被充分利用起来。

测量 GPU 利用率通常有以下几种方式:

NVIDIA System Management Interface (NSMI): nvidia-smi 包含一个设备监视器，在执行程序的同时通过 nvidia-smi dmon 可以查看许多 GPU 硬件数据，默认每一秒钟更新一次，其中的 sm % 列就是 GPU 利用率;

$ nvidia-smi dmon -o T -i 0
#Time        gpu    pwr  gtemp  mtemp     sm    mem    enc    dec    jpg    ofa   mclk   pclk
#HH:MM:SS    Idx      W      C      C      %      %      %      %      %      %    MHz    MHz
 21:10:26      0    360     40     46     86     44      0      0      0      0   2619   1980
 21:10:27      0    364     40     48     87     53      0      0      0      0   2619   1980
 21:10:28      0    363     39     49     86     47      0      0      0      0   2619   1980

NVIDIA Nsight Systems (NSYS): 作为 GPU 的性能分析器，NSight System 的功能十分强大，虽然 NSight System 可以用来测量 GPU 利用率，但它的运行时开销并不小，导致测量 GPU 利用率这种简单指标时，会有测不准的现象。想要测准需要一定的技巧，这里不推荐初学者使用 NSight System 测量 GPU 利用率;
NVIDIA Data Center GPU Manager (DCGM): 数据中心级 GPU 标配，通过 GPU Telemetry 可以实时观测到 GPU 的许多性能指标，例如 GPU 利用率、Tensor Core 利用率等等，目前许多数据中心上都已经配备了 DCGM，特别是在云上;
NVIDIA Management Library (NVML): NVML 为 GPU 硬件数据提供了编程接口，开发者可以通过编程的方式访问 GPU 的各项数据，其中就包含 GPU 利用率，nvidia-smi 和 DCGM 的背后就是 NVML，推荐高级开发者使用;

这里我们最推荐的查看 GPU 利用率的方式是 nvidia-smi dmon，它非常简单易用。

适用场景 Link to heading

GPU 利用率的高低依赖于用户手中的软硬件系统情况，还依赖于用户手中的应用程序特性。

这里以 NGC PyTorch 23.09 (nvcr.io/nvidia/pytorch:23.09-py3) 为例，用经 ColossalAI 优化过的 Stable Diffusion V2 和 Stable Diffusion XL 的训练过程作为测试模型，前者被用作 MLPerf Stable Diffusion 的参考实现 (mlcommons/training/stable_diffusion)。GPU 为 NVIDIA H100 80GB HBM3，使用两种 CPU:

Intel(R) Xeon(R) Platinum 8480C，DGXH100 标配，最大主频 3.8 GHz，2 sockets，共 112 个物理核;
AMD EPYC 7413 24-Core Processor，最大主频 3.6 GHz，1 socket，共 24 个物理核;

测得 GPU 利用率如下:

Model	CPU	Batch Size	GPU Utilization
Stable Diffusion 2.0	Intel(R) Xeon(R) Platinum 8480C	1	76%
Stable Diffusion 2.0	Intel(R) Xeon(R) Platinum 8480C	32	95%
Stable Diffusion 2.0	AMD EPYC 7413 24-Core Processor	1	55%
Stable Diffusion 2.0	AMD EPYC 7413 24-Core Processor	32	93%
Stable Diffusion XL 1.0	AMD EPYC 7413 24-Core Processor	1	38%

可见，软硬件系统和应用程序配置不同，GPU 利用率通常会有比较大的差异。仔细观察以上数据，GPU 利用率从 38% 到 95% 不等，我们不难看到 GPU 利用率受以下因素影响：

系统硬件配置: 在 batch size 为 1 时，Stable Diffusion 2.0 在 AMD EPYC 7413 的 GPU 利用率为 55%，在 Intel(R) Xeon(R) Platinum 8480C 上的 GPU 利用率为76%，一个好的 CPU 更能充分发挥 GPU 的能力；
GPU 上负载的大小: 相同的软硬件和相同的模型，大 batch size 的 GPU 利用率就远高于小 batch size 的 GPU 利用率，例如 Stable Diffusion 2.0 在 8480C 上，batch size 为 32 时的 GPU 利用率是 95%，batch size 为 1 时的 GPU 利用率是 76%；
应用程序特性: 同样是 Stable Diffusion，更大尺寸的 SD XL 的 GPU 利用率比 SD 2.0 还要低;

实际上，影响 GPU 利用率的因素还有很多，而且往往是多种因素相互作用的结果。GPU 利用率低的最常见情况是每个 GPU 上的 batch size 过小，在以下 3 个场景中频发:

大模型: 随着模型参数的增多，需要的 GPU 内存也随之增长，当没有足够的内存存放 activation 时，就不得不通过减小 batch size 的方法来限制模型的内存使用量;
大规模训练 (scale-out): 如果模型在单 GPU 或单个 DGX 上的训练时间过长，就需要使用大量的 GPU 训练模型。模型的 global batch size 通常有一个极大值，受此影响，在大规模训练的场景中通常也需要降低单 GPU 上的 batch size;
模型本身特性: 一些模型本身的特性决定了其 batch size 不能过大，否则难以收敛，例如 SSD。当然，存在一些增大 batch size 的手段，但往往需要额外的代价;

根本原因 Link to heading

造成 GPU 利用率低的表象虽然五花八门，但其深层次的原因却是 CPU 与 GPU 不协调: CPU 负载太多，或者 GPU 负载太少。GPU 负载好理解，主要指的是 CUDA kernel 以及 CPU 和 GPU 之间必要的内存拷贝。这里的 CPU 负载有两层含义:

一是指 CPU 参与的数据运算和逻辑控制: 数据运算如 dataloader 中的数据增强，或者把输入的文本划分为 token 并转为 tensor，逻辑控制如判定要不要梯度裁剪。它们的共同特征是 GPU 需要等来自 CPU 的数据或者决策完全就绪才能继续执行，一旦 GPU 需要这些数据或者决策，而 CPU 没有及时处理完，那么 GPU 就处在空闲状态，降低了 GPU 利用率。

二是指 CPU 为调用 CUDA API 所做的准备和清理工作: 比如一个 nn.Linear，它最终调用了 cuBLAS 中的 GEMM，但 PyTorch 需要从 Python 前端执行到 C++ 后端，需要为输出 tensor 和必要的 workspace 分配内存，需要根据 device、合适的精度、正向传播/反向传播等一些列信息决定执行哪个算子，cuBLAS 需要根据输入参数的信息查询 heuristics 决定调用哪一个 CUDA kernel，PyTorch 还需要在调用 kernel 后维护计算图信息、执行必要的清理工作。调用一个 CUDA kernel 需要准备一系列复杂的运行时环境，这些都由 CPU 负责完成。当模型中存在大量的小 CUDA kernel 时，CPU 准备运行时环境的时间就会超过 CUDA kernel 在 GPU 上的执行时间，造成 GPU 需要等待 CPU 完成运行时环境的准备，从而降低了GPU 利用率。

不幸的是，PyTorch 的设计以灵活易用为主，牺牲了一些性能方面的考虑。时至今日，PyTorch 的运行时环境已经非常复杂，同时一些不好的用户习惯带来了额外的运行时负担，造成 PyTorch 模型在使用过程中通常不能把 GPU 充分利用起来。

解决方案 Link to heading

知道了 GPU 利用率低的根本原因是 CPU 与 GPU 不协调，我们就可以对症下药，主要从两方面入手：降低 CPU 负载，和增加 GPU 负载。

首先是硬件解决方案，如果你的硬件系统配置不合理，请及时升级硬件系统。好马配好鞍，好的 GPU 需要好的 CPU 驱动。合理的硬件系统配置，能够充分发挥 GPU 的能力。硬件系统升级并不是必须的（土豪例外），硬件系统升级也不一定能解决所有的问题，但有助于提升GPU利用率。

从增加 GPU 负载的角度，以下手段可以提升 GPU 利用率:

增加每个 GPU 上的 batch size: 不是所有的场景都适用，GPU 数目不多且模型本身的 global batch size 较大时可以一试。除了直接增加单 GPU 上的 batch size，一些通过削减 GPU 内存使用量从而间接提升 batch size 的手段也可以考虑，包括但不限于 gradient/activation checkpointing，ZeRO，model parallelism，等等;
共享 GPU: 当一个应用无法充分利用 GPU 时，也可以考虑通过多进程、Multi-Process Service (MPS)、Multi-Instance GPU(MIG) 等手段让多个应用共享同一个GPU，量化投资领域的机器学习模型就有此类案例;

从减小 CPU 负载的角度，以下手段可以提升 GPU 利用率:

预处理/缓存: 针对需要 CPU 重复处理的数据，可以在 training 开始前通过预处理把它保存到文件系统中，例如把 training 数据集中的文本转为 token，在预处理阶段把 token 保存在文件系统中，避免在 training 中实时处理数据;
多线程: 现如今的 CPU 基本上都是多核处理器，把主线程上的重负载任务分配给其他 CPU 核，从而降低主线程的工作量，避免主线程因任务繁重而无法及时调用 CUDA kernel;
CPU 与 GPU 流水线: 把 CPU 和 GPU 组成软件流水线，确保 GPU 需要的数据及时就绪，例如在 GPU 执行当前 iteration 时，CPU 上的 dataloader 预取并处理下一个 iteration 的数据，从而避免由于 CPU 数据没有及时就绪而造成的 GPU 等待;
迁移到 GPU 上: 针对 CPU 负责的数据处理和逻辑控制部分，可以把他们迁移到 GPU 上，例如使用 DALI 在 GPU 上进行图像数据预处理，又比如梯度裁剪，原本是否进行梯度裁剪的 if/else 是在 CPU 上执行的，现在把整个梯度裁剪用 GPU 实现，包含其中的 if/else;
减少 CUDA kernel 数目: 通过算子融合，可以把多个小 CUDA kernel 合并成一个大 CUDA kernel，原来 CPU 需要为每一个 CUDA kernel 准备运行时环境，融合后则只需要准备一次，这会显著降低 CPU 的运行时开销。能够有效减少 CUDA kernel 数目的手段包括但不限于 vertical fusion，horizontal fusion，multi tensor apply 等等;
消除 CPU 与 GPU 之间的同步: CPU 与 GPU 之间的每一次同步，GPU 都会因为 CPU 在准备下一个 CUDA kernel 的运行时环境而进入空闲状态，甚至可能引起后续 CUDA kernel 的连锁反应，使得 CPU 没有足够的时间依次调用 CUDA kernel。移除 CPU 与 GPU 之间的所有同步点，可以把 PyTorch 程序变为 异步(asynchronous, sync-free) 程序，GPU 被 CPU 阻塞的概率大大降低。现阶段，不是所有 PyTorch 中的同步都可以被移除，个别同步点也比较难移除，好在绝大多数同步点都可以避免，例如 .cuda() 会导致一次同步，而 .cuda(non_blocking=True) 则不会有同步;
CUDA Graph: CUDA Graph 是解决运行时开销的终极武器，它可以基本消除 PyTorch 的所有运行时开销，威力十分强大，但在 PyTorch 中有一定的使用难度，目前也仅限于 static shape 的模型，需要有一定的经验，可以参考 CUDA Graphs。在 PyTorch 中，CUDA Graph 又分为 Partial-network capture 和 Whole-network capture，前者要容易很多，但效果不如后者，Torch 2.0 中的 TorchInductor 也可以自动给一些子图自动应用 CUDA Graph;

在我们所遇到的案例中，上述方法基本都有被用到，但如果追求极致性能，异步和 Whole-network CUDA Graph 是必不可少的手段。针对 Stable Diffusion，当我们应用了这两个优化后，即使单个 GPU 的batch size 为 1，GPU 利用率也是 100%:

$ nvidia-smi dmon -o T -i 0
#Time        gpu    pwr  gtemp  mtemp     sm    mem    enc    dec    jpg    ofa   mclk   pclk
#HH:MM:SS    Idx      W      C      C      %      %      %      %      %      %    MHz    MHz
 21:05:30      0    295     37     44    100     44      0      0      0      0   2619   1980
 21:05:31      0    356     38     46    100     43      0      0      0      0   2619   1980
 21:05:32      0    369     38     46    100     43      0      0      0      0   2619   1980

总结 Link to heading

总而言之，人工智能领域的竞争是分秒必争，GPU 等硬件资源十分稀缺且昂贵，提升 GPU 利用率十分有必要。一点小结:

使用设备监视器查看 GPU 利用率;
软硬件系统和应用程序配置不同，GPU 利用率通常会有比较大的差异;
常见的 GPU 利用率低的案例有单 GPU 上 batch size 小、大规模训练等等;
GPU 利用率低的根本原因是 CPU 与 GPU 不协调;
提升 GPU 利用率主要从降低 CPU 负载和增加 GPU 负载入手;
异步和 CUDA Graph 提升 GPU 利用率的终极手段;

虽然本文所述内容主要针对 PyTorch，但提升 GPU 利用率的原则和手段同样适用于其他框架和 CUDA 程序。

最后，把 GPU 利用率提升到 100%，只是端到端场景中优化深度学习模型性能的第一步，只有 GPU 利用率上来了，谈 GPU 的执行效率 (efficiency) 才是有意义的事情。