如何解决 Amazon EC2 中的 NVIDIA GPU 和 GPU 驱动程序通信问题？

解决方法

根据您用来与 GPU 或 GPU 驱动程序通信的命令或工具，您可能会遇到以下错误消息。

对于 nvidia-smi，您会收到以下错误消息：

"NVIDIA-SMI has failed because it couldn't communicate with the NVIDIA driver. Make sure that the latest NVIDIA driver is installed and running."

对于 jax.devices()，您会在终端中收到以下错误消息：

"WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.) [CpuDevice(id=0)]"

通常，由于硬件、驱动程序和库之间的兼容性问题，您会收到上述错误消息。要使用带有 GPU 的实例类型，必须在操作系统 (OS) 中安装 GPU 驱动程序和相关库。此外，与 GPU 和 GPU 驱动程序通信的工具和库必须与 GPU 兼容。

要检查您是否收到了 jax.devices() 错误，请运行以下命令：

python
import jax
jax.devices()

输出示例：

$ python
Python 3.9.22 (main, Apr 29 2025, 00:00:00)
[GCC 11.5.0 20240719 (Red Hat 11.5.0-5)] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import jax
>>> jax.devices()
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
[CpuDevice(id=0)]

验证您的 AMI 是否支持您使用的实例类型

最佳做法是将 AWS Deep Learning AMI (DLAMI) 用于 NVDIA GPU 驱动程序。查看要使用的 DLAMI 发行说明，以确保它与您的配置兼容。

注意：NVIDIA 驱动程序有两种类型的 DLAMI，它们使用专有驱动程序或开源驱动程序。每个 DLAMI 都支持特定的实例类型。

确认是否已安装兼容 CUDA 的库

如果您尝试运行 jax.devices() 但遇到了错误，则您可能尚未安装兼容 CUDA 的 JAX 库。

要检查您是否已安装该库，请运行以下命令：

pip list|grep jax

输出示例：

$ pip list|grep jax
jax                              0.4.30
jaxlib                           0.4.30

如果命令的输出中没有与 CUDA 兼容的库，请运行以下命令来安装 JAX 库：

pip install -U "jax[cuda12]"

**注意：**将 jax[cuda12] 替换为您的 CUDA 版本。

然后，重新运行 pip list 命令以验证是否已正确安装该库。

输出示例：

$ pip list|grep jax
jax                              0.4.30
jax-cuda12-pjrt                  0.4.30
jax-cuda12-plugin                0.4.30
jaxlib                           0.4.30

验证 "nvcc --version" 和 "nvidia-smi" 中的 CUDA 版本是否相同

首先，运行以下命令：

nvcc --version

输出示例：

$ nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2024 NVIDIA Corporation
Built on Wed_Apr_17_19:19:55_PDT_2024
Cuda compilation tools, release 12.5, V12.5.40
Build cuda_12.5.r12.5/compiler.34177558_0

然后，运行以下命令：

nvidia-smi

输出示例：

$ nvidia-smi
Mon Apr 21 10:52:36 2025
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.183.01             Driver Version: 535.183.01   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+

比较命令输出中的 CUDA 版本。在前面的示例中，出现问题是因为您只能在 NVIDIA 驱动程序中使用 CUDA 版本 12.2 或更早版本： 535.183.01。但是，该配置使用的是 CUDA 版本 12.5。

要解决此问题，请根据您的操作系统运行以下命令来卸载 CUDA 12.5 并安装 CUDA 12.2。

Amazon Linux 2023 (AL2023)、Red Hat Enterprise Linux (RHEL) 8 或 RHEL 9：

sudo dnf remove cuda-toolkit-12-5
sudo dnf install cuda-toolkit-12-2

Amazon Linux 2 (AL2) 或 RHEL 7：

sudo yum remove cuda-toolkit-12-5
sudo yum install cuda-toolkit-12-2

Ubuntu：

sudo apt remove cuda-toolkit-12-5
sudo apt install cuda-toolkit-12-2

安装 CUDA 12.2 后，运行以下命令以验证是否可以使用 jax.devices() 与 GPU 通信：

import jax
jax.devices()
[cuda(id=0)]

为避免 GPU 兼容性问题，最佳做法是使用具有优化的 NVIDIA 驱动程序、已配置的 CUDA 工具包以及更好的支持和兼容性的 DLAMI。

（可选）升级您的内核和驱动程序

**重要事项：**当您使用支持的亚马逊机器映像 (AMI) 时，无需手动安装 NVIDIA 驱动程序，因为它们已经安装在 AMI 上。

更新内核版本以保持与已安装驱动程序和软件包版本的兼容性并不是最佳做法。但是，如果由于安全补丁而必须更新内核版本，请运行以下命令来更新内核：

sudo dnf versionlock delete kernel*
sudo dnf update -y

**注意：**前面的示例命令适用于 AWS 深度学习基础 AMI (Amazon Linux 2023)。

要手动安装或更新 NVIDIA 驱动程序，请参阅安装选项。

验证通信是否按预期进行

要测试您的实例是否可以与 GPU 通信，请运行以下命令之一：

nvidia-smi

-或-

import jax
jax.devices()

输出示例：

$ nvidia-smi
Wed May 21 11:04:43 2025
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 570.133.20 Driver Version: 570.133.20 CUDA Version: 12.8 |
|-----------------------------------------+------------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+========================+======================|
| 0 NVIDIA A10G On | 00000000:00:1E.0 Off | 0 |
| 0% 32C P8 10W / 300W | 0MiB / 23028MiB | 0% Default |
| | | N/A |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=========================================================================================|
| No running processes found |
+-----------------------------------------------------------------------------------------+

>>> import jax
>>> jax.devices()
[cuda(id=0)]

相关信息

Amazon EC2 实例的 NVIDIA 驱动程序