How to detect CUDA out of memory errors?

[nalzok]

Let's say you want to SVD on a skinny matrix.

>>> import jax
>>> import jax.numpy as jnp
>>>
>>> A = jnp.zeros((786432, 5), dtype=jnp.float32)
>>> jnp.linalg.svd(A, full_matrices=False)
(Array([[1., 0., 0., 0., 0.],
       [0., 1., 0., 0., 0.],
       [0., 0., 1., 0., 0.],
       ...,
       [0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]], dtype=float32), Array([0., 0., 0., 0., 0.], dtype=float32), Array([[1., 0., 0., 0., 0.],
       [0., 1., 0., 0., 0.],
       [0., 0., 1., 0., 0.],
       [0., 0., 0., 1., 0.],
       [0., 0., 0., 0., 1.]], dtype=float32))

That works as expected, but if we make the matrix sufficiently tall

>>> A = jnp.zeros((38597376, 5), dtype=jnp.float32)
>>> jnp.linalg.svd(A, full_matrices=False)
2023-08-03 14:08:20.842664: W external/xla/xla/service/gpu/runtime/support.cc:58] Intercepted XLA runtime error:
INTERNAL: jaxlib/gpu/solver_kernels.cc:680: operation gpusolverDnSgesvd( handle.get(), d.jobu, d.jobvt, d.m, d.n, a, d.m, s, u, d.m, vt, d.n, static_cast<float*>(work), d.lwork, nullptr, info) failed: cuSolver invalid value error
2023-08-03 14:08:20.842796: E external/xla/xla/pjrt/pjrt_stream_executor_client.cc:2593] Execution of replica 0 failed: INTERNAL: Failed to execute XLA Runtime executable: run time error: custom call 'xla.gpu.custom_call' failed: jaxlib/gpu/solver_kernels.cc:680: operation gpusolverDnSgesvd( handle.get(), d.jobu, d.jobvt, d.m, d.n, a, d.m, s, u, d.m, vt, d.n, static_cast<float*>(work), d.lwork, nullptr, info) failed: cuSolver invalid value error; current tracing scope: custom-call.9; current profiling annotation: XlaModule:#hlo_module=jit_svd,program_id=11#.
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/data/research/smolGPT/.venv/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/data/research/smolGPT/.venv/lib/python3.10/site-packages/jax/_src/pjit.py", line 253, in cache_miss
    outs, out_flat, out_tree, args_flat, jaxpr = _python_pjit_helper(
  File "/data/research/smolGPT/.venv/lib/python3.10/site-packages/jax/_src/pjit.py", line 166, in _python_pjit_helper
    out_flat = pjit_p.bind(*args_flat, **params)
  File "/data/research/smolGPT/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 2596, in bind
    return self.bind_with_trace(top_trace, args, params)
  File "/data/research/smolGPT/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 389, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
  File "/data/research/smolGPT/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 821, in process_primitive
    return primitive.impl(*tracers, **params)
  File "/data/research/smolGPT/.venv/lib/python3.10/site-packages/jax/_src/pjit.py", line 1209, in _pjit_call_impl
    return xc._xla.pjit(name, f, call_impl_cache_miss, [], [], donated_argnums,
  File "/data/research/smolGPT/.venv/lib/python3.10/site-packages/jax/_src/pjit.py", line 1192, in call_impl_cache_miss
    out_flat, compiled = _pjit_call_impl_python(
  File "/data/research/smolGPT/.venv/lib/python3.10/site-packages/jax/_src/pjit.py", line 1148, in _pjit_call_impl_python
    return compiled.unsafe_call(*args), compiled
  File "/data/research/smolGPT/.venv/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
  File "/data/research/smolGPT/.venv/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py", line 1229, in __call__
    results = self.xla_executable.execute_sharded(input_bufs)
jax._src.traceback_util.UnfilteredStackTrace: jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Failed to execute XLA Runtime executable: run time error: custom call 'xla.gpu.custom_call' failed: jaxlib/gpu/solver_kernels.cc:680: operation gpusolverDnSgesvd( handle.get(), d.jobu, d.jobvt, d.m, d.n, a, d.m, s, u, d.m, vt, d.n, static_cast<float*>(work), d.lwork, nullptr, info) failed: cuSolver invalid value error; current tracing scope: custom-call.9; current profiling annotation: XlaModule:#hlo_module=jit_svd,program_id=11#.

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Failed to execute XLA Runtime executable: run time error: custom call 'xla.gpu.custom_call' failed: jaxlib/gpu/solver_kernels.cc:680: operation gpusolverDnSgesvd( handle.get(), d.jobu, d.jobvt, d.m, d.n, a, d.m, s, u, d.m, vt, d.n, static_cast<float*>(work), d.lwork, nullptr, info) failed: cuSolver invalid value error; current tracing scope: custom-call.9; current profiling annotation: XlaModule:#hlo_module=jit_svd,program_id=11#.

How can we check in advance whether an SVD calculation can fit into the memory, so that I can fall back to jax.pure_callback and call np.linalg.svd in NumPy? Trying to catch an exception isn't very convenient (if not possible) in a JIT'ed function.