Enabling the graph mode in PyTorch

The Ascend platform supports the execution of the PyTorch network in graph mode. The following figure shows the software architecture for execution in graph mode.

Figure 1 Software architecture for graph mode of PyTorch

TorchAir (Torch Ascend Intermediate Representation) is an extended library that supports the graph mode capability in Ascend Extension for PyTorch (torch_npu). It supports graph mode training and inference of the PyTorch network on the Ascend device. TorchAir interconnects with the Dynamo feature of PyTorch to convert the FX (functionalization) graph of PyTorch into an Ascend IR graph, compile and optimize the computational graph through GE, and deliver the graph to the Ascend hardware for execution.

Before enabling the graph mode for the PyTorch network, you are advised to learn about Ascend Extension for PyTorchs. For details, see Ascend Extension for PyTorch Quick Start s.

The Ascend compilation backend provided by TorchAir can be passed to the torch.compile API of PyTorch as a parameter to enable the graph mode. An example is shown below:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

# Import torch_npu and then torchair.
import torch
import torch_npu
import torchair

# Define a model.
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x, y):
        return torch.add(x, y)

# Instantiate the model.
model = Model()

# Obtain the default backend provided by the NPU from TorchAir.
config = torchair.CompilerConfig()
npu_backend = torchair.get_npu_backend(compiler_config=config)

# Use the backend of TorchAir to call the compile API to compile the model.
opt_model = torch.compile(model, backend=npu_backend)

# Execute the compiled model.
x = torch.randn(2, 2)
y = torch.randn(2, 2)
opt_model(x, y)

For details about how to use the TorchAir-based PyTorch graph mode, see PyTorch Graph Mode Instructions (TorchAir) .