参数超过阈值后Python多进程运行异常缓慢的原因排查

问题重现 #

脚本代码如下：

import torch
import multiprocessing as mp
from tqdm import tqdm

def simulate_condition(args):
    for i in range(args['n_steps']):
        # 基于PyTorch的随机矩阵马尔可夫链模拟逻辑
        ...

def proxy(args):
    return simulate_condition(args)

if __name__ == "__main__":
    num_cores = mp.cpu_count()-2  # 22核心
    args = [{'n_steps': n} for n in ...]  # 共40个任务
    with mp.Pool(num_cores) as pool:
        res = list(tqdm(pool.imap(proxy, args), desc="Simulate", total=len(args)))

测试不同n_steps的运行时间：

环境：Ubuntu 22.04 LTS、Ryzen 5900X、Python 3.10.4，22个进程仅占用6GB内存（总64GB），htop显示核心满负载但任务无进展。已尝试切换pool.map、spawn启动方式，均无效；但MacOS运行正常，替换为Numpy计算后恢复线性扩展。

可能的原因及排查方向 #

1. PyTorch CPU线程与多进程资源竞争 #

PyTorch默认会为CPU运算启用多线程（依赖OpenBLAS/MKL后端），当启动22个多进程后，每个进程再占用多个CPU线程，总线程数远超物理线程数，导致核心上下文切换过载，任务陷入低效循环。

• 验证方法：在simulate_condition开头添加torch.set_num_threads(1)，强制每个PyTorch进程只用1个线程，避免线程抢占。

2. Ubuntu内存分配器的碎片问题 #

Ubuntu默认的glibc内存分配器与MacOS的malloc实现不同，PyTorch在多进程中反复分配/释放小张量时，易产生内存碎片，导致后续分配效率急剧下降。

• 验证方法：
  • 在子进程中预先分配大张量复用，避免循环内频繁创建小张量；
  • 设置环境变量export MALLOC_TRIM_THRESHOLD_=0，或切换使用jemalloc内存分配器。

3. PyTorch内部全局状态的隐性同步 #

PyTorch CPU运算时可能存在全局状态（如随机数生成器共享、后端初始化锁），多进程高负载下触发隐性同步等待，导致进程阻塞。

• 验证方法：
  • 在simulate_condition开头为每个任务初始化独立的随机数生成器：torch.manual_seed(args['seed'])（每个任务分配唯一seed）；
  • 主进程避免提前初始化PyTorch资源，让子进程独立初始化上下文。

4. Linux内核调度或软锁问题 #

即便内存占用低，Linux内核可能因进程内存突发需求触发软锁，或调度器对密集CPU任务的策略导致进程饥饿。

• 验证方法：
  • 查看dmesg或/var/log/syslog，检查是否有OOM Killer记录；
  • 降低进程数至物理核心数（12），避免超线程过载；
  • 用nice命令降低脚本优先级，减少调度冲突。

优先验证的解决建议 #

优先尝试限制PyTorch单进程线程数，这是多进程+PyTorch CPU运算最常见的性能瓶颈：
修改子进程函数：

def simulate_condition(args):
    torch.set_num_threads(1)  # 强制单线程
    torch.set_num_interop_threads(1)
    for i in range(args['n_steps']):
        ...

若无效，再依次尝试内存分配器调整、进程数优化等方案。

内容的提问来源于stack exchange，提问作者Carol Eisen