Installing Bagua

Wheels (precompiled binary packages) are available for Linux (x86_64). Package names are different depending on your CUDA Toolkit version (CUDA Toolkit version is shown in nvcc --version).

CUDA Toolkit versionInstallation command
>= v10.2pip install bagua-cuda102
>= v11.1pip install bagua-cuda111
>= v11.3pip install bagua-cuda113
>= v11.5pip install bagua-cuda115

Add --pre to pip install commands to install pre-release (development) versions.

Install from source

To install Bagua by compiling from source code on your machine, you need the following dependencies installed on your system:

We provide an automatic installation script for Ubuntu. Just run the following command to install Bagua and above libraries (except for CUDA, you should always install CUDA by yourself):

curl -Ls | sudo bash

Run the following commands to install Bagua (source code packages, which will be compiled on your machine).

# release version
python3 -m pip install bagua --upgrade

# develop version (git master)
python3 -m pip install --pre bagua --upgrade

Use Docker image

We provide Docker image with Bagua installed based on official PyTorch images. You can find them on DockerHub.