JAX suddenly 2x slower on V100 GPUs

[minqi]

Seemingly overnight, my JAX project (in v. 0.4.1) became suddenly ~2x slower in terms operations per second when running on V100 and Quadro GP100s. The project git commit hash is identical and there were no changes to my conda environment. Does anyone have any ideas for what might cause this performance degradation?

[astanziola]

I am experiencing also increasing slowness of JAX since a few releases, expecially when working with very large array (on the order of a Gb or more)

[mattjj]

Hrm, this is weird...

The project git commit hash is identical and there were no changes to my conda environment.

@minqi do you think your jaxlib changed? Could you try installing an older vs current jaxlib?

Any chance at sharing a repro?

[yashk2810]

Also, what change did you make when you noticed this slowness? Was it a jax/jaxlib upgrade or was it just shifting from 1 GPU to another GPU or something else?

If you upgraded JAX, don't forget to bump your jaxlib to 0.4.1 too!

[minqi]

Thanks for the responses! I ended up reinstalling both jax + jaxlib and now see the prior training speeds.

[mattjj]

That's great! Perhaps that's the first line of defense: if you come across this discussion page, update your jax and jaxlib first!