在大模型部署过程中，GPU驱动版本不兼容是一个常见但棘手的问题。本文将通过对比不同解决方案，帮助ML工程师快速定位并解决该问题。
问题现象
当部署如LLaMA、ChatGLM等大模型时，若GPU驱动版本与CUDA版本不匹配，会出现如下错误：
RuntimeError: CUDA out of memory. Tried to allocate ...
AssertionError: CUDA error: invalid device function
问题排查步骤
1. 检查当前环境
# 查看CUDA版本
nvcc --version
# 查看GPU驱动版本
nvidia-smi
# 查看PyTorch CUDA支持
python -c "import torch; print(torch.version.cuda)"
2. 对比兼容性
常见兼容性对照表： | PyTorch版本 | CUDA版本 | 驱动版本 | |-------------|----------|----------| | 2.0.x | 11.8 | >=470.82 | | 2.1.x | 12.1 | >=535.54 |
3. 解决方案对比
方案A：升级驱动
# Ubuntu系统
sudo apt update
sudo apt install nvidia-driver-535
方案B：降级CUDA版本
# 使用conda环境
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia
方案C：使用Docker容器
FROM nvidia/cuda:11.8-runtime-ubuntu20.04
RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
最佳实践建议
推荐使用Docker方案进行生产环境部署，可有效避免驱动版本冲突问题。
讨论