Version of NCCL for jax cuda #15233

sh0416 · 2023-03-27T07:37:44Z

sh0416
Mar 27, 2023

I encounter the following error.

2023-03-27 07:31:43.753066: E external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:2389] Execution of replica 0 failed: INTERNAL: Failed to execute XLA Runtime executable: run time error: custom call 'xla.gpu.all_reduce_start' failed: external/org_tensorflow/tensorflow/compiler/xla/service/gpu/nccl_utils.cc:138: NCCL operation ncclGetUniqueId(&id) failed: internal error.
2023-03-27 07:31:43.753138: E external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:2389] Execution of replica 0 failed: INTERNAL: Failed to execute XLA Runtime executable: run time error: custom call 'xla.gpu.all_reduce_start' failed: external/org_tensorflow/tensorflow/compiler/xla/service/gpu/nccl_utils.cc:138: NCCL operation ncclGetUniqueId(&id) failed: internal error.
Traceback (most recent call last):
  File "train_mlp.py", line 56, in <module>
    loss_jit(params, batch)
  File "/home/seonghyeon_scatterlab_co_kr/anaconda3/envs/scatterlab-jax/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/seonghyeon_scatterlab_co_kr/anaconda3/envs/scatterlab-jax/lib/python3.8/site-packages/jax/_src/api.py", line 565, in cache_miss
    out_flat = call_bind_continuation(execute(*args_flat))
  File "/home/seonghyeon_scatterlab_co_kr/anaconda3/envs/scatterlab-jax/lib/python3.8/site-packages/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
  File "/home/seonghyeon_scatterlab_co_kr/anaconda3/envs/scatterlab-jax/lib/python3.8/site-packages/jax/interpreters/pxla.py", line 2113, in __call__
    out_bufs = self.xla_executable.execute_sharded_on_local_devices(
jax._src.traceback_util.UnfilteredStackTrace: jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Failed to execute XLA Runtime executable: run time error: custom call 'xla.gpu.all_reduce_start' failed: external/org_tensorflow/tensorflow/compiler/xla/service/gpu/nccl_utils.cc:138: NCCL operation ncclGetUniqueId(&id) failed: internal error.: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

I think nccl needs to be updated. What version are you recommended?