Is it possible to distribute a LatitudeLongitudeGrid over multiple GPUs?

[mbareford]

Hello,

I've been experimenting with a particular Oceananigans example from the z-star-coordinate branch of the repo.

https://github.com/CliMA/Oceananigans.jl/blob/as/z-star-coordinate/double_gyre_5.jl

I can get this double gyre model to run on a single NVIDIA A100 GPU and was wondering how to run the model across multiple GPUs using one MPI rank per GPU.

1. Should I start with a ImplicitGlobalGrid then determine the Nx/Lx for the local mpi rank and create a corresponding LatitudeLongitudeGrid?

2. Should I instead use Oceananigans Distributed, assuming that that module is compatible with LatitudeLongitudeGrid.
   If not, would representing the problem using a MultiRegion.ConformalCubedSphereGrid work?

Best wishes,
Michael

[mbareford]

The multi GPU Oceananigans benchmark shows how to combine LatitudeLongitudeGrid with MultiRegionGrid.
Basically, start with the LatLon grid than distribute across GPU devices using MultiRegionGrid and XPartition.

https://github.com/CliMA/Oceananigans.jl/blob/main/benchmark/benchmark_multi_GPU.jl

I'm not sure if this code actually works as XPartition takes a parameter called regions which I assume should be a list of GPU IDs,
however, that info is conveyed via the dev parameter.

And could this code be made to work if running across GPUs on multiple nodes?

[glwagner]

Don't use MultiRegionGrid. Sorry for the misleading script! I am removing it in #4787

For multi-GPU simulations, use Distributed(GPU()) architecture (see the docstring for Distributed for more information -- you can also control how the simulation is distributed across the GPUs). Everything should work as-is by only changing the architecture.

There is also more information here:

https://clima.github.io/OceananigansDocumentation/stable/grids/#Distributed-grids

If your script does not work after changing the architecture, please do report it (either open an issue or discuss here).

[mbareford]

Thank you, I'll read up on Distributed Grids and get back to you.