No GPU found, despite installing Jax with GPU support #14404

jamesheald · 2023-02-10T15:58:35Z

jamesheald
Feb 10, 2023

I have followed the 'pip installation: GPU (CUDA)' instructions here https://github.com/google/jax#installation but jax does not find my GPU. I implemented many potential solutions after reading numerous posts related to this issue online, to no avail. Any help regarding how I can make Jax detect my GPU would be greatly appreciated. Details of what I have installed are provided below.

I have CUDA 11.4 installed:

$ nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2021 NVIDIA Corporation
Built on Wed_Jun__2_19:15:15_PDT_2021
Cuda compilation tools, release 11.4, V11.4.48
Build cuda_11.4.r11.4/compiler.30033411_0

CUDA 11.4 is located at /usr/local/cuda-11.4. Before I import jax into python I also specify this path using

os.environ['XLA_FLAGS'] = '--xla_gpu_cuda_data_dir=/usr/local/cuda-11.4'

CuDNN version is 8.7:

$ cat /usr/local/cuda-11.4/include/cudnn_version.h | grep CUDNN_MAJOR -A 2
#define CUDNN_MAJOR 8
#define CUDNN_MINOR 7
#define CUDNN_PATCHLEVEL 0
--
#define CUDNN_VERSION (CUDNN_MAJOR * 1000 + CUDNN_MINOR * 100 + CUDNN_PATCHLEVEL)

/* cannot use constexpr here since this is a C-only file */

I have a CUDA-compatible GPU:

$ lspci | grep NVIDIA

2d:00.0 VGA compatible controller: NVIDIA Corporation TU106 [GeForce RTX 2070] (rev a1)
2d:00.1 Audio device: NVIDIA Corporation TU106 High Definition Audio Controller (rev a1)
2d:00.2 USB controller: NVIDIA Corporation TU106 USB 3.1 Host Controller (rev a1)
2d:00.3 Serial bus controller [0c80]: NVIDIA Corporation TU106 USB Type-C UCSI Controller (rev a1)

I believe I have installed the appropriate jaxlib version:

# Installs the wheel compatible with Cuda >= 11.4 and cudnn >= 8.2
pip install "jax[cuda11_cudnn82]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

NVIDIA driver is version 470.161.03:

$ find /usr/lib/modules -name nvidia.ko -exec modinfo {} \;
filename:       /usr/lib/modules/5.14.0-1057-oem/updates/dkms/nvidia.ko
firmware:       nvidia/470.161.03/gsp.bin
alias:          char-major-195-*
version:        470.161.03
supported:      external
license:        NVIDIA
srcversion:     6651300FD4B9D974BB30104
alias:          pci:v000010DEd*sv*sd*bc03sc02i00*
alias:          pci:v000010DEd*sv*sd*bc03sc00i00*
depends:        drm
retpoline:      Y
name:           nvidia
vermagic:       5.14.0-1057-oem SMP mod_unload modversions 
sig_id:         PKCS#7
signer:         u433c Secure Boot Module Signature key
sig_key:        77:96:C5:ED:5D:89:79:34:2F:51:66:26:F7:CB:A0:BC:50:6B:6B:15
sig_hashalgo:   sha512
signature:      51:7A:13:6C:39:7F:08:61:7C:26:3B:15:DB:A8:C7:75:D1:3A:4B:B5:
		ED:0F:2E:C6:EA:9B:09:52:9A:F5:32:C7:DF:CA:86:35:9F:16:94:83:
		60:E4:51:8F:A8:52:41:68:C8:8D:2B:16:7B:6D:D9:45:B9:7F:47:CF:
		9D:13:7C:BE:19:DE:5C:A0:F4:99:AF:89:66:10:3E:13:EE:B1:AB:88:
		2C:D2:29:03:12:D1:F7:C1:3E:F9:44:43:FA:C7:45:E4:53:84:CD:C4:
		5A:43:58:85:B6:BD:19:F6:FD:7D:C9:89:21:B0:59:AD:94:3F:8F:70:
		5E:93:8D:6C:FA:CF:BC:1C:1F:3F:E6:42:F0:01:A7:09:A3:EE:00:81:
		53:95:93:6C:83:F7:C2:DC:D6:65:0C:A3:1C:52:F5:80:88:4B:1F:80:
		66:FB:4E:DB:96:83:C6:04:5D:76:78:6A:4B:CE:42:91:93:24:4A:3B:
		11:79:04:FF:BE:FD:96:ED:0C:D4:C2:06:00:7D:2B:A0:E1:1F:33:72:
		BA:2C:DA:7F:57:B7:98:76:15:66:D3:D1:83:A6:35:32:35:25:B3:96:
		18:E9:C3:88:2B:A6:4D:40:4B:F4:3F:63:8A:1D:38:2A:18:A5:0C:30:
		3D:1C:FF:68:22:E0:CB:63:C8:22:A5:60:DC:9D:7A:F1
