You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
I have an issue with a custom partionned primitive.
Everything works well in its own but when used inside a bigger function I am getting this error
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Failed to load in-memory CUBIN (compiled for a different GPU?).: CUDA_ERROR_ILLEGAL_ADDRESS: an illegal memory access was encountered
My HPC cluster's driver version is 12.2, and I want to use the lates JAX so with prebuilt CUDA 12.3 .
I have also tried a local JAX with a CUDA 12.2 module .. same error
E0427 12:54:10.576806 3139367 pjrt_stream_executor_client.cc:2809] Execution of replica 0 failed: INTERNAL: Failed to load in-memory CUBIN (compiled for a different GPU?).: CUDA_ERROR_ILLEGAL_ADDRESS: an illegal memory access was encountered
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
I will try to provide a MWE but I can point you to my implementation
reacted with thumbs up emoji reacted with thumbs down emoji reacted with laugh emoji reacted with hooray emoji reacted with confused emoji reacted with heart emoji reacted with rocket emoji reacted with eyes emoji
-
Hello,
I have an issue with a custom partionned primitive.
Everything works well in its own but when used inside a bigger function I am getting this error
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Failed to load in-memory CUBIN (compiled for a different GPU?).: CUDA_ERROR_ILLEGAL_ADDRESS: an illegal memory access was encountered
My HPC cluster's driver version is 12.2, and I want to use the lates JAX so with prebuilt CUDA 12.3 .
I have also tried a local JAX with a CUDA 12.2 module .. same error
I will try to provide a MWE but I can point you to my implementation
This works
This fails
Thank you
Beta Was this translation helpful? Give feedback.
All reactions