GPU 环境配置

Note

本文已按 Ubuntu 26.04 LTS 2026 年 4 月发布说明做初步复核，AI/GPU 驱动和 CUDA/ROCm 版本仍需随供应商支持矩阵继续验证。

GPU 是运行大多数 AI 工作负载的关键硬件。本文将详细介绍在 Ubuntu 26.04 上配置 NVIDIA GPU 环境的完整流程，包括驱动安装、CUDA Toolkit、cuDNN 以及验证步骤。同时也简要介绍 AMD GPU 的 ROCm 配置。

26.04 仓库状态

Ubuntu 26.04 的仓库已提供 NVIDIA CUDA Toolkit（sudo apt install cuda-toolkit）和 AMD ROCm 7.1.0（sudo apt install rocm）计算栈包。它们不会默认预装；如果用于生产训练、特定 PyTorch/TensorFlow 版本或多 GPU 服务器，仍应以 NVIDIA/AMD 官方支持矩阵和框架发布说明为最终依据。

NVIDIA 驱动安装

Ubuntu 26.04 提供了便捷的 NVIDIA 驱动安装工具 ubuntu-drivers。

检查硬件


# 查看系统中的 NVIDIA GPU
lspci | grep -i nvidia
 
# 查看推荐的驱动版本
sudo ubuntu-drivers devices

输出示例：


== /sys/devices/pci0000:00/0000:00:01.0/0000:01:00.0 ==
modalias : pci:v000010DEd00002684sv...
vendor   : NVIDIA Corporation
model    : AD102 [GeForce RTX 4090]
driver   : nvidia-driver-560 - third-party non-free recommended
driver   : nvidia-driver-550 - third-party non-free
driver   : nvidia-driver-555 - third-party non-free

自动安装(推荐)


# 自动安装推荐的驱动版本
sudo ubuntu-drivers autoinstall
 
# 重启系统使驱动生效
sudo reboot

手动安装指定版本


# 查看可用驱动版本
apt list nvidia-driver-* 2>/dev/null | grep -v "Listing"
 
# 安装指定版本
sudo apt install -y nvidia-driver-560
 
# 重启系统
sudo reboot

使用 NVIDIA 官方仓库安装(最新版本)


# 添加 NVIDIA 官方仓库（Ubuntu 26.04 对应 ubuntu2604）
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2604/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt update
 
# 安装最新驱动
sudo apt install -y nvidia-open
 
# 重启
sudo reboot

验证驱动安装


# 查看驱动版本和 GPU 状态
nvidia-smi

正常输出示例：


+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 560.35.03              Driver Version: 560.35.03      CUDA Version: 12.6     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 4090       Off  | 00000000:01:00.0  On  |                  Off |
|  0%   35C    P8              18W / 450W |     512MiB / 24564MiB  |      0%      Default |
+-----------------------------------------+------------------------+----------------------+

Tip

nvidia-smi 输出中的 “CUDA Version” 显示的是驱动支持的最高 CUDA 版本，不是实际安装的 CUDA Toolkit 版本。你需要单独安装 CUDA Toolkit。

CUDA Toolkit 安装

CUDA Toolkit 是 NVIDIA 的 GPU 计算平台和编程模型，是绝大多数 AI 框架(PyTorch、TensorFlow 等)的依赖。

驱动与 CUDA 版本的关系

CUDA Toolkit 对最低 NVIDIA 驱动版本有要求，新版 CUDA 需要更新的驱动。安装前请以 NVIDIA CUDA Toolkit 发布说明和 PyTorch / TensorFlow 官方的版本支持矩阵为准，选择三者都兼容的组合。一般做法是：先用 nvidia-smi 查看驱动支持的最高 CUDA 版本，再据此选择 CUDA Toolkit 与框架轮子（wheel）。

方法一（推荐）：通过 Ubuntu 26.04 仓库安装

Ubuntu 26.04 的仓库已经收录 CUDA Toolkit，对大多数本地开发来说这是最省事的方式：


# 直接安装 Ubuntu 仓库中的 CUDA Toolkit
sudo apt install -y cuda-toolkit
 
# 验证
nvcc --version

方法二：通过 NVIDIA 官方仓库安装指定版本

如果你需要特定版本的 CUDA Toolkit（例如匹配某个 PyTorch 轮子）：


# 添加 NVIDIA 官方仓库（Ubuntu 26.04 对应 ubuntu2604）
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2604/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt update
 
# 查看仓库中可用的 CUDA Toolkit 版本
apt-cache search cuda-toolkit
 
