Memory profiling questions

[dbalabka]

I'm trying to migrate some code from Numba to JAX. As a result, I've spent much time figuring out why following code example memory usage spikes up to ~1.5 GB while the Numba version of the same algorithm has only a few megabytes memory footprint:

In the following example, I've implemented a bootstrap algorithm that estimates mean distribution:

import jax
import jax.random as random
import numpy as np
from memory_profiler import profile

@profile
def sampling(x, key):
    return random.choice(random.split(key), x, (x.shape[0], )).mean()

@profile
def main():
    size = 100_000
    data = np.random.normal(0, 0.5, size)
    n_samples = 1000

    random_key = jax.random.PRNGKey(42)
    jax.vmap(jax.partial(sampling, data))(jax.random.split(random_key, n_samples)).block_until_ready()

    jax.profiler.save_device_memory_profile("memory.prof")

main()

Profiling with Memory Profiler:

Line #    Mem usage    Increment  Occurences   Line Contents
============================================================
     6    202.3 MiB    202.3 MiB           1   @profile
     7                                         def sampling(x, key):
     8    207.7 MiB      5.4 MiB           1       return random.choice(random.split(key), x, (x.shape[0], )).mean()


Filename: /home/torinaki/src/ab-testing/batch_test.py

Line #    Mem usage    Increment  Occurences   Line Contents
============================================================
    10    151.5 MiB    151.5 MiB           1   @profile
    11                                         def main():
    12    151.5 MiB      0.0 MiB           1       size = 100_000
    13    152.4 MiB      0.9 MiB           1       data = np.random.normal(0, 0.5, size)
    14    152.4 MiB      0.0 MiB           1       n_samples = 1000
    15                                         
    16    195.3 MiB     42.9 MiB           1       random_key = jax.random.PRNGKey(42)
    17    207.7 MiB     12.4 MiB           1       jax.vmap(jax.partial(sampling, data))(jax.random.split(random_key, n_samples))

Also, I've tried to profile the code with JAX's Device Memory Profiler. Unfortunately, both profilers fail to spot the memory consumption problem.

A similar discussion can be found here: #7509

It would be highly helpful if you can help me to answer these questions:

1. Could you please explain what is the reason for high memory usage? Probably, it is somehow related to autobatching in vmap and the explanation can be found in the documentation but it isn't clear for me.
2. Is it possible to spot this problem using JAX's memory profiler somehow or any other profiler? It seems that compiled code can not be profiled so easily.
3. Is it correct to pass random keys split by jax.random.split into vmap or there is any other better way to do it?

Thanks in advance.

[ThomasWarford]

Hey, did you find a good solution to this issue? I'm having a similar problem with high memory usage.

I'm thinking perhaps vmap isn't the best option for running the same computation with different data many times in the fastest way possible,

[dbalabka]

Unfortunately, I had to use Numba instead and didn't have a chance to get back to this issue. IMO there should be a simple answer to why high memory consumption is happening. Probably, Jax creates all possible batches and only then starts the computing via vmap. If it is the case, then IMO JAX should have a pool with a size of vCPUs count which should fill with gradually created batches to not reach OOM. Probably, it is worth trying to simulate the pooling on pure Python.