Utkarsh Utkarsh, Valentin Churavy, Yingbo Ma
Tim Besard, Prakitr Srisuma, Tim Gymnich, Adam R. Gerlach, Alan Edelman, George Barbastathis, Richard D. Braatz, Christopher Rackauckas
2023 · arXiv
We demonstrate a high-performance vendor-agnostic method for massively parallel solving of ensembles of ordinary differential equations (ODEs) and stochastic differential equations (SDEs) on GPUs. The method is integrated with a widely used differential equation solver library in a high-level language (Julia's DifferentialEquations.jl) and enables GPU acceleration without requiring code changes by the user. Our approach achieves state-of-the-art performance compared to hand-optimized CUDA-C++ kernels while performing 20--100$\times$ faster than the vectorizing map (vmap) approach implemented in JAX and PyTorch. Performance evaluation on NVIDIA, AMD, Intel, and Apple GPUs demonstrates performance portability and vendor-agnosticism. We show composability with MPI to enable distributed multi-GPU workflows. The implemented solvers are fully featured -- supporting event handling, automatic differentiation, and incorporation of datasets via the GPU's texture memory -- allowing scientists to take advantage of GPU acceleration on all major current architectures without changing their model code and without loss of performance. We distribute the software as an open-source library https://github.com/SciML/DiffEqGPU.jl
Metadata: arXiv. Abstracts retain their original authors’ rights.
Original HN comments, with labels from simple text rules. A label describes the whole comment and may miss context or irony. Read the discussion to make up your own mind. About the labels
5 selected excerpts · 5 total mentions
One of the really nice things about Julia for GPU programming is that you can write your own kernels. CUDA.jl isn't just C kernels. This is why (for example) DiffEQGPU.jl is able to be a lot faster than the other GPU based ODE solvers (see https://arxiv.org/abs/2304.06835 for details).
No, there are two types of GPU usage: the ensemble form and the array form. In the array form, you simply make your state variable a GPU and can run all internal operations on the GPU. This is used for GPU acceleration of `f` functions where you have a large state and a relatively "regular" function in `f`, such as neural networks or PDE discretizations. Thus you can see this in action for example on things like the DeepEquilibriumNetworks.jl https://docs.sciml.ai/DeepEquilibriumNetworks/stable/tutoria... use cases, where it's big neural networks with nonlinear solvers in them. The other use case is if you're solving tons of f(u,p)=0 equations for many different p, but…
all the same `f`. This shows up for example when doing parameter sweeps. You might think that you can just use the former for this: if you originally had `u` a vector, then you can make `u` a matrix and solve `f(u[:,i],p[:,i])` all simultaneously. This is in fact what machine learning libraries tend to do with their `vmap` approach. However, that hides a choice: how do you change `f` to build the composite `f`? You can either do the matrix operations on each operation of `f`, or you can build a specific GPU kernel for `f`. Of course you can do the former, that is exactly the same as the use case above, reusing the same machinery. However, we also offer the latter. And the reason for that is because calling GPU kernels has a lot of overhead and inhibits interprocedural optimizations, and therefore the former is a lot slower than the latter when you can exploit this known repeated structure. We were able to demonstrate this in a recent publication which showed using the latter approach for ODEs was about 20x-100x faster than PyTorch, Jax, and our own implementation of the former, while our approach to the latter is just about as fast as MPGOS which is a CUDA kernel, indicating that it's some limitation of the design (see https://www.sciencedirect.com/science/article/abs/pii/S00457..., or the arxiv version https://arxiv.org/abs/2304.06835). The kernel building approach is thus simply more efficient for things like parameter sweeps, and it's a pretty unique way of using the GPU as most machine learning libraries are not built around kernel generation but instead only give you the linear algebra kernel calling approaches. tl;dr, DiffEqGPU.jl's documentation is very clear on the two types of GPU usage https://docs.sciml.ai/DiffEqGPU/stable/getting_started/, but is specifically for ODEs. We need to write some similar tutorials for GPU usage on the nonlinear solvers to this effect.
https://ar5iv.org/abs/2304.06835 Thank you! I guess that's what: > If you are familiar with ar5iv, an arXivLabs collaboration, our HTML offering is essentially bringing this impactful project fully “in-house”. Our ultimate goal is to backfill arXiv’s entire corpus so that every paper will have an HTML version, but for now this feature is reserved for new papers. Refer to.
Indeed, and this year we created a system for compiling ODE code not just optimized CUDA kernels but also OneAPI kernels, AMD GPU kernels, and Metal. Peer reviewed version is here (https://www.sciencedirect.com/science/article/abs/pii/S00457...), open access is here (https://arxiv.org/abs/2304.06835), and the open source code is at https://github.com/SciML/DiffEqGPU.jl. The key that the paper describes is that in this case kernel generation is about 20x-100x faster than PyTorch and Jax (see the Jax compilation in multiple ways in this notebook https://colab.research.google.com/drive/1d7G-O5JX31lHbg7jTzz..., extra overhead though from calling Julia from Python but still…
shows a 10x). The point really is that while deep learning libraries are amazing, at the end of the day they are DSL and really pull towards one specific way of computing and parallelization. It turns out that way of parallelizing is good for deep learning, but not for all things you may want to accelerate. Sometimes (i.e. cases that aren't dominated by large linear algebra) building problem-specific kernels is a major win, and it's over-extrapolating to see ML frameworks do well with GPUs and think that's the only thing that's required. There are many ways to parallelize a code, ML libraries hardcode a very specific way, and it's good for what they are used for but not every problem that can arise.
We recently showed DiffEqGPU.jl generating customized ODE solver kernels for NVIDIA CUDA, AMD GPUs, Intel OneAPI, and Apple Metal, where for CUDA it matches state of the art (MPGOS) which is about 10x-100x something like Jax/PyTorch (where the performance difference comes from inefficiencies of using vmap vs actually writing and calling a kernel). It's all in https://arxiv.org/abs/2304.06835. So this stuff exists and people are using it. Of course the caveat here is this is the context of engineering applications so someone would need to do similar for LLMs to fully relate back to the article, but it shows the tools are ready to a large extent for someone to step up in the ML space.
More from the same topics.
Martin Kleppmann
The CAP Theorem is a frequently cited impossibility result in distributed systems, especially among NoSQL distributed databases.
Anders Sandberg, Eric Drexler, Toby Ord
The Fermi paradox is the conflict between an expectation of a high {\em ex ante} probability of intelligent life elsewhere in the universe and the apparently lifeless universe we in fact observe.
Chiyuan Zhang, Samy Bengio, Moritz Hardt +2
Despite their massive size, successful deep artificial neural networks can exhibit a remarkably small difference between training and test performance.
Scott Aaronson
In honor of Alan Turing's hundredth birthday, I unwisely set out some thoughts about one of Turing's obsessions throughout his life, the question of physics and free will.