Ubuntu上 PyTorch 与 CUDA 的配合指南
一 核心概念与版本匹配
- PyTorch 并不要求与系统 CUDA 完全一致，关键是让 PyTorch 的 CUDA 运行时与系统驱动能力匹配。用命令查看：
- 驱动与最高可用 CUDA：运行 nvidia-smi（右上显示 Supported/Runtime CUDA）
- 工具链版本：运行 nvcc --version
- 选择策略：优先让 PyTorch 的 cu12x/cu11x 与系统的 CUDA 12.x/11.x 大版本一致；多数情况下，PyTorch 的预编译包会随带合适的 CUDA 运行时（cudatoolkit），无需与系统 CUDA 完全相同。若遇到如 flash-attn 等需要编译的库，版本匹配更为关键。
二 安装与配置步骤
- 准备环境
- 更新系统并安装基础工具：sudo apt update && sudo apt install -y build-essential cmake git python3-dev python3-pip
- 建议使用 Miniconda 管理隔离环境：bash Miniconda3-latest-Linux-x86_64.sh，然后创建环境 conda create -n torch-gpu python=3.10
- 安装驱动与 CUDA（两种思路，二选一或混合）
- 思路A（系统级 Toolkit）：安装与你 GPU 匹配的 NVIDIA 驱动 与 CUDA Toolkit，并配置环境变量（见下节）。适合需要系统级 CUDA 工具链的场景。
- 思路B（Conda 自带 cudatoolkit）：在 Conda 环境中直接安装 cudatoolkit，PyTorch 使用自带的 CUDA 运行时，通常无需系统级 CUDA。示例：conda install pytorch torchvision torchaudio cudatoolkit=12.1 -c pytorch
- 安装 PyTorch（与 CUDA 对齐）
- Pip 方式（按系统 CUDA 选择对应 cu 版本）：pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu121
- Conda 方式（与上一步 cudatoolkit 保持一致）：conda install pytorch torchvision torchaudio cudatoolkit=12.1 -c pytorch
- 环境变量（仅在采用系统级 CUDA 时必要）
- 在 ~/.bashrc 或 ~/.zshrc 中添加并生效：
- export CUDA_HOME=/usr/local/cuda
- export PATH=$CUDA_HOME/bin:$PATH
- export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
- 生效：source ~/.bashrc
三 验证与常见故障排查
- 快速验证
- 查看 PyTorch、CUDA 可用性：
- python - <<‘PY’
import torch
print(“PyTorch:”, torch.version)
print(“CUDA available:”, torch.cuda.is_available())
print(“CUDA runtime:”, torch.version.cuda)
print(“Device:”, torch.cuda.get_device_name(0) if torch.cuda.is_available() else “None”)
PY
- 常见问题
- nvidia-smi 正常但 torch.cuda.is_available() 为 False
- 可能安装了 CPU 版 PyTorch；卸载后按上节重装带 CUDA 的版本（pip 选 cu12x/cu11x，或 conda 指定 cudatoolkit）。
- 版本不匹配（如安装 flash-attn 编译失败）
- 先确认 nvidia-smi 与 nvcc 的版本，再选择与系统 12.x/11.x 对齐的 PyTorch（如 cu121/cu118），然后重装依赖。
- 驱动问题
- 确保已安装兼容驱动；如遇到 Secure Boot 导致驱动未生效，按提示处理或临时关闭 Secure Boot 后重装驱动。
四 实践建议
- 优先使用 Conda 环境 + cudatoolkit，减少与系统 CUDA 的耦合；需要系统级 CUDA 工具链（如 nsight、自定义 CUDA 扩展）时，再采用系统级安装并正确设置 CUDA_HOME/LD_LIBRARY_PATH。
- 安装命令务必与 CUDA 大版本对齐（如本地为 12.x，选择 cu12x 的 PyTorch 预编译包），这是保证稳定性的关键。
- 记录环境信息（PyTorch 版本、CUDA runtime、驱动版本、是否 conda 安装），便于复现与排障。