6 ms·

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

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?

fnbr 5y ago
(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 5y ago
  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?
- sivakon 5y ago
  The first benchmark is wrong. Numpy random by default uses float64, and jax uses float32. In [9]: x = np.random.randn(10000,10000).astype('f') In [10]: %timeit -n5 fn(x) 623 ms ± 9.31 ms per loop (mean ± std. dev. of 7 runs, 5 loops each) With Jax In [17]: %timeit -n5 jax_fn(x).block_until_ready() 98.4 ms ± 1.55 ms per loop (mean ± std. dev. of 7 runs, 5 loops each) It's still 6x improvement but not as large that the article claims. I am on the latest Intel Mac
nestorD 5y ago
> 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.
deleted 5y ago
[deleted]
chillee 5y ago
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 5y ago
  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 5y ago
    re: leading comparison, that makes sense. I recommend labeling it a bit better: JAX on TPU vs NumPy on CPU
    - SleekEagle 5y ago
      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)"
      - p1esk 5y ago
        Comparing Jax on TPU vs Numpy on CPU does not make any sense. Of course a GEMM hardware accelerator will be much faster than a general purpose CPU. What is the point of this comparison? You either run the same code on different hardware, or different code on the same hardware. A much more interesting comparison would be CuPy or Pytorch code vs Jax code running on A100.
SleekEagle 5y ago
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 5y ago
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 5y ago
  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!

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?

fnbr 5y ago

(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 5y ago
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?
sivakon 5y ago
The first benchmark is wrong. Numpy random by default uses float64, and jax uses float32. In [9]: x = np.random.randn(10000,10000).astype('f') In [10]: %timeit -n5 fn(x) 623 ms ± 9.31 ms per loop (mean ± std. dev. of 7 runs, 5 loops each) With Jax In [17]: %timeit -n5 jax_fn(x).block_until_ready() 98.4 ms ± 1.55 ms per loop (mean ± std. dev. of 7 runs, 5 loops each) It's still 6x improvement but not as large that the article claims. I am on the latest Intel Mac

nestorD 5y ago

> 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.

deleted 5y ago

[deleted]

chillee 5y ago

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 5y ago
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 5y ago
  re: leading comparison, that makes sense. I recommend labeling it a bit better: JAX on TPU vs NumPy on CPU
  - SleekEagle 5y ago
    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)"
    - p1esk 5y ago
      Comparing Jax on TPU vs Numpy on CPU does not make any sense. Of course a GEMM hardware accelerator will be much faster than a general purpose CPU. What is the point of this comparison? You either run the same code on different hardware, or different code on the same hardware. A much more interesting comparison would be CuPy or Pytorch code vs Jax code running on A100.

SleekEagle 5y ago

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 5y ago

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 5y ago
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!