torch_pg
Description
torch_pg is a collection of PyTorch Distributed ProcessGroup Backends.
Supports
- tested on PyTorch 1.6.0/1.7.0 and NCCL 2.7.8
Installation
MPI_HOME=/path/to/mpi NCCL_HOME=/path/to/nccl pip install -e .Examples
MPI
import torch.distributed as dist
import torch_pg
# Install or replace built-in MPI backend.
torch_pg.init_mpi()
dist.init_process_group("mpi")NCCL
import torch.distributed as dist
import torch_pg
# Install or replace built-in NCCL backend.
torch_pg.init_nccl()
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "29501"
dist.init_process_group("nccl", rank=rank, world_size=world_size)License
torch_pg is licensed under the BSD-3-Clause License.