GPU memory

[mariogeiger]

Hi,

I am worrying about my code design choices with respect to the potential OOM it might cause.
I design my code in such a way that it contain some dead code that I count on jax.jit to get rid of during the optimization.
However I am worrying it might cause troubles during the compilation.

Q1: Is there a way to print the current GPU memory usage (without setting XLA_PYTHON_CLIENT_ALLOCATOR=platform )? Ideally something like print(jax.memory_currently_allocated())

Q2: Imagine I have a code with some dead code parts

def f(x):
    y = g()    # some dead code
    return x

I bet the dead code will be removed by the compilation (jax.jit) but does the dead code could cause OOM error during the compilation?

[YouJiacheng]

Q1: I think there isn't a pytorch-like API that can interactively check device memory usage. There is a non-interactive but comprehensive method for device memory profiling https://jax.readthedocs.io/en/latest/device_memory_profiling.html
Q2: I think default behavior of jit is lazy evaluation, only using abstract array(without data) during compilation. see https://jax.readthedocs.io/en/latest/_autosummary/jax.ensure_compile_time_eval.html