4 ms·
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
by yablak 3y ago
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/ef807198584e4634c82bfd6a82bb9241c51d56d1/GPU_ODE_JAX/bench_diffrax.py#L87 https://github.com/utkarsh530/GPUODEBenchmarks/blob/ef807198...
- yablak 3y agoWhat they should do is build the vmap and jit that, then run timing on calling the resulting function.
- geysersam 3y agoSo instead of jitting main they should do something like @jax.jit @jax.vmap def main(... ?
- yablak 3y agoYes, that's right
- sroussey 3y agoHere? https://github.com/utkarsh530/GPUODEBenchmarks/blob/ef807198584e4634c82bfd6a82bb9241c51d56d1/GPU_ODE_JAX/bench_diffrax.py#L45 https://github.com/utkarsh530/GPUODEBenchmarks/blob/ef807198...
- deleted 3y ago[deleted]
- ChrisRackauckas 3y agoThis collab notebook shows effectively no difference from doing this: https://colab.research.google.com/drive/1d7G-O5JX31lHbg7jTzzozbo5-Gp7DBEv?usp=sharing https://colab.research.google.com/drive/1d7G-O5JX31lHbg7jTzz.... The average for diffrax on this collab machine goes from 20.5 to something like 20.3 seconds. You can see DiffEqGPU.jl running from Python via diffeqpy at around 2.3 seconds. This is a very rough benchmark of course since using DiffEqGPU has a fairly high (documented) overhead, and the free tier T4 GPU is not great, but it shows the ballpark of an order of magnitude or so. Note that you can also see that compile times are pretty negligible even at this scale (and the paper benchmarks are a few orders of magnitude larger than this, so at that point it's really negligible). That shouldn't be surprising though since we're talking about operations involving hundreds of thousands or millions of ODEs. At this scale, micro optimizations tend to have a much more minor effect. And the paper describes in detail that we developed two versions of the methods in Julia, one that was an array-based approach like Jax and PyTorch vmaps (EnsembleGPUArray), and another that was a kernel generating approach like MPGOS (EnsembleGPUKernel). Jax, PyTorch, and EnsembleGPUArray all performed similarly while MPGOS performed similarly to EnsembleGPUKernel. To us, this was a pretty strong indicator that the performance difference comes from the fact that the way EnsembleGPUKernel is performing the parallelism is very different from the approach that an ML library takes. And yes, there's small differences in the groups, but those are like 2x-3x or so, while the paper benchmarks are in log-scale because the difference between the two classes of designs are much larger.
- yablak 3y agoSame for pytorch. I don't know enough pyt, but guessing they didn't jit anything. https://github.com/utkarsh530/GPUODEBenchmarks/blob/ef807198584e4634c82bfd6a82bb9241c51d56d1/GPU_ODE_PyTorch/bench_torchdiffeq.py#L82 https://github.com/utkarsh530/GPUODEBenchmarks/blob/ef807198...
- geysersam 3y agoThis 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?
- yablak 3y agoThere 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 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.
- ubj 3y agoOn your last point, as long as you jit the topmost level, it doesn't matter whether or not you have inner jitted functions. The end result should be the same. Source: https://github.com/google/jax/discussions/5199#discussioncomment-219142 https://github.com/google/jax/discussions/5199#discussioncom...
- ChrisRackauckas 3y agoConfirmed in this collab notebook it doesn't make a tangible difference: https://colab.research.google.com/drive/1d7G-O5JX31lHbg7jTzzozbo5-Gp7DBEv#scrollTo=ExZj1Ak9UTuj https://colab.research.google.com/drive/1d7G-O5JX31lHbg7jTzz... .
- npalli 3y agoSeems like the benchmarking code is a small script and what you are suggesting might be a few lines of code. Might be worthwhile to take a stab and see if there is a difference.