4 ms·

Yes, we discuss this in the post, and cases where this is not possible. For example, in the adjoints of universal differential equations you have to take a Jaco

Yes, we discuss this in the post, and cases where this is not possible. For example, in the adjoints of universal differential equations you have to take a Jacobian-vector product. You can make that a matrix by batching, but then you are making a trade-off because you're effectively solving multiple ODEs simultaneously. When doing this, you might be able to fill the GPU kernels more, but you also have to take more time steps! This is because adaptive time stepping is only correct if you take the minimum of all of the steps (average would make some above error tolerance, or worse unstable). If you have some ODEs stiff, or worse have stochasticity that can cause random times to shrink the time steps at different points for different solvers, making a solver of the union of equations something that is extremely wasteful. And that becomes a dilemma: we started trying to solve multiple ODEs at the same time to fill a matrix-matrix multiplication kernel, but by doing so we now have to take 100x more matrix-matrix multiplications! So that ended up nowhere, and we focused on matrix-vector performance via direct SIMD and effective use of mutation (as the blog post describes) and got very good results. We are also redoing some of the GPU ODE solvers and seeing about 100x improvement by avoiding batching in a similar way, but that's a post for another time. All in all, some applications can just slam a big net. Others can't. The two cases need different optimizations.

patrickkidger 5y ago
Taking a union of differential equations is a pretty bad idea for exactly the reasons you describe. But it's absolutely possible to parallelise multiple diffeq solves without needing to use the same time steps for each solve. That is, the time steps can be vectorised, rather than broadcast, over the batch axis. So the only thing you actually need are the same number of steps for each solve, which can be easily accomplished by just padding out the solves that finish with slightly fewer steps. In practice this ends up introducing negligible overhead whilst solving the above issue very neatly. For example this is precisely what Diffrax (https://github.com/patrick-kidger/diffrax https://github.com/patrick-kidger/diffrax) does under `jax.vmap`. I've not dug into what Julia does here; is this not already done when broadcasting `DifferentialEquations.solve`?
- adgjlsfhk1 5y ago
  How does this work with stiffness adaptive solvers? The computation you need to do for when an equation is stiff looks pretty different from when it's not stiff, and I'm having a hard time seeing how that would work if you vectorize the timesteps. Do you need to synchronize after every timestep to split into separate batches for stiff and non-stiff?
  - patrickkidger 5y ago
    You're discussing switching algorithms like LSODA? This is a really good question that I don't have a neat answer to. You can actually also hit similar issues when naively vectorising some stiff algorithms that detect when to recalculate Jacobians: at each timestep, some batch elements will want the recalculation whilst some won't. In both cases, something like what you suggest might be a possible solution. And thankfully support for such custom batching will Soon (tm) be coming to JAX. In the mean time, Diffrax actually side-steps the problem by not implementing those kinds of solvers. New solvers usually get added on a by-request basis, so this just hasn't been an issue yet. So: 1. Do you know what DifferentialEquations.jl does in this scenario when running on the GPU? (On the CPU needing to sychronise batch elements usually isn't such a concern.) 2. If the price of using JAX is giving up a bit of performance in this one use case, then I think I'm okay with that. The clear contenders in this space are JAX and Julia, and right now it's a concious choice to be using JAX as a better fit for the problems I'm tackling. This comparison is something I've written about before: https://discourse.julialang.org/t/state-of-machine-learning-in-julia/74385/4 https://discourse.julialang.org/t/state-of-machine-learning-....
    - adgjlsfhk1 5y ago
      One of the best general purpose solvers (which DifferentialEquations.jl often uses by default) is AutoTsit5. This is a solver that detect stiffness per timestep and uses either Tsit5 (a 5th order RK solver) or Rosenbrock23 depending on the stiffness. These solvers are great because they typically give really good results for a wide variety of problems. The main way DifferentialEquations.jl deals with this currently is being faster on the CPU than other solvers are on GPU. For small models, GPU has too much overhead, and for large models, you can do linear algebra within timesteps on the GPU without compromising solver efficiency.

patrickkidger 5y ago

Taking a union of differential equations is a pretty bad idea for exactly the reasons you describe. But it's absolutely possible to parallelise multiple diffeq solves without needing to use the same time steps for each solve. That is, the time steps can be vectorised, rather than broadcast, over the batch axis. So the only thing you actually need are the same number of steps for each solve, which can be easily accomplished by just padding out the solves that finish with slightly fewer steps. In practice this ends up introducing negligible overhead whilst solving the above issue very neatly. For example this is precisely what Diffrax (https://github.com/patrick-kidger/diffrax https://github.com/patrick-kidger/diffrax) does under `jax.vmap`. I've not dug into what Julia does here; is this not already done when broadcasting `DifferentialEquations.solve`?

adgjlsfhk1 5y ago
How does this work with stiffness adaptive solvers? The computation you need to do for when an equation is stiff looks pretty different from when it's not stiff, and I'm having a hard time seeing how that would work if you vectorize the timesteps. Do you need to synchronize after every timestep to split into separate batches for stiff and non-stiff?
- patrickkidger 5y ago
  You're discussing switching algorithms like LSODA? This is a really good question that I don't have a neat answer to. You can actually also hit similar issues when naively vectorising some stiff algorithms that detect when to recalculate Jacobians: at each timestep, some batch elements will want the recalculation whilst some won't. In both cases, something like what you suggest might be a possible solution. And thankfully support for such custom batching will Soon (tm) be coming to JAX. In the mean time, Diffrax actually side-steps the problem by not implementing those kinds of solvers. New solvers usually get added on a by-request basis, so this just hasn't been an issue yet. So: 1. Do you know what DifferentialEquations.jl does in this scenario when running on the GPU? (On the CPU needing to sychronise batch elements usually isn't such a concern.) 2. If the price of using JAX is giving up a bit of performance in this one use case, then I think I'm okay with that. The clear contenders in this space are JAX and Julia, and right now it's a concious choice to be using JAX as a better fit for the problems I'm tackling. This comparison is something I've written about before: https://discourse.julialang.org/t/state-of-machine-learning-in-julia/74385/4 https://discourse.julialang.org/t/state-of-machine-learning-....
  - adgjlsfhk1 5y ago
    One of the best general purpose solvers (which DifferentialEquations.jl often uses by default) is AutoTsit5. This is a solver that detect stiffness per timestep and uses either Tsit5 (a 5th order RK solver) or Rosenbrock23 depending on the stiffness. These solvers are great because they typically give really good results for a wide variety of problems. The main way DifferentialEquations.jl deals with this currently is being faster on the CPU than other solvers are on GPU. For small models, GPU has too much overhead, and for large models, you can do linear algebra within timesteps on the GPU without compromising solver efficiency.