Supercharged high-resolution ocean simulation with Jax • Hacker News Board

Re: Supercharged high-resolution ocean simulation with Jax

#41

Post by dekhn » Mon, Dec 06, 2021, 6:40 PM UTC

ChrisRackauckas wrote: post #39

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

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

IIUC people have already run MD (which is the field I used to work in) on XLA, https://twitter.com/sschoenholz/status/1334997741185814530 In these cases it's almost always better (unless you are a numerical genius) to port to the engine, than to try to make a better algorithm that runs on a smaller engine.

Re: Supercharged high-resolution ocean simulation with Jax

#42

Post by sundarurfriend » Mon, Dec 06, 2021, 6:49 PM UTC

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…

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

Re: Supercharged high-resolution ocean simulation with Jax

#43

Post by dionhaefner » Mon, Dec 06, 2021, 6:51 PM UTC

adgjlsfhk1 wrote: post #38

dionhaefner wrote: post #36

Earlier quoted context omitted.

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?

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

Re: Supercharged high-resolution ocean simulation with Jax

#44

Post by ChrisRackauckas » Mon, Dec 06, 2021, 8:52 PM UTC

dekhn wrote: post #41

ChrisRackauckas wrote: post #39

Earlier quoted context omitted.

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

IIUC people have already run MD (which is the field I used to work in) on XLA, https://twitter.com/sschoenholz/status/1334997741185814530 In these cases it's almost always better (unless you are a numerical genius) to port to the engine, than to try to make a better algorithm that runs on a smaller engine.

Yes, that has nothing to do with what I just said though. Of course MD is fine because symplectic ODE solvers cannot generally have adaptivity (without tricky and very expensive handling of `t` inside of the Hamiltonian which nobody does because it's still an active research topic how to make it computationally viable). So MD gets a quasi-static code which XLA is fine with optimizing. I was explicitly talking about the non-quasi-static cases.

Re: Supercharged high-resolution ocean simulation with Jax

#45

Post by dekhn » Mon, Dec 06, 2021, 8:57 PM UTC

ChrisRackauckas wrote: post #44

dekhn wrote: post #41

Earlier quoted context omitted.

IIUC people have already run MD (which is the field I used to work in) on XLA, https://twitter.com/sschoenholz/status/1334997741185814530 In these cases it's almost always better (unless you are a numerical genius) to port to the engine, than to try to make a better algorithm that runs on a smaller engine.

Yes, that has nothing to do with what I just said though. Of course MD is fine because symplectic ODE solvers cannot generally have adaptivity (without tricky and very expensive handling of `t` inside of the Hamiltonian which nobody does because it's still an active research topic how to make it computationally viable). So MD gets a quasi-static code which XLA is fine with optimizing. I was explicitly talking about t…

I've worked in ODEs for 20+ years and I don't think that non-quasi-static solvers have really ever come up. Are these commonly used? IE, how much CPU/GPU/TPU time is spent on them globally and how useful are they?

Re: Supercharged high-resolution ocean simulation with Jax

#46

Post by ChrisRackauckas » Mon, Dec 06, 2021, 8:59 PM UTC

dionhaefner wrote: post #43

adgjlsfhk1 wrote: post #38

Earlier quoted context omitted.

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?

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

That's precisely what composability solves. We're seeing in CLIMA that using more general highly optimized solvers can greatly decrease the `f` cost count moreso than focusing on really low level optimizations. Especially in things like the land model where you can have many stability issues (such as large complex eigenvalues which happen to work very poorly with multistep methods, even BDF), the ability to split the develop of the time stepping to a huge community of 100's of developers without losing performance gives something where more optimal methods for a domain arise and are found. Yes, the standard is to use something simpler. No, it's not even close to optimal and that is something that is being made very clear.

Re: Supercharged high-resolution ocean simulation with Jax

#47

Post by PhilipVinc » Mon, Dec 06, 2021, 9:35 PM UTC

adgjlsfhk1 wrote: post #40

fault1 wrote: post #34

Earlier quoted context omitted.

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.

That's incorrect. If you work with mid-sized neural networks and MCMC sampling, allocations start to play a significant role (And Flux.jl is bad at preallocation). Prealloc.jl does not work properly. Zygote.jl adds even more allocations to the mix...

Jax/XLA completely solves this problem. Yes, it's annoying that you have to work with a static graph but if your problem fits the description... it's great.

Re: Supercharged high-resolution ocean simulation with Jax

#48

Post by dklend122 » Mon, Dec 06, 2021, 10:12 PM UTC

PhilipVinc wrote: post #47

adgjlsfhk1 wrote: post #40

Earlier quoted context omitted.

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

That's incorrect. If you work with mid-sized neural networks and MCMC sampling, allocations start to play a significant role (And Flux.jl is bad at preallocation). Prealloc.jl does not work properly. Zygote.jl adds even more allocations to the mix... Jax/XLA completely solves this problem. Yes, it's annoying that you have to work with a static graph but if your problem fits the description... it's great.

There's work being done to solve this in Julia. See escapeanalysis.jl and the immutable array pr in base

Re: Supercharged high-resolution ocean simulation with Jax

#49

Post by ChrisRackauckas » Tue, Dec 07, 2021, 3:10 AM UTC

dekhn wrote: post #45

ChrisRackauckas wrote: post #44

Earlier quoted context omitted.

Yes, that has nothing to do with what I just said though. Of course MD is fine because symplectic ODE solvers cannot generally have adaptivity (without tricky and very expensive handling of `t` inside of the Hamiltonian which nobody does because it's still an active research topic how to make it computationally viable). So MD gets a quasi-static code which XLA is fine with optimizing. I was explicitly talking about t…

I've worked in ODEs for 20+ years and I don't think that non-quasi-static solvers have really ever come up. Are these commonly used? IE, how much CPU/GPU/TPU time is spent on them globally and how useful are they?

Have you used almost any ODE solver? Almost every single one uses embedded methods to adapt time steps. ode23, ode45, ode23t, ode23tb, ode15s, LSODE, LSODA, radau, rodas, VODE, CVODE, ... even for DAEs you have DASSL, IDA, ... I can keep going but it's just listing every ODE solver code out there. Once you do that then the computation is dependent on values and thus the full compute is not determined by the input sizes, which is something known to be blocking the full usage in Jax because of XLA limitations (for example the implementation of dense output).