Hacker News Reader

Comment by patrickkidger | Hacker News Reader

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/68bf7b99ba02c246b20ea...)

View on HN

"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/68bf7b99ba02c246b20ea...

Nope, this was on v0.2 of SimpleChains.jl; I did check.

Anyway, something to investigate if I ever get the time.

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.

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.

Just to confirm, v0.2.1 is the fast version.

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.

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.

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.

Reply on news.ycombinator.com