MPI Tools

Core MPI Utilities

spinup.utils.mpi_tools.mpi_avg(x)[source]

Average a scalar or vector over MPI processes.

spinup.utils.mpi_tools.mpi_fork(n, bind_to_core=False)[source]

Re-launches the current script with workers linked by MPI.

Also, terminates the original process that launched it.

Taken almost without modification from the Baselines function of the same name.

Parameters:
  • n (int) – Number of process to split into.
  • bind_to_core (bool) – Bind each MPI process to a core.
spinup.utils.mpi_tools.mpi_statistics_scalar(x, with_min_and_max=False)[source]

Get mean/std and optional min/max of scalar x across MPI processes.

Parameters:
  • x – An array containing samples of the scalar to produce statistics for.
  • with_min_and_max (bool) – If true, return min and max of x in addition to mean and std.
spinup.utils.mpi_tools.num_procs()[source]

Count active MPI processes.

spinup.utils.mpi_tools.proc_id()[source]

Get rank of calling process.

MPI + PyTorch Utilities

spinup.utils.mpi_pytorch contains a few tools to make it easy to do data-parallel PyTorch optimization across MPI processes. The two main ingredients are syncing parameters and averaging gradients before they are used by the adaptive optimizer. Also there’s a hacky fix for a problem where the PyTorch instance in each separate process tries to get too many threads, and they start to clobber each other.

The pattern for using these tools looks something like this:

  1. At the beginning of the training script, call setup_pytorch_for_mpi(). (Avoids clobbering problem.)
  2. After you’ve constructed a PyTorch module, call sync_params(module).
  3. Then, during gradient descent, call mpi_avg_grads after the backward pass, like so:
optimizer.zero_grad()
loss = compute_loss(module)
loss.backward()
mpi_avg_grads(module)   # averages gradient buffers across MPI processes!
optimizer.step()
spinup.utils.mpi_pytorch.mpi_avg_grads(module)[source]

Average contents of gradient buffers across MPI processes.

spinup.utils.mpi_pytorch.setup_pytorch_for_mpi()[source]

Avoid slowdowns caused by each separate process’s PyTorch using more than its fair share of CPU resources.

spinup.utils.mpi_pytorch.sync_params(module)[source]

Sync all parameters of module across all MPI processes.

MPI + Tensorflow Utilities

The spinup.utils.mpi_tf contains a a few tools to make it easy to use the AdamOptimizer across many MPI processes. This is a bit hacky—if you’re looking for something more sophisticated and general-purpose, consider horovod.

class spinup.utils.mpi_tf.MpiAdamOptimizer(**kwargs)[source]

Adam optimizer that averages gradients across MPI processes.

The compute_gradients method is taken from Baselines MpiAdamOptimizer. For documentation on method arguments, see the Tensorflow docs page for the base AdamOptimizer.

apply_gradients(grads_and_vars, global_step=None, name=None)[source]

Same as normal apply_gradients, except sync params after update.

compute_gradients(loss, var_list, **kwargs)[source]

Same as normal compute_gradients, except average grads over processes.

spinup.utils.mpi_tf.sync_all_params()[source]

Sync all tf variables across MPI processes.