How to avoid doubled memory consumption when updating a huge state?

[ZhengKeli]

Consider the following code, where I want to:

1. Load a state from a file (non-pure function)
2. Update it without changing its shape (pure function)
3. Save it back to a file (non-pure function)

import jax
import jax.numpy as jnp


@jax.jit
def update(state: jax.Array) -> jax.Array:
    return state + 1.0


# load state from a file
state = jnp.zeros([1024, 1024, 1024], dtype=jnp.float32)

# update the state
state = compute_new_state(state)

# save state to a file
print(state)

The problem is that the state is too huge that the memory can only barely hold it. Theoretically this is ok if all the operations in update() can be executed "in-place". Since jax requires update() to be a pure function, the code is not writen in a "in-place" manner. Hopefully, jax is smart enough to optimize the code to be in-place.

However, when it comes to the assignment state = compute_new_state(state), there is a moment when the result of compute_new_state(state) has been computed but the state is not assigned yet. At this moment the memory has to hold 2 instances of state - the newlly computed one and the old one, which is impossible because the memory can barely hold only one.

I can immediately think of approaches like:

• Give the state a flag so that it can be garbage-collected once needed before assignment.
• Deferring computation until after the assignment, when the old state object has already been garbage-collected.

But I haven't found any way to do that yet. Is there any feasible API or tools?

[yashk2810]

Have you tried partial(jit, donate_argnums=0)?

[ZhengKeli]

Great! This is exactly what I'm looking for! Thank you very much! 😀

By the way, after a series of experiments with memory profiling I found that:

• When using cpu, the computation is oftenly schduled after the assignment, so it is actully alright.
• When using gpu, adding a "donate_argnums" becomes necessary, maybe because the interaction with GPU is not instant.
• When using gpu, in some extreme cases, it needs to set XLA_PYTHON_CLIENT_ALLOCATOR=platform to make it run.