SuperAPI AI Directory
Partner with us中文AI newsFavorites
Skillstorch_npu

torch_npu

Deployment & operations

View repository

Updated 10-15 17:20

About this skill

torch_npu capabilities and usage guidance

When to use this Skill

  • The user @torch_npu, mentions Ascend NPU, CANN, Ascend, or PyTorch running on NPU.
  • The user needs environment check/deployment: check or deploy PyTorch, check whether the environment supports NPU.
  • When the user uses @torch_npu_doc : based on this skill's reference.md provide the project's Chinese documentation capability description and operation steps.
  • The user needs code hints for NPU format conversion: when writing or completing torch_npu.npu_format_cast、npu_format_cast_、torch_npu.Format、get_npu_format , provide parameter and enum hints according to the sub-capability in §2.1; if this function is not understood, it may be suggested to use this skill's MCP to fetch version documentation as supplementary explanation.

1. Environment check and deployment

1.1 Automatically check the PyTorch and NPU environment

When answering or generating scripts, perform as needed or suggest that the user perform the following checks:

Check whether the PyTorch and Python versions are within the compatible range(see the "PyTorch and Python Version Compatibility Table" in README.zh.md):

  • Supports multiple versions such as PyTorch 1.11.0–2.9.0, corresponding to Python 3.7–3.11 (depending on the specific PyTorch version).

Check whether the environment supports NPU:

import torch
import torch_npu  # 2.5.1 及以后可不显式 import，仍建议写以便兼容

# 是否可用 NPU、设备数量
if torch.npu.is_available():
    count = torch.npu.device_count()
    # 使用 device='npu' 或 .npu()
else:
    # 未安装 CANN / 未 source set_env.sh / 无 NPU 设备
    pass

Check the CANN environment variables (verify after installation):

  • Before use, run:source /usr/local/Ascend/ascend-toolkit/set_env.sh(the path is subject to the actual CANN installation).
  • If ASCEND_HOME_PATH、ASCEND_OPP_PATH is not set or the path does not exist, torch_npu will report an error and prompt to run source set_env.sh。

1.2 Deployment step summary

  1. Install CANN: according to CANN Installation Guide install it, and check the CANN and PyTorch/Extension versions against the "Ascend Auxiliary Software" table in README.zh.md.
  2. Install PyTorch: for x86 use pip3 install torch==2.7.1+cpu --index-url https://download.pytorch.org/whl/cpu, for aarch64 use pip3 install torch==2.7.1(example versions, replace as needed).
  3. Install dependencies：pip3 install pyyaml setuptools。
  4. Install torch_npu：pip3 install torch-npu==2.7.1(the version must be compatible with PyTorch and CANN).
  5. Verify：source set_env.sh then run the quick verification code (see "Post-installation Verification" in README.zh.md).

2. torch_npu capability catalog (brief)

Category Capability description
Device and memory torch.npu: device management,device_count、current_device、set_device、synchronize、Stream/Event, memory statistics and allocation (memory_allocated、empty_cache、MemPool , etc.).
Tensor/storage tensor.npu()、tensor.is_npu, NPU Storage, serialization torch.save/load support NPU, DDP/multiprocess reductions.
Training/optimization torch.npu.amp mixed precision,torch_npu.optim, FSDP patch (ShardedGradScaler), gradient checkpointing defaults to NPU.
Distributed torch_npu.distributed: HCCL/LCCL backend,is_hccl_available、reinit_process_group, RPC, symmetric memory, DTensor rules.
Extended API torch_npu.contrib: NMS, IoU series, ROIAlign, DCN, FusedAttention, custom modules (such as DropoutWithByteMask), etc.
Graph and Compilation NPU Graph（npugraphify), Dynamo, Inductor, and torch.compile support the NPU.
Inference/ONNX ONNX export and NPU custom operator wrappers (such as OneHot, RoiAlign, NMS, FastGelu, MultiHeadAttention, etc.).
** profiling** torch_npu.profiler, MSTX patch, performance dump.
Other HiFloat8Tensor, erase_stream, matmul_checksum, transfer_to_npu (optional), op_plugin.

For detailed APIs, refer to Ascend Extension for PyTorch Custom API Reference and the project README.zh.md .

2.1 Sub-capability: torch_npu.npu_format_cast code hints

When a user writes in code or asks about torch_npu.npu_format_cast、npu_format_cast_, NPU tensor format conversion, the followingcode hints and completion guidanceshould be provided so that all hints can be completed in the IDE.

