torch_npu (Ascend PyTorch)

TorchNPU is a deep learning adapter plugin that enables PyTorch to invoke Ascend NPUs. Its architecture consists of several key functional modules:

• Basic Compute: Provides support for PyTorch native and custom APIs for mainstream AI scenarios.
• Distributed: Accelerates distributed training via FSDP2 and supports DTensor, collective communication (AllGather, AllReduce, AllToAll), and point-to-point primitives (Send, Recv).
• Graph Mode: Accelerates training/inference via dynamic graph capture, static graph optimization, and efficient code generation. Supports offloading via NPUGraph (v2.6.0+).
• Debug & Tuning: Provides profiling for compute, communication, and memory, plus real-time anomaly monitoring via WatchDog.
• TorchNPU Core: Manages virtual memory to reduce fragmentation, optimizes cross-stream memory reuse in distributed settings, and integrates resources via PrivateUse1.

When using run_and_get_code, the returned codes contains the captured MLIR output code snippets. These snippets typically include a call function and a benchmark_compiled_model function.

The call function encapsulates the MLIR-compiled operator logic. Within this function, you will find the core execution interface, such as mlir_fused_xxx.run. This interface is responsible for executing the fused operator on a specific NPU device and compute stream, taking input tensors and writing results to the output tensor.

def call(args):
    arg0_1, arg1_1 = args
    args.clear()
    with torch.npu.utils.device(0):
        torch.npu.set_device(0)
        buf0 = empty_strided((3, ), (1, ), device='npu', dtype=torch.float32)
        stream0 = get_raw_stream(0)
        # This is the core execution interface for the fused operator
        mlir_fused_mul_0.run(arg0_1, arg1_1, buf0, stream=stream0)
        del arg0_1
        del arg1_1
    return (buf0, )

The profiler offers different modes for data collection. Note: You cannot use two or more modes simultaneously in the same process.

Data Profiling Modes

• torch_npu.profiler.profile (Recommended): Provides complete profiling APIs. You manually add these APIs to your code to freely choose specific content to profile.
• dynamic_profile (Recommended): Supports starting profiling at any time during training and allows starting profiling without modifying user code, offering higher flexibility.

Extended Profiling (Optional)

If you need to profile additional data, the following functions are available:

• Profile MSTX data
• Profile environment variables
• Mark performance data profiling process
• Device memory visualization
• Create Profiler sub-thread profiling

Data Parsing

• Automatic Parsing: The standard profiling APIs perform automatic parsing during the profiling process.
• Offline Parsing: Use this if you have unparsed raw performance data that needs to be processed after the fact.

The HCCL AllReduce example demonstrates the standard lifecycle for performing collective communication on NPU devices using Libtorch:

1. NPU Device Initialization: Initialize the device using torch_npu::init_npu("npu:X").
2. Create ProcessGroupHCCL: Instantiate the process group specifically for NPU collective communication.
3. Tensor Creation: Create tensors directly on the NPU device.
4. Execute AllReduce: Perform the global reduction operation on the tensors.
5. Verification: Validate that the reduction results match expected mathematical outcomes.
6. Resource Cleanup: Crucial step—you must call torch_npu::finalize_npu() to release NPU resources after use.

Important Requirements:

• Ensure NPU devices are available and CANN is correctly installed.
• All participating processes must be started simultaneously for the collective communication to function.