Skip to content

Latest commit

 

History

History
41 lines (33 loc) · 1.23 KB

README.md

File metadata and controls

41 lines (33 loc) · 1.23 KB
Raw

Torch-UCC

The Torch-UCC plugin is a research prototype that enables collective communication over XCCL for distributed PyTorch applications that load the plugin at application runtime. The XCCL interface is a non-standard API with a corresponding reference implementation used to guide design decisions for the UCC project.

Licenses

The torch-ucc plugin is licensed as:

Contributor Agreement and Guidelines

In order to contribute to torch-ucc, please sign up with an appropriate Contributor Agreement. Follow these instructions when submitting contributions and changes.

Build

Required packages:

  • PyTorch
  • UCX
  • XCCL
UCX_HOME=<PATH_TO_UCX> UCC_HOME=<PATH_TO_XCCL> python setup.py clean --all install

Run

export LD_LIBRARY_PATH=<PATH_TO_UCX>/lib:<PATH_TO_XCCL>/lib:$LD_LIBRARY_PATH
python example.py
import torch
import torch.distributed as dist
import torch_ucc

....
dist.init_process_group('ucc', rank=comm_rank, world_size=comm_size)
....
dist.all_to_all_single(recv_tensor, send_tensor)