jax on SLURM with CUDA 11.6 #25373

ashok-arora · 2024-12-10T07:48:12Z

ashok-arora
Dec 10, 2024

Hey,
I am unable to find the installation instructions for jax versions that support gpu compatibility and would work with cuda 11.6, I wish I could upgrade the driver version but since I am on a SLURM cluster, it is fixed by the university provider. Moreover, I can change the cuda-toolkit version using the spack manager but I am not sure how to find the right combination of cuda and cudnn that would work. I tried to install different versions but they all gave errors. Is there a guide on older versions?

Answered by ashok-arora

Dec 10, 2024

Finally worked, heres my setup for future reference:

in the slurm script:

module load cuda/11.6
module load spack
source /home/apps/spack/share/spack/setup-env.sh
spack load cudnn@8.7.0.84-11.8

code to check:

import jax
import logging

  logger = logging.getLogger(__name__)
  logger.setLevel(logging.INFO)

  logger.info("Starting Job for Config:\n"+str(OmegaConf.to_yaml(config)))
  logger.info("Available Backends:"+str(jax.devices()))
  jax.print_environment_info()

Output of the code:

[2024-12-10 16:44:32,071][jax._src.xla_bridge][INFO] - Unable to initialize backend 'tpu_driver': NOT_FOUND: Unable to find driver in registry given worker: 
[2024-12-10 16:44:34,106][jax._src.xla_br…

View full answer

ashok-arora · 2024-12-10T10:52:24Z

ashok-arora
Dec 10, 2024
Author

Here's a stack trace from one of my recent endeavors:

/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/flax/struct.py:132: FutureWarning: jax.tree_util.register_keypaths is deprecated, and will be removed in a future release. Please use `register_pytree_with_keys()` instead.
  jax.tree_util.register_keypaths(data_clz, keypaths)
/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/flax/struct.py:132: FutureWarning: jax.tree_util.register_keypaths is deprecated, and will be removed in a future release. Please use `register_pytree_with_keys()` instead.
  jax.tree_util.register_keypaths(data_clz, keypaths)
2024-12-10 16:19:17.297366: E external/org_tensorflow/tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:429] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERROR
2024-12-10 16:19:17.306907: E external/org_tensorflow/tensorflow/compiler/xla/status_macros.cc:54] INTERNAL: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:641) dnn != nullptr 
*** Begin stack trace ***
	
	
	
	
	
	
	
	
	
	
	
	
	
	
	
	
	
	_PyObject_MakeTpCall
	
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	PyObject_Call
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	PyObject_Call
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	
	PyObject_Call
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	
	
	_PyEval_EvalFrameDefault
	
	
	
	_PyEval_EvalFrameDefault
	
	PyEval_EvalCode
	
	
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	
	_PyObject_CallMethodIdObjArgs
	PyImport_ImportModuleLevelObject
	_PyEval_EvalFrameDefault
	
	PyEval_EvalCode
	
	
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	
	_PyObject_CallMethodIdObjArgs
	PyImport_ImportModuleLevelObject
	_PyEval_EvalFrameDefault
	
	PyEval_EvalCode
	
	
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	
	_PyObject_CallMethodIdObjArgs
	PyImport_ImportModuleLevelObject
	_PyEval_EvalFrameDefault
	
	PyEval_EvalCode
	
	
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	_PyEval_EvalFrameDefault
	_PyFunction_Vectorcall
	
	_PyObject_CallMethodIdObjArgs
	PyImport_ImportModuleLevelObject
	_PyEval_EvalFrameDefault
	
*** End stack trace ***

