7 ms·

Supercharged high-resolution ocean simulation with Jax

90 points by dionhaefner 5y ago

geoalchimista 5y ago
What made you choose JAX over Julia? I'm interested in this question, because I have been thinking about transitioning to Julia but have always hesitated to make the move, since overall the Python ecosystem still seems way ahead in terms of visualization and toolchain. Also, would you expect JAX acceleration to work well with other types of discretization, such as spectral methods?
- krasin 5y ago
  From the post: > JAX on GPU outperforms everything I've only skimmed through the blog post, but it feels that the GPU acceleration without a need to write any custom code, was the primary reason to choose JAX.
  - moelf 5y ago
    >without a need to write any custom code, was the primary reason to choose JAX. this is even more "free" in Julia, JAX at least need to worry when foreign call happens (library not derived from Numpy/JAX ecosystem, or outright C/C++ binding without JAX rules).
    - krasin 5y ago
      For someone like me who is familiar with JAX, but only recently starting to consider to pay attention to Julia (it does have a momentum), would it be possible to provide any good examples of using GPU / multiple GPUs from Julia? I've tried to search on my own, but only a way to write CUDA-dependent code: https://juliagpu.gitlab.io/CUDA.jl/usage/multigpu/ https://juliagpu.gitlab.io/CUDA.jl/usage/multigpu/
      - moelf 5y ago
        I'm not familiar with multi-GPU setup in general. GPU programming in Julia has the advantage that naive operation doesn't even need to be GPU-aware (for writers), since GPU arrays (of any vendor backend) conforms the AbstractArray interface. If you're advanced library writer, you can leverage: https://juliagpu.github.io/KernelAbstractions.jl/stable/#Writing-your-first-kernel-1 https://juliagpu.github.io/KernelAbstractions.jl/stable/#Wri... which allows you to write kernel, in Julia, that compiles efficiently with rest of native Julia code, and that works cross-vendor! Back to multi GPU, it seems there's: https://clima.github.io/OceananigansDocumentation/stable/appendix/benchmarks/#Multi-GPU-Shallow-Water-Model https://clima.github.io/OceananigansDocumentation/stable/app... which is MPI based?
      - adgjlsfhk1 5y ago
        https://www.juliapackages.com/p/pencilarrays https://www.juliapackages.com/p/pencilarrays is a really good tool to do this type of stuff automatically for some applications.
- 6gvONxR4sf7o 5y ago
  If numpy is a good fit for it, JAX is a good fit for accelerating it, basically. I think of it as numpy plus program transformations, such as differentiation, JIT, parallelization, compiling to XLA, compiling to TPU/GPU, etc. The magic of JAX is that it keeps all that stuff about as simple as writing numpy code.
