Hacker News Board Hacker News, browsed like a classic message board
Live data from Hacker News

Supercharged high-resolution ocean simulation with Jax

dionhaefner.github.io

Post reply on HN
Print view View original

21–30 of 49 posts

Re: Supercharged high-resolution ocean simulation with Jax

#21

Post by dionhaefner »

civilized wrote: post #2

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.

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

Re: Supercharged high-resolution ocean simulation with Jax

#22

Post by dionhaefner »

adgjlsfhk1 wrote: post #9
mjn wrote: post #3

Earlier quoted context omitted.

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

"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 b…

I was mostly referring to the millions (billions?) of dollars getting poured into Python library development by tech companies. With the effect that Python stays relevant and has a thriving library ecosystem. Maybe I'm wrong and Julia is just that good that it doesn't matter - I guess time will tell.

Re: Supercharged high-resolution ocean simulation with Jax

#23

Post by dionhaefner »

bertr4nd wrote: post #4

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/upda…

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

Re: Supercharged high-resolution ocean simulation with Jax

#24

Post by dionhaefner »

oakridge wrote: post #6

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.

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.

Re: Supercharged high-resolution ocean simulation with Jax

#25

Post by bertr4nd »

dionhaefner wrote: post #23
bertr4nd wrote: post #4

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/upda…

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

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!

Re: Supercharged high-resolution ocean simulation with Jax

#26

Post by civilized »

dionhaefner wrote: post #21
civilized wrote: post #2

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.

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 who…

Those are very good reasons!

Re: Supercharged high-resolution ocean simulation with Jax

#27

Post by adgjlsfhk1 »

dionhaefner wrote: post #24
oakridge wrote: post #6

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.

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),…

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.

Re: Supercharged high-resolution ocean simulation with Jax

#28

Post by dekhn »

adgjlsfhk1 wrote: post #9
mjn wrote: post #3

Earlier quoted context omitted.

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

"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 b…

Jax is just a tool to generate XLA, which produces extremely high performance computational graphs which can map to arbitrarily fast hardware, so I'm very skeptical of the utility of the conclusions of thelink you provided (which seems to be comparing single process CPU linear algebra?)

Re: Supercharged high-resolution ocean simulation with Jax

#29

Post by adgjlsfhk1 »

dekhn wrote: post #28
adgjlsfhk1 wrote: post #9

Earlier quoted context omitted.

"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 b…

Jax is just a tool to generate XLA, which produces extremely high performance computational graphs which can map to arbitrarily fast hardware, so I'm very skeptical of the utility of the conclusions of thelink you provided (which seems to be comparing single process CPU linear algebra?)

Single thread CPU Linear algebra is the bottleneck of most small systems, so if you can't do that right, you are going to have problems. If you don't believe the benchmark, feel free to run them yourself.

That said, Jax also has bigger issues in it's handling of higher derivatives. Currently, it only supports a few types of jacobians, and the ones it is missing include all the sparse methods that can make your code orders of magnitude faster. https://jax.readthedocs.io/en/latest/notebooks/autodiff_cook.... DifferentialEquations, on the other hand can do automatic sparsity detection https://diffeq.sciml.ai/stable/tutorials/advanced_ode_exampl....

Re: Supercharged high-resolution ocean simulation with Jax

#30

Post by ChrisRackauckas »

adgjlsfhk1 wrote: post #9
mjn wrote: post #3

Earlier quoted context omitted.

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

"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 b…

That's an old example. It will now default to Enzyme and should do quite a bit faster. I should update that.
Post reply on HN