Pleasingly, this was actually very easy to do! Basically:
There's a new class of things called 'communicators' which let you get rank, number of ranks etc and perform basic communications - I've made an mpi4py implementation for now as that's what my old code used
CompiledNetwork does some basic stuff if a communicator is provided:
Only building on the 1st rank
Waiting on a barrier before loading
Doing the NCCL initialisation
Compiler subdivides batches across ranks if a communicator is provided and turns on the magic NCCL flag so GeNN generates the additional bits of code (https://github.com/genn-team/genn/pull/449)
Metrics like SparseCategoricalAccuracy get passed the communicator and use it to
Other than that, it's all just passing the communicator around and a few places where the 'full' batch size is used rather than the scaled down one e.g. in the EventProp compiler to scale stuff. I've also added a couple of additional examples (at some point I need to tidy the examples up a bit) which demonstrate how you need to change your code to run across multiple GPUs - mostly just splitting the dataset and turning off progress bars etc apart from on the first rank.
Pleasingly, this was actually very easy to do! Basically:
mpi4py
implementation for now as that's what my old code usedCompiledNetwork
does some basic stuff if a communicator is provided:Compiler
subdivides batches across ranks if a communicator is provided and turns on the magic NCCL flag so GeNN generates the additional bits of code (https://github.com/genn-team/genn/pull/449)SparseCategoricalAccuracy
get passed the communicator and use it toOther than that, it's all just passing the communicator around and a few places where the 'full' batch size is used rather than the scaled down one e.g. in the EventProp compiler to scale stuff. I've also added a couple of additional examples (at some point I need to tidy the examples up a bit) which demonstrate how you need to change your code to run across multiple GPUs - mostly just splitting the dataset and turning off progress bars etc apart from on the first rank.