Memory leak when using @partial(jax.jit, static_argnums=(0,)) in class

[pipme]

From the activity monitor, the process's memory keeps growing during the for loop. How to deal with it? static_argnums would cause memory leak? Thanks in advance.

import jax
import jax.numpy as jnp
import numpy as np
from time import sleep
from functools import partial

class A:
    def __init__(self, name):
        self.name = name
        self.X = np.ones((5000, 5000))

    @partial(jax.jit, static_argnums=(0,))
    def f(self, Y):
        return jnp.sum(Y**2)

Y = np.ones(10)
for i in range(100):
    a = A(i)
    a.f(Y)
    sleep(2)

[jakevdp]

Hi - thanks for the question! In some senses this is intended behavior, in that the JAX jit cache uses strong references to track static arguments, so that if an argument goes out of scope the cache entry is not discarded. You can see why this would be useful, because, e.g. you wouldn't want to have to recompile a function just because you used a non-global string value as a static argument!

The side effect, then, is that if you have very large temporary static arguments, a strong reference will be created within the jit cache that will keep them in scope. This isn't really a leak per-se – the object is deliberately referenced in order to keep cached versions of compiled functions and prevent having to re-compile. You could address this by manually clearing the JIT cache of the function, in which case the temporary objects are garbage collected and the memory use remains constant:

import os, psutil

for i in range(41):
  A.f.clear_cache()
  a = A(i)
  a.f(Y)
  if i % 10 == 0:
    process = psutil.Process(os.getpid())
    print(f"Iteration {i}: {process.memory_info().rss} bytes")
# Iteration 0: 471412736 bytes
# Iteration 10: 873562112 bytes
# Iteration 20: 873660416 bytes
# Iteration 30: 873660416 bytes
# Iteration 40: 873672704 bytes

The jit cache is an LRU cache of a fixed size, so eventually unused static arguments will be discarded, but if your static arguments are large, you'll probably run out of memory before that point.

A better solution here would be to use static argnums as JAX intends: i.e. pass small hashable truly static values. Note that your use of static_argnums here goes against the recommendations in the jax.jit documentation; namely:

Static arguments should be hashable, meaning both __hash__ and __eq__ are implemented, and immutable.

Your class does not implement __hash__ or __eq__, and is not immutable (the value of self.X could change).

If you're working with classes that contain arrays, marking them as static is bad practice as it can easily result in surprises; here's a quick example:

class A:
  def __init__(self):
    self.x = jnp.arange(5)

  @partial(jax.jit, static_argnums=0)
  def f(self):
    return self.x.sum()

a = A()
print(a.f())
# 10

a.x += 1 # mutation that doesn't change the hash
print(a.f())
# 10

The class has violated the contract of static_argnums – it is not immutable – and so the output of the second call is incorrect! It would be nice if JAX could catch this and warn users that it's happening, but that's a somewhat tricky thing to do well (see related discussion in #9024).

A better approach for this kind of thing would be to register your class as a custom pytree, which is effectively a way to tell jax which aspects of the class are static, and which are dynamic. For your class it might look like this:

from jax.tree_util import register_pytree_node

class A:
  def __init__(self, name):
    self.name = name
    self.X = np.ones((5000, 5000))

  @jax.jit  # note: no static argnums
  def f(self, Y):
    return jnp.sum(Y**2)

def _flatten_A(a):
  children = (a.X,)  # non-static values
  aux_data = (a.name,)  # static values
  return (children, aux_data)

def _unflatten_A(aux_data, children):
  X, = children
  name, = aux_data
  a = A(name)
  a.X = X
  return A

register_pytree_node(A, _flatten_A, _unflatten_A)

a = A(1)
print(a.f(Y))
# 10.0

here a.X is no longer treated as a static argument, so no long-lived reference to it is created, and you'll no longer see the blow-up in memory on repeated calls.

Does that make sense?

[pipme]

Oh A.f._clear_cache() doesn't work but A.f.clear_cache() does work. There was no A.f.clear_cache() in 0.3.10 so I thought it was a typo 😅.

[pipme]

I will have a closer look at your other comments tomorrow. Thanks again!

[pipme]

from jax.tree_util import register_pytree_node


class A:
    def __init__(self, name):
        self.name = name
        self.X = np.ones((5000, 5000))

    @jax.jit  # note: no static argnums
    def f(self, Y):
        return jnp.sum(Y ** 2) * jnp.sum(self.X ** 2)


def _flatten_A(a):
    children = (a.X,)  # non-static values
    aux_data = (a.name,)  # static values
    return (children, aux_data)


def _unflatten_A(aux_data, children):
    (X,) = children
    (name,) = aux_data
    a = A(name)
    a.X = X
    return A

register_pytree_node(A, _flatten_A, _unflatten_A)

Y = jnp.ones(2)
a = A(1)
print(a.X.shape)  # output: (5000, 5000)
a.f(Y)  # AttributeError

I might still not fully understand. The above raises AttributeError: type object 'A' has no attribute 'X'.

[jakevdp]

Sorry, there's a typo. In _unflatten_A, it should be return a rather than return A (i.e. you return the unflattened object)

[pipme]

Oops get it. Thanks!

[patrick-kidger]

For context, cross-referencing patrick-kidger/diffrax#142 here as another example of the same issue. (Where the solution is the same -- clear the JIT cache.)

In addition, @pipme - if you want to do "JAX with classes", and in particular have the register-as-custom-pytree handled for you, then have a look at Equinox.

[pipme]

@patrick-kidger Thanks for the info! I was looking at treeo and will definitely have a look at Equinox too : )

