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 でインスタンスタイプを使用するには、GPU ドライバーおよび関連するライブラリをオペレーティングシステム (OS) にインストールする必要があります。また、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 が使用するインスタンスタイプをサポートしていることを確認する

NVDIA GPU ドライバーには、AWS Deep Learning AMI (DLAMI) の使用をお勧めします。使用する DLAMI に関するリリースノートを参照し、現在の構成と互換性があることを確認してください。

注: NVIDIA ドライバーには 2 種類の 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 ドライバー 535.183.01 では、バージョン 12.2 以前の CUDA のみを使用できるものの、この構成では CUDA バージョン 12.5 を使用したことが原因で問題が発生します。

この問題を解決するには、CUDA 12.5 をアンインストールし、CUDA 12.2 をインストールするために、使用する OS に応じて次のコマンドを実行します。

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 マシンイメージ (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 ドライバー