本地运行Hugging Face Code Llama无法调用GPU，报Torch未编译CUDA

本地部署Hugging Face的Code Llama模型时，已安装Nvidia驱动和CUDA Toolkit，但模型仅用CPU运行；设置device_map="cuda:0"后报错：

AssertionError: Torch not compiled with CUDA enabled

核心问题是当前环境的PyTorch未编译CUDA支持，与Nvidia驱动/CUDA Toolkit无关，按以下步骤修复：

• 验证PyTorch的CUDA支持状态
  运行以下代码确认：

  import torch
  print(torch.cuda.is_available())  # 输出False则说明无CUDA支持
  print(torch.version.cuda)         # 若无CUDA支持，可能输出None
  

• 卸载现有PyTorch
  使用pip卸载：

  pip uninstall torch torchvision torchaudio -y
  

  若用conda环境，执行：

  conda uninstall torch torchvision torchaudio -y
  

• 安装匹配CUDA版本的PyTorch
  先通过nvcc --version查看本地CUDA Toolkit版本，选择对应版本的PyTorch安装命令：

  • 若CUDA版本为11.8：

    pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
    

  • 若CUDA版本为12.1：

    pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
    

  注：确保PyTorch版本与本地CUDA版本兼容，无需强行追求最新版本。

• 验证安装成功
  再次运行之前的验证代码，若torch.cuda.is_available()返回True，说明CUDA支持已生效。

• 优化代码显存占用
  为7B模型添加torch_dtype=torch.float16参数，减少GPU显存消耗（约14G显存即可运行）：

  from transformers import AutoTokenizer
  import transformers
  import torch
  
  model = "codellama/CodeLlama-7b-hf"
  
  tokenizer = AutoTokenizer.from_pretrained(model)
  pipeline = transformers.pipeline(
      "text-generation",
      model=model,
      torch_dtype=torch.float16,
      device_map="cuda:0"
  )
  
  prompt = "Write python code to reverse a string"
  
  sequences = pipeline(
      prompt,
      do_sample=True,
      top_k=10,
      temperature=0.1,
      top_p=0.95,
      num_return_sequences=1,
      eos_token_id=tokenizer.eos_token_id,
      max_length=200,
  )
  for seq in sequences:
      print(f"Result: {seq['generated_text']}")
  

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