Conda 版 PyTorch 的 CPU affinity 问题

问题表现

在部分 Conda 环境中,即使把 DataLoader 的 num_workers 设为大于 1,数据加载进程仍可能集中在少数 CPU 上,GPU 利用率也随之偏低。原记录中的 GPU 利用率曾低于 20%,但这个数值只是当时环境下的表现。

DataLoader 还可能提示:

This DataLoader will create 8 worker processes in total. Our suggested max number of worker in current system is 1, which is smaller than what this DataLoader is going to create.

这里需要区分两个概念:CPU affinity 是允许进程运行的逻辑 CPU 集合,CPU 编号则是集合中的元素。 例如,{0} 表示只能在编号为 0 的逻辑 CPU 上运行,不是“有 0 个 CPU 可用”。

检查亲和性

在 Linux 上,可以在主进程和 worker 中分别打印:

import os

cpus = os.sched_getaffinity(0)
print(f"pid={os.getpid()}, cpus={sorted(cpus)}, count={len(cpus)}")

其中,参数 0 指当前调用线程,返回值是它允许使用的 CPU 集合,具体语义见 Python 文档。

如果主进程可以使用多个 CPU,而 fork 后的 worker 只剩一个或少数几个,就需要继续检查运行时是否修改了亲和性。若主进程本身就受限,则应先核对任务调度器、容器或启动命令分配的 CPU 范围。

与 LLVM OpenMP 的关系

PyTorch 社区曾报告,某些 Conda 依赖组合中的 llvm-openmp 会导致 fork 子进程的亲和性异常。相关修复 pytorch/builder#1389 通过添加 llvm-openmp < 16 的依赖约束,规避当时的问题。

原笔记还记录了 17.x 下的异常,并提到 18.x 仍有同类表现。这些是旧环境记录,不能据此判断所有较新版本或所有渠道的 PyTorch 都存在相同问题。

临时处理

调整环境变量或依赖

可以先尝试在启动 Python 前设置:

export KMP_AFFINITY=disabled

该变量控制 OpenMP 运行时的线程亲和性行为,需在运行时初始化前设置,详见 LLVM OpenMP 文档。设置后,应重新检查 worker 的 CPU 集合,而不是仅凭警告是否消失判断效果。

对需要复现的旧环境,也可以参考上述 PR,将 llvm-openmp 约束为 <16。是否需要调整,应结合实际安装渠道和依赖版本判断。

在 worker 中恢复允许的 CPU 集合

另一种方法是通过 worker_init_fn 恢复亲和性。比起直接使用 range(os.cpu_count()),更合适的是先保存任务实际获准使用的 CPU 集合,再传给 worker:

import os
from functools import partial

# 在导入相关运行时、创建 worker 之前保存任务允许的 CPU 集合。
allowed_cpus = os.sched_getaffinity(0)


def restore_affinity(worker_id, *, cpus):
    os.sched_setaffinity(0, cpus)


worker_init_fn = partial(restore_affinity, cpus=allowed_cpus)

将这里的 worker_init_fn 传给 DataLoader 即可。如果保存时主进程的亲和性已经异常,应先修复主进程的设置;否则只是把错误范围重新传给 worker。

相关讨论

本文采用 CC BY-NC-SA 4.0 协议