API signature and parameters

  • torch_npu.npu_format_cast(tensor, acl_format, customize_dtype=None)

    • tensor: on the NPU, torch.Tensor(requires first .npu()）。
    • acl_format: target storage format, which can be int or torch_npu.Format enum member.
    • customize_dtype: optional, for custom dtype in scenarios such as ONNX.
    • Returns: a new tensor (does not modify the original tensor).

  • torch_npu.npu_format_cast_(tensor, acl_format)

    • Same as above, but the in-place version, which directly modifies the format of tensor .

  • torch_npu.get_npu_format(tensor)

    • Returns the current NPU storage format of the tensor (torch_npu.Format or an integer).

Common Format enums (torch_npu.Format)

In code hints, the following common values can be suggested first (from torch_npu.npu._format.Format）：

Enum name Value Common use
Format.NCHW 0 Default 4D convolution layout
Format.NHWC 1 4D layout with channels last
Format.ND 2 General ND layout
Format.NC1HWC0 3 Commonly used by operators such as Conv/BatchNorm
Format.FRACTAL_Z 4 3D convolution, etc.
Format.FRACTAL_NZ 29 Linear/matrix multiplication, Attention weights, etc.
Format.NDC1HWC0 32 5D
Format.FRACTAL_Z_3D 33 3D convolution
Format.UNDEFINED -1 Undefined

Other optional:NC1HWC0_C04(12)、HWCN(16)、NDHWC(27)、NCDHW(30)、NC(35)、NCL(47)、FRACTAL_NZ_C0_*(50–54), etc.

Code hint and completion rules

  1. Complete the second parameter: when the user enters torch_npu.npu_format_cast(x, , suggest acl_format can be int or torch_npu.Format.xxx, and list common enums (such as Format.NCHW、Format.NHWC、Format.FRACTAL_NZ、Format.NC1HWC0）。
  2. Complete Format enums: when the user enters torch_npu.Format. , suggest the above list of enum members.
  3. Paired use: if the code already has get_npu_format(t), when it needs to be converted to the same format, you can suggest torch_npu.npu_format_cast(other, torch_npu.get_npu_format(t))。
  4. Common scenarios：
    • Linear layer weight quantization/migration to NPU:torch_npu.npu_format_cast(weight.npu(), 29)（FRACTAL_NZ）；
    • Gradient consistent with the parameter format:torch_npu.npu_format_cast(p.grad, torch_npu.get_npu_format(p))；
    • NC1HWC0 of BN/Conv during module migration:torch_npu.npu_format_cast(tensor, 3) or Format.NC1HWC0。

Documentation source description

If a more authoritative format description or the correspondence with CANN is needed, the user can be instructed to use the MCPof this skill (see mcp/README.md) to use fetch_torch_npu_doc to fetch the version documentation under the corresponding branch of GitCode Ascend/pytorch (such as docs/zh for documents related to format or operator adaptation), or look up entries related to "Framework Feature Guide" in the reference.

3. How to automatically provide the capabilities in the documentation (@torch_npu_doc)

  • When the user @torch_npu_doc or explicitly requests to consult the torch_npu Chinese documentation:
    • Prefer to use this skill's reference documentation：reference.md. It contains the index of all Chinese documentation in the project and summaries of key content.
    • Based on the user's question, locate the corresponding document path in the "Documentation Index" of reference.md (such as torch_npu/README.zh.md、torch_npu/docs/zh/quick_start/quick_start.md、torch_npu/docs/zh/framework_feature_guide_pytorch/、torch_npu/docs/zh/troubleshooting/ , etc.), read that file in the workspace as needed, and cite the relevant passages.
    • For installation, verification, uninstallation, and version compatibility: use the summaries in README.zh.md or reference; for migration and quick start: use quick_start; for operator adaptation, memory, performance, and graph mode: use the corresponding documents under framework_feature_guide_pytorch; for errors and troubleshooting: use the corresponding documents under troubleshooting.
  • When answering, give a brief explanation in Chinese, and note the cited document paths so the user can consult them on their own.

4. Reference Resources

  • This skill's documentation reference：reference.md. Summarizes the index and key points of all Chinese documentation in the torch_npu project; when the user @torch_npu_doc , use it to locate and cite specific documents.
  • Project Chinese Documentation: in the workspace torch_npu/README.zh.md(version table, deployment, verification, reference documentation links).
  • Official Documentation：Ascend Community Ascend Extension for PyTorch, software installation guide, model migration and training, operator adaptation, PyTorch native API list, custom API reference (see the table at the bottom of README.zh.md).
  • Fetch versioned documentation from GitCode: this skill provides an MCP service, which can fetch documentation from GitCode Ascend/pytorch by branch and save it to fetched_docs/. For installation and configuration, see mcp/README.md; the project has already configured .cursor/mcp.json, after installing dependencies, restart Cursor to use the tool fetch_torch_npu_doc、fetch_torch_npu_docs_batch、list_torch_npu_doc_paths。