Supercharged high-resolution ocean simulation with Jax

https://dionhaefner.github.io/2021/12/supercharged-high-resolution-ocean-simulation-with-jax/

by dionhaefner • 5 years ago

90 48 5 years ago

dionhaefner.github.io

civilized4• 5 years 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.

mjn2• 5 years 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-geophysic...

adgjlsfhk18• 5 years 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/62a063f23cccf3a55a4a... shows a pretty simple case where DifferentialEquations.JL is 6x faster at gradient calculations than Jax.

dionhaefner6• 5 years ago

A bit late to the party, but here are some reasons:

- When we started Veros (~4 years ago) Julia was very new on our radar and we didn't know whether it would stick. And to be frank, I'm still not convinced whether it will stick. Yes it seems like a fantastic language, but we all know how long it took Python to gain traction.

- Climate scientists and students already do their post-processing in Python. Having the whole stack in the same language makes things a lot easier for domain experts whose first priority is physics, not coding.

- Python skills translate better to other jobs, which I think is important for young academics.

- The Python library ecosystem is so good. Need to use PETSc? `import petsc4py`. Simplify postprocessing? Export your model state as `xarray` dataset. Julia is great for bleeding edge autodiff through everything stuff, but the bread and butter libraries are just so polished and battle tested in Python.

- I don't know Julia :)

civilized4• 5 years ago

Those are very good reasons!

bertr4nd2• 5 years 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:68640106214579...

dionhaefner6• 5 years ago

True, but unfortunately Pytorch is not quite there yet when it comes to more complex benchmarks:

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

bertr4nd2• 5 years 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!

geoalchimista• 5 years 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?

krasin2• 5 years 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.

moelf4• 5 years 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).

6gvONxR4sf7o• 5 years 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.

oakridge• 5 years 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.

adgjlsfhk18• 5 years ago

I would recommend checking out https://www.youtube.com/watch?v=GiSsoA1udUk. It shows that you can can do climate models with 16 bit numbers.

dionhaefner6• 5 years 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.

adgjlsfhk18• 5 years 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.

moelf4• 5 years ago

like some other commenters here, 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!

ChrisRackauckas8• 5 years 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.

dionhaefner6• 5 years 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.

ChrisRackauckas8• 5 years ago

Oh another thing, in the chaotic regime which is of interest, standard automatic differentiation schemes don't even apply as you require shadow adjoints given the shadow trajectory leads to inaccurate calculations for the gradient. Julia's system is the only one that I know of that has shadow adjoints for differentiation of ergodic properties.

https://frankschae.github.io/post/shadowing/

So unless the purpose is to only differentiate the simulator for short time periods or in the absence of chaos, I cannot see differentiation as a good justification because AD will not give a stable algorithm on that type of problem.

sundarurfriend• 5 years ago

https://github.com/CliMA/Oceananigans.jl

Off-topic, but I have to say that is one of my favourite package names in Julia. (A more recent one is [MATDaemon](https://github.com/jondeuce/MATDaemon.jl))

Supercharged high-resolution ocean simulation with Jax

Related Stories