Strategies for dealing with XlaRuntimeError: RESOURCE_EXHAUSTED: Error loading program: Attempting to reserve

[dlwh]

(on tpu v3)

any suggestions for figuring out memory use for errors like:

  File "/home/kml11/venv310/lib/python3.10/site-packages/jax/_src/pjit.py", line 1148, in _pjit_call_impl_python
    return compiled.unsafe_call(*args), compiled
  File "/home/kml11/venv310/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
  File "/home/kml11/venv310/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: RESOURCE_EXHAUSTED: Error loading program: Attempting to reserve 4.67G at the bottom of memory. That was not possible. There are 432.95M free, 0B reserved, and 432.95M reservable.: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).
The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

usually we get the more helpful (if still somewhat cryptic) errors showing allocated buffers sorted by max size, but some times we just get this trace with nothing really to help guide us.

What are the best strategies for understanding the memory use here? The model is only 125M params (and we're using fsdp), so memory shouldn't be an issue, but it is for some reason

[steveepreston]

Same issue