如何让LlamaForCausalLM的model.generate()占用2个以上CPU核心？

针对Llama模型在CPU上生成文本仅占用少量核心的问题，可通过调整PyTorch线程设置、模型加载参数来最大化CPU使用率，具体操作如下：

核心配置步骤 #

1. 强制PyTorch使用全部CPU核心
   PyTorch默认会限制线程数，需通过环境变量和API设置让其匹配你的CPU核心数（此处以12核为例）：

   • 设置底层BLAS库的线程数（适配OpenBLAS/MKL等后端）
   • 显式指定PyTorch的运算线程数

2. 优化模型加载参数
   CPU对float32 dtype的运算优化更充分，加载模型时指定该 dtype 并开启低内存模式，能更好地利用多线程资源。

修改后的完整代码 #

import torch
import os
from transformers import LlamaTokenizer, LlamaForCausalLM

# 配置CPU线程数，替换为你的实际核心数
cpu_core_count = 12
os.environ["OMP_NUM_THREADS"] = str(cpu_core_count)
os.environ["MKL_NUM_THREADS"] = str(cpu_core_count)  # 若用OpenBLAS则替换为OPENBLAS_NUM_THREADS
torch.set_num_threads(cpu_core_count)
torch.set_num_interop_threads(cpu_core_count)

# 加载tokenizer与模型
tokenizer = LlamaTokenizer.from_pretrained("/path/to/model")
model = LlamaForCausalLM.from_pretrained(
    "/path/to/model",
    torch_dtype=torch.float32,
    low_cpu_mem_usage=True,
    device_map="cpu"
)

prompt = "prompt text"
inputs = tokenizer(prompt, return_tensors="pt")

# 生成文本
generate_ids = model.generate(
    inputs.input_ids,
    max_length=1500,
    temperature=0.7,
    do_sample=True
)

output = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
print(output)

补充说明 #

• 自回归生成的特性决定了每一步计算依赖上一步结果，无法做到完全并行，但上述配置能让单步计算充分利用全部CPU核心，显著提升整体使用率。
• 若调整后仍未占满核心，可检查系统是否有其他进程占用资源，或确认PyTorch是否正确关联了优化的BLAS库（如MKL）。

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