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

31–40 of 49 posts

Re: Supercharged high-resolution ocean simulation with Jax

#31

Post by ChrisRackauckas »

moelf wrote: post #10

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

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.

Re: Supercharged high-resolution ocean simulation with Jax

#32

Post by ChrisRackauckas »

moelf wrote: post #10

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

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.

Re: Supercharged high-resolution ocean simulation with Jax

#33

Post by ChrisRackauckas »

adgjlsfhk1 wrote: post #27
dionhaefner wrote: post #24

Earlier quoted context omitted.

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.

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)

Re: Supercharged high-resolution ocean simulation with Jax

#34

Post by fault1 »

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

I'm also surprised that XLA.jl doesn't seem to have had continued development: https://github.com/FluxML/XLA.jl

When in doubt, piggybacking on (or at least interoperating with) what the large technology companies are investing in is probably savvy, sort of what the OP did.

Re: Supercharged high-resolution ocean simulation with Jax

#35

Post by dekhn »

adgjlsfhk1 wrote: post #29
dekhn wrote: post #28

Earlier quoted context omitted.

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

This post is about Big Simulations, not small systems. Like, hundreds to thousands of cores wiht parameters that don't fit in RAM on a single machine.

I am sure the benchmark produces the numbers the author says, but it's not measuring something useful to the posters of this simulation.

Re: Supercharged high-resolution ocean simulation with Jax

#36

Post by dionhaefner »

ChrisRackauckas wrote: post #31
moelf wrote: post #10

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

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.

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.

Re: Supercharged high-resolution ocean simulation with Jax

#37

Post by civilized »

adgjlsfhk1 wrote: post #20
fault1 wrote: post #19

Earlier quoted context omitted.

It may take a while, however. 15-20 years ago, you kind of had to use Python on the sly in the scientific setting vs the incumbents (MATLAB, C++, Fortran). Julia seems to be in a similar phase. That being said, Python does have some structural advantages since it positions itself as a universal glue. It's much easier to gain a critical mass in that regard vs a niche area like scientific or numerical computing. That b…

I think Julia has a much better path to wide adoption for numerical computing/HPC. It is a much better language for package developers (you pretty much never have to go to a lower level language and everything can compose together with much less work). If you look at Julia and Python packages with similar functionality, the Julia one will typically be much more general and 1/10th the lines of code. This is a pretty p…

I'm very excited about this, but my guess is it will take years for the packages to be so broadly and clearly superior that there's a mass migration to Julia. And even then people may prefer to just call Julia from a language they're more comfortable with. Still, it would be amazing for Julia to become the single go-to high-performance language of ML/DL/AI, advanced statistical modeling, HPC, etc.

Re: Supercharged high-resolution ocean simulation with Jax

#38

Post by adgjlsfhk1 »

dionhaefner wrote: post #36
ChrisRackauckas wrote: post #31

Earlier quoted context omitted.

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.

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.

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?

Re: Supercharged high-resolution ocean simulation with Jax

#39

Post by ChrisRackauckas »

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

XLA only optimizes quasi-static code, which does not include adaptive numerical solvers like those for ODEs. It's a generally good assumption for ML though, but there are ways to break it. I wrote a piece showcasing some ideas around that: https://www.stochasticlifestyle.com/useful-algorithms-that-a...

Re: Supercharged high-resolution ocean simulation with Jax

#40

Post by adgjlsfhk1 »

fault1 wrote: post #34
dekhn wrote: post #28

Earlier quoted context omitted.

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

I'm also surprised that XLA.jl doesn't seem to have had continued development: https://github.com/FluxML/XLA.jl When in doubt, piggybacking on (or at least interoperating with) what the large technology companies are investing in is probably savvy, sort of what the OP did.

XLA.jl was kind of a solution looking for a problem. If you want fast code in Julia, you can just write Julia.
Post reply on HN