Guide to integrating FluxMPI into your code

There are essentially 6 main steps to remember:

1. Initialize FluxMPI FluxMPI.Init().
2. Sync Model Parameters and States FluxMPI.synchronize!(ps; root_rank). (Remember to use FluxMPIFluxModel for Flux models.)
3. Use DistributedDataContainer to distribute your data evenly across the processes. (Of course an alternative is to just manually partition your data.)
4. Wrap the optimizer in DistributedOptimizer or call allreduce_gradients(gs::NamedTuple) before eveery Optimisers.update.
5. Sync the optimizer state across the processes FluxMPI.synchronize!(opt_state; root_rank).
6. Change logging code to check for local_rank == 0.

Finally, start the code using mpiexecjl -n <np> julia --project=. <filename>.jl