# 安装指定版本（将 XX-Y 替换为实际可用版本，如 cuda-toolkit-13-0）
sudo apt install -y cuda-toolkit-XX-Y

Note

Ubuntu 26.04 已将 NVIDIA CUDA Toolkit 纳入官方仓库，可直接 sudo apt install cuda-toolkit 安装。NVIDIA 的 ubuntu2604 仓库仍可用于需要特定 CUDA 版本（例如匹配某个 PyTorch 轮子）的场景。

配置环境变量


# 添加 CUDA 到 PATH
cat >> ~/.bashrc << 'EOF'
 
# CUDA Toolkit
export CUDA_HOME=/usr/local/cuda
export PATH=$CUDA_HOME/bin:$PATH
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
EOF
 
source ~/.bashrc

验证 CUDA 安装


# 检查 CUDA 编译器版本
nvcc --version
 
# 输出示例：
# nvcc: NVIDIA (R) Cuda compiler driver
# Cuda compilation tools, release X.Y, VX.Y.xxx
 
# 运行 CUDA 示例(可选)
# 较新版本的 CUDA 不再随 Toolkit 分发示例代码，需从 GitHub 克隆
sudo apt install -y cmake build-essential git
git clone https://github.com/NVIDIA/cuda-samples.git
cd cuda-samples
cmake -B build && cmake --build build -j$(nproc)
 
# 运行设备查询（路径随版本可能不同）
./build/Samples/1_Utilities/deviceQuery/deviceQuery
 
# 如果输出 "Result = PASS" 则安装成功

安装多个 CUDA 版本


# 通过 NVIDIA 官方仓库可以同时安装多个版本（将 XX-Y 替换为实际版本）
sudo apt install -y cuda-toolkit-XX-Y
sudo apt install -y cuda-toolkit-ZZ-W
 
# 查看已安装的版本
ls /usr/local/ | grep cuda
 
# 切换默认版本(通过更新软链接，将 X.Y 替换为目标版本)
sudo rm /usr/local/cuda
sudo ln -s /usr/local/cuda-X.Y /usr/local/cuda
 
# 验证当前版本
nvcc --version

cuDNN 安装

cuDNN(CUDA Deep Neural Network library)是 NVIDIA 的深度学习加速库，PyTorch 和 TensorFlow 都依赖它。

通过 apt 安装(推荐)


# 确保已添加 NVIDIA 仓库
# 安装 cuDNN for CUDA 12
sudo apt install -y cudnn9-cuda-12
 
# 或安装特定版本
sudo apt install -y libcudnn9-cuda-12 libcudnn9-dev-cuda-12
 
# 验证安装
dpkg -l | grep cudnn

通过 tar 包安装


# 从 NVIDIA 开发者网站下载 cuDNN
# https://developer.nvidia.com/cudnn-downloads
# 需要免费的 NVIDIA 开发者账号
 
