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.

解决的问题

使 PyTorch 用户能够在 Moore Threads 显卡上运行其机器学习工作负载。通过提供一个模仿 PyTorch API 的后端，无需为不同硬件重写代码，开发者只需进行最少的修改，即可从 CUDA 或 CPU 切换到 MUSA。

工作原理

torch_musa 是 PyTorch 的一个扩展。通过将设备字符串从 cuda 或 cpu 更改为 musa，用户即可利用 Moore Threads 硬件。它与 MUSA-SDK 集成，包括 muDNN 和 MCCL（用于分布式训练）等工具，并提供 C/C++ 扩展 API，可编写自定义高性能层，且无需编写大量样板代码。

适用人群

使用 PyTorch 并希望将其模型迁移到 Moore Threads GPU 的 AI 开发者和研究人员。

主要亮点

  • API 一致性：使用与 PyTorch 相同的格式，使习惯标准库的用户迁移过程顺畅。
  • 硬件加速：通过 musa 后端提供对 Moore Threads GPU 性能的访问。
  • C/C++ 扩展：提供高效的 API，可编写自定义层，无需编写大量包装代码。
  • 分布式支持：支持使用 mccl 后端进行分布式训练和推理。
  • 优化工具：提供动态双精度转换和统一内存管理等高级功能。

相关

  • 项目
    aiqm/torchani

    TorchANI 2.0 是一个 PyTorch 库，用于训练和使用 ANI 风格的神经网络原子间势能，提供 GPU 加速的能量/力预测以及分子动力学和 ML/MM 模拟工具。

  • 项目
    ankane/torch.rb

    一个用于 LibTorch 的 Ruby 接口，让深度学习和神经网络开发能直接在 Ruby 中进行，并镜像 PyTorch API。

  • 项目
    facebookresearch/multimodal

    一个用于大规模训练最先进多模态模型的 PyTorch 库，为 CLIP 和 BLIP-2 等模型提供模块化构建块和预训练权重。

  • 项目
    iot-salzburg/gpu-jupyter

    GPU‑Jupyter 提供基于 Docker 的 GPU 支持 JupyterLab 环境，预装 Python、PyTorch、TensorFlow、Julia 和 R。支持版本镜像、易于自定义和 Docker‑Compose，可在任何配备 NVIDIA GPU 的机器上实现可复现的深度学习实验。

  • 项目
    pytorch/benchmark

    PyTorch Benchmarks 是一组标准化、轻量级的流行深度学习模型（如 BERT、ResNet、Stable Diffusion 等）的集合，可用于衡量不同 PyTorch 构建、CUDA 版本和后端的性能。它提供统一 API、安装脚本、多种运行器（简单测试、pytest-benchmark、自定义用户基准）以及用于在 AWS g4dn.metal 实例上进行低噪声调优的实用工具。