mpi4torch is an automatic-differentiable wrapper of MPI functions for the pytorch tensor library.
MPI stands for Message Passing Interface and is the de facto standard communication interface on high-performance computing resources. To facilitate the usage of pytorch on these resources an MPI wrapper that is transparent to pytorch's automatic differentiation (AD) engine is much in need. This library tries to bridge this gap.
mpi4torch is also hosted on PyPI. However, due to the ABI-incompatibility of the different MPI implementations it is not provided as a binary wheel and needs to be built locally. Hence, you should have an appropriate C++ compiler installed, as well as the development files of your MPI library be present. The latter are usually provided through the module system of your local cluster, and you should consult the manuals of your cluster for this, or through the package manager of your Linux distribution.
Once the dependencies have been satisfied the installation can be triggered by the usual
pip install mpi4torch
It is highly advised to first read the basic usage chapter of the documentation before jumping into action, since there are some implications of the pytorch AD design on the usage of mpi4torch. In other words, there are some footguns lurking!
You have been warned, but if you insist on an easy usage example, consider the following code snippet, which is an excerpt from examples/simple_linear_regression.py
comm = mpi4torch.COMM_WORLD
def lossfunction(params):
# average initial params to bring all ranks on the same page
params = comm.Allreduce(params, mpi4torch.MPI_SUM) / comm.size
# compute local loss
localloss = torch.sum(torch.square(youtput - some_parametrized_function(xinput, params)))
# sum up the loss among all ranks
return comm.Allreduce(localloss, mpi4torch.MPI_SUM)Here we have parallelized a loss function simply by adding two calls to Allreduce. For a more thorough
discussion of the example see here.
Running tests is as easy as
mpirun -np 2 nose2
