¿Cómo soluciono los problemas de comunicación entre la GPU y los controladores de GPU de NVIDIA en Amazon EC2?

Resolución

Es posible que aparezcan los siguientes mensajes de error según el comando o la herramienta que utilices para comunicarte con la GPU o los controladores de la GPU.

En el caso de nvidia-smi, se mostrará el siguiente mensaje de error:

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

En el caso de jax.devices(), se mostrará el siguiente mensaje de error en la terminal:

"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)]"

Por lo general, se muestran los mensajes de error anteriores debido a problemas de compatibilidad entre el hardware, los controladores y las bibliotecas. Para usar tipos de instancias con GPU, debes instalar los controladores de GPU y las bibliotecas relacionadas en tu sistema operativo (SO). Además, las herramientas y bibliotecas que se comunican con la GPU y los controladores de la GPU deben ser compatibles con la GPU.

Para comprobar si se ha producido un error de jax.devices(), ejecuta los siguientes comandos:

python
import jax
jax.devices()

Resultado de ejemplo:

$ 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)]

Verificación de que la AMI sea compatible con el tipo de instancia que se ha utilizado

Se recomienda utilizar las AMI de aprendizaje profundo de AWS (DLAMI) para los controladores de GPU de NVIDIA. Consulta las notas de la versión de DLAMI que quieres usar para asegurarte de que es compatible con tu configuración.

Nota: Los controladores de NVIDIA tienen dos tipos de DLAMI que utilizan controladores de origen propietarios o de código abierto. Cada DLAMI admite tipos de instancias específicos.

Verificación de que se ha instalado la biblioteca compatible con CUDA

Si has intentado ejecutar jax.devices() y se ha producido un error, es posible que no hayas instalado la biblioteca JAX compatible con CUDA.

Para comprobar si has instalado la biblioteca, ejecuta el siguiente comando:

pip list|grep jax

Resultado de ejemplo:

$ pip list|grep jax
jax                              0.4.30
jaxlib                           0.4.30

Si las bibliotecas compatibles con CUDA no están en el resultado del comando, ejecuta el siguiente comando para instalar la biblioteca JAX:

pip install -U "jax[cuda12]"

Nota: Sustituye jax[cuda12] por tu versión de CUDA.

A continuación, vuelve a ejecutar el comando pip list para comprobar que has instalado correctamente la biblioteca.

Resultado de ejemplo:

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

Verifica que las versiones de CUDA en "nvcc --version" y "nvidia-smi" sean las mismas

Ejecuta el siguiente comando:

nvcc --version

Resultado de ejemplo:

$ 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

A continuación, ejecuta el siguiente comando:

nvidia-smi

Resultado de ejemplo:

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

Compara las versiones de CUDA en los resultados de los comandos. En el ejemplo anterior, el problema se ha producido porque solo puedes usar la versión 12.2 o anterior de CUDA en el controlador de NVIDIA: 535.183.01. Sin embargo, la configuración usa la versión 12.5 de CUDA.

Para resolver el problema, ejecuta los siguientes comandos según tu sistema operativo para desinstalar CUDA 12.5 e instalar CUDA 12.2.

Amazon Linux 2023 (AL2023), Red Hat Enterprise Linux (RHEL) 8 o RHEL 9:

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

Amazon Linux 2 (AL2) o 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

Después de instalar CUDA 12.2, ejecuta el siguiente comando para comprobar que puedes usar jax.devices() para comunicarte con la GPU:

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

Para evitar problemas de compatibilidad con la GPU, se recomienda usar DLAMI con controladores de NVIDIA optimizados, un conjunto de herramientas de CUDA configurado y mejor soporte y compatibilidad.

(Opcional) Actualización del kernel y el controlador

Importante: Cuando usas una imagen de máquina de Amazon (AMI) compatible, no necesitas instalar manualmente los controladores de NVIDIA porque ya están instalados en la AMI.

No se recomienda actualizar la versión del kernel para mantener la compatibilidad con la versión del paquete y del controlador instalados. Sin embargo, si debes actualizar la versión del kernel debido a un parche de seguridad, ejecuta los siguientes comandos para actualizarlo:

sudo dnf versionlock delete kernel*
sudo dnf update -y

Nota: Los comandos de ejemplo anteriores son para la AMI de la base de aprendizaje profundo de AWS (Amazon Linux 2023).

Para instalar o actualizar manualmente los controladores de NVIDIA, consulta Opciones de instalación.

Verificación de que la comunicación funcione según lo esperado

Para comprobar que la instancia se puede comunicar con la GPU, ejecuta uno de los siguientes comandos:

nvidia-smi

Alternativa:

import jax
jax.devices()

Resultados de ejemplo:

$ 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)]

Información relacionada

Controladores de NVIDIA para la instancia de Amazon EC2