[HN Gopher] Using JAX in 2022
___________________________________________________________________
Using JAX in 2022
Author : SleekEagle
Score : 54 points
Date : 2022-02-15 17:32 UTC (5 hours ago)
(HTM) web link (www.assemblyai.com)
(TXT) w3m dump (www.assemblyai.com)
| azinman2 wrote:
| These benchmarks are pretty crazy, especially as I presumed NumPy
| to do far better. Isn't tensorflow etc already doing GPU
| acceleration? Is this actually a fair comparison?
| [deleted]
| nestorD wrote:
| > especially as I presumed NumPy to do far better
|
| If you are running your benchmark on a single Numpy operation
| then I would expect Numpy to have equal or better speed (you
| are not paying for JIT compilation). However, when you are
| doing several operations, Numpy will do a loop on array
| elements for each operation while JAX can fuse everything, that
| can end up making a big difference.
| chillee wrote:
| It's not quite a fair comparison, since Numpy is running with
| float64 while Jax is running with float32.
|
| If you fix the benchmarks then looks like this
|
| 5 loops, best of 5: 99.2 ms per loop
|
| 10 loops, best of 5: 114 ms per loop
|
| 10 loops, best of 5: 20.2 ms per loop
|
| 5x faster is to be expected as there are 5 pointwise operations
| (that are bandwidth bound) that can be fused.
|
| The leading comparison is also quite misleading, imo, since I
| think it's comparing Numpy on CPU vs. Jax on an accelerator.
| SleekEagle wrote:
| Thanks very much for posting this - I forgot JAX defaults to
| float32 - I'll fix that soon.
|
| As for the other part about the leading comparison - I was
| trying to highlight just how much faster JAX could be in the
| best-case scenario. Beyond the accelerator and JIT, the
| function itself lends to being expedited significantly when
| JITted. I posted benchmarks with a comparison of JAX vs NumPy
| both on CPU, and then with JAX on TPU further down to control
| more variables. (reposted from reddit)
| brrrrrm wrote:
| re: leading comparison, that makes sense. I recommend
| labeling it a bit better: JAX on TPU vs NumPy on CPU
| SleekEagle wrote:
| Updated! Thanks for the feedback. Added this note in the
| figure description
|
| "(n.b. JAX is using TPU and NumPy is using CPU in order
| to highlight that JAX's speed ceiling is much higher than
| NumPy's)"
| fnbr wrote:
| (I'm an engineer at DeepMind, and I work with JAX daily)
|
| It's a somewhat fair comparison; in my experience, highly
| optimized JAX matches highly optimized Tensorflow. However,
| non-optimized (but JITted) JAX beats non-optimized Tensorflow,
| as Tensorflow requires a lot of architectural changes to make
| it perform well. JAX, on the other hand, tends to perform well
| as long as you just JIT it. So it's much easier to get to, say,
| 90% of optimal performance. In Tensorflow, it's much harder (in
| my experience- maybe I'm just bad at Tensorflow).
|
| The JIT compilation that JAX does is really, really good, as it
| combines operations together in a highly performant way.
| SleekEagle wrote:
| Thanks for this comment, I probably could've stressed JIT
| more. Random question - I noted in the article that you all
| at DeepMind announced that you're using JAX to accelerate
| your research.
|
| IIRC, you guys standardized TensorFlow several years back.
| What does the current split look like between JAX and TF
| internally? Do some people use TF and some use JAX, or do you
| use JAX for specific tasks?
| SleekEagle wrote:
| NumPy can't run on accelerators, which is part of where the
| difference stems from. I implemented the same operation first
| with both NP/JAX on CPU, then with JAX on TPU and NP again on
| CPU, not sure if you saw that part.
|
| Obviously, the ability to use an accelerator makes JAX faster,
| but even without that on CPU it was faster. This is in part
| because of JIT, and in fairness the calculation in question
| does lend itself well to being expedited by JIT.
|
| I actually have some preliminary benchmarks for a follow up
| specifically on just NumPy vs JAX, and it has become clear so
| far that NumPy is better in certain cases, especially for small
| operations where the overhead of JAX is not worth it.
|
| In the article I mention this briefly, along with how JAX
| hasn't been focused on being optimized on CPU because they have
| bigger fish to fry, so to speak. I also link to the JAX
| documentation that has some comments comparing the two!
|
| Relating to TF - I don't actually use TF at any point, but I
| did use PyTorch for Hessian calculation. TF and PT obviously do
| both work on GPU, but JAX has the benefit of being able to JIT
| more and implement everything in terms of XLA (although TF
| obviously has XLA support as well, and PT kind of does but just
| to get PT working on TPU).
|
| Thinking about doing another article on a direct comparison of
| JAX with PT and TF - let me know if that's something you'd like
| to see!
| brilee wrote:
| NumPy does not automatically use GPU, so my guess is that the
| speedup is due to the CPU/GPU differential. This is
| specifically noted later on in the "Should I use JAX" section -
| if you just want numpy for accelerators, JAX is a no-brainer.
| SleekEagle wrote:
| Thanks for your comment - in the JIT section under "Jax
| Transformations" I implement the same calculation on NP and
| JAX on CPU (and TPU) to get a direct comparison for CPU. I
| mentioned in another comment that the calculation lends
| itself well to being JITted, but I have some preliminary data
| for a follow-up JAX vs NumPy performance comparison.
|
| It's clear so far that NumPy can outperform JAX on CPU for
| small computations (unsurprisingly). Please let me know if
| you'd be interested in seeing a more thorough analysis!
| skrebbel wrote:
| With the year so prominently in the title, i thought this was
| going to be about a technology that is obviously not trendy
| anymore, ie, JAX The Terrible Java XML Parser. I'm disappointed
| that it was just blog title spam about some non-controversial
| modern technology.
| jszymborski wrote:
| Wait, is the "The Terrible Java XML Parser" moniker tongue-and-
| cheek or sincere?
|
| I remember JAX (the XML parser) really saving my bacon when I
| was parsing larger-than-memory XML files ages ago.
| otsaloma wrote:
| How does JAX's JIT compare to Numba?
| SleekEagle wrote:
| I didn't get a chance to run experiments myself, but from other
| people I've heard JAX is superior.
| lmeyerov wrote:
| Am curious esp wrt numba gpu jit...
| brilee wrote:
| This is a really great overview of JAX - the best I've seen
| outside of primary Google/DeepMind sources. Glad to see people
| besides Googlers getting familiar with its capabilities.
| SleekEagle wrote:
| Thanks for your feedback! I wanted to introduce it for people
| who haven't worked with it before giving recommendations, but I
| got a little carried away because I think JAX is really cool!
|
| I was planning on doing a more thorough introductory tutorial
| or deep dive into Transformations, so let me know if you (all)
| think that would be instructive!
| time_to_smile wrote:
| > Why Should I Care About JAX?
|
| > In short - speed.
|
| For me personally, the magic of JAX is that it able to have this
| performance, while being as close as possible to having _first
| class differentiation_ in Python. The latter is a far more
| important reason to use JAX. It can really change how you think
| about programming and ML. Rather than implementing a specific
| model, you can write up the parameterized solution to a problem
| then _solve it_.
|
| However first class differentiation ultimately isn't really
| useful unless you happen to also solve the speed problem. That is
| what makes JAX incredible. From the programming perspective JAX
| is to differentiable programming what Prolog is to logic
| programming, however Prolog has always been limited ultimately by
| performance problems where JAX is not.
| SleekEagle wrote:
| Thanks for your comment - I wanted to keep it as general as
| possible, and I think speed touches basically every possible
| use of JAX, but I definitely hear what you are saying. For DL,
| jvps and quick/easy Hessians (along with speed) are really the
| reasons to use JAX imo.
|
| I like your comment on thinking about implementing
| models/thinking about them as parameterized solutions
| baldeagle wrote:
| I've been really struggling to see the benefit from autodiff...
| I work with DAGs, where the data processes are batches that
| could be processed in parallel and I would love to speed that
| up, or to enable some gradient descent style optimization....
| but I can't figure out how autodiff plays into it. Do you know
| of a good place to help learn how it could apply to my domain?
| SleekEagle wrote:
| Firstly vmap() might make it easier for you from a
| programming standpoint. You can think of operations in terms
| of single instances and then autobatch them.
|
| Also, jit() is where you'll likely see your speed increases.
|
| As for autodiff - not every application will use it. What
| exactly are you trying to optimize? If you can parameterize
| your model somehow (e.g. parameterize edge weights) and then
| figure out some way to measure how "bad" your model is, you
| can use autodiff to tune your edge weights to minimize that
| metric. Not too familiar with DAGs but the first step is
| figuring out how to parameterize your model in a way that
| tuning the params can lead to your goal (and how to measure
| how close you are to that goal)
| paperwork wrote:
| I wish JAX worked with windows natively (without using wsl). I
| teach a very high level intro to numpy and would _love_ to have
| my students try jax. These students are relatively new to
| programming and the idea of using a linux shell or having to
| compile anything themselves just wouldn't work.
| SleekEagle wrote:
| VM with Ubuntu? If you write a guide they can get set up with
| VMWare pretty easily!
___________________________________________________________________
(page generated 2022-02-15 23:02 UTC)