MooreThreads/torch_musa

torch_musa – PyTorch for Moore Threads MUSA GPUs

What it is

• An official PyTorch extension that lets you run PyTorch code on Moore Threads’ MUSA graphics cards. It mirrors the regular PyTorch API, so you only need to replace device strings like cuda with musa.

Why it matters

• Enables the large PyTorch ecosystem (models, libraries, tutorials) to run on a hardware platform that is not CUDA‑based, opening up a new class of accelerators for deep‑learning training and inference.

Key capabilities

• Drop‑in backend – torch.musa.is_available(), tensor.to('musa'), torch.musa.Stream(), etc., work just like their CUDA counterparts.
• Distributed training – supports the mccl backend (MCCL library) for multi‑GPU scaling.
• C/C++ extensions – provides a low‑boilerplate API for writing custom ops in native code.
• Tooling – utilities for porting CUDA code, building extensions, and debugging.
• Ecosystem support – pre‑built MUSA‑compatible wheels for torchvision, torchaudio, and a curated list of other projects (e.g., Hugging Face Transformers, PyTorch Lightning, PyTorch3D).

Getting started

1. Prerequisites – install the MUSA SDK (driver, toolkit, muDNN, MCCL) and math libraries (muThrust, muAlg). Verify with musaInfo.
2. Installation – either download pre‑compiled wheels from the GitHub releases page or build from source with the provided build.sh script (options for building only PyTorch, only torch_musa, debug/ASAN builds, wheel generation, etc.).
3. Docker – ready‑made images are hosted on registry.mthreads.com/mcconline. Pull and run with --privileged and MTHREADS_VISIBLE_DEVICES=all to expose the GPUs inside the container.
4. Code change – import torch as usual; the library auto‑loads. Switch device strings:

   import torch
   a = torch.tensor([1., 2.], device='musa')
   torch.backends.mudnn.allow_tf32 = True
   

   For distributed jobs:

   import torch.distributed as dist
   dist.init_process_group('mccl', rank=..., world_size=...)
   

Supported companion libraries

• torchvision – install from the Moore Threads vision repo (v0.26.0.post1 for torch_musa ≥ 2.7.0) or from upstream PyTorch if versions match.
• torchaudio – built from the official PyTorch audio source matching the torch version.
• Other projects – a list of “musified” repos (pytorch3d, pytorch_sparse, pytorch_scatter, pytorch_cluster, pytorch_lightning, etc.) is maintained in the README; they can be installed with pip install . or python setup.py install.

Troubleshooting highlights

• torch.musa.is_available() returns False → check MUSA driver/SKD installation and LD_LIBRARY_PATH.
• Missing shared libraries (libmusa.so, libmudnn.so, libmccl.so) → ensure MUSA_HOME points to the SDK and the libraries are on the library path.
• Distributed init fails → verify MCCL is installed (only on S4000) or build with USE_MCCL=0 for other chips.
• Container sees no device → run the Docker command shown in the README and set MTHREADS_VISIBLE_DEVICES.

License

• BSD‑style license (see LICENSE).

---

Bottom line: torch_musa lets developers leverage the familiar PyTorch workflow on Moore Threads’ MUSA GPUs, with full support for tensors, autograd, CUDA‑like extensions, and distributed training, backed by official Docker images and a growing set of compatible libraries.