parm:           NvSwitchRegDwords:NvSwitch regkey (charp)
parm:           NvSwitchBlacklist:NvSwitchBlacklist=uuid[,uuid...] (charp)
parm:           NVreg_ResmanDebugLevel:int
parm:           NVreg_RmLogonRC:int
parm:           NVreg_ModifyDeviceFiles:int
parm:           NVreg_DeviceFileUID:int
parm:           NVreg_DeviceFileGID:int
parm:           NVreg_DeviceFileMode:int
parm:           NVreg_InitializeSystemMemoryAllocations:int
parm:           NVreg_UsePageAttributeTable:int
parm:           NVreg_RegisterForACPIEvents:int
parm:           NVreg_EnablePCIeGen3:int
parm:           NVreg_EnableMSI:int
parm:           NVreg_TCEBypassMode:int
parm:           NVreg_EnableStreamMemOPs:int
parm:           NVreg_RestrictProfilingToAdminUsers:int
parm:           NVreg_PreserveVideoMemoryAllocations:int
parm:           NVreg_EnableS0ixPowerManagement:int
parm:           NVreg_S0ixPowerManagementVideoMemoryThreshold:int
parm:           NVreg_DynamicPowerManagement:int
parm:           NVreg_DynamicPowerManagementVideoMemoryThreshold:int
parm:           NVreg_EnableGpuFirmware:int
parm:           NVreg_EnableUserNUMAManagement:int
parm:           NVreg_MemoryPoolSize:int
parm:           NVreg_KMallocHeapMaxSize:int
parm:           NVreg_VMallocHeapMaxSize:int
parm:           NVreg_IgnoreMMIOCheck:int
parm:           NVreg_NvLinkDisable:int
parm:           NVreg_EnablePCIERelaxedOrderingMode:int
parm:           NVreg_RegisterPCIDriver:int
parm:           NVreg_RegistryDwords:charp
parm:           NVreg_RegistryDwordsPerDevice:charp
parm:           NVreg_RmMsg:charp
parm:           NVreg_GpuBlacklist:charp
parm:           NVreg_TemporaryFilePath:charp
parm:           NVreg_ExcludedGpus:charp
parm:           rm_firmware_active:charp
filename:       /usr/lib/modules/5.15.0-58-generic/updates/dkms/nvidia.ko
firmware:       nvidia/470.161.03/gsp.bin
alias:          char-major-195-*
version:        470.161.03
supported:      external
license:        NVIDIA
srcversion:     6651300FD4B9D974BB30104
alias:          pci:v000010DEd*sv*sd*bc03sc02i00*
alias:          pci:v000010DEd*sv*sd*bc03sc00i00*
depends:        drm
retpoline:      Y
name:           nvidia
vermagic:       5.15.0-58-generic SMP mod_unload modversions 
sig_id:         PKCS#7
signer:         u433c Secure Boot Module Signature key
sig_key:        77:96:C5:ED:5D:89:79:34:2F:51:66:26:F7:CB:A0:BC:50:6B:6B:15
sig_hashalgo:   sha512
signature:      7A:C5:AC:E1:41:7D:F3:8D:F9:C0:BF:57:F2:F0:54:F8:71:23:73:79:
		FC:B3:65:D3:78:32:21:DE:9B:0F:F4:F2:3B:AF:7C:1A:3A:DB:87:C5:
		3F:F1:B6:38:63:39:12:83:E4:F9:E0:A7:6D:C3:D7:FD:E0:09:C7:7E:
		D5:8D:4A:D1:BB:FB:D6:D0:B1:BB:3E:72:76:C6:97:45:A1:A9:16:E3:
		D3:69:0E:A6:DD:81:0E:D7:D7:03:7D:BB:44:A5:30:8C:D0:EE:13:43:
		66:49:D4:57:1A:85:AB:E1:2C:D0:8E:6C:F8:9E:12:E0:50:C4:0F:3C:
		5E:BA:E1:B4:A0:04:B4:CA:CF:E5:23:C1:2E:FA:80:4E:D7:41:38:57:
		99:71:B8:D7:2A:F9:C8:DF:65:9D:05:C3:F2:84:70:43:1C:CB:D8:81:
		36:8D:00:3A:05:A1:9B:A5:27:C6:9F:24:D7:94:0D:1A:FF:F9:49:AF:
		BE:1B:AE:62:07:BA:36:B7:D5:F2:53:35:D2:25:6E:2D:21:A3:7C:63:
		2F:4E:43:CB:8E:F8:BF:0A:27:3B:F2:FA:0B:C4:95:51:61:A3:6E:92:
		DD:9B:61:F3:BF:B5:80:52:FF:70:0C:A7:BE:79:E0:F9:76:F5:B8:0A:
		B7:CE:3E:00:4C:5C:88:1F:9A:26:F6:95:9E:66:4A:C2
