Parallel simulations using multiple nodes/MPI

[fspereira1]

Dear All,

I am new to oceananigans, so my apologies for any basic questions. I want to run oceananigans on an HPC machine using several nodes (MPI). I wonder if anyone could help me with this. What is the best way of doing this?

One of your reference cases is quite close to my problem:

[https://github.com/CliMA/Oceananigans.jl/blob/main/examples/ocean_wind_mixing_and_convection.jl]

Best,
Filipe

[glwagner]

@fspereira1, do you mind if I convert this to a discussion? We like to use discussions for usage and syntax questions (in part to make it easier to find for future users, and in part because then we won't need to "close" the discussion).

[fspereira1]

Sure.

[fspereira1]

and I can provide the .jl file I am trying to run in parallel

[glwagner]

The short answer is that distributed support is experimental, which is why there is little documentation or information about it.

Here's a short summary:

• The HydrostaticFreeSurfaceModel can be run on multiple GPUs via MultiRegion, but only in a single process / on one node. With 4 V100s, for example, the maximum simulation size will be something like 1024^2 x 512.

• The NonhydrostaticModel support distributed CPU simulations, but not distributed GPU simulations. I believe there are some usability issues there (see Usability concerns and possible improvements for MultiArch #2349). Some prototypes can be found in https://github.com/CliMA/Oceananigans.jl/tree/main/validation/distributed_simulations.

• Note that on one V100 with 32 GB, you can run NonhydrostaticModel with one tracer up to 512 x 512 x 384 (with more tracers, the size will have to be reduced). I think @tomchor has done some exploration running the absolutely largest simulations that can be afforded on one GPU. Slightly larger models can be run if you have access to an A100 with 80 GB (which may also be ~3x faster).

• There is also a PR (Distributed mixed FFT / vertical tridiagonal solver #2538, getting a bit ancient now), which extends the distributed FFT algorithm to cover all cases for NonhydrostaticModel. Last I checked the code on this PR seemed to work, but the tests were not passing for some reason.

• You'll likely get best results on GPUs, since we haven't focused much on CPU performance. However, if you don't have access to GPUs, then you of course have to consider CPU solutions.

Hope that helps.

[fspereira1]

I am trying to run

pressure_solver = FourierTridiagonalPoissonSolver(grid, FFTW.ESTIMATE)
model = NonhydrostaticModel(; grid, pressure_solver, buoyancy,
advection = UpwindBiasedFifthOrder(),
timestepper = :RungeKutta3,
tracers = (:T, :S),
coriolis = FPlane(f=1e-4),
closure = AnisotropicMinimumDissipation(),
boundary_conditions = (u=u_bcs, T=T_bcs, S=S_bcs))

just on CPUs

[glwagner]

Check out this script for inspiration:

https://github.com/CliMA/Oceananigans.jl/blob/main/validation/distributed_simulations/mpi_nonhydrostatic_two_dimensional_turbulence.jl

I believe on main you can only distribute grids that are regular in all directions. I think you'll need #2538 to distribute a stretched vertical grid.

[tomchor]

• The NonhydrostaticModel support distributed CPU simulations, but not distributed GPU simulations. I believe there are some usability issues there (see Usability concerns and possible improvements for MultiArch #2349). Some prototypes can be found in https://github.com/CliMA/Oceananigans.jl/tree/main/validation/distributed_simulations.

It's a shame we haven't been able to merge that PR. IIRC there's only one test that hangs and never completes, forcing us to kill the test manually. I tried debugging it at some point, but since it doesn't throw any errors I couldn't quite figure it out. Maybe we should try to revive that PR since it's such a nice addition.

[glwagner]

Contributions very welcome on #2538 !

[tomchor]

@fspereira1 if you're trying to run large models on a GPU let me know and I can give you some tips. There's some things you can do when defining diagnostics that can save quite a bit of space in my experience (i.e. reusing scratch space).

[iuryt]

Hi all!
I know that it is being a while since this discussion started and maybe things have changed.
I noticed that in the example Nranks is set to the y direction,

Oceananigans.jl/validation/distributed_simulations/distributed_nonhydrostatic_turbulence.jl

Line 29 in 2946027

arch = Distributed(CPU(); topology, ranks=(1, Nranks, 1), communicator=comm)

but each rank is running a chunk in the x direction

Oceananigans.jl/validation/distributed_simulations/distributed_nonhydrostatic_turbulence.jl

Line 30 in 2946027

grid = RectilinearGrid(arch; topology, size=(Nx ÷ Nranks, Ny), halo=(3, 3), x=(0, 2π), y=(0, 2π))

Is this correct?

[iuryt]

hm.. I will check that, but I was having the following error while trying to run
Attempting to use an MPI routine before initializing MPICH

[iuryt]

Okay, so, MPI needs to be initialized before accessing its communication variables.

Thus, despite calling Distributed initializes MPI, we need to access Nranks = MPI.Comm_size(comm) before to define Partition.

Oceananigans.jl/validation/distributed_simulations/distributed_nonhydrostatic_turbulence.jl

Lines 20 to 22 in 2946027

	comm = MPI.COMM_WORLD
	rank = MPI.Comm_rank(comm)
	Nranks = MPI.Comm_size(comm)

Should we initialize MPI when importing DistributedComputations?

[iuryt]

I will transfer this discussion to a PR.

[iuryt]

We can continue the discussion on #3429

[glwagner]

Taking a step back, we should not have to write using MPI to do distributed simulations.

If we do require it for some specialized purpose (or to demonstrate something, as in a validation example), then perhaps we can right a comment about this in the script.