如何解決 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 版本。在上述範例中，問題發生的原因是，CUDA 版本 12.2 或更舊版本只能搭配 NVIDIA 驅動程式版本： 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。

(選用) 升級核心和驅動程式

**重要：**當您使用支援的 Amazon Machine Image (AMI) 時，無需手動安裝 NVIDIA 驅動程式，因為其已安裝在 AMI 上。

為了保持與已安裝的驅動程式和套件版本的相容性，而更新核心版本並不是最佳實務。但若因安全性修補而必須更新核心版本，請執行以下命令來更新核心：

sudo dnf versionlock delete kernel*
sudo dnf update -y

**注意：**上述範例命令適用於 AWS Deep Learning Base 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 驅動程式