Understanding Memory Buffers and donate_argnums

[Jeevesh8]

I have a setup of the form:

@partial(jax.pmap, donate_argnums=(0))
def fn(x):
    new_x = 2+x
    return new_x, jnp.mean(x)

@jax.pmap
def fn2(x):
    return x, jnp.mean(x)

x = jnp.array([[1,2,3,4,5,6,7,8]*800000]*8)

for i in range(1,4):
    print(x)
    x, y = fn(x)
    if i%2==0:
        for j in range(3):
            print(x)
            x, y = fn2(x)

Now, my problem is, when we enter the second loop, I end up getting a RuntimeError: Invalid argument: DeviceArray has been deleted., when the first call to fn2 occurs. I am able to print x just before entering fn2. That means, that x has not been deleted, but still it says that it has been deleted. I want to know why?

Note: The above example doesn't reproduce the error, and is just representative of the flow of actual computation I am doing.

Any suggestions as to how to debug this are also welcome.

[jakevdp]

There is some information on donate_argnums in the pmap docstring. To quote:

Specify which arguments are “donated” to the computation. It is safe to donate arguments if you no longer need them once the computation has finished. In some cases XLA can make use of donated buffers to reduce the amount of memory needed to perform a computation, for example recycling one of your input buffers to store a result. You should not reuse buffers that you donate to a computation, JAX will raise an error if you try to. By default, no arguments are donated.

So to paraphrase, arguments marked for donation may be used by XLA to make the computation more efficient, in which case the buffers will be marked deleted. The reason you can print the contents of a deleted buffer is I think because of internal cacheing used for printing device arrays that ends up bypassing the deletion check; I would probably consider this a bug.

Edit: I just looked and the code path I suspected would lead to this behavior does not. Can you give an example of a short code snippet that reproduces the issue you're seeing? The one in your question does not seem to result in any error.

[Jeevesh8]

Ohh.. but I see that the async problem may even occur even if the second function has print or other blocking statements. For if the first function doesn't have any, it will go on executing in a non-blocking way, and when the control reaches for calling to the function being pmap-ed(the "second function") for the first time, will the second function wait for the output from the first to become available ? And then continue, to compile? I think it will compile from just the tracer values(or whatever the first function outputs instead of the actual computed outputs) output from the first function, but just after it compiles, the state it will try to access, will it wait for it to become available?

[Jeevesh8]

@jakevdp I am not using the same array, but trying to use the array returned in place of it. But I see that that too gets donated in the next step.. so probably not safe to use it, unless I block_until_ready() and stop it from getting donated in the next iteration, and only after using it in the internal function, donate it for the next iteration.

[jakevdp]

Oh, I see – sorry I missed that aspect of it.

I think it would be worth taking a deeper look at what's going on here – if you're able to put together a minimal repro, it would be great if you could open an issue with the report.

[Jeevesh8]

How can I check if a device array has been deleted or not?

[jakevdp]

import jax.numpy as jnp
x = jnp.arange(10)
print(x.is_deleted())
# False