moolib - a communications library for distributed ML training
moolib offers general purpose RPC with automatic transport selection (shared memory, TCP/IP, Infiniband) allowing models to data-parallelise their training and synchronize gradients and model weights across many nodes.
moolib is an RPC library to help you perform distributed machine learning research, particularly reinforcement learning. It is designed to be highly flexible and highly performant.
It is flexible because it allows researchers to define their own training loops and data-collection policies with minimal interference or abstractions -
moolib gets out of the way of research code.
It is performant because it gives researchers the power of efficient data-parallelization across GPUs with minimal overhead, in a manner that is highly scalable.
moolib aims to provide researchers with the freedom to implement whatever experiment loop they desire, and the freedom to scale it up from single GPUs to hundreds at will (with no additional code). It ships with a reference implementations IMPALA on Atari that can easily be adapted to other environments or algorithms.
moolib without CUDA support
To install from GitHub:
pip install git+https://github.com/facebookresearch/moolib
To build from source:
git clone --recursive [email protected]:facebookresearch/moolib cd moolib && pip install .
How to host docs (after installation):
pip install sphinx==4.1.2 cd docs && ./run_docs.sh
Run an Example
To run the example agent on a given Atari level:
First, start the broker:
python -m moolib.broker
It will output something like
Broker listening at 0.0.0.0:4431.
Note that a single broker is enough for all your experiments.
Now take the IP address of your computer. If you ssh'd into your machine, this should work (in a new shell):
export BROKER_IP=$(echo $SSH_CONNECTION | cut -d' ' -f3) # Should give your machine's IP. export BROKER_PORT=4431
To start an experiment with a single peer:
python -m examples.vtrace.experiment connect=BROKER_IP:BROKER_PORT \ savedir=/tmp/moolib-atari/savedir \ project=moolib-atari \ group=Zaxxon-Breakout \ env.name=ALE/Breakout-v5
To add more peers to this experiment, start more processes with the same
group settings, using a different setting for
moolib is licensed under the MIT License. See
LICENSE for details.