3 ms·

If your network is small, use a larger batch size to take advantage of GPUs

Sukera 5y ago
The article addresses that: > This problem is far too small to saturate the GPU, even with such a large batch size. Time is dominated by moving batches from the CPU to the GPU. Unfortunately, as the batch sizes get larger, we need more epochs to reach the same accuracy, so we can hit a limit in terms of maximizing accuracy/time. The batch size in their example is 2048.
- lostmsu 5y ago
  You could also preload your datasets to GPU, if its RAM is large enough.
  - celrod 5y ago
    The Flux.jl example did this. A PR to the PyTorch example to do this would be welcome: https://github.com/chriselrod/LeNetTorch https://github.com/chriselrod/LeNetTorch
ChrisRackauckas 5y ago
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-....
paradite 5y ago
For reinforcement learning using mini batch updates, batch size beyond 32/64 is shown to decrease performance in various problems.

The article addresses that: > This problem is far too small to saturate the GPU, even with such a large batch size. Time is dominated by moving batches from the CPU to the GPU. Unfortunately, as the batch sizes get larger, we need more epochs to reach the same accuracy, so we can hit a limit in terms of maximizing accuracy/time. The batch size in their example is 2048.

lostmsu 5y ago
You could also preload your datasets to GPU, if its RAM is large enough.
- celrod 5y ago
  The Flux.jl example did this. A PR to the PyTorch example to do this would be welcome: https://github.com/chriselrod/LeNetTorch https://github.com/chriselrod/LeNetTorch

ChrisRackauckas 5y ago

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

paradite 5y ago

For reinforcement learning using mini batch updates, batch size beyond 32/64 is shown to decrease performance in various problems.