CPU 最佳实践

本章重点介绍如何在 CPU 上进行训练和推理时,通过环境设置获得最佳性能的最佳实践。

Intel

超线程

对于 GNN 领域等特定工作负载,为了获得最佳性能,建议的默认设置是关闭超线程。关闭超线程功能可以在 BIOS [1] 或操作系统级别 [2] [3] 完成。

替代内存分配器

替代内存分配器,例如 tcmalloc,可以通过更高效的内存使用显著提高性能,减少不必要的内存分配或释放开销。 tcmalloc 使用线程局部缓存来减少线程同步开销,分别通过自旋锁和每线程竞技场减少锁竞争,并按大小对内存分配进行分类以减少内存碎片化开销。

为了利用 tcmalloc 提供的优化,请在您的系统上安装它(在 Ubuntu 上,tcmalloc 包含在 libgoogle-perftools4 包中),并将共享库添加到 LD_PRELOAD 环境变量中

export LD_PRELOAD=/lib/x86_64-linux-gnu/libtcmalloc.so.4:$LD_PRELOAD

OpenMP 设置

由于 OpenMP 是默认的并行后端,我们可以通过 dgl.utils.set_num_threads() 控制包括采样和训练在内的性能。

如果未设置 OpenMP 线程数且数据加载器中的 num_workers 设置为 0,则 OpenMP 运行时通常默认使用可用的 CPU 核心数。这在大多数情况下效果良好,也是 DGL 中的默认行为。

如果数据加载器中的 num_workers 设置为大于 0,则每个工作进程的 OpenMP 线程数将设置为 1。这是 PyTorch 中的默认行为。在这种情况下,我们可以将主进程的 OpenMP 线程数设置为 CPU 核心数。

性能调优高度依赖于工作负载和硬件配置。我们建议用户尝试不同的设置,并为自己的情况选择最佳设置。

Dataloader CPU 亲和性

注意

此功能仅适用于 dgl.dataloading.DataLoader。尚不支持 dgl.graphbolt 中的数据加载器。

如果数据加载器工作进程数大于 0,请考虑使用 DGL Dataloader 类的 use_cpu_affinity() 方法,这通常会显著提高训练性能。

use_cpu_affinity 将设置适当的 OpenMP 线程数(等于分配给主进程的 CPU 核心数),将数据加载器工作进程绑定到单独的 CPU 核心,并将主进程限制在剩余核心。

在多个 NUMA 节点设置中,use_cpu_affinity 默认只使用 NUMA 节点 0 的核心,并假设工作负载在多个 NUMA 节点上扩展性较差。如果您认为您的工作负载在使用多个 NUMA 节点时会获得更好的性能,您可以传递用于数据加载的核心列表 (loader_cores) 和用于计算的核心列表 (compute_cores)。

loader_cores 和 compute_cores 参数(CPU 核心列表)可以传递给 enable_cpu_affinity,以便更精细地控制使用哪些核心,例如在工作负载在多个 NUMA 节点上良好扩展的情况下。

用法
dataloader = dgl.dataloading.DataLoader(...)
...
with dataloader.enable_cpu_affinity():
    <training loop or inferencing>

手动控制

关于 OpenMP 设置的更高级和更精细控制,请参考文章 Maximize Performance of Intel® Optimization for PyTorch* on CPU [4]

注脚

脚本总运行时间: (0 minutes 0.000 seconds)

由 Sphinx-Gallery 生成的图库