Multi-GPU and distributed demonstrations of PINNs

[ChrisRackauckas]

Multi-GPU should "just work" through distributed, so we should make sure to have a demo of this.

[ChrisRackauckas]

@leios would it be a good idea to use Clima tools here, or build something off of distributed arrays of CuArrays?

[killah-t-cell]

What is the status of multi-GPU support for PINNs?

[ChrisRackauckas]

It should just work. Try it.

[udemirezen]

@ChrisRackauckas Do you have a specific example? I am sorry I can not find any tutorial or example code in NeuralPDE.jl repository. Have I been missing something?

[ChrisRackauckas]

It's not a NeuralPDE thing, it's just a Flux thing. If the neural network is multi-GPU then the PINN will be multi-GPU. We can make a tutorial here to help people out, but you won't need a tutorial to do it.

CUDA.jl describes how to setup multiple GPUs here: https://fluxml.ai/Flux.jl/stable/gpu/ Multiple GPU nets are defined in: https://github.com/DhairyaLGandhi/DaggerFlux.jl

Stick those into NeuralPDE and you got gold.

[udemirezen]

Thank you very much for the information

[jpsamaroo]

Note that performance on multiple GPUs and/or multiple processes will be poor until https://github.com/JuliaGPU/DaggerGPU.jl/pull/17 is merged (planning that for the end of this week), but things should still function correctly without it.