vmap seems to replicate memory in newer version of jaxlib

[helange23]

I just upgraded jaxlib from 0.1.57 to 0.1.68 and that basically broke all of my code.

I have a large vmap (over 1m) and within the vmap I am accessing a fairly large matrix (1GB+). I realized that in the newer version of jaxlib, it attempts to replicate this matrix 1m times which obviously doesn't fit.

Downgrading to the older jaxlib version got rid of the problem. I would like to release my code working with latest version of JAX.

Is this intended behavior? If so, can I somehow tell vmap not to replicate the matrix?

[jakevdp]

Hi - thanks for the question and sorry for the problems you're facing! Jumping from 0.1.57 to 0.1.68 reflects over 7 months of active development in JAX and XLA, so it's not easy to guess what changes might have affected your use-case. Is it possible to give a small example of a script that demonstrates the issue?

[helange23]

So my code runs smoothly with jax 0.2.7 and jaxlib 0.1.57 but crashes with jax 0.2.16 and jaxlib 0.1.68.

The error I get is: RuntimeError: Unimplemented: Kernel launch needs more blocks (28672000000) than allowed by hardware (2147483647).

If I reduce something akin to a batch size I get an out of memory error and when I don't access the large array during the vmap everything seems to run fine.

I couldnt just downgrade jaxlib because I wasn't able to import jax anymore. I will try to get a minimum example running.

[mattjj]

I suspect the issue is with the jax version and not the jaxlib version. Are you able to run jax==0.2.7 against jaxlib==0.1.68? If so that'd be something to try.

A minimal repro will let us get to the bottom of it!

[helange23]

So I stepped through the different version of jax and jaxlib and it seems to go bad going from jax 0.2.9 to 0.2.10 (both tried with jaxlib 0.1.60).

The error is completely different now though. Jitting my loss now never finishes and I get a "very slow compile?" warning. I killed the process after about 5min. So I suspect that I am overwhelming the JIT compiler?

In the code I have a large vmap that has a while in it that does a bunch of polynomial arithmetic. I found that unrolling the poly arithmetic was much faster than using the jax.numpy implementations. So instead of doing:

c = jnp.polymul(a,b)

I do:

c = [a0*b0, a0*b1 + a1*b0, a0*b2 + a1*b1 + a2*b0, a0*b3 + a1*b2 + a2*b1,
       a0*b4 + a1*b3 + a2*b2, a1*b4 + a2*b3, a2*b4]

At least for the old version of JAX that I was using this was sometimes 3x faster but led to fairly slow JIT compile times (about 40s). But I was happy because I only pay that price once. (Poly arithmetic is probably about 5.000-10.000 operations).

Was something like a setting changed for the JIT compiler that makes it optimization more aggressive at potentially increasing its runtime? If so, is there a way to tell it to chill like its 0.2.9?

[mattjj]

Here's the changelog entry for 0.2.10. Nothing really stands out to me.

All the aggressiveness of optimizations would be in XLA (i.e. in jaxlib), not the jax version.

This feels like we're gathering more info, but we may still need a repro to really pin it down.

[helange23]

Here is the smallest example I could find that produces the behavior. It's unfortunately not that minimal. Calling rf_render_bg works fine with version <0.2.9 and is actually pretty fast, never terminates with version 0.2.10 and throws an Unimplemented RuntimeError with version 0.2.16. Any pointers towards how I can make the code work for most (or the latest) versions of JAX would be greatly appreciated.

https://github.com/helange23/fast_taylor_transform/blob/main/this_breaks.py

Also, I noticed that if you change the 512 in the code (which is a number representing resolution so everything scales quadratically with it) to 16, then JAX doesn't throw the "not enough blocks" error but runs out of memory. It tries to allocate the large matrix exactly 16**2 times. Not sure if that helps.