AllenNLP解码时使用multiprocessing比单进程运行迟缓问题咨询

问题根因 #

你观察到的torch._C._nn.linear耗时随进程数升高，核心是CPU计算资源超卖+缓存命中率下降，和共享内存的锁竞争没有关系，你之前做的模型共享内存、spawn启动的配置都是正确的，不需要调整：

• PyTorch CPU版本默认启用OpenMP进行算子内多线程并行，默认会占用所有可用CPU核。你启动N个进程时，每个进程都会启动等于CPU核数的OpenMP线程，总线程数远超物理核数，大量CPU时间浪费在线程上下文切换上。你测试的6核设备开12个进程时，总线程数可达72个，是物理核数的12倍，自然单步算子耗时飙升。
• 多进程同时访问共享内存中的模型参数时，会频繁冲毁CPU缓存，单进程推理时参数可以长期留在L3缓存中，多进程场景下每次linear算子调用都需要从主内存重新加载参数，访问延迟提升一个数量级。

解决方案 #

按优先级执行以下优化即可：

• 限制单进程算子并行数
  在导入PyTorch/AllenNLP之前，先设置环境变量关闭进程内的多线程：

import os
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"
# 再导入torch和allennlp相关库

调整后每个进程只用1个线程做计算，进程数设置为物理核数减1即可（6核设备开5个进程最优，不要开超过物理核的进程数，超线程对PyTorch CPU推理没有收益）。

• 调整任务提交逻辑
  不要用apply_async单句提交任务，改为按小批量提交：每次给子进程传入16~32句的批次，子进程调用predict_batch接口批量推理，大幅降低跨进程通信开销，同时提升算子计算效率。
• 显式关闭梯度计算
  在子进程的预测逻辑外层包裹torch.no_grad()，避免推理过程中生成不必要的计算图和梯度记录，进一步降低计算开销。
• MacOS适配调整
  MacOS对多进程CPU密集任务的调度效率远低于Linux，若有条件切换到Linux环境，相同配置下性能可以再提升30%以上。

优化后预期效果 #

按上述方案调整后，6核设备开5个进程，总耗时可以降到140秒左右，单进程的处理耗时和单跑时基本一致，不会出现随进程数升高的情况。

内容的提问来源于stack exchange，提问作者Sam Bayer