mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-15 19:36:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add /dev/nvidiactl to the list of NVIDIA GPU devices. This is to cover the use case where a subset of GPUs are exposed to Docker/Kubernetes; the container might not necessarily see /dev/nvidia0.

Browse Source 
Testing:

Tested on a A100 16 GPU VM with only 8 GPU exposed to the container.

PiperOrigin-RevId: 657801141
This commit is contained in:
Gleb Pobudzey 2024-07-30 18:24:42 -07:00 committed by jax authors
parent bd14d6ab64
commit 35ba6f78bb
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
jax/_src/hardware_utils.py
View File

@ -38,6 +38,7 @@ _TPU_ENHANCED_BARRIER_SUPPORTED = [
_NVIDIA_GPU_DEVICES = [
'/dev/nvidia0',
'/dev/nvidiactl', # Docker/Kubernetes
'/dev/dxg', # WSL2
]