Skip to content

Running and debugging with PyTorch RBLN

Overview

In this tutorial, we demonstrate how to run and debug with PyTorch RBLN.

Setup & Installation

Note

Please note that torch-rbln requires an RBLN Portal account.

Running a simple example (torch.add)

The following add.py is an example of verifying fp16 tensor ops on the Rebellions NPU. As you can see, the below example uses the same code as GPU or CPU except using rbln instead of cuda or cpu.

add.py
1
2
3
4
5
6
7
8
9
import torch

device = 'rbln'

a = torch.tensor([1, 2, 3], dtype=torch.float16, device=device)
b = torch.tensor([4, 5, 6], dtype=torch.float16, device=device)

c = a + b
print(f"{c}")  # expect: [5., 7., 9.] on rbln:0 (fp16)

If you run the above script, we expect the following:

1
2
$ python3 add.py
tensor([5., 7., 9.], device='rbln:0', dtype=torch.float16)

As we can see, fp16 element-wise torch.add on rbln matches the expected value.

Using PDB in PyTorch RBLN for further debugging

The following shows how to confirm the result tensor and indexed values in PDB.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
$ python3 -m pdb add.py
> add.py(1)<module>()
-> import torch
(Pdb) b 6
Breakpoint 1 at add.py:6
(Pdb) c
> add.py(6)<module>()
-> print(c)
(Pdb) p c
tensor([5., 7., 9.], device='rbln:0', dtype=torch.float16)
(Pdb) p c[0], c[1]
(tensor(5., device='rbln:0', dtype=torch.float16), tensor(7., device='rbln:0', dtype=torch.float16))

In the above example, we can see that both the full tensor and indexed tensor match expectations.