Skip to content

Running and debugging with PyTorch RBLN

Overview

This tutorial demonstrates how to run and debug with PyTorch RBLN.

Setup & Installation

1
2
3
$ pip install \
  --extra-index-url https://download.pytorch.org/whl/cpu \
  torch-rbln==0.1.8

1
2
3
$ pip install \
  --extra-index-url https://pypi.rbln.ai/simple \
  rebel-compiler==0.10.2

Note

Please note that rebel-compiler requires an RBLN Portal account.

Running a simple example (torch.add)

The following add.py example verifies fp16 tensor ops on the RBLN NPUs. It uses the same code you would use on a GPU or CPU, except that it uses 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 9
Breakpoint 1 at add.py:9
(Pdb) c
> add.py(9)<module>()
-> print(f"{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.