Hacker News Reader

Comment by brrrrrm | Hacker News Reader

here's an impl: https://jott.live/code/train_matrix_exp_estim_jax.py

it works but surprisingly slow on my devices

    Initial Train Loss: 5.9414
    Initial Test Loss: 5.9414
    Took: 175.90 seconds
    Train Loss: 0.0108
    Test Loss: 0.0108
    Took: 183.89 seconds
    Train Loss: 0.0016
    Test Loss: 0.0016
    Took: 188.95 seconds
    Train Loss: 0.0010
    Test Loss: 0.0010

SimpleChains is about 0.5 seconds on this computer, about 5 seconds to compile. So you're right, it is slow, too slow, like 300x faster than Jax, and I suspect something is wrong with the Jax script. It should probably be closer to like 2.5x or so because of memory handling and vectorization stuff, but not any slower than PyTorch.

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

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.

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

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.

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.

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.

Thanks!

Reply on news.ycombinator.com