civilized 5y ago
Awesome! One question immediately comes to mind. Any interest in doing this stuff with Julia? You're basically the epitome of their target audience: a scientific computing expert who does HPC with differentiable programs.
- mjn 5y ago
  There's an earlier blog post by the same author where they discuss three possible ways of moving away from the Fortran/C status quo towards higher-level models. They mention Julia as one of the routes, but not the one they decided to take: https://dionhaefner.github.io/2021/04/higher-level-geophysical-modelling/#the-taxonomy-of-high-level-modelling https://dionhaefner.github.io/2021/04/higher-level-geophysic...
  - adgjlsfhk1 5y ago
    "On the other hand, Julia’s focus on scientific applications is both blessing and curse. In this day and age, a lot of the progress in computing is driven by applications outside academia (mostly through machine learning)." This seems like a crazy mis-read to me. Julia is probably the language that has the best integration of differential equations and machine learning. Jax closes the gap a little, but is still way behind. For example https://gist.github.com/ChrisRackauckas/62a063f23cccf3a55a4ac9f6e497739a https://gist.github.com/ChrisRackauckas/62a063f23cccf3a55a4a... shows a pretty simple case where DifferentialEquations.JL is 6x faster at gradient calculations than Jax.
    - ChrisRackauckas 5y ago
      That's an old example. It will now default to Enzyme and should do quite a bit faster. I should update that.
    - mjn 5y ago
      I read that as being about what language industry uses to write ML applications, not about technical feasibility of integrating machine learning methods into a codebase. Put differently: industry most often uses Python (especially in ML), therefore the author wants to target Python in order to maximize uptake outside of academia. They even admit that doing it in Python is technically harder than doing it in Julia ("Unfortunately, this type [Type III] is also the hardest to get right"), but consider it worth the trouble for the broader accessibility. (That's more or less the direction I've been going with research code lately too, so I can sympathize, although I'm not entirely happy with the situation and definitely also sympathize with the Julia folks being unhappy about it.)
bertr4nd 5y ago
In a shameless plug, I want to note that running these sorts of workloads on CPU using Pytorch got much faster (some results on a benchmark from this post’s author’s suite in [0]) in the most recent torch release thanks to the addition of a JIT compiler. Obviously there’s much to recommend Jax (the XLA compiler is quite excellent), but it’s nice to have some choice in the space. [0] https://www.linkedin.com/feed/update/activity:6864010621457977344 https://www.linkedin.com/feed/update/activity:68640106214579...
- dionhaefner 5y ago
  True, but unfortunately Pytorch is not quite there yet when it comes to more complex benchmarks: https://github.com/dionhaefner/pyhpc-benchmarks#example-results https://github.com/dionhaefner/pyhpc-benchmarks#example-resu... JAX really is the only library that comes close to low-level code on CPU, almost always (that I've tried).
  - bertr4nd 5y ago
    Interesting, I thought pytorch was a bit more competitive on those other benchmarks (but admittedly it’s been a while since I looked). Slicing shouldn’t be a fundamental problem, but perhaps there are some important details that have been overlooked. Thanks for pointing it out!
oakridge 5y ago
My question is how much of the operations in JAX here can be done with reduced precision and can utilize training accelerators i.e. TPUs. I've noticed a lot of research coming out in physics, where everything is simulated in at least double float, being augmented with ML approaches where precision is traded for dynamic range.
- adgjlsfhk1 5y ago
  I would recommend checking out https://www.youtube.com/watch?v=GiSsoA1udUk https://www.youtube.com/watch?v=GiSsoA1udUk. It shows that you can can do climate models with 16 bit numbers.
- dionhaefner 5y ago
  The thing with reduced precision is that things may look fine at first, but then you eventually notice unphysical features in your solution (like additional wave modes after very long simulation times, or energy conservation issues). So we really don't know as a community yet how far we can venture from float64, but it looks like float32 may be viable. Veros works OK on TPUs (about the same speed as a high-end GPU), but since you can't buy TPUs that's an immediate no for most academic users of climate models. Renting hardware doesn't really make sense when you keep it busy for months at a time and the HPC infrastructure is already in place.
  - adgjlsfhk1 5y ago
    can't you fix a lot of the nonphysical issues by using better integration schemes? that might be hard in Jax though. From what I know, it's options for better numerical stability are pretty limited.
    - ChrisRackauckas 5y ago
      No, in fact, you want to go lower order with lower precision. The real answer is that if the solution is in the chaotic regime then maybe Float16 is fine because you'll be dominated by other numerical errors anyways (if you're also making sure you have adequate conservation so the solution doesn't explode in some way), but if you're not in the chaotic regime then even Float32 is pushing it in many cases (i.e. it better be non-stiff as stiffness pretty much guernetees operations which span beyond Float32 relative epsilon). So it's a case-dependent topic and not something that has an easy answer, though the case for Float16 is rather small. (We had some small tests generating TPU ODE solver code from Julia and showcased some rather bizarre stuff back when Keno was working on it, but never wrote a post summarizing all of it)
moelf 5y ago
like some other commenters here, https://github.com/CliMA/Oceananigans.jl https://github.com/CliMA/Oceananigans.jl immediately comes to mind, maybe it would be fun to compare projects on this scale between JAX/Julia. > JAX offers more than just a JIT compiler: JAX functions are also differentiable if the downstream library is completely implemented in JAX (numba) ecosystem. Similar for Julia, except implementing fast code in Julia is natural, doesn't involve debugging 3 compilers (Cpython, Numba, Jax). Many python library is only differentiable because the 100x more effort were put in writing C/C++ backend, binding to python, and writing chain rules for foreign functions. I would imagine Julia to be a good fit for this direction in the future!
- ChrisRackauckas 5y ago
  The real problem with the Jax code is that the non-composable programming language setup put it into a corner where it's using an extremely inefficient time stepping method that it has "optimized", but how is it optimized if you're doing 100 times more function calls than you have to? Algorithms matter, and "optimizing Adams-Bashforth 2" is a pretty silly idea.
  - dionhaefner 5y ago
    I agree with your point regarding non-composability and ”algorithm lock in” (which may or may not be solvable woth better abstractions), but explicit time stepping schemes are still the main workhorse of global ocean modelling, so I’m not sure whether ”silly” is the right label here.
    - adgjlsfhk1 5y ago
      Why are explicit time stepping schemes the main tool used? Is it because the languages that these models are written in aren't flexible enough, or is there a math reason why dynamic time-stepping isn't better?
      - dionhaefner 5y ago
        Climate models are vastly complex, and you need to bring together many experts from many disciplines to write and maintain one, and analyze the output. This seems to lead to the simplest methods coming out on top. Perhaps it could be solved with better abstractions (a lot of very smart people are trying).

90 points by dionhaefner 5y ago

What made you choose JAX over Julia? I'm interested in this question, because I have been thinking about transitioning to Julia but have always hesitated to make the move, since overall the Python ecosystem still seems way ahead in terms of visualization and toolchain. Also, would you expect JAX acceleration to work well with other types of discretization, such as spectral methods?

krasin 5y ago
From the post: > JAX on GPU outperforms everything I've only skimmed through the blog post, but it feels that the GPU acceleration without a need to write any custom code, was the primary reason to choose JAX.
- moelf 5y ago
  >without a need to write any custom code, was the primary reason to choose JAX. this is even more "free" in Julia, JAX at least need to worry when foreign call happens (library not derived from Numpy/JAX ecosystem, or outright C/C++ binding without JAX rules).
  - krasin 5y ago
    For someone like me who is familiar with JAX, but only recently starting to consider to pay attention to Julia (it does have a momentum), would it be possible to provide any good examples of using GPU / multiple GPUs from Julia? I've tried to search on my own, but only a way to write CUDA-dependent code: https://juliagpu.gitlab.io/CUDA.jl/usage/multigpu/ https://juliagpu.gitlab.io/CUDA.jl/usage/multigpu/
    - moelf 5y ago
      I'm not familiar with multi-GPU setup in general. GPU programming in Julia has the advantage that naive operation doesn't even need to be GPU-aware (for writers), since GPU arrays (of any vendor backend) conforms the AbstractArray interface. If you're advanced library writer, you can leverage: https://juliagpu.github.io/KernelAbstractions.jl/stable/#Writing-your-first-kernel-1 https://juliagpu.github.io/KernelAbstractions.jl/stable/#Wri... which allows you to write kernel, in Julia, that compiles efficiently with rest of native Julia code, and that works cross-vendor! Back to multi GPU, it seems there's: https://clima.github.io/OceananigansDocumentation/stable/appendix/benchmarks/#Multi-GPU-Shallow-Water-Model https://clima.github.io/OceananigansDocumentation/stable/app... which is MPI based?
    - adgjlsfhk1 5y ago
      https://www.juliapackages.com/p/pencilarrays https://www.juliapackages.com/p/pencilarrays is a really good tool to do this type of stuff automatically for some applications.
6gvONxR4sf7o 5y ago
If numpy is a good fit for it, JAX is a good fit for accelerating it, basically. I think of it as numpy plus program transformations, such as differentiation, JIT, parallelization, compiling to XLA, compiling to TPU/GPU, etc. The magic of JAX is that it keeps all that stuff about as simple as writing numpy code.

civilized 5y ago

Awesome! One question immediately comes to mind. Any interest in doing this stuff with Julia? You're basically the epitome of their target audience: a scientific computing expert who does HPC with differentiable programs.

mjn 5y ago
There's an earlier blog post by the same author where they discuss three possible ways of moving away from the Fortran/C status quo towards higher-level models. They mention Julia as one of the routes, but not the one they decided to take: https://dionhaefner.github.io/2021/04/higher-level-geophysical-modelling/#the-taxonomy-of-high-level-modelling https://dionhaefner.github.io/2021/04/higher-level-geophysic...
- adgjlsfhk1 5y ago
  "On the other hand, Julia’s focus on scientific applications is both blessing and curse. In this day and age, a lot of the progress in computing is driven by applications outside academia (mostly through machine learning)." This seems like a crazy mis-read to me. Julia is probably the language that has the best integration of differential equations and machine learning. Jax closes the gap a little, but is still way behind. For example https://gist.github.com/ChrisRackauckas/62a063f23cccf3a55a4ac9f6e497739a https://gist.github.com/ChrisRackauckas/62a063f23cccf3a55a4a... shows a pretty simple case where DifferentialEquations.JL is 6x faster at gradient calculations than Jax.
  - ChrisRackauckas 5y ago
    That's an old example. It will now default to Enzyme and should do quite a bit faster. I should update that.
  - mjn 5y ago
    I read that as being about what language industry uses to write ML applications, not about technical feasibility of integrating machine learning methods into a codebase. Put differently: industry most often uses Python (especially in ML), therefore the author wants to target Python in order to maximize uptake outside of academia. They even admit that doing it in Python is technically harder than doing it in Julia ("Unfortunately, this type [Type III] is also the hardest to get right"), but consider it worth the trouble for the broader accessibility. (That's more or less the direction I've been going with research code lately too, so I can sympathize, although I'm not entirely happy with the situation and definitely also sympathize with the Julia folks being unhappy about it.)

bertr4nd 5y ago

In a shameless plug, I want to note that running these sorts of workloads on CPU using Pytorch got much faster (some results on a benchmark from this post’s author’s suite in [0]) in the most recent torch release thanks to the addition of a JIT compiler. Obviously there’s much to recommend Jax (the XLA compiler is quite excellent), but it’s nice to have some choice in the space. [0] https://www.linkedin.com/feed/update/activity:6864010621457977344 https://www.linkedin.com/feed/update/activity:68640106214579...

dionhaefner 5y ago
True, but unfortunately Pytorch is not quite there yet when it comes to more complex benchmarks: https://github.com/dionhaefner/pyhpc-benchmarks#example-results https://github.com/dionhaefner/pyhpc-benchmarks#example-resu... JAX really is the only library that comes close to low-level code on CPU, almost always (that I've tried).
- bertr4nd 5y ago
  Interesting, I thought pytorch was a bit more competitive on those other benchmarks (but admittedly it’s been a while since I looked). Slicing shouldn’t be a fundamental problem, but perhaps there are some important details that have been overlooked. Thanks for pointing it out!

oakridge 5y ago

My question is how much of the operations in JAX here can be done with reduced precision and can utilize training accelerators i.e. TPUs. I've noticed a lot of research coming out in physics, where everything is simulated in at least double float, being augmented with ML approaches where precision is traded for dynamic range.

adgjlsfhk1 5y ago
I would recommend checking out https://www.youtube.com/watch?v=GiSsoA1udUk https://www.youtube.com/watch?v=GiSsoA1udUk. It shows that you can can do climate models with 16 bit numbers.
dionhaefner 5y ago
The thing with reduced precision is that things may look fine at first, but then you eventually notice unphysical features in your solution (like additional wave modes after very long simulation times, or energy conservation issues). So we really don't know as a community yet how far we can venture from float64, but it looks like float32 may be viable. Veros works OK on TPUs (about the same speed as a high-end GPU), but since you can't buy TPUs that's an immediate no for most academic users of climate models. Renting hardware doesn't really make sense when you keep it busy for months at a time and the HPC infrastructure is already in place.
- adgjlsfhk1 5y ago
  can't you fix a lot of the nonphysical issues by using better integration schemes? that might be hard in Jax though. From what I know, it's options for better numerical stability are pretty limited.
  - ChrisRackauckas 5y ago
    No, in fact, you want to go lower order with lower precision. The real answer is that if the solution is in the chaotic regime then maybe Float16 is fine because you'll be dominated by other numerical errors anyways (if you're also making sure you have adequate conservation so the solution doesn't explode in some way), but if you're not in the chaotic regime then even Float32 is pushing it in many cases (i.e. it better be non-stiff as stiffness pretty much guernetees operations which span beyond Float32 relative epsilon). So it's a case-dependent topic and not something that has an easy answer, though the case for Float16 is rather small. (We had some small tests generating TPU ODE solver code from Julia and showcased some rather bizarre stuff back when Keno was working on it, but never wrote a post summarizing all of it)

moelf 5y ago

like some other commenters here, https://github.com/CliMA/Oceananigans.jl https://github.com/CliMA/Oceananigans.jl immediately comes to mind, maybe it would be fun to compare projects on this scale between JAX/Julia. > JAX offers more than just a JIT compiler: JAX functions are also differentiable if the downstream library is completely implemented in JAX (numba) ecosystem. Similar for Julia, except implementing fast code in Julia is natural, doesn't involve debugging 3 compilers (Cpython, Numba, Jax). Many python library is only differentiable because the 100x more effort were put in writing C/C++ backend, binding to python, and writing chain rules for foreign functions. I would imagine Julia to be a good fit for this direction in the future!

ChrisRackauckas 5y ago
The real problem with the Jax code is that the non-composable programming language setup put it into a corner where it's using an extremely inefficient time stepping method that it has "optimized", but how is it optimized if you're doing 100 times more function calls than you have to? Algorithms matter, and "optimizing Adams-Bashforth 2" is a pretty silly idea.
- dionhaefner 5y ago
  I agree with your point regarding non-composability and ”algorithm lock in” (which may or may not be solvable woth better abstractions), but explicit time stepping schemes are still the main workhorse of global ocean modelling, so I’m not sure whether ”silly” is the right label here.
  - adgjlsfhk1 5y ago
    Why are explicit time stepping schemes the main tool used? Is it because the languages that these models are written in aren't flexible enough, or is there a math reason why dynamic time-stepping isn't better?
    - dionhaefner 5y ago
      Climate models are vastly complex, and you need to bring together many experts from many disciplines to write and maintain one, and analyze the output. This seems to lead to the simplest methods coming out on top. Perhaps it could be solved with better abstractions (a lot of very smart people are trying).