• Overview
• Installation
  • Prerequisites
  • Docker Image
  • From Python wheels
  • From Source
• MUSA Supported Repositories
  • torchvision and torchaudio
  • Other Repositories
• Usage
  • Key Changes
  • Codegen
• License

Overview

torch_musa is an extended Python package based on PyTorch. Combined with PyTorch, users can take advantage of the strong power of Moore Threads graphics cards through torch_musa.

torch_musa's APIs are consistent with PyTorch in format, which allows users accustomed to PyTorch to migrate smoothly to torch_musa, so for the usage users can refer to PyTorch Official Doc, all you need is just switch the backend string from "cpu" or "cuda" to "musa".

torch_musa also provides a bundle of tools for users to conduct cuda-porting, building musa extension and debugging. Please refer to README.md.

For some customize optimizations, like Dynamic Double Casting and Unified Memory Management, please refer to README.md.

If you want to write your layers in C/C++, we provide a convenient extension API that is efficient and with minimal boilerplate. No wrapper code needs to be written. You can see a ResNet50 example here.

---

Installation

Prerequisites

Before installing torch_musa, here are some software packages that need to be installed on your machine first:

• (For Docker users) KUAE Cloud Native Toolkits, including Container-Toolkit, MTML, sGPU;
• MUSA-SDK, including MUSA Driver, musa_toolkit, muDNN and MCCL(S4000 only);
• Other libraries, including muThrust and muAlg

Docker Image

We provide several released docker images and they can be easily found in mcconline.

Pull the docker image with docker pull, create a container with command below and there you go.

# For example, start a S80 docker with Python3.10
docker run -it --privileged \
  --pull always --network=host \
  --env MTHREADS_VISIBLE_DEVICES=all \
  registry.mthreads.com/mcconline/musa-pytorch-release-public:latest /bin/bash

During its initial startup, it performs a self-check, so you can see the MUSA environment is ready or not.

From Python wheels

Download torch & torch_musa wheels from our Release Page, please make sure you have all prerequisites installed.

From Source

First, clone the torch_musa repository

git clone https://github.com/MooreThreads/torch_musa.git
cd torch_musa

then, we need to set MUSA_HOME, LD_LIBRARY_PATH and PYTORCH_REPO_PATH for building torch_musa:

export MUSA_HOME=path/to/musa_libraries(including mudnn and musa_toolkits) # defalut value is /usr/local/musa/
export LD_LIBRARY_PATH=$MUSA_HOME/lib:$LD_LIBRARY_PATH
# if PYTORCH_REPO_PATH is not set, PyTorch will be downloaded outside this directory when building with build.sh
export PYTORCH_REPO_PATH=/path/to/PyTorch

To building torch_musa, run:

bash build.sh -c  # clean cache then build PyTorch and torch_musa from scratch

Some important building parameters are as follows:

• --torch/-t: build original PyTorch only
• --musa/-m: build torch_musa only
• --debug: build in debug mode
• --asan: build in asan mode
• --clean/-c: clean everything built and build
• --wheel/-w: generate wheels

For example, if one has built PyTorch and only needs to build torch_musa with wheel, run:

bash build.sh -m -w

For S80/S3000 users, since the MCCL library is not provided for such architectures, please add USE_MCCL=0 whilst building torch_musa:

USE_MCCL=0 bash build.sh -c

MUSA Supported Repositories

We provide some widely used PyTorch environment repositories, which have all adapted with our MUSA platform. Besides, many repositories have supported MUSA backend upstream, like Transformers, Accelerate, you can install them from PyPi with pip install [repo-name].

torchvision

For torch_musa v2.7.0 and later, install torchvision from our musified one:

git clone https://github.com/MooreThreads/vision -b v0.22.1-musa --depth 1
cd vision && python setup.py install

For torch_musa v2.5.0 and earlier, install torchvision from source with:

git clone https://github.com/pytorch/vision -b ${version} --depth 1
cd vision && python setup.py install

the version is depend on torch version, for example you have torch_musa v2.5.0 with torch v2.5.0, install torchvision==0.20.0.

torchaudio

Install torchaudio from source: them from source:

git clone https://github.com/pytorch/audio.git -b ${version} --depth 1
cd audio && python setup.py install

the version is same as the torch version.

Other Repositories

There are many widely used pytorch-related repositories, and we musified some of them and put them into our GitHub, here's the list:

Repo	Branch	Link	Build command
pytorch3d	musa-dev	https://github.com/MooreThreads/pytorch3d	python setup.py install
pytorch_sparse	master	https://github.com/MooreThreads/pytorch_sparse	python setup.py install
pytorch_scatter	master	https://github.com/MooreThreads/pytorch_scatter	python setup.py install
torchvision	v0.22.1-musa	https://github.com/MooreThreads/vision	python setup.py install
More to come...

If users find any question about these repos, please file issues in torch_musa, and if anyone musify a repository, you can submit a Pull Request that helping us to expand this list.

Usage

Key Changes

The following two key changes are required when using torch_musa:

• Import torch_musa package

  import torch
  import torch_musa

• Change the device to musa

  import torch
  import torch_musa
  
  a = torch.tensor([1.2, 2.3], dtype=torch.float32, device='musa')
  b = torch.tensor([1.2, 2.3], dtype=torch.float32, device='cpu').to('musa')
  c = torch.tensor([1.2, 2.3], dtype=torch.float32).musa()

torch musa has integrated torchvision ops in the musa backend. Please do the following if torchvision is not installed:

• Install torchvision package via building from source

  # ensure torchvision is not installed
  pip uninstall torchvision
  
  git clone https://github.com/pytorch/vision.git
  cd vision
  python setup.py install
  

• Use torchvision musa backend:

  import torch
  import torch_musa
  import torchvision
  
  def get_forge_data(num_boxes):
      boxes = torch.cat((torch.rand(num_boxes, 2), torch.rand(num_boxes, 2) + 10), dim=1)
      assert max(boxes[:, 0]) < min(boxes[:, 2])  # x1 < x2
      assert max(boxes[:, 1]) < min(boxes[:, 3])  # y1 < y2
      scores = torch.rand(num_boxes)
      return boxes, scores
  
  num_boxes = 10
  boxes, scores = get_forge_data(num_boxes)
  iou_threshold = 0.5
  print(torchvision.ops.nms(boxes=boxes.to("musa"), scores=scores.to("musa"), iou_threshold=iou_threshold))
  

Codegen

In torch_musa, we provide the codegen module to implement bindings and registrations of customized MUSA kernels, see link.

License

torch_musa has a BSD-style license, as found in the LICENSE file.