Open mehranmirramezani opened 1 month ago
I suspect your matrix is very close to singular.
On CPU, the solver is implemented using float64 intermediates (via scipy's spsolve
), while on GPU the solver is implemented using float32 intermediates (via cuda's csrlsvqr
).
If that's the case, there's not much we can do here, unfortunately, aside from recommend that you do not pass near-singular matrices to spsolve
.
Many thanks for the details on the solver. I am wondering if there is a way to switch to cpu just for the solver part. I tried the below code but it seems still running on gpu and giving same error.
jax.device_put(spsolve(), device=jax.devices('cpu')[0])
device-putting the output of spsolve
on CPU will not cause the operation to be executed on CPU. But if you device-put the inputs to spsolve
onto the CPU, then the operation should execute on CPU.
Hi! I can confirm that putting arguments onto CPU triggers solution on CPU. That's nice! It seems solving on CPU is much faster on many occasions.
I have additional question however. Is there a way to reuse the factorized matrix for many right hand sides? With scipy
one can do:
factorized_K = scipy.sparse.linalg.splu(K)
x = factorized_K.solve(b)
and solve against many b
vectors. Are there some obstacles to implement something like that with jax
?
There's not currently any method of pre-factorization available in JAX. JAX's spsolve
was implemented to wrap a CUDA routine that a user needed; the CPU implementation is basically a reference implementation based on pure_callback
. If you're on CPU and need further funtionality, the best approach would probably be to use pure_callback
directly to call that functionality in scipy
.
Thanks, pure_callback
is the way. However, what i do not see immediately is how to preserve an object to be reused between subsequent callbacks. In other words how to perform such a sequence:
K
(on GPU or CPU)pure_callback
to factorize K
on CPUb
(on GPU or CPU)pure_callback
to get x
using the factorized K
on CPUb
pure_callback
againFactorized K must be 'saved' somehow, maybe as global variable?...
factorized_K
is just an object with a few arrays and static attributes. You could probably register it as a PyTree and return it to jax, and then pass it to the next pure_callback
Description
I am developing an FEM solver using JAX using sparse matrices. The code works perfectly fine on cpu but the exact same code fails on gpu by giving this error:
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Failed to execute XLA Runtime executable: run time error: custom call 'xla.gpu.custom_call' failed: Singular matrix in linear solve.
Attached is the simplified version of the code that was used:
System info (python version, jaxlib version, accelerator, etc.)
jax: 0.4.26
jaxlib: 0.4.23.dev20240502
numpy: 1.26.4
python: 3.10.14 (main, May 6 2024, 19:42:50) [GCC 11.2.0]
jax.devices (1 total, 1 local): [cuda(id=0)]
$ nvidia-smi Wed Jul 17 17:43:15 2024
+-----------------------------------------------------------------------------------------+ | NVIDIA-SMI 555.42.02 Driver Version: 555.42.02 CUDA Version: 12.5 | |-----------------------------------------+------------------------+----------------------+ | 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 NVIDIA GeForce RTX 2080 Ti On | 00000000:B1:00.0 Off | N/A | | 0% 33C P0 64W / 260W | 160MiB / 11264MiB | 0% Default | | | | N/A | +-----------------------------------------+------------------------+----------------------+
+-----------------------------------------------------------------------------------------+ | Processes: | | GPU GI CI PID Type Process name GPU Memory | | ID ID Usage | |=========================================================================================| | 0 N/A N/A 2131480 C python 156MiB | +-----------------------------------------------------------------------------------------+