[pipme]

import jax
import jax.numpy as jnp
from jax.tree_util import register_pytree_node

class A:
    def __init__(self, name):
        self.name = name
        self.X = jnp.ones((5000, 5000))

    @jax.jit  # note: no static argnums
    def f(self, Y):
        return jnp.sum(Y ** 2) * jnp.sum(self.X ** 2)


def _flatten_A(a):
    children = (a.X,)  # non-static values
    aux_data = (a.name,)  # static values
    return (children, aux_data)


def _unflatten_A(aux_data, children):
    (X,) = children
    (name,) = aux_data
    a = A(name)
    a.X = X
    return a

register_pytree_node(A, _flatten_A, _unflatten_A)

a = A(1)
for i in range(1,5):
    a.f(jnp.ones(i))
print(A.f._cache_size())  # 4
del a
print(A.f._cache_size())  # 4

@jakevdp After registering the class as custom pytree, I would like the jit cache to be automatically cleared while it seems not. I am aware of #7930 (comment). I need to create a lot of objects like a and call a.f on dynamic shape inputs, resulting non-negligible increasing memory due to jit cache. How could I deal with it?

Is calling from xxx import A; A.f.clear_cache() periodically the only way? I also had a look on #10828. Ideally we'd like to treat jitting class bound method in a different way and clear jit cache once the underlying self object is deleted. Am I missing something? Thanks for help!

[jakevdp]

It's hard to guess why you might have a memory leak without seeing the code. Can you create a self-contained minimal example that demonstrates the issue?

[pipme]

Thanks. It's a bit complex to create a minimal example for my concrete problem. In the mean time, a simple snippet below increases memory by ~20MB at each print, until 4096 maxmium cache size. It starts with 112MB and finalize with 637MB, which is quite big and non-negligible as far as I can see. As in my problem, the array shape keeps changing. Is there a way to clear all cache programmatically? I saw #2072 (comment) but it looks like xla._xla_callable.cache_clear() has been removed. Or a method to config lru's maxsize can also be useful?

@jax.jit
@jax.value_and_grad
def f(x):
    return jnp.sum(x)


print(f"{psutil.Process().memory_info().rss / 1e6} MB")
for i in tqdm(range(5000)):
    f(jnp.ones(i + 1))
    if i % 100 == 0:
        print(f"{psutil.Process().memory_info().rss / 1e6} MB")
print(f"{psutil.Process().memory_info().rss / 1e6} MB")

[pipme]

Now I am trying with a hack #10828 (comment) provided by @patrick-kidger . It reduces the memory usage quite a lot during the long run!

[pipme]

Ok I see clear_backend(). Will have a try.
https://github.com/google/jax/blob/7721579700ee5c4a951e72156a0bdac4c9768f43/jax/_src/api.py#L3250-L3264

[jakevdp]

Thanks for the example - it looks like each cached compilation is using ~200KB of memory. If this is problematic, then the new clear_backends API introduced in #11462 may be helpful...

At this point the cache size is fixed at 4096 (I believe) but perhaps this makes clear that we should allow this to be user-configurable?