Ascend/pytorch

Ascend PyTorch adapter (torch_npu). Mirror of https://gitcode.com/Ascend/pytorch

What it solves

TorchNPU is a deep learning adapter plugin that allows the PyTorch framework to run directly on Ascend AI processors (NPUs). It bridges the gap between PyTorch's high-level APIs and Huawei's Ascend hardware, enabling developers to leverage NPU compute capabilities without leaving the PyTorch ecosystem.

How it works

It functions as an extension library that integrates Ascend NPU operators and device resources into PyTorch via a mechanism called PrivateUse1. It provides a set of core compute APIs, distributed training support (including FSDP2 and DTensor), and a "Graph Mode" that uses dynamic graph capture and static graph optimization to accelerate training and inference.

Who it’s for

AI developers and researchers who are using PyTorch and want to migrate their models to Ascend NPU hardware for improved performance and compute power.

Highlights

  • PyTorch Native Experience: Reuses upstream PyTorch capabilities while providing deep adaptation for NPUs.
  • Distributed Training: Supports FSDP2, DTensor, and collective communication primitives (AllGather, AllReduce, AllToAll).
  • Graph Mode Acceleration: Uses dynamic graph capture and efficient code generation to speed up execution.
  • Performance Tuning: Includes profiling tools for compute, communication, and memory usage, plus a WatchDog for communication anomalies.
  • Memory Optimization: Features virtual memory management to reduce fragmentation and cross-stream memory reuse.

Related

  • Project
    pytorch/pytorch

    A Python-based machine learning framework that provides GPU-accelerated tensor computation and a dynamic autograd system for building deep neural networks.

  • Project
    Lightning-AI/pytorch-lightning

    A deep learning framework that abstracts away the engineering boilerplate of PyTorch, enabling researchers to pretrain and finetune AI models at scale across various hardware.

  • Project
    pytorch/torchtitan

    torchtitan is a PyTorch‑native library for large‑scale generative‑AI training. It bundles modern parallelism (FSDP2, tensor, pipeline, context), memory‑saving tricks (activation checkpointing, Float8/MXFP8 quantization), distributed checkpointing, and extensive logging. Ready‑to‑use for Llama 3.1 (8 B‑405 B) and extensible to new models, it runs on single‑node or multi‑node clusters via simple CLI/config files, and is released under a BSD‑3‑Clause license.

  • Project
    pytorch/TensorRT

    Torch-TensorRT integrates NVIDIA TensorRT into PyTorch to accelerate inference latency by up to 5x on NVIDIA GPUs.

  • Project
    huggingface/accelerate

    A PyTorch wrapper that abstracts the boilerplate code for distributed training, allowing users to run the same script on any device configuration including multi-GPU and TPU.