parm:           NvSwitchRegDwords:NvSwitch regkey (charp)
parm:           NvSwitchBlacklist:NvSwitchBlacklist=uuid[,uuid...] (charp)
parm:           NVreg_ResmanDebugLevel:int
parm:           NVreg_RmLogonRC:int
parm:           NVreg_ModifyDeviceFiles:int
parm:           NVreg_DeviceFileUID:int
parm:           NVreg_DeviceFileGID:int
parm:           NVreg_DeviceFileMode:int
parm:           NVreg_InitializeSystemMemoryAllocations:int
parm:           NVreg_UsePageAttributeTable:int
parm:           NVreg_RegisterForACPIEvents:int
parm:           NVreg_EnablePCIeGen3:int
parm:           NVreg_EnableMSI:int
parm:           NVreg_TCEBypassMode:int
parm:           NVreg_EnableStreamMemOPs:int
parm:           NVreg_RestrictProfilingToAdminUsers:int
parm:           NVreg_PreserveVideoMemoryAllocations:int
parm:           NVreg_EnableS0ixPowerManagement:int
parm:           NVreg_S0ixPowerManagementVideoMemoryThreshold:int
parm:           NVreg_DynamicPowerManagement:int
parm:           NVreg_DynamicPowerManagementVideoMemoryThreshold:int
parm:           NVreg_EnableGpuFirmware:int
parm:           NVreg_EnableUserNUMAManagement:int
parm:           NVreg_MemoryPoolSize:int
parm:           NVreg_KMallocHeapMaxSize:int
parm:           NVreg_VMallocHeapMaxSize:int
parm:           NVreg_IgnoreMMIOCheck:int
parm:           NVreg_NvLinkDisable:int
parm:           NVreg_EnablePCIERelaxedOrderingMode:int
parm:           NVreg_RegisterPCIDriver:int
parm:           NVreg_RegistryDwords:charp
parm:           NVreg_RegistryDwordsPerDevice:charp
parm:           NVreg_RmMsg:charp
parm:           NVreg_GpuBlacklist:charp
parm:           NVreg_TemporaryFilePath:charp
parm:           NVreg_ExcludedGpus:charp
parm:           rm_firmware_active:charp
filename:       /usr/lib/modules/5.15.0-60-generic/updates/dkms/nvidia.ko
firmware:       nvidia/470.161.03/gsp.bin
alias:          char-major-195-*
version:        470.161.03
supported:      external
license:        NVIDIA
srcversion:     6651300FD4B9D974BB30104
alias:          pci:v000010DEd*sv*sd*bc03sc02i00*
alias:          pci:v000010DEd*sv*sd*bc03sc00i00*
depends:        drm
retpoline:      Y
name:           nvidia
vermagic:       5.15.0-60-generic SMP mod_unload modversions 
sig_id:         PKCS#7
signer:         u433c Secure Boot Module Signature key
sig_key:        77:96:C5:ED:5D:89:79:34:2F:51:66:26:F7:CB:A0:BC:50:6B:6B:15
sig_hashalgo:   sha512
signature:      68:E2:DF:ED:97:4E:A2:67:93:CA:68:BD:C1:10:B6:EE:CF:E7:F2:5F:
		95:F6:B7:99:29:56:F3:1A:35:84:4C:76:0C:05:2F:15:86:2D:52:46:
		91:6F:E9:BC:4B:B8:B7:23:4F:7B:FD:D4:D5:21:6D:2C:A7:5F:DA:FD:
		C9:CE:BE:5E:14:F5:7B:4E:5B:97:FD:AE:A0:F6:4A:A2:C5:0B:89:8D:
		F6:28:C3:20:B0:4D:74:CD:D7:6A:E5:50:65:1A:2F:76:E3:77:0D:D0:
		9A:C2:6C:40:4C:3F:F4:70:75:95:0B:12:2B:1C:03:0D:0E:93:C3:5C:
		02:8A:A7:00:F7:72:6B:94:F5:B3:80:6D:6F:88:EC:F5:1A:9E:C0:AC:
		82:B9:24:FB:3D:38:72:92:E9:5B:2E:6B:CC:87:06:E7:78:30:14:ED:
		DC:7A:47:CE:3F:08:AC:78:94:EC:FE:DB:F4:1B:77:D7:2B:13:B2:97:
		14:65:4C:02:F9:89:4B:F2:76:EF:59:B5:23:99:AE:47:5E:FF:A3:6A:
		DE:C8:C8:C7:BD:10:31:B4:ED:6A:2C:AD:66:3E:AA:B9:13:40:55:FA:
		4D:2A:83:6C:80:45:AB:4B:85:F1:4E:D3:8B:BD:A5:D6:A8:AD:0B:EC:
		BE:4E:0E:B3:C4:31:D4:12:84:54:98:CD:D7:CD:E0:35
