llama-cpp-python模型无法调用NVIDIA GPU，默认使用CPU计算求助

在搭载可用NVIDIA GPU的RHEL节点上，通过Docker运行llama-2-7b-chat.Q3_K_L.gguf模型时，llama-cpp-python始终采用CPU计算，未调用GPU资源。已验证该GPU可正常运行其他模型。

使用的Docker启动命令 #

docker run -it --rm -p 8888:8888 --runtime=nvidia --gpus all -v /users/jupyter/data:/data -v /users/jupyter/notebooks:/project/notebooks llama-gpu

相关代码 #

环境配置与安装代码 #

!export FORCE_CMAKE=1
!export CMAKE_ARGS="-DLLAMA_CUBLAS=on"
!export LLAMA_CPP_LIB=/azureml-envs/tensorflow-2.12-cuda11/lib/python3.8/site-packages/llama_cpp_cuda/libllama.so
pip install llama-cpp-python

模型加载与推理代码 #

from llama_cpp import Llama

def question_generator(context):
    prompt = """[INST] <<SYS>>
        You are a helpful, respectful and honest assistant.
        Always respond as helpfully as possible, while being safe.
        Please ensure you generate the question based on the given context only
        <</SYS>>
        generate 3 questions based on the given content:-{}.
        """.format(context)
    
    llm = Llama(
        model_path="llama-2-7b-chat.Q3_K_L.gguf",
        n_ctx=8192,
        n_batch=512,
        use_mlock=True,
        n_gpu_layers=248,
        n_threads=8
    )
    
    output = llm(prompt,
               max_tokens=-1,
               echo=False,
               temperature=0.2,
               top_p=0.1)
    
    return output['choices'][0]['text']

批量推理代码 #

df["questions"]=""

for i in range(len(df)):
    df["questions"].iloc[i]=question_generator(df["text"].iloc[i])

已尝试的解决方法 #

CMAKE_ARGS="-DLLAMA_CUBLAS=on" FORCE_CMAKE=1 pip install llama-cpp-python --force-reinstall --upgrade --no-cache-dir 

排查与解决建议 #

• 验证CUDA编译状态
  运行以下代码检查llama-cpp-python是否启用CUDA支持：

  import llama_cpp
  print(llama_cpp.__version__)
  print(llama_cpp._llama.lib._params)
  

  输出中若包含cublas相关字段则编译正常，否则说明编译时未正确引入CUDA环境。

• 确认环境变量生效时机
  在Docker容器内直接执行echo $CMAKE_ARGS和echo $FORCE_CMAKE，确保安装llama-cpp-python时这些变量已正确设置。Jupyter Notebook中!export仅在当前shell会话生效，若后续重启内核或重新安装会失效，建议在Dockerfile中预先配置环境变量。

• 修正n_gpu_layers参数
  7B模型总层数约为32层，设置n_gpu_layers=248超出实际层数，建议改为n_gpu_layers=-1（将所有层加载到GPU）或n_gpu_layers=32，避免参数无效导致GPU未被调用。

• 检查模型与版本兼容性
  确认llama-2-7b-chat.Q3_K_L.gguf模型与当前llama-cpp-python版本兼容，部分旧版本对新GGUF格式支持不完善，可升级到最新稳定版尝试。

• 验证容器内GPU环境
  在容器内执行nvidia-smi确认GPU被识别，执行ldconfig -p | grep cublas确认CUDA库存在，确保容器内CUDA环境完整。

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