MooreThreads/torch_musa

torch_musa is an open source repository based on PyTorch, which can make full use of the super computing power of MooreThreads graphics cards.

What it solves

It enables PyTorch users to run their machine learning workloads on Moore Threads graphics cards. It removes the need to rewrite code for different hardware by providing a backend that mimics the PyTorch API, allowing developers to switch from CUDA or CPU to MUSA with minimal changes.

How it works

torch_musa acts as an extension to PyTorch. By changing the device string from cuda or cpu to musa, users can leverage Moore Threads hardware. It integrates with the MUSA-SDK, including tools like muDNN and MCCL (for distributed training), and provides a C/C++ extension API for writing custom high-performance layers with minimal boilerplate.

Who it’s for

AI developers and researchers who use PyTorch and want to migrate their models to Moore Threads GPUs.

Highlights

  • API Consistency: Uses the same format as PyTorch, making migration smooth for those accustomed to the standard library.
  • Hardware Acceleration: Provides access to Moore Threads GPU power through the musa backend.
  • C/C++ Extensions: Includes an efficient API for writing custom layers without needing extensive wrapper code.
  • Distributed Support: Supports distributed training and inference using the mccl backend.
  • Optimization Tools: Offers advanced features like Dynamic Double Casting and Unified Memory Management.

Related

  • Project
    aiqm/torchani

    TorchANI 2.0 is a PyTorch library for training and using ANI‑style neural‑network interatomic potentials, providing GPU‑accelerated energy/force predictions and tools for molecular dynamics and ML/MM simulations.

  • Project
    ankane/torch.rb

    A Ruby interface for LibTorch that enables deep learning and neural network development directly in Ruby, mirroring the PyTorch API.

  • Project
    facebookresearch/multimodal

    A PyTorch library for training state-of-the-art multimodal models at scale, providing modular building blocks and pretrained weights for models like CLIP and BLIP-2.

  • Project
    iot-salzburg/gpu-jupyter

    GPU‑Jupyter provides a Docker‑based, GPU‑enabled JupyterLab environment pre‑installed with Python, PyTorch, TensorFlow, Julia, and R. It offers versioned images, easy customization, and Docker‑Compose support, enabling reproducible deep‑learning experiments on any machine with an NVIDIA GPU.

  • Project
    pytorch/benchmark

    PyTorch Benchmarks is a collection of standardized, lightweight versions of popular deep‑learning models (BERT, ResNet, Stable Diffusion, etc.) that can be run to measure the performance of different PyTorch builds, CUDA versions, and back‑ends. It offers a uniform API, install scripts, and multiple runners (simple test, pytest‑benchmark, custom userbenchmark) plus utilities for low‑noise tuning on AWS g4dn.metal instances.