Pytorch性能调优指南[模型性能调优系列1]

对于深度学习模型的性能调优，基本出发点有两个：

  • 提高计算速度和效率
  • 减少IO，包括磁盘和内存数据的读写

以下是对 PyTorch 官方文档《性能调优指南》（Performance Tuning Guide）的中文总结，涵盖了提升训练和推理效率的关键策略，在阅读时可以参照以上两点思考一下：

1. 通用优化策略

1.1. 异步数据加载与增强

  • torch.utils.data.DataLoader的num_workers参数默认为0，此时数据加载和训练都是在主进程中同步执行的，效率比较低；通过设置num_workers > 0可以实现数据加载与训练的并行处理。
    • 注意虽然由于python的GIL（Global Interpreter Lock）存在，其多线程并发并不是真正的并发，但并不会影响加载数据这类IO密集型任务：在等待某个worker的IO任务完成式会释放GIL，使得其他worker继续其数据加载任务
  • 启用 pin_memory=True 可加快数据从主机到 GPU 的传输速度。其原理在于为GPU在内存中专门分配一块区域用于存储训练数据，该内存区域不可被操作系统内核分页，且是连续的。GPU侧知道该内存区域的物理地址，可无需cpu介入通过DMA（Direct Memmory Access）直接访问，实现了训练+数据加载异步流式执行

1.2. 验证和推理时禁用梯度计算

  • 在执行验证或推理任务时，使用 torch.no_grad() 上下文管理器（也可以作为函数装饰器），避免不必要的梯度计算，减少内存占用并加快执行速度。

1.3. 卷积后接 BatchNorm 时禁用偏置

  • nn.Conv2d 层初始化时默认bias=True，后直接连接 nn.BatchNorm2d 时，可设置 bias=False，因为 BatchNorm 会抵消偏置的影响。

1.4. 使用 param.grad = None 替代 zero_grad()

  • pytorch会默认累加梯度，所以在每轮训练任务迭代开始前需要将上轮迭代的梯度置零，以上两种方式均可以达到目的。但通过将梯度设置为 None，可以减少内存操作，提高效率（对于参数量较大的大模型来说效果更好）。
    • 在下一轮迭代中，optimizer认为该tensor还没有执行过bp，会复用内存直接assign一个新的梯度变量给它
    • pytorch 1.7之后接口optimizer.zero_grad(set_to_none=True)默认为true，无需区分

1.5. 融合操作以减少内存访问

  • 在PyTorch的eager模式下，会为每个算子生成一个kernel并执行，如此每个算子都要执行：1)加载显存数据到GPU；2)执行对应的操作；3)将结果写回显存。其中1和3往往是性能瓶颈，通过将多个逐元素操作（如加法、乘法、激活函数等）融合为单个内核，减少内存读取和写入次数(参考算子融合)。示例如下：
@torch.compile
def gelu(x):
   return x * 0.5 * (1.0 + torch.erf(x / 1.41421))

1.6. 启用 channels_last 内存格式

  • 对于计算机视觉模型，使用 channels_last 内存格式可提高内存访问效率，尤其在使用混合精度训练时效果显著。

1.7. 中间结果检查点（Checkpointing）

  • 该策略通过只存储部分tensor中间结果以缓解模型训练时的内存压力，通过 torch.utils.checkpoint，在反向传播时重新计算部分前向传播结果，适用于深层模型或大批量训练。注意该策略对那些计算量小的，存储量大的tensor尤为有效，如各种激活函数(ReLU, Sigmoid, Tanh)， 上/下采样等

1.8. 禁用调试 API

  • 在常规训练中，禁用如 torch.autograd.detect_anomaly 等调试工具，以减少开销。

2. CPU 优化策略

2.1. 利用非一致性内存访问（NUMA）控制

在多插槽服务器上，使用 numactl 将进程绑定到特定的 CPU 节点，减少跨节点内存访问延迟。

2.2. 配置 OpenMP 线程和亲和性

通过设置环境变量（如 OMP_NUM_THREADS、GOMP_CPU_AFFINITY）优化线程使用和 CPU 亲和性，提高并行计算性能。

2.3. 使用 Intel OpenMP 运行时库（libiomp）

在 Intel 平台上，使用 libiomp 替代默认的 GNU OpenMP，可能获得更好的性能。

2.4. 切换内存分配器

使用 jemalloc 或 tcmalloc 替代默认的 malloc，提高内存分配和释放的效率。

2.5. 结合 TorchScript 使用 oneDNN Graph 进行推理优化

通过 torch.jit.trace 和 torch.jit.freeze，利用 oneDNN Graph 进行操作融合，提升推理性能，特别适用于 Float32 和 BFloat16 数据类型。

2.6. 在 CPU 上使用分布式数据并行（DDP）训练模型

对于小规模或内存受限的模型，使用 torch-ccl 和 DDP 在多核 CPU 上进行高效训练。

3. GPU 优化策略

3.1. 启用 Tensor Cores

在支持 Tensor Cores 的 GPU 上，使用混合精度训练（AMP）和适当的张量尺寸（通常为 8 的倍数）以充分利用硬件加速。

3.2. 使用 CUDA Graphs

通过 torch.compile(model, mode="reduce-overhead")，减少 CPU 与 GPU 之间的上下文切换，提高执行效率。

3.3. 启用 cuDNN 自动调优器

设置 torch.backends.cudnn.benchmark = True，让 cuDNN 自动选择最优的卷积算法，提高性能。

3.4. 避免不必要的 CPU-GPU 同步

尽量减少如 tensor.cpu()、tensor.item() 等操作，以避免阻塞 GPU 的执行。

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

使用 torch.rand(size, device='cuda') 直接在 GPU 上创建张量，避免额外的数据传输。

3.6. 使用混合精度训练（AMP）

启用 AMP，可在 Volta 及更新的 GPU 架构上获得高达 3 倍的速度提升，同时减少内存使用。

3.7. 预分配内存以处理可变长度输入

对于语音识别或 NLP 模型，预先分配最大序列长度的内存，避免训练过程中频繁的内存分配和释放。

4. 分布式训练优化

4.1. 使用高效的数据并行后端

推荐使用 torch.nn.parallel.DistributedDataParallel（DDP）替代 DataParallel，以获得更好的性能和扩展性。

4.2. 在梯度累积时跳过不必要的 all-reduce 操作

在使用 DDP 进行梯度累积时，利用 no_sync() 上下文管理器，避免在每次反向传播后进行 all-reduce 操作，仅在最后一次反向传播后执行。

4.3. 确保构造函数和执行过程中层的顺序一致

在使用 DistributedDataParallel(find_unused_parameters=True) 时，确保模型中层的定义顺序与实际执行顺序一致，以避免梯度同步问题。

参考链接

posted @ 2025-06-01 02:08  beanmoon  阅读(94)  评论(0)    收藏  举报