4 ms·

What they should do is build the vmap and jit that, then run timing on calling the resulting function.

by yablak 3y ago

geysersam 3y ago
So instead of jitting main they should do something like @jax.jit @jax.vmap def main(... ?
- yablak 3y ago
  Yes, that's right
  - sroussey 3y ago
    Here? 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 ago
  This 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.

What they should do is build the vmap and jit that, then run timing on calling the resulting function.

by yablak 3y ago

What they should do is build the vmap and jit that, then run timing on calling the resulting function.

geysersam 3y ago

So instead of jitting main they should do something like @jax.jit @jax.vmap def main(... ?

yablak 3y ago
Yes, that's right
- sroussey 3y ago
  Here? 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 ago
This 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.