测试环境:
龙芯3C5000L,KVM虚拟化环境,8核/4核,皆配置8GB内存
Loongnix-20 Server,UOS v20,内核均为4.19
Python版本3.8/3.9都是从官方源下载的
- Loongnix-20 Server:Python 3.8
- UOS v20:Python 3.9
torch版本,皆来自Loongnix源 https://pypi.loongnix.cn/loongson/pypi/torch/
- Python 3.8:torch-2.0.0a0+git0bd6be9
- Python 3.9:torch-1.13.0a0+git7c98e70
尝试推理Phi 1.5语言模型和Tiny SD绘图模型
https://huggingface.co/microsoft/phi-1_5
https://huggingface.co/segmind/tiny-sd
但发现
- Python 3.8+Torch 2.0+Phi 1.5模型推理时仅有一个线程活跃
- Python 3.9+torch-1.13+tiny-sd模型推理时仅有两个线程活跃
查看torch.__config__.parallel info()
可以确定核心数被正确识别,OpenMP有启用,使用top -H可以发现对应核心数量的python线程已建立,但大部分线程在不活跃的状态,更改环境变量,使用numactl
命令指定CPU核心范围,设置set_num_threads
和set_num_interop_threads
都没有明显效果,示例代码片段如下
os.environ ['OMP_NUM_THREADS'] = str(cpu_num)
os.environ ['OPENBLAS_NUM_THREADS'] = str(cpu_num)
os.environ ['MKL_NUM_THREADS'] = str(cpu_num)
os.environ ['VECLIB_MAXIMUM_THREADS'] = str(cpu_num)
os.environ ['NUMEXPR_NUM_THREADS'] = str(cpu_num)
torch.set_num_threads(cpu_num)
torch.set_num_interop_threads(cpu_num)
实测Phi 1.5生成200个token耗时需要1000秒左右,tiny-sd大概100s/Iteration,执行20步推理生成一张512x512图像需要约30分钟。
作为对比,i5-7200u笔记本上四个线程可以跑满Phi 1.5生成200个token约为76秒,tiny-sd大概16s/Iteration,分别约为13倍和6.2倍,与Phi 1.5推理仅调用单核,tiny-sd仅调用双核的情况相符。
请问是否有小伙伴有相关的经验,能够让Pytorch在3C5000L上调用更多的核心?