Error executing job with overrides: []
Traceback (most recent call last):
  File "/scratch/ashok/agalite/check_jax.py", line 39, in <module>
    main()
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/hydra/main.py", line 94, in decorated_main
    _run_hydra(
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
    _run_app(
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/hydra/_internal/utils.py", line 457, in _run_app
    run_and_report(
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/hydra/_internal/utils.py", line 223, in run_and_report
    raise ex
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/hydra/_internal/utils.py", line 220, in run_and_report
    return func()
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/hydra/_internal/utils.py", line 458, in <lambda>
    lambda: hydra.run(
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/hydra/_internal/hydra.py", line 132, in run
    _ = ret.return_value
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/hydra/core/utils.py", line 260, in return_value
    raise self._return_value
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
  File "/scratch/ashok/agalite/check_jax.py", line 20, in main
    from src.trainers.trainers_control import ControlTrainer
  File "/scratch/ashok/agalite/src/trainers/trainers_control.py", line 19, in <module>
    from src.model_fns import *
  File "/scratch/ashok/agalite/src/model_fns/__init__.py", line 3, in <module>
    from .seq_fns import *
  File "/scratch/ashok/agalite/src/model_fns/seq_fns.py", line 8, in <module>
    from src.models.agalite.galite import GaLiTe
  File "/scratch/ashok/agalite/src/models/agalite/galite.py", line 7, in <module>
    from src.models.agalite.layers import *
  File "/scratch/ashok/agalite/src/models/agalite/layers.py", line 55, in <module>
    class ParameteredProjection(nn.Module):
  File "/scratch/ashok/agalite/src/models/agalite/layers.py", line 61, in ParameteredProjection
    def __call__(self, inputs,kernel_init=orthogonal(jnp.sqrt(2)),bias_init=constant(0.0)):
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/jax/_src/pjit.py", line 235, in cache_miss
    outs, out_flat, out_tree, args_flat = _python_pjit_helper(
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/jax/_src/pjit.py", line 184, in _python_pjit_helper
    out_flat = pjit_p.bind(*args_flat, **params)
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/jax/_src/core.py", line 2577, in bind
    return self.bind_with_trace(top_trace, args, params)
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/jax/_src/core.py", line 363, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/jax/_src/core.py", line 807, in process_primitive
    return primitive.impl(*tracers, **params)
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/jax/_src/pjit.py", line 1291, in _pjit_call_impl
    compiled = _pjit_lower(
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py", line 3373, in compile
    executable = self._compile_unloaded(
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py", line 3344, in _compile_unloaded
    return UnloadedMeshExecutable.from_hlo(
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py", line 3626, in from_hlo
    xla_executable = dispatch.compile_or_get_cached(
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/jax/_src/dispatch.py", line 1091, in compile_or_get_cached
    return backend_compile(backend, serialized_computation, compile_options,
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
  File "/home/ashok/.conda/envs/agalite/lib/python3.10/site-packages/jax/_src/dispatch.py", line 1036, in backend_compile
    return backend.compile(built_c, compile_options=options)
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:641) dnn != nullptr

0 replies

ashok-arora · 2024-12-10T11:16:14Z

ashok-arora
Dec 10, 2024
Author

Finally worked, heres my setup for future reference:

in the slurm script:

module load cuda/11.6
module load spack
source /home/apps/spack/share/spack/setup-env.sh
spack load cudnn@8.7.0.84-11.8

code to check:

import jax
import logging

  logger = logging.getLogger(__name__)
  logger.setLevel(logging.INFO)

  logger.info("Starting Job for Config:\n"+str(OmegaConf.to_yaml(config)))
  logger.info("Available Backends:"+str(jax.devices()))
  jax.print_environment_info()

Output of the code:

[2024-12-10 16:44:32,071][jax._src.xla_bridge][INFO] - Unable to initialize backend 'tpu_driver': NOT_FOUND: Unable to find driver in registry given worker: 
[2024-12-10 16:44:34,106][jax._src.xla_bridge][INFO] - Unable to initialize backend 'rocm': NOT_FOUND: Could not find registered platform with name: "rocm". Available platform names are: CUDA Host Interpreter
[2024-12-10 16:44:34,107][jax._src.xla_bridge][INFO] - Unable to initialize backend 'tpu': module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'
[2024-12-10 16:44:34,107][jax._src.xla_bridge][INFO] - Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.
[2024-12-10 16:44:36,466][__main__][INFO] - Starting Job for Config:
{}

[2024-12-10 16:44:36,466][__main__][INFO] - Available Backends:[StreamExecutorGpuDevice(id=0, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=1, process_index=0, slice_index=0)]
jax:    0.4.6
jaxlib: 0.4.4
numpy:  1.26.3
python: 3.10.15 (main, Oct  3 2024, 07:27:34) [GCC 11.2.0]
jax.devices (2 total, 2 local): [StreamExecutorGpuDevice(id=0, process_index=0, slice_index=0) StreamExecutorGpuDevice(id=1, process_index=0, slice_index=0)]
process_count: 1

$ nvidia-smi
Tue Dec 10 16:44:36 2024       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 510.47.03    Driver Version: 510.47.03    CUDA Version: 11.6     |
|-------------------------------+----------------------+----------------------+
| 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  Tesla V100-SXM2...  Off  | 00000000:61:00.0 Off |                    0 |
| N/A   44C    P0    40W / 300W |  12655MiB / 16384MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  Tesla V100-SXM2...  Off  | 00000000:89:00.0 Off |                    0 |
| N/A   44C    P0    39W / 300W |    307MiB / 16384MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A     39875      C   .../envs/agalite/bin/python3    12653MiB |
|    1   N/A  N/A     39875      C   .../envs/agalite/bin/python3      305MiB |
+-----------------------------------------------------------------------------+

0 replies

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

jax on SLURM with CUDA 11.6 #25373

{{title}}

Replies: 2 comments

{{title}}

{{title}}

Select a reply

jax on SLURM with CUDA 11.6 #25373

ashok-arora Dec 10, 2024

Replies: 2 comments

ashok-arora Dec 10, 2024 Author

ashok-arora Dec 10, 2024 Author

ashok-arora
Dec 10, 2024

ashok-arora
Dec 10, 2024
Author

ashok-arora
Dec 10, 2024
Author