评价此页
★ ★ ★ ★ ★
recipes/recipes/tuning_guide

注意

前往文末 下载完整的示例代码。

性能调优指南#

创建日期：2020年9月21日 | 最后更新：2025年7月9日 | 最后验证：2024年11月5日

作者：Szymon Migacz

性能调优指南是一系列优化和最佳实践，可以加速 PyTorch 中深度学习模型的训练和推理。所呈现的技术通常只需更改几行代码即可实现，并可应用于所有领域的各种深度学习模型。

您将学到什么

  • PyTorch 模型通用优化技术

  • CPU 特定性能优化

  • GPU 加速策略

  • 分布式训练优化

先决条件

  • PyTorch 2.0 或更高版本

  • Python 3.8 或更高版本

  • 支持 CUDA 的 GPU（推荐用于 GPU 优化）

  • Linux、macOS 或 Windows 操作系统

概述#

性能优化对于高效的深度学习模型训练和推理至关重要。本教程涵盖了一套全面的技术，用于加速不同硬件配置和用例下的 PyTorch 工作负载。

通用优化#

import torch
import torchvision

启用异步数据加载和增强#

torch.utils.data.DataLoader 支持在单独的 worker 子进程中进行异步数据加载和数据增强。`DataLoader` 的默认设置为 `num_workers=0`，这意味着数据加载是同步的，并在主进程中完成。因此，主训练进程必须等待数据可用才能继续执行。

将 `num_workers` 设置为 `> 0` 可以启用异步数据加载，并使训练和数据加载之间重叠。`num_workers` 的值应根据工作负载、CPU、GPU 和训练数据的位置进行调整。

`DataLoader` 接受 `pin_memory` 参数，默认为 `False`。在使用 GPU 时，最好将 `pin_memory` 设置为 `True`，这将指示 `DataLoader` 使用 pinned 内存，并实现从主机到 GPU 更快、异步的内存复制。

禁用验证或推理的梯度计算#

PyTorch 会保存所有涉及需要梯度的张量的操作的中间缓冲区。通常，验证或推理不需要梯度。`torch.no_grad()` 上下文管理器可以应用于禁用指定代码块内的梯度计算，这可以加快执行速度并减少所需内存量。`torch.no_grad()` 也可以用作函数装饰器。

当卷积层直接后跟 batch norm 时，禁用卷积的偏置#

`torch.nn.Conv2d()` 具有 `bias` 参数，默认为 `True`（`Conv1d` 和 `Conv3d` 也是如此）。

如果 `nn.Conv2d` 层后面直接跟着一个 `nn.BatchNorm2d` 层，则卷积中的偏置是不需要的，可以使用 `nn.Conv2d(..., bias=False, ....)` 代替。偏置是不需要的，因为在第一步 `BatchNorm` 会减去均值，这有效地抵消了偏置的影响。

这同样适用于 1d 和 3d 卷积，只要 `BatchNorm`（或其他归一化层）在与卷积偏置相同的维度上进行归一化。

来自 `torchvision` 的模型已实现了此优化。

使用 `parameter.grad = None` 而不是 `model.zero_grad()` 或 `optimizer.zero_grad()`#

与其调用

model.zero_grad()
# or
optimizer.zero_grad()

来清零梯度，不如使用以下方法：

for param in model.parameters():
    param.grad = None

第二个代码片段不会清零每个单独参数的内存，并且后续的反向传播使用赋值而不是加法来存储梯度，这减少了内存操作的数量。

