如何确定多线程NumPy程序的MPI进程数？

问题背景 #

查阅mpi4py文档后，未找到关于MPI多进程与多线程NumPy结合的实践指导。文档中涉及NumPy的内容多集中在Bcast、Scatter这类数据移动操作，但对于并行化计算密集型任务（比如np.linalg.solve这类依赖BLAS/LAPACK的线性求解），缺少两个关键信息：

• 当每个MPI进程本身会占用多线程/核心时，如何选择合适的MPI进程数
• 如何确定当前NumPy程序实际使用的线程/核心数

测试场景与代码 #

在12线程笔记本上做了基础测试，场景为每个MPI进程独立求解不同的线性系统（各进程的(A, b)数据互不相同），测试代码如下：

import time 

from mpi4py import MPI 
import numpy as np 
import numpy.random as npr 

world = MPI.COMM_WORLD
world_size: int = world.Get_size()
rank: int = world.Get_rank()
master: bool = (rank == 0)

n: int = 8_000 # 单进程运行耗时约2.5s
A: np.ndarray = npr.randint(0, 10, (n, n))
b: np.ndarray = npr.randint(0, 10, (n,))
start: float = time.perf_counter() 
_ = np.linalg.solve(A, b) 
runtime: float = time.perf_counter() - start 
if master: 
    print(f"{runtime=:.3f}s")

测试结果 #

• 启动的MPI进程数越多，运行时间基本呈线性增长：比如mpiexec -n 2的耗时约为mpiexec -n 1的两倍
• 若将np.linalg.solve替换为time.sleep，则进程数增加不会导致运行时间增长

推测原因：BLAS/LAPACK例程默认会占用全部可用核心/线程，多MPI进程同时运行时，会引发CPU资源竞争，导致实际串行执行。

核心问题 #

当目标是并行化本身使用多线程的NumPy程序时，如何确定可使用/应使用的MPI进程数？

解决方案 #

1. 明确NumPy依赖的线程库及当前线程数 #

• 查看NumPy链接的BLAS/LAPACK库：执行np.__config__.show()，输出中会显示当前使用的库（比如OpenBLAS、MKL、ATLAS）
• 对应库的线程数控制环境变量：
  • OpenBLAS: OPENBLAS_NUM_THREADS
  • MKL: MKL_NUM_THREADS
  • ATLAS: ATLAS_NUM_THREADS
  • 通用兼容：OMP_NUM_THREADS

2. 计算最优MPI进程数 #

核心原则是总CPU线程数 ≈ MPI进程数 × 每个进程的BLAS线程数，同时预留1-2个线程给系统避免资源耗尽。
举例：12线程的机器，若设置每个进程的BLAS线程数为2，则最优MPI进程数为(12-1)/2 ≈ 5；若设置每个进程用1个BLAS线程，则MPI进程数可设为11。

3. 显式控制线程数并验证 #

• 方式一：运行前通过环境变量设置

  # 以OpenBLAS为例，设置每个进程用1个线程
  export OPENBLAS_NUM_THREADS=1
  mpiexec -n 11 python your_script.py
  

• 方式二：代码开头设置（需在导入numpy之前）

  import os
  os.environ['OPENBLAS_NUM_THREADS'] = '1'
  import numpy as np
  

• 验证：测试不同进程数+线程数的组合，找到耗时最短的配置

内容的提问来源于stack exchange，提问作者Nick Richardson