5 ms·

I've just run this myself. Tidying up the JAX script, I find that it runs in ~15 seconds on my laptop CPU. Likewise, fixing the crash error in the blog post (U

I've just run this myself. Tidying up the JAX script, I find that it runs in ~15 seconds on my laptop CPU. Likewise, fixing the crash error in the blog post (UndefVarError: alloc_threaded_grad not defined) then I find that the Julia implementation also runs in ~15 seconds on my laptop CPU. So, 1x faster than JAX? Clearly we're getting very different results here, so (a) I'd like to get to the bottom of this, but (b) perhaps we should be more cautious about making performance claims. (EDIT: updated script in this GitHub gist: https://gist.github.com/patrick-kidger/68bf7b99ba02c246b20eaa38f2ad3d38 https://gist.github.com/patrick-kidger/68bf7b99ba02c246b20ea...)

ChrisRackauckas 4y ago
"UndefVarError: alloc_threaded_grad not defined" means that you weren't running the released version that is discussed in the blog post but the one of the much older slower ones. Any anyways, removing threaded calls and then benchmarking is clearly going to make it slower. With the updated ones, we seem to see 2x, 10x, and 22x, with increasing performance improvements based on the availability of threads and the existence of AVX512. For anyone who wants the details, just check that same gist: https://gist.github.com/patrick-kidger/68bf7b99ba02c246b20eaa38f2ad3d38 https://gist.github.com/patrick-kidger/68bf7b99ba02c246b20ea...
- patrickkidger 4y ago
  Nope, this was on v0.2 of SimpleChains.jl; I did check. Anyway, something to investigate if I ever get the time.
  - adgjlsfhk1 4y ago
    Just to confirm, v0.2.1 is the fast version.
    - ChrisRackauckas 4y ago
      Hmm... looks like he never came back to properly benchmark it with the right version, never did a version with the multithreaded call, and has a 1x number that is suspiciously different from everyone else who ran the benchmark in the Gist. Peculiar.
  - celrod 4y ago
    We probably should have released the blog post version as 0.3 instead of 0.2.1. There were not any breaking changes, but enough fixes and additions (e.g., the convolutional layer and threading) that 0.2.0 isn't really comparable. Odds are there is a version conflict in 0.2.1 (which also bumped dependency requirements) that does not exist in 0.2.0, and that this is preventing you from upgrading. You could start a smaller environment, or try `] add SimpleChains@0.2.1` to force it/get an error message telling you what the conflict is.
  - sundarurfriend 4y ago
    Do you mean v0.2.0 or v0.2.2? From a cursory look at the repo, it looks like alloc_threaded_grad was added in 0.2.1, so the UndefVarError makes sense if you were using the older 0.2.0 version.
brrrrrm 4y ago
Thanks for the updated script! I added a comment to my code pointing to yours. I’m curious what the key slowdown in my attempt was.
- patrickkidger 4y ago
  Not at all -- I was very happy to see some Equinox code out in the wild! I think the main slowdown was doing the model updates out of the JIT region. JAX-without-JIT is ridiculously slow.

ChrisRackauckas 4y ago

"UndefVarError: alloc_threaded_grad not defined" means that you weren't running the released version that is discussed in the blog post but the one of the much older slower ones. Any anyways, removing threaded calls and then benchmarking is clearly going to make it slower. With the updated ones, we seem to see 2x, 10x, and 22x, with increasing performance improvements based on the availability of threads and the existence of AVX512. For anyone who wants the details, just check that same gist: https://gist.github.com/patrick-kidger/68bf7b99ba02c246b20eaa38f2ad3d38 https://gist.github.com/patrick-kidger/68bf7b99ba02c246b20ea...

patrickkidger 4y ago
Nope, this was on v0.2 of SimpleChains.jl; I did check. Anyway, something to investigate if I ever get the time.
- adgjlsfhk1 4y ago
  Just to confirm, v0.2.1 is the fast version.
  - ChrisRackauckas 4y ago
    Hmm... looks like he never came back to properly benchmark it with the right version, never did a version with the multithreaded call, and has a 1x number that is suspiciously different from everyone else who ran the benchmark in the Gist. Peculiar.
- celrod 4y ago
  We probably should have released the blog post version as 0.3 instead of 0.2.1. There were not any breaking changes, but enough fixes and additions (e.g., the convolutional layer and threading) that 0.2.0 isn't really comparable. Odds are there is a version conflict in 0.2.1 (which also bumped dependency requirements) that does not exist in 0.2.0, and that this is preventing you from upgrading. You could start a smaller environment, or try `] add SimpleChains@0.2.1` to force it/get an error message telling you what the conflict is.
- sundarurfriend 4y ago
  Do you mean v0.2.0 or v0.2.2? From a cursory look at the repo, it looks like alloc_threaded_grad was added in 0.2.1, so the UndefVarError makes sense if you were using the older 0.2.0 version.

brrrrrm 4y ago

Thanks for the updated script! I added a comment to my code pointing to yours. I’m curious what the key slowdown in my attempt was.

patrickkidger 4y ago
Not at all -- I was very happy to see some Equinox code out in the wild! I think the main slowdown was doing the model updates out of the JIT region. JAX-without-JIT is ridiculously slow.