将梯度设置为 `None` 的数值行为与将其设置为零略有不同，更多细节请参阅[文档](https://pytorch.ac.cn/docs/stable/optim.html#torch.optim.Optimizer.zero_grad)。

或者，调用 `model` 或 `optimizer.zero_grad(set_to_none=True)`。

融合操作#

点运算（如元素加法、乘法以及 `sin()`、`cos()`、`sigmoid()` 等数学函数）可以合并为单个内核。这种融合有助于减少内存访问和内核启动时间。通常，点运算是内存绑定的；PyTorch 动态模式会为每个操作启动一个单独的内核，这涉及从内存加载数据、执行操作（通常不是最耗时的步骤）以及将结果写回内存。

通过使用融合算子，多个点运算只启动一个内核，数据只加载和存储一次。这种效率对于激活函数、优化器和自定义 RNN 单元等特别有益。

PyTorch 2 引入了由 TorchInductor 提供的编译模式，TorchInductor 是一个底层的编译器，可以自动融合内核。TorchInductor 将其功能扩展到简单的逐元素操作之外，能够对符合条件的逐元素和归约操作进行高级融合，以优化性能。

在最简单的情况下，可以通过将 `torch.compile` 装饰器应用于函数定义来启用融合，例如：

@torch.compile
def gelu(x):
    return x * 0.5 * (1.0 + torch.erf(x / 1.41421))

有关更高级的用例，请参阅[torch.compile 简介](https://pytorch.ac.cn/tutorials/intermediate/torch_compile_tutorial.html)。

为计算机视觉模型启用 `channels_last` 内存格式#

PyTorch 支持卷积网络的 `channels_last` 内存格式。该格式旨在与 AMP 结合使用，以使用 Tensor Cores 进一步加速卷积神经网络。

对 `channels_last` 的支持是实验性的，但预计适用于标准的计算机视觉模型（例如 ResNet-50、SSD）。要将模型转换为 `channels_last` 格式，请遵循[Channels Last Memory Format Tutorial](https://pytorch.ac.cn/tutorials/intermediate/memory_format_tutorial.html)。该教程包含一个关于[转换现有模型](https://pytorch.ac.cn/tutorials/intermediate/memory_format_tutorial.html#converting-existing-models)的部分。

检查点中间缓冲区#

缓冲区检查点是一种减轻模型训练内存容量负担的技术。它不是在反向传播中存储所有层的输入以计算上游梯度，而是存储少量层的输入，其余的在反向传播过程中重新计算。减少的内存需求允许增加批次大小，从而提高利用率。

应仔细选择检查点目标。最好的方法是不存储计算成本很小的、但输出很大的层。示例目标层包括激活函数（如 `ReLU`、`Sigmoid`、`Tanh`）、上/下采样以及具有小累积深度的矩阵向量运算。

PyTorch 支持本地 `torch.utils.checkpoint` API 来自动执行检查点和重新计算。

禁用调试 API#

许多 PyTorch API 旨在用于调试，应在常规训练运行时禁用它们：

  • 异常检测：`torch.autograd.detect_anomaly` 或 `torch.autograd.set_detect_anomaly(True)`

  • 剖析器相关：`torch.autograd.profiler.emit_nvtx`、`torch.autograd.profiler.profile`

  • autograd `gradcheck`：`torch.autograd.gradcheck` 或 `torch.autograd.gradgradcheck`

CPU 特定优化#

利用非统一内存访问 (NUMA) 控制#

NUMA 或非统一内存访问是一种内存布局设计，用于数据中心机器，旨在利用多处理器机器中具有多个内存控制器和块的内存局部性。总的来说，所有深度学习工作负载（训练或推理）在不访问跨 NUMA 节点的硬件资源的情况下都能获得更好的性能。因此，推理可以运行多个实例，每个实例运行在一个处理器上，以提高吞吐量。对于单节点上的训练任务，建议使用分布式训练，使每个训练进程运行在一个处理器上。

在一般情况下，以下命令仅在第 N 个节点的内核上执行 PyTorch 脚本，并避免跨处理器内存访问以减少内存访问开销。

numactl --cpunodebind=N --membind=N python <pytorch_script>

更详细的描述可以在[这里](https://intel.github.io/intel-extension-for-pytorch/cpu/latest/tutorials/performance_tuning/tuning_guide.html)找到。

利用 OpenMP#

OpenMP 用于提高并行计算任务的性能。`OMP_NUM_THREADS` 是可用于加速计算的最简单的开关。它决定了用于 OpenMP 计算的线程数。CPU 亲和性设置控制工作负载在多个核心上的分布。它会影响通信开销、缓存行失效开销或页面抖动，因此正确设置 CPU 亲和性可以带来性能优势。`GOMP_CPU_AFFINITY` 或 `KMP_AFFINITY` 决定如何将 OpenMP* 线程绑定到物理处理单元。详细信息可以在[这里](https://intel.github.io/intel-extension-for-pytorch/cpu/latest/tutorials/performance_tuning/tuning_guide.html)找到。

使用以下命令，PyTorch 在 N 个 OpenMP 线程上运行任务。

export OMP_NUM_THREADS=N

通常，以下环境变量用于使用 GNU OpenMP 实现设置 CPU 亲和性。`OMP_PROC_BIND` 指定线程是否可以跨处理器移动。将其设置为 `CLOSE` 会将 OpenMP 线程保持在连续分区内的主线程附近。`OMP_SCHEDULE` 决定 OpenMP 线程的调度方式。`GOMP_CPU_AFFINITY` 将线程绑定到特定的 CPU。一个重要的调优参数是核心固定（core pinning），它防止线程在多个 CPU 之间迁移，从而增强数据局部性并最小化跨核心通信。

export OMP_SCHEDULE=STATIC
export OMP_PROC_BIND=CLOSE
export GOMP_CPU_AFFINITY="N-M"

Intel OpenMP 运行时库 (`libiomp`)#

默认情况下，PyTorch 使用 GNU OpenMP (GNU `libgomp`) 进行并行计算。在 Intel 平台上，Intel OpenMP 运行时库 (`libiomp`) 提供 OpenMP API 规范支持。与 `libgomp` 相比，它有时能带来更高的性能优势。利用环境变量 `LD_PRELOAD` 可以将 OpenMP 库切换到 `libiomp`。

export LD_PRELOAD=<path>/libiomp5.so:$LD_PRELOAD

与 GNU OpenMP 中的 CPU 亲和性设置类似，`libiomp` 中提供了环境变量来控制 CPU 亲和性设置。`KMP_AFFINITY` 将 OpenMP 线程绑定到物理处理单元。`KMP_BLOCKTIME` 设置线程在完成并行区域执行后等待睡眠的时间（以毫秒为单位）。在大多数情况下，将 `KMP_BLOCKTIME` 设置为 1 或 0 会产生良好的性能。以下命令显示了使用 Intel OpenMP 运行时库的常见设置。

export KMP_AFFINITY=granularity=fine,compact,1,0
export KMP_BLOCKTIME=1

切换内存分配器#

对于深度学习工作负载，`Jemalloc` 或 `TCMalloc` 可以通过尽可能多地重用内存来比默认的 `malloc` 函数获得更好的性能。[Jemalloc](https://github.com/jemalloc/jemalloc) 是一个通用 `malloc` 实现，它侧重于避免碎片和可扩展的并发支持。[TCMalloc](https://ggdocs.cn/tcmalloc/overview.html) 也具有一些优化功能，可以加快程序执行速度。其中一项是持有内存缓存以加快对常用对象的访问。即使在释放内存后也持有这些缓存有助于避免昂贵的系统调用，如果这些内存稍后被重新分配的话。使用环境变量 `LD_PRELOAD` 来利用其中一个。

export LD_PRELOAD=<jemalloc.so/tcmalloc.so>:$LD_PRELOAD

使用 PyTorch ``DistributedDataParallel``(DDP) 功能在 CPU 上训练模型#

对于小型模型或内存绑定的模型（如 DLRM），在 CPU 上训练也是一个不错的选择。在具有多个处理器的机器上，分布式训练可以实现高效的硬件资源利用，以加速训练过程。`Torch-ccl` 针对 Intel(R) `oneCCL`（集合通信库）进行了优化，用于高效的分布式深度学习训练，实现了 `allreduce`、`allgather`、`alltoall` 等集合通信操作，实现了 PyTorch C10D `ProcessGroup` API，并且可以动态加载为外部 `ProcessGroup`。基于 PyTorch DDP 模块实现的优化，`torch-ccl` 加速了通信操作。除了对通信内核进行的优化外，`torch-ccl` 还具有同步计算-通信功能。

GPU 特定优化#

启用 Tensor Cores#

Tensor Cores 是专门用于计算矩阵-矩阵乘法运算的硬件，主要用于深度学习和人工智能工作负载。Tensor Cores 具有特定的精度要求，可以通过手动调整或通过自动混合精度 API 进行调整。

特别是，张量运算利用了较低精度的工作负载。这可以通过 `torch.set_float32_matmul_precision` 进行控制。默认格式设置为“highest”，它利用了张量数据类型。但是，PyTorch 提供了替代的精度设置：“high”和“medium”。这些选项优先考虑计算速度而非数值精度。

使用 CUDA Graphs#

在使用 GPU 时，工作首先必须从 CPU 启动，在某些情况下，CPU 和 GPU 之间的上下文切换可能导致资源利用率低下。CUDA Graphs 是一种在不产生额外的内核启动和主机同步成本的情况下，在 GPU 内部进行计算的方法。

# It can be enabled using
torch.compile(m, "reduce-overhead")
# or
torch.compile(m, "max-autotune")

CUDA Graph 的支持尚在开发中，其使用可能会导致设备内存消耗增加，并且某些模型可能无法编译。

启用 cuDNN 自动调优器#

[NVIDIA cuDNN](https://developer.nvidia.com/cudnn) 支持多种算法来计算卷积。自动调优器运行一个简短的基准测试，并为给定输入大小的给定硬件选择性能最佳的内核。

对于卷积网络（其他类型目前不支持），在启动训练循环之前，通过设置

torch.backends.cudnn.benchmark = True

  • 自动调优器的决策可能是非确定性的；不同的运行可能会选择不同的算法。有关更多详细信息，请参阅[PyTorch：可复现性](https://pytorch.ac.cn/docs/stable/notes/randomness.html?highlight=determinism)。

  • 在某些罕见情况下，例如输入大小变化很大的情况，最好禁用自动调优器来运行卷积网络，以避免为每个输入大小选择算法所带来的开销。

避免不必要的 CPU-GPU 同步#

避免不必要的同步，尽可能让 CPU 超前于加速器运行，以确保加速器工作队列中包含许多操作。

在可能的情况下，避免需要同步的操作，例如：

  • print(cuda_tensor)

  • cuda_tensor.item()

  • 内存复制：`tensor.cuda()`、`cuda_tensor.cpu()` 和等效的 `tensor.to(device)` 调用

  • cuda_tensor.nonzero()

  • 依赖 CUDA 张量操作结果的 Python 控制流（例如 `if (cuda_tensor != 0).all()`）

直接在目标设备上创建张量#

与其调用 `torch.rand(size).cuda()` 来生成随机张量，不如直接在目标设备上生成输出：`torch.rand(size, device='cuda')`。

这适用于接受 `device` 参数的所有创建新张量的函数：`torch.rand()`、`torch.zeros()`、`torch.full()` 等。

使用混合精度和 AMP#

混合精度利用 [Tensor Cores](https://www.nvidia.com/en-us/data-center/tensor-cores/)，并在 Volta 及更新的 GPU 架构上提供高达 3 倍的整体加速。要使用 Tensor Cores，应启用 AMP，并且矩阵/张量维度应满足调用使用 Tensor Cores 的内核的要求。

要使用 Tensor Cores：

  • 将尺寸设置为 8 的倍数（以映射到 Tensor Cores 的维度）

    • 有关更多详细信息和特定于层类型的指南，请参阅[深度学习性能文档](https://docs.nvda.net.cn/deeplearning/performance/index.html#optimizing-performance)。

    • 如果层大小派生自其他参数而不是固定的，也可以显式填充（例如 NLP 模型中的词汇大小）。

  • 启用 AMP

    • 混合精度训练和 AMP 简介：[幻灯片](https://nvlabs.github.io/eccv2020-mixed-precision-tutorial/files/dusan_stosic-training-neural-networks-with-tensor-cores.pdf)

    • 提供原生 PyTorch AMP：[文档](https://pytorch.ac.cn/docs/stable/amp.html)，[示例](https://pytorch.ac.cn/docs/stable/notes/amp_examples.html#amp-examples)，[教程](https://pytorch.ac.cn/tutorials/recipes/recipes/amp_recipe.html)

在输入长度可变的情况下预分配内存#

语音识别或 NLP 模型通常在具有可变序列长度的输入张量上进行训练。可变长度可能会对 PyTorch 缓存分配器造成问题，并可能导致性能下降或意外的内存不足错误。如果一个批次包含短序列长度，后面跟着另一个批次包含长序列长度，那么 PyTorch 将被迫释放前一次迭代的中间缓冲区并重新分配新的缓冲区。这个过程耗时，并导致缓存分配器中的碎片，这可能导致内存不足错误。

一种典型的解决方案是实现预分配。它包括以下步骤：

  1. 生成一个（通常是随机的）最大序列长度的输入批次（对应于训练数据集中的最大长度或某个预定义阈值）。

  2. 使用生成的批次执行前向和后向传播，不要执行优化器或学习率调度器，此步骤将预分配最大尺寸的缓冲区，这些缓冲区可以在后续训练迭代中重用。

  3. 清零梯度。

  4. 继续进行常规训练。

分布式优化#

使用高效的数据并行后端#

PyTorch 有两种实现数据并行训练的方法：

`DistributedDataParallel` 提供了更好的性能和多 GPU 扩展性。更多信息请参阅 PyTorch 文档中[CUDA 最佳实践的相关部分](https://pytorch.ac.cn/docs/stable/notes/cuda.html#use-nn-parallel-distributeddataparallel-instead-of-multiprocessing-or-nn-dataparallel)。

在使用 `DistributedDataParallel` 和梯度累积进行训练时，跳过不必要的 all-reduce#

默认情况下，`torch.nn.parallel.DistributedDataParallel` 在每次反向传播后执行梯度 all-reduce，以计算参与训练的所有工作节点上的平均梯度。如果训练使用 N 步的梯度累积，那么在每次训练步骤后都不需要 all-reduce，仅在最后一次调用 backward 之后、优化器执行之前执行 all-reduce 是必要的。

`DistributedDataParallel` 提供了 `no_sync()` 上下文管理器，它会禁用特定迭代的梯度 all-reduce。`no_sync()` 应应用于梯度累积的前 `N-1` 次迭代，最后一次迭代应遵循默认执行并执行所需的梯度 all-reduce。

如果使用 `DistributedDataParallel(find_unused_parameters=True)`，请匹配构造函数和执行过程中的层顺序#

`torch.nn.parallel.DistributedDataParallel` 结合 `find_unused_parameters=True` 使用模型构造函数中的层和参数顺序来构建用于 `DistributedDataParallel` 梯度 all-reduce 的桶。`DistributedDataParallel` 将 all-reduce 与反向传播重叠。特定桶的 all-reduce 仅在给定桶中所有参数的梯度都可用时才被异步触发。

为了最大化重叠量，模型构造函数中的顺序应大致匹配执行过程中的顺序。如果顺序不匹配，则整个桶的 all-reduce 将等待最后到达的梯度，这可能会减少反向传播和 all-reduce 之间的重叠，all-reduce 最终可能会暴露出来，从而减慢训练速度。

`DistributedDataParallel` 结合 `find_unused_parameters=False`（这是默认设置）依赖于在反向传播过程中遇到的操作顺序的自动桶形成。使用 `find_unused_parameters=False` 时，不需要重新排序层或参数即可实现最佳性能。

在分布式环境中平衡工作负载#

工作负载不平衡通常发生在处理顺序数据的模型中（语音识别、翻译、语言模型等）。如果一个设备收到的数据批次中的序列长度比其他设备的序列长度长，那么所有设备都将等待最后一个完成的工作节点。在具有 `DistributedDataParallel` 后端的分布式环境中，反向传播充当隐式同步点。

有多种方法可以解决工作负载不平衡问题。核心思想是在每个全局批次中尽可能均匀地将工作负载分配给所有工作节点。例如，Transformer 通过形成具有近似恒定数量 token 的批次（以及可变数量的序列）来解决不平衡问题，其他模型通过对具有相似序列长度的样本进行分桶，甚至通过对数据集按序列长度排序来解决不平衡问题。

结论#

本教程涵盖了一套全面的 PyTorch 模型性能优化技术。关键要点包括：

  • 通用优化：启用异步数据加载，禁用推理的梯度，使用 `torch.compile` 融合操作，以及使用高效的内存格式。

  • CPU 优化：利用 NUMA 控制，优化 OpenMP 设置，以及使用高效的内存分配器。

  • GPU 优化：启用 Tensor Cores，使用 CUDA Graphs，启用 cuDNN 自动调优器，以及实现混合精度训练。

  • 分布式优化：使用 DistributedDataParallel，优化梯度同步，以及平衡跨设备的负载。

其中许多优化可以通过最小的代码更改来实现，并在各种深度学习模型中提供显著的性能提升。

进一步阅读#

在此页面上
PyTorch 库