Starting point for multi-GPU simulations?

[BrodiePearson]

Are there any ClimaOcean examples/branches that would provide a template for multi-GPU (near-)global simulations on several 2-4 GPU systems? [A100s (80GB), H100s (80GB) and GH200s (95GB or 480GB if we can figure out how to access the extra memory...)]

It looks like examples/distributed_perfect_one_degree_model_calibration.jl has some of this capability, but I think the multi-GPU part is doing the inverse problem only.

Our group has been benchmarking speed/efficiency of Oceananigans LES on these single GPUs (stop by Alana Kihn's poster if you are at Ocean Sciences!), but we'd love to extend this to global and GPU-distributed simulations if its possible.

@glwagner @simone-silvestri @navidcy

[navidcy]

It's definitely on the list to put together such examples!

[simone-silvestri]

Should be quite simple, if you have a setup on single GPU, the only oceananigans-related change you need is to modify the architecture:

arch = Distributed(GPU(); partition = Partition(2))

runs on 2 gpus.
The only problem for the moment is the distributed output.
For benchmarking on clusters, you need to make sure all other mpi-related settings are enabled:

• loading the correct cuda-aware mpi modules
• use MPI local binaries through MPIPreferences (https://juliaparallel.org/MPI.jl/stable/configuration/).

[glwagner]

Nice. Output should be pretty easy to figure out too, I think. A good exercise might be to go through the examples and make sure they all run merely by changing the architecture. We went through that exercise way back when to help us line up the interface for CPU->GPU transition.

[iuryt]

Hi all,
Nothing to add here, but just want to explain that despite I have paused the implementation of the validation code because of OSM24, I am definitely coming back to this asap.

[BrodiePearson]

Thanks for all your responses! I think the architecture modification was the key piece I was missing to get the ball rolling, although I still have some MPI issues to smooth out on our system...

[iuryt]

Let us know how it works out and create an issue if you encounter any problem.