parm:           NvSwitchRegDwords:NvSwitch regkey (charp)
parm:           NvSwitchBlacklist:NvSwitchBlacklist=uuid[,uuid...] (charp)
parm:           NVreg_ResmanDebugLevel:int
parm:           NVreg_RmLogonRC:int
parm:           NVreg_ModifyDeviceFiles:int
parm:           NVreg_DeviceFileUID:int
parm:           NVreg_DeviceFileGID:int
parm:           NVreg_DeviceFileMode:int
parm:           NVreg_InitializeSystemMemoryAllocations:int
parm:           NVreg_UsePageAttributeTable:int
parm:           NVreg_RegisterForACPIEvents:int
parm:           NVreg_EnablePCIeGen3:int
parm:           NVreg_EnableMSI:int
parm:           NVreg_TCEBypassMode:int
parm:           NVreg_EnableStreamMemOPs:int
parm:           NVreg_RestrictProfilingToAdminUsers:int
parm:           NVreg_PreserveVideoMemoryAllocations:int
parm:           NVreg_EnableS0ixPowerManagement:int
parm:           NVreg_S0ixPowerManagementVideoMemoryThreshold:int
parm:           NVreg_DynamicPowerManagement:int
parm:           NVreg_DynamicPowerManagementVideoMemoryThreshold:int
parm:           NVreg_EnableGpuFirmware:int
parm:           NVreg_EnableUserNUMAManagement:int
parm:           NVreg_MemoryPoolSize:int
parm:           NVreg_KMallocHeapMaxSize:int
parm:           NVreg_VMallocHeapMaxSize:int
parm:           NVreg_IgnoreMMIOCheck:int
parm:           NVreg_NvLinkDisable:int
parm:           NVreg_EnablePCIERelaxedOrderingMode:int
parm:           NVreg_RegisterPCIDriver:int
parm:           NVreg_RegistryDwords:charp
parm:           NVreg_RegistryDwordsPerDevice:charp
parm:           NVreg_RmMsg:charp
parm:           NVreg_GpuBlacklist:charp
parm:           NVreg_TemporaryFilePath:charp
parm:           NVreg_ExcludedGpus:charp
parm:           rm_firmware_active:charp

When I open python I get the following message:

>>> os.environ['XLA_FLAGS'] = '--xla_gpu_cuda_data_dir=/usr/local/cuda-11.4'
>>> import jax
>>> jax.devices()
2023-02-09 18:13:51.106666: E external/org_tensorflow/tensorflow/compiler/xla/stream_executor/cuda/cuda_driver.cc:267] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected
No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
[CpuDevice(id=0)]

I am using Ubuntu version 20.04.3:

$ lsb_release -a
No LSB modules are available.
Distributor ID:	Ubuntu
Description:	Ubuntu 20.04.3 LTS
Release:	20.04
Codename:	focal

hawkinsp · 2023-02-10T16:00:01Z

hawkinsp
Feb 10, 2023
Maintainer

What does nvidia-smi print?

3 replies

jamesheald Feb 10, 2023
Author

$ 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.

hawkinsp Feb 10, 2023
Maintainer

I'd definitely fix that problem first! Until nvidia-smi is happy, JAX won't be happy.

Rebooting might work.

jamesheald Feb 14, 2023
Author

I needed to do a machine owner key (MOK) secure boot. GPU is now found. Cheers!

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

No GPU found, despite installing Jax with GPU support #14404

Uh oh!

{{title}}

Uh oh!

Replies: 1 comment 3 replies

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Select a reply

Uh oh!

No GPU found, despite installing Jax with GPU support #14404

Uh oh!

jamesheald Feb 10, 2023

Replies: 1 comment · 3 replies

Uh oh!

hawkinsp Feb 10, 2023 Maintainer

Uh oh!

jamesheald Feb 10, 2023 Author

Uh oh!

hawkinsp Feb 10, 2023 Maintainer

Uh oh!

jamesheald Feb 14, 2023 Author

jamesheald
Feb 10, 2023

Replies: 1 comment 3 replies

hawkinsp
Feb 10, 2023
Maintainer

jamesheald Feb 10, 2023
Author

hawkinsp Feb 10, 2023
Maintainer

jamesheald Feb 14, 2023
Author