评价此页
★ ★ ★ ★ ★
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 可以启用异步数据加载，从而实现训练与数据加载的重叠。应根据工作负载、CPU、GPU 以及训练数据的位置来调整 num_workers。

DataLoader 接受 pin_memory 参数，默认为 False。当使用 GPU 时，最好设置 pin_memory=True。这会指示 DataLoader 使用固定内存（pinned memory），从而实现从主机到 GPU 更快且异步的内存拷贝。

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

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

对紧跟 Batch Norm 的卷积层禁用偏置（bias）#

torch.nn.Conv2d() 具有 bias 参数，默认值为 True（Conv1d 和 Conv3d 同理）。

如果 nn.Conv2d 层之后直接紧跟一个 nn.BatchNorm2d 层，那么卷积层中的偏置就不需要了。此时应使用 nn.Conv2d(..., bias=False, ....)。这是因为 BatchNorm 的第一步会减去平均值，这有效地抵消了偏置的影响。

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

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 与将其设置为零在数值行为上略有不同，详情请参阅文档。

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

融合操作#

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

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

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

在最简单的情况下，通过在函数定义上应用 torch.compile 装饰器即可启用融合，例如

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

有关更高级的用例，请参阅 torch.compile 入门。

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

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

channels_last 的支持目前处于实验阶段，但预计在标准计算机视觉模型（如 ResNet-50, SSD）上可以使用。要将模型转换为 channels_last 格式，请参考 Channels Last 内存格式教程。该教程包含关于 转换现有模型 的章节。

检查点（Checkpoint）中间缓冲区#

缓冲区检查点是一种缓解模型训练中内存容量负担的技术。它不再存储所有层的输入以用于反向传播中的梯度计算，而是仅存储少数层的输入，其余层则在反向传播期间重新计算。这种减少的内存需求使得增加批处理大小成为可能，从而提高利用率。

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

PyTorch 支持原生的 torch.utils.checkpoint API，可自动执行检查点和重新计算。

禁用调试 API#

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

CPU 专用优化#

利用非统一内存访问（NUMA）控制#

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

通常情况下，以下命令仅在第 N 个节点的内核上执行 PyTorch 脚本，避免跨插槽内存访问，以减少内存访问开销。

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

更详细的描述可以在此处找到。

利用 OpenMP#

OpenMP 用于为并行计算任务带来更好的性能。OMP_NUM_THREADS 是最简单的加速计算的开关。它决定了用于 OpenMP 计算的线程数。CPU 亲和性设置控制工作负载如何在多个内核上分布。它会影响通信开销、缓存行失效开销或页面抖动，因此正确设置 CPU 亲和性会带来性能收益。GOMP_CPU_AFFINITY 或 KMP_AFFINITY 决定了如何将 OpenMP 线程绑定到物理处理单元。详细信息可以在此处找到。

使用以下命令，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"

英特尔 OpenMP 运行时库（libiomp）#

默认情况下，PyTorch 使用 GNU OpenMP（GNU libgomp）进行并行计算。在英特尔平台上，英特尔 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 可获得良好的性能。以下命令显示了英特尔 OpenMP 运行时库的常见设置。

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

切换内存分配器#

对于深度学习工作负载，Jemalloc 或 TCMalloc 可以通过尽可能多地复用内存来获得比默认 malloc 函数更好的性能。Jemalloc 是一个通用 malloc 实现，强调避免碎片化和可扩展的并发支持。TCMalloc 也具有许多优化功能以加快程序执行。其中之一是在缓存中保留内存，以加快对常用对象的访问。即使在释放后仍保留此类缓存，也有助于避免如果稍后重新分配内存时产生的昂贵的系统调用。使用环境变量 LD_PRELOAD 来利用它们中的一个。

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

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

对于小规模模型或内存受限的模型（如 DLRM），在 CPU 上训练也是一个不错的选择。在多插槽机器上，分布式训练可以带来高效的硬件资源利用，从而加速训练过程。Torch-ccl 使用英特尔(R) oneCCL（集体通信库）进行了优化，用于高效的分布式深度学习训练，实现了如 allreduce, allgather, alltoall 等集体通信操作，实现了 PyTorch C10D ProcessGroup API，并可以作为外部 ProcessGroup 动态加载。基于 PyTorch DDP 模块中实现的优化，torch-ccl 加速了通信操作。除了对通信内核进行的优化外，torch-ccl 还具有同步计算-通信功能。

GPU 专用优化#

启用 Tensor Cores#

Tensor Cores 是专为计算矩阵-矩阵乘法运算而设计的专业硬件，主要用于深度学习和 AI 工作负载。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 自动调优器（auto-tuner）#

NVIDIA cuDNN 支持许多算法来计算卷积。自动调优器运行一个简短的基准测试，并针对给定的硬件和输入大小选择性能最好的内核。

对于卷积网络（目前不支持其他类型），请在启动训练循环之前通过设置以下内容启用 cuDNN 自动调优器：

torch.backends.cudnn.benchmark = True

  • 自动调优器的决策可能不是确定性的；对于不同的运行，可能会选择不同的算法。有关详细信息，请参阅 PyTorch：可复现性

  • 在极少数情况下，例如具有高度可变输入大小的模型，最好在禁用自动调优器的情况下运行卷积网络，以避免为每个输入大小选择算法所带来的开销。

避免不必要的 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，在 Volta 及更新的 GPU 架构上提供高达 3 倍的总体速度提升。要使用 Tensor Cores，必须启用 AMP，并且矩阵/张量维度应满足调用使用 Tensor Cores 的内核的要求。

要使用 Tensor Cores：

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

    • 参阅 深度学习性能文档 以获取更多详细信息和针对层类型的具体准则

    • 如果层大小是从其他参数推导出来的而不是固定的，仍然可以显式填充，例如 NLP 模型中的词汇表大小

  • 启用 AMP

针对变长输入预分配内存#

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

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

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

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

  3. 将梯度清零

  4. 继续常规训练

分布式优化#

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

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

DistributedDataParallel 提供了更好的性能和扩展到多 GPU 的能力。有关更多信息，请参阅 PyTorch 文档中 CUDA 最佳实践的相关章节。

如果在 DistributedDataParallel 和梯度累积下训练，请跳过不必要的 all-reduce#

默认情况下，torch.nn.parallel.DistributedDataParallel 在每次反向传播后执行梯度 all-reduce，以计算参与训练的所有工作节点上的平均梯度。如果训练使用 N 步的梯度累积，则不需要在每个训练步之后进行 all-reduce，只需在反向传播的最后一次调用后，在执行优化器之前进行一次 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 的存储桶（buckets）。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 库