11–20 of 34 posts
Uhh they time the vmap of the jit on Jax, basically skipping a ton of optimizations,.esp if there is any linear algebra in there. They also include the cost of building the vmap functional. Not a valid comparison. https://github.com/utkarsh530/GPUODEBenchmarks/blob/ef807198...
https://github.com/utkarsh530/GPUODEBenchmarks/blob/ef807198...
Uhh they time the vmap of the jit on Jax, basically skipping a ton of optimizations,.esp if there is any linear algebra in there. They also include the cost of building the vmap functional. Not a valid comparison. https://github.com/utkarsh530/GPUODEBenchmarks/blob/ef807198...
Uhh they time the vmap of the jit on Jax, basically skipping a ton of optimizations,.esp if there is any linear algebra in there. They also include the cost of building the vmap functional. Not a valid comparison. https://github.com/utkarsh530/GPUODEBenchmarks/blob/ef807198...
What they should do is build the vmap and jit that , then run timing on calling the resulting function.
@jax.jit
@jax.vmap
def main(...
?If you want to say what you think is important about an article, that's fine, but do it by adding a comment to the thread. Then your view will be on a level playing field with everyone else's: https://hn.algolia.com/?dateRange=all&page=0&prefix=false&so...
(Submitted title was "Julia GPU-based ODE solver 20-100 x faster than those in Jax and PyTorch". We've changed that to a shortened version of the paper title, to fit HN's 80 char limit.)
Uhh they time the vmap of the jit on Jax, basically skipping a ton of optimizations,.esp if there is any linear algebra in there. They also include the cost of building the vmap functional. Not a valid comparison. https://github.com/utkarsh530/GPUODEBenchmarks/blob/ef807198...
Uhh they time the vmap of the jit on Jax, basically skipping a ton of optimizations,.esp if there is any linear algebra in there. They also include the cost of building the vmap functional. Not a valid comparison. https://github.com/utkarsh530/GPUODEBenchmarks/blob/ef807198...
This is very interesting. The claim does sound too good to be true. Am I understanding you correctly that you are saying including the vmap operation in the timing is wrong because it involves compilation time that could have been amortized over all the runs, and that the compilation time is considerable compared to the ode-solve itself?
There are many reasons for this. First, the former will loop iterations on the GPU in a serial fashion. Second, the internal jit makes optimization options opaque to Jax. For example, if there's a loop of matmuls inside main, that loop can be converted to a loop of einsums if you vmap first. It can also be fused into sometimes into a bigger operation that doesn't jump control variables back and forth between CPU and GPU between time steps. Between the two you both increase throughput and decrease latency.
I think in Jax, jit(vmap(jit(...))) will also reoptimize the same way as jit(vmap(...)) but I'm not 100% certain.
Earlier quoted context omitted.
What is direct kernel generation?
In this context I would imagine it's constructing source code for a kernel- the engine that implements a step in a neural network- that is closer to optimal. See https://cuda.juliagpu.org/stable/tutorials/performance/ for related work
> One can still batch together evaluations through the ODE solver by concatenating the states of each batch element together, creating a combined ODE with dimension D × K. In some cases, controlling error on all batch elements together might require evaluating the combined system K times more often than if each system was solved individually. However, in practice the number of evaluations did not increase substantially when using minibatches.
I don't understand why someone wrote a 30-page, obfuscated paper on just... parallelizing it the obvious way.
Earlier quoted context omitted.
This is very interesting. The claim does sound too good to be true. Am I understanding you correctly that you are saying including the vmap operation in the timing is wrong because it involves compilation time that could have been amortized over all the runs, and that the compilation time is considerable compared to the ode-solve itself?
There are two things going on. First, you're right that the vmap should have been done once outside the timing. But equally important, vmap(jit(...)) speed will generally be lower than jit(vmap(...)). There are many reasons for this. First, the former will loop iterations on the GPU in a serial fashion. Second, the internal jit makes optimization options opaque to Jax. For example, if there's a loop of matmuls inside…
Source: https://github.com/google/jax/discussions/5199#discussioncom...