在Kaggle Notebook部署Oobabooga TextGen API时遇CUDA版本过低问题

问题描述 #

我是新手，尝试在Kaggle Notebook上部署Oobabooga TextGen WebUI，作为Autogen的开源LLM替代OpenAI API，但一直遇到CUDA版本过低的错误。相同脚本在Google Colab可正常运行，两者Python版本均为3.10，仅CUDA版本不同（Kaggle为11.4，Colab为12）。

报错后应用无法使用GPU，Notebook持续出现内存不足问题（Kaggle内存为Colab两倍）。

错误提示 #

CUDA initialization: The NVIDIA driver on your system is too old (found version 11040). Please update your GPU driver by downloading and installing a new version from the URL: http://www.nvidia.com/Download/index.aspx Alternatively, go to: https://pytorch.org to install a PyTorch version that has been compiled with your version of the CUDA driver.

解决方案 #

核心是强制安装适配Kaggle CUDA 11.4的PyTorch，并调整Oobabooga的依赖包版本，修改后的完整部署脚本如下：

import torch
from pathlib import Path

# 先卸载不兼容的PyTorch组件，适配CUDA 11.4
!pip uninstall -y torch torchvision torchaudio
!pip install torch==2.0.1+cu117 torchvision==0.15.2+cu117 torchaudio==2.0.2+cu117 --index-url https://download.pytorch.org/whl/cu117
# 注：CUDA 11.4兼容cu117的PyTorch包，无需安装cu114专属版本

if Path.cwd().name != 'text-generation-webui':
  print("Installing the webui...")

  !git clone https://github.com/oobabooga/text-generation-webui
  %cd text-generation-webui

  torver = torch.__version__
  print(f"TORCH: {torver}")
  is_cuda117 = '+cu117' in torver  # 适配Kaggle CUDA 11.4的兼容版本
  is_cuda118 = '+cu118' in torver
  is_cuda114 = '+cu114' in torver  # 新增CUDA 11.4判断

  textgen_requirements = open('requirements.txt').read().splitlines()
  if is_cuda117 or is_cuda114:
      # 替换所有高版本CUDA依赖为cu117，适配CUDA 11.4
      textgen_requirements = [req.replace('+cu121', '+cu117').replace('+cu122', '+cu117').replace('torch2.1', 'torch2.0') for req in textgen_requirements]
  elif is_cuda118:
      textgen_requirements = [req.replace('+cu121', '+cu118').replace('+cu122', '+cu118') for req in textgen_requirements]
  with open('temp_requirements.txt', 'w') as file:
      file.write('\n'.join(textgen_requirements))

  !pip install -r requirements.txt --upgrade
  !pip install -r extensions/openai/requirements.txt --upgrade
  !pip install -r temp_requirements.txt --upgrade

  print("\033[1;32;1m\n --> 如果看到\"previously imported packages\"的警告，忽略即可。\033[0;37;0m")
  print("\033[1;32;1m\n --> 无需重启运行时。\n\033[0;37;0m")

  try:
    import flash_attn
  except:
    !pip uninstall -y flash_attn

# Parameters
model_url ="typeof/dolphin-2.2.1-mistral-7b-sharded"
branch = ""
command_line_flags = "--n-gpu-layers 128 --load-in-4bit --use_double_quant"
api = True  # 开启API供Autogen调用

if api:
  for param in ['--api', '--public-api']:
    if param not in command_line_flags:
      command_line_flags += f" {param}"

model_url = model_url.strip()
if model_url != "":
    if not model_url.startswith('http'):
        model_url = 'https://huggingface.co/' + model_url

    # Download the model
    url_parts = model_url.strip('/').strip().split('/')
    output_folder = f"{url_parts[-2]}_{url_parts[-1]}"
    branch = branch.strip('"\' ')
    if branch.strip() != '':
        output_folder += f"_{branch}"
        !python download-model.py {model_url} --branch {branch}
    else:
        !python download-model.py {model_url}
else:
    output_folder = ""

# Start the web UI
cmd = f"python server.py --extensions openai --share --public-api"
if output_folder != "":
    cmd += f" --model {output_folder}"
cmd += f" {command_line_flags}"
print(cmd)
!$cmd

关键修改说明 #

• 强制安装兼容CUDA 11.4的PyTorch：先卸载原有PyTorch，安装torch==2.0.1+cu117版本，该版本可向下兼容CUDA 11.4环境
• 扩展CUDA版本判断逻辑：新增对cu114的检测，并统一将高版本CUDA依赖替换为cu117，确保所有库适配当前环境
• 开启API模式：默认设置api=True，自动添加--api和--public-api参数，方便Autogen调用

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