Skip to content

Latest commit

 

History

History
47 lines (30 loc) · 856 Bytes

README.md

File metadata and controls

47 lines (30 loc) · 856 Bytes
Raw

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.