Best practices for updating array parameters and buffer management

[hamzamerzic]

Hello,

I am currently working with a model and trying to understand the most efficient methods to manage and update arrays, specifically model parameters, in JAX without unnecessarily duplicating parameters on the device. In addition, I would like to gain a deeper understanding of the technical aspects behind these methods as I believe my current comprehension might be incomplete.

For context, let's consider the following class:

class Inference:
  forward_fn: ...  # assuming jitted
  params: ...

  def serve(self, inputs):
    return self.forward_fn(inputs, self.params)

  def set_params(self, new_params):
    self.params = new_params

My primary concern revolves around the set_params function. The new_params are nested numpy arrays which need to be transferred from host to device (H2D) which happens during the next execution of forward_fn.

However, this raises several questions and concerns:

1. When does self.params get deallocated from its current device?
2. How does the garbage collector (gc) interact with self.params? Is there a dependency on gc for initiating memory cleanup on the device?
3. Does the H2D transfer of new_params occur only upon the next execution of forward_fn or is it dependent on gc?

I am curious if there's a way to manually update the buffer at our chosen time, taking into account that forward_fn could potentially be executing concurrently. I've looked into jax.device_put and jax.device_put_replicated, but as they don't seem to accept an existing buffer to update, I'm concerned this could lead to unnecessary device-side copying and dependency on background processes for memory cleanup.

Any advice, best practices, or insights into how JAX handles these situations would be greatly appreciated.

Thank you.

[jakevdp]

1. When does self.params get deallocated from its current device?
2. How does the garbage collector (gc) interact with self.params? Is there a dependency on gc for initiating memory cleanup on the device?

In short, assuming your code is not wrapped in jit, the way to think of it is like this,

• self.params references some Python object that we might call obj – Python maintains a refcount on obj that has a value of at least 1, becuse self.params currently references it.
• when you write self.params = new_params, the refcount of obj is decreased by 1. If self.params is the only reference to obj, then the refcount will become zero.
• at some point later (not necessarily immediately) Python's garbage collection process will detect that obj is no longer referenced in the program, and will deallocate obj
• if obj is a JAX array, this deallocation will tell XLA that the associated memory can be freed. At some point later (not necessarily immediately), XLA will deallocate and free that memory buffer from the device

3. Does the H2D transfer of new_params occur only upon the next execution of forward_fn or is it dependent on gc?

The operation self.params = new_params does not cause any device transfer. All it does is add a reference named self.params to the pre-existing object referenced by new_params. If this object is a JAX array, it will already be on a device. If this object is a numpy array, then device transfer will happen if and when it is passed to a JAX function. If you'd like to explicitly control when some object is transferred to a device, you can use jax.device_put or similar.

they don't seem to accept an existing buffer to update

Outside of JIT, there is no way to update the buffer pointed to by self.params with new data, because JAX arrays are immutable. Inside JIT, the XLA compiler may utilize in-place updates when the compiler considers it advantageous, so long as such updates don't change the semantics of the program (for example, when the original buffer is not referenced later in the program).

Hope that helps!

[hamzamerzic]

Thanks a lot for the detailed explanation Jake.

assuming your code is not wrapped in jit

Yes, that's a reasonable assumption. I'm only assuming self.forward_fn is jitted.

Just to summarize, reassigning the reference to self.params triggers gc which then at some point sends the signal to XLA to free the buffer. We cannot in-place update the buffer that self.params is referencing. Is that correct?

This means that to update self.params without using any additional memory we first need to free all the references to it and then manually put the new params to device. So then something like this should work?

self.params = jax.tree_map(lambda x, y: jax.device_put(x, device=y.sharding), new_params, self.params)

But, this still risks doubling the memory in case gc does not immediately clear the references, sends the signal to XLA and XLA deletes the buffer. Is there a more controlled way to do this? Is it enough to explicitly del the params?

[jakevdp]

When you explicitly del the params, it is a Python operation that decrements the reference count, which leads to (asynchronous) Python garbage collection followed by (asynchronous) XLA buffer deletion.

If you wish, you can call arr.delete() on any JAX array instance to manually trigger deletion of its XLA buffer, after which any attempt to use arr will result in an error.