# 解压到 CUDA 目录
sudo tar -xzf cudnn-linux-x86_64-9.x.x.x_cuda12-archive.tar.xz
sudo cp cudnn-linux-x86_64-9.x.x.x_cuda12-archive/include/* /usr/local/cuda/include/
sudo cp cudnn-linux-x86_64-9.x.x.x_cuda12-archive/lib/* /usr/local/cuda/lib64/
sudo ldconfig

验证 cuDNN


# 检查 cuDNN 是否可以被找到
ldconfig -p | grep cudnn
 
# 通过 Python 验证(需要先安装 PyTorch)
python3 -c "import torch; print('cuDNN version:', torch.backends.cudnn.version()); print('cuDNN enabled:', torch.backends.cudnn.enabled)"

验证 GPU 可用性

安装完所有组件后，进行全面验证。

系统级验证


# 1. 检查驱动
nvidia-smi
 
# 2. 检查 CUDA
nvcc --version
 
# 3. 检查 cuDNN
ldconfig -p | grep cudnn
 
# 4. 查看 GPU 详细信息
nvidia-smi -q | head -50
 
# 5. 监控 GPU 实时状态
watch -n 1 nvidia-smi

PyTorch GPU 验证


# 安装 PyTorch(在已激活的虚拟环境中)
# pip install torch 默认安装当前 CUDA 版本的轮子；
# 旧驱动可显式指定较低 CUDA，如 cu126
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126
 
# 验证 GPU 是否可用
python3 << 'EOF'
import torch
 
print(f"PyTorch 版本: {torch.__version__}")
print(f"CUDA 可用: {torch.cuda.is_available()}")
print(f"CUDA 版本: {torch.version.cuda}")
print(f"cuDNN 版本: {torch.backends.cudnn.version()}")
print(f"GPU 数量: {torch.cuda.device_count()}")
 
if torch.cuda.is_available():
    print(f"当前 GPU: {torch.cuda.get_device_name(0)}")
    print(f"GPU 显存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
 
    # 简单的 GPU 计算测试
    x = torch.randn(1000, 1000, device='cuda')
    y = torch.randn(1000, 1000, device='cuda')
    z = torch.mm(x, y)
    print(f"GPU 计算测试通过！矩阵乘法结果形状: {z.shape}")
else:
    print("警告: CUDA 不可用，请检查驱动和 CUDA 安装")
EOF

TensorFlow GPU 验证


# 安装 TensorFlow
pip install tensorflow
 
# 验证
python3 << 'EOF'
import tensorflow as tf
 
print(f"TensorFlow 版本: {tf.__version__}")
gpus = tf.config.list_physical_devices('GPU')
print(f"可用 GPU 数量: {len(gpus)}")
 
for gpu in gpus:
    print(f"  GPU: {gpu.name}")
 
if gpus:
    # 简单测试
    with tf.device('/GPU:0'):
        a = tf.random.normal([1000, 1000])
        b = tf.random.normal([1000, 1000])
        c = tf.matmul(a, b)
        print(f"GPU 计算测试通过！矩阵乘法结果形状: {c.shape}")
EOF

ROCm(AMD GPU)简介

ROCm(Radeon Open Compute)是 AMD 的开源 GPU 计算平台，类似于 NVIDIA 的 CUDA。

支持的 AMD GPU

ROCm 主要支持以下 AMD GPU 系列：

Radeon RX 7900 XTX / 7900 XT(RDNA 3)
Radeon RX 7800 XT / 7700 XT(RDNA 3)
Radeon PRO W7900 / W7800
Instinct MI300X / MI250X / MI210(数据中心)

安装 ROCm

Ubuntu 26.04 是首个可以直接用 apt 一行命令安装 ROCm 的版本，Universe 仓库已包含 ROCm 7.1.0，推荐优先使用 26.04 仓库自带的包：


# 方法一（推荐）：使用 Ubuntu 26.04 仓库自带的 ROCm
# 完整运行时栈
sudo apt install -y rocm
 
# 仅安装开发所需的库和头文件
# sudo apt install -y rocm-dev
 
# 将用户加入 render 和 video 组
sudo usermod -aG render,video $USER
 
# 重启
sudo reboot

Note

Ubuntu 26.04 仓库收录的是 ROCm 7.1.0。如需更新的版本，请参考下面 AMD 官方仓库的方式，并以 AMD 官方文档的当前版本号为准。


# 方法二：使用 AMD 官方仓库安装较新版本
# 请到 https://repo.radeon.com/rocm/apt/ 查看当前版本号，替换下方的 <version>
# Ubuntu 26.04（resolute）的源在 AMD 上线前，可暂用 24.04（noble）的源作为回退
wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | \
  gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
 
echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/<version> noble main" | \
  sudo tee /etc/apt/sources.list.d/rocm.list
 
sudo apt update
sudo apt install -y rocm
sudo usermod -aG render,video $USER
sudo reboot

验证 ROCm


# 查看 AMD GPU 信息
rocm-smi
 
# 检查 ROCm 版本
rocminfo | head -20
 
# 安装 PyTorch ROCm 版本
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.2
 
# 验证
python3 -c "import torch; print('ROCm available:', torch.cuda.is_available()); print('Device:', torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'N/A')"

Note

AMD ROCm 的软件生态不如 NVIDIA CUDA 成熟，部分 AI 工具可能不支持或支持有限。如果你主要用于 AI 开发，NVIDIA GPU 仍然是更稳妥的选择。但 ROCm 在持续改进，PyTorch 和部分主流框架已有良好支持。

常见问题排查

问题 1：nvidia-smi 无输出或报错


# 检查驱动是否加载
lsmod | grep nvidia
 
# 如果没有输出，尝试重新加载驱动
sudo modprobe nvidia
 
# 检查是否有 Secure Boot 问题
mokutil --sb-state
# 如果 Secure Boot 开启，可能需要签名驱动或关闭 Secure Boot
 
# 检查内核日志
dmesg | grep -i nvidia
sudo journalctl -b | grep -i nvidia

问题 2：驱动版本冲突


# 卸载所有 NVIDIA 驱动
sudo apt purge -y 'nvidia-*'
sudo apt purge -y 'libnvidia-*'
sudo apt autoremove -y
 
# 重启
sudo reboot
 
# 重新安装
sudo ubuntu-drivers autoinstall
sudo reboot

问题 3：CUDA 版本与 PyTorch 不匹配


# 查看当前 CUDA 版本
nvcc --version
 
# 查看 PyTorch 期望的 CUDA 版本
python3 -c "import torch; print(torch.version.cuda)"
 
# 如果不匹配，安装对应 CUDA 版本的 PyTorch 轮子
# 默认（当前 CUDA）
pip install torch
 
# CUDA 12.6
pip install torch --index-url https://download.pytorch.org/whl/cu126
 
# CUDA 12.8
pip install torch --index-url https://download.pytorch.org/whl/cu128
 
# 可用的 cuXXX 索引以 https://pytorch.org/get-started/locally/ 为准

问题 4：GPU 显存不足(OOM)


# 查看当前 GPU 显存使用情况
nvidia-smi
 
# 找出占用 GPU 的进程
nvidia-smi --query-compute-apps=pid,name,used_memory --format=csv
 
# 释放 GPU 显存(杀死占用进程)
# 先确认进程 PID
kill -9 <PID>
 
# 清理 PyTorch 缓存
python3 -c "import torch; torch.cuda.empty_cache()"

问题 5：Nouveau 驱动冲突


# Nouveau 是开源 NVIDIA 驱动，可能与官方驱动冲突
# 检查是否加载了 nouveau
lsmod | grep nouveau
 
# 如果有输出，禁用 nouveau
sudo bash -c 'cat > /etc/modprobe.d/blacklist-nouveau.conf << EOF
blacklist nouveau
options nouveau modeset=0
EOF'
 
# 重新生成 initramfs
sudo update-initramfs -u
 
# 重启
sudo reboot

问题 6：Docker 容器中使用 GPU


# 安装 NVIDIA Container Toolkit
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \
  sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
 
curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
  sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
  sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
 
sudo apt update
sudo apt install -y nvidia-container-toolkit
 
# 配置 Docker
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker
 
# 测试 Docker GPU 访问
docker run --rm --gpus all nvidia/cuda:12.6.0-base-ubuntu24.04 nvidia-smi
# 提示：nvidia/cuda 的镜像标签按 CUDA 版本与基础系统组织，
# 请到 https://hub.docker.com/r/nvidia/cuda/tags 选择存在的标签

Tip

安装 NVIDIA 驱动后建议立即安装 NVIDIA Container Toolkit。这样在 Docker 容器中运行 AI 工具(如 Open WebUI、ComfyUI 等)时可以直接使用 GPU 加速，无需在容器内单独配置驱动。

完整安装脚本

以下脚本汇总了完整的 GPU 环境配置流程：


#!/bin/bash
# gpu-setup.sh — Ubuntu 26.04 NVIDIA GPU 环境一键配置
set -e
 
echo "=== 1. 安装 NVIDIA 驱动 ==="
sudo apt update
sudo ubuntu-drivers autoinstall
 
echo "=== 2. 安装 CUDA Toolkit（Ubuntu 26.04 仓库版本）==="
sudo apt install -y cuda-toolkit
 
echo "=== 3.（可选）如需 NVIDIA 官方仓库的较新版本，可改用以下方式 ==="
# wget -q https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2604/x86_64/cuda-keyring_1.1-1_all.deb
# sudo dpkg -i cuda-keyring_1.1-1_all.deb
# sudo apt update
# sudo apt install -y cuda-toolkit-XX-Y   # 替换为实际可用版本
 
echo "=== 4. 安装 cuDNN ==="
sudo apt install -y cudnn9-cuda-12
 
echo "=== 5. 配置环境变量 ==="
if ! grep -q "CUDA_HOME" ~/.bashrc; then
  cat >> ~/.bashrc << 'ENVEOF'
 
# CUDA Environment
export CUDA_HOME=/usr/local/cuda
export PATH=$CUDA_HOME/bin:$PATH
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
ENVEOF
fi
 
echo "=== 6. 安装 NVIDIA Container Toolkit ==="
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \
  sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
  sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
  sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
sudo apt update
sudo apt install -y nvidia-container-toolkit
 
echo ""
echo "=== 安装完成！==="
echo "请重启系统后运行 nvidia-smi 和 nvcc --version 验证安装。"
echo "重启命令: sudo reboot"


# 保存并运行脚本
chmod +x gpu-setup.sh
./gpu-setup.sh