Skip to content

Running and debugging with RBLN PyTorch

This tutorial runs a first operation and a first module on an RBLN NPU, then shows the two tools you use to understand what happened: torch.rbln.explain() and the CPU fallback log. You will learn how to tell whether an operation ran on the device, and where to look when it did not.

Prerequisites

  • An RBLN NPU host with the RBLN Driver installed
  • Python: 3.10–3.14
  • torch, torch-rbln, and rebel-compiler installed as described in Installation

If you have not installed the packages yet:

1
2
3
4
5
$ pip install \
  --index-url https://download.pytorch.org/whl/cpu \
  torch==2.13.0+cpu
$ pip install \
  torch-rbln==0.12.0

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

Note

rebel-compiler requires an RBLN Portal account.

Confirm the environment before starting:

1
python -m torch_rbln.diagnose

The rebel ABI handshake: block of the report ends with verdict: OK. If it does not, follow Installation › Verification first.

Step 1: Run a first operation

The following add.py adds two 64-element fp16 tensors on the NPU. It is the same code you would write for a GPU or CPU, with rbln in place of cuda or cpu.

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

device = "rbln"

a = torch.arange(64, dtype=torch.float16, device=device)
b = torch.ones(64, dtype=torch.float16, device=device)

c = a + b
print(c.device, c.dtype, c[:4].cpu())  # expect: rbln:0, fp16, [1., 2., 3., 4.]

$ python3 add.py
rbln:0 torch.float16 tensor([1., 2., 3., 4.], dtype=torch.float16)

Two things happened that are invisible in the output:

  • a and b were created lazily. Creating a tensor on rbln reserves a virtual address only; the physical device memory was allocated, and the data uploaded, when a + b first needed them on the device.
  • a + b compiled aten::add for this shape and dtype on the first call. Later calls in the same process with the same shape and dtype reuse the compiled program. The cache lives in the process, so a new process compiles again.

Two details of the script are deliberate. To run eager addition on the NPU instead of the CPU, the last dimension of each input tensor must be a multiple of 64. This example uses one-dimensional tensors, so each has 64 elements (see Execution Semantics › Pre-dispatch checks). And the script prints c[:4].cpu() rather than c itself, because formatting an rbln tensor for printing runs host-side operators on it; reading the values back with .cpu() first keeps the fallback log in Step 3 free of that noise.

The dtype matters. Device kernels dispatch float16 and bfloat16; any other dtype runs on the host. Change torch.float16 to torch.float32 and the result is still correct, but the addition runs on CPU. Step 3 shows how to see that.

Step 2: Run a module on the device

A module works the same way as a single operator: convert it to a device dtype, move it to rbln, and call it. Create module.py:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
import torch

net = (
    torch.nn.Sequential(
        torch.nn.Linear(64, 128),
        torch.nn.SiLU(),
        torch.nn.Linear(128, 64),
    )
    .to(torch.float16)
    .to("rbln")
    .eval()
)

x = torch.ones(4, 64, dtype=torch.float16, device="rbln")

with torch.inference_mode():
    y = net(x)

print(y.shape, y.dtype, y.device)

$ python3 module.py
torch.Size([4, 64]) torch.float16 rbln:0

Two rules carry over to every model you port:

  • Cast before moving. .to(torch.float16) first, then .to("rbln"). A float32 module placed on rbln runs, but every operator falls back to CPU.
  • Call model.eval() to set the model to evaluation mode. Use torch.inference_mode() during inference to avoid unnecessary memory overhead.

Each Linear and the SiLU compiled on the first call. Run the forward again inside the same process and it replays the compiled programs.

Step 3: Find operations that ran on CPU

When an operator or dtype is not supported on the device, RBLN PyTorch runs it on CPU so execution can continue. The result is still an rbln tensor, so tensor.device says where the tensor is placed, not where the computation ran. The tool that tells you is torch.rbln.explain(): a context manager that counts what the backend did inside a region that you did not write and cannot see, namely CPU fallbacks, recompilations, and host round-trips, and attributes each to a cause.

Create add_fp32.py. It is Step 1 with float32 tensors, wrapped in an explain() region:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
import torch

device = "rbln"

a = torch.arange(64, dtype=torch.float32, device=device)
b = torch.ones(64, dtype=torch.float32, device=device)

with torch.rbln.explain(with_stack=True) as p:
    c = a + b

print(c.device, c.dtype, c[:4].cpu())
print(p.report())

$ python3 add_fp32.py
rbln:0 torch.float32 tensor([1., 2., 3., 4.])
[overhead: 1 signal]  RBLN EXPLAIN   (region wall 989.001us | device mem 256.00 MB peak, reserved)
  ...
---------------------  -----  -----  -------------------------------------
Signal                 Count  Bytes  Note
---------------------  -----  -----  -------------------------------------
dispatch/cpu_fallback      1     --  try: graph mode, or a supported dtype
---------------------  -----  -----  -------------------------------------
  dispatch/cpu_fallback:  (sum 850.661us wall)
    aten::add.out  1  at add_fp32.py:9(<module>)
    why: dtype-not-fp16 1
    candidates (no fast-path handler): add.out
  (detail: p.help(signal) | raw: p.dump())

The result is correct and reports rbln:0, but the addition ran on the CPU. How to read the report:

  • The marker [overhead: N signals] is a fact, not a severity grade. It says that N kinds of hidden event fired; cost is judged from the Count, Bytes, and Note columns.
  • dispatch/cpu_fallback groups the fallen-back operators with their counts, and why: names the reason. Here it is dtype-not-fp16: the inputs are float32, which has no device kernel. The candidates line lists the fallen-back operators that had no host fast path.
  • with_stack=True adds the at file:line(function) call site per operator. It is off by default so a plain region costs nothing.
  • Note leads with what to try. Rows without a fix are runtime-internal notifications.

Fallbacks come from two places, and only one of them can be turned into an error:

  • An operator with no RBLN registration at all goes through the generic fallback. With TORCH_RBLN_DISABLE_FALLBACK=unsupported_op set, executing an unregistered operator raises an error naming the operator instead of running it on the CPU. This environment variable is read once when the extension loads, so set it before starting the process.
  • A registered operator that fails a pre-dispatch check, as aten::add does here because its inputs are float32, falls back regardless of that variable. The why: line is the signal; the fix is the dtype.

See Environment Variables › Execution and fallback for the categories.

explain() covers one region. For a whole run, RBLN PyTorch also logs fallbacks at INFO level, which the default WARNING level hides. The log names every operator that ran on the CPU, including those sent there by a pre-dispatch check such as the dtype check that sent aten::add to the CPU above, together with the source line when it can be traced, so it points at the layer of a model that produced the fallback. It does not give the reason. explain() gives one for the pre-dispatch checks, but it does not count an operator that has no device kernel or is pinned to the CPU, so for those the log is the only signal. Enable the log with:

1
export TORCH_RBLN_LOG_LEVEL=INFO

For example, torch.sin has no device kernel, so torch.sin(x) on an rbln tensor logs:

[2026-09-14 19:49:00.145][I] `aten::sin` op ran on CPU instead of RBLN
[2026-09-14 19:49:00.145][W] TRACE
sin_demo.py:5: UserWarning: TRACE (Triggered internally at c10/rbln/RBLNLogging.cpp:266.)
  y = torch.sin(x)

Warning

While a Python tracer such as pdb or coverage is active, RBLN PyTorch runs every operator on CPU rather than compiling under the tracer. Values inspected under a debugger are correct, but the device placement and timing you observe there do not represent a normal run. Read fallback logs and explain() reports from a run outside the debugger.

Restore the default log level when you are done:

1
export TORCH_RBLN_LOG_LEVEL=WARNING

Step 4: Explain hidden overhead

Step 3 used explain() to catch a CPU fallback. The same report catches the two other kinds of hidden work, recompilation and copies that travel through the host, and it can separate the one-time cost of a first call from what recurs on every call. It is not a timing profiler; a region can be [clean] and still slow because of device compute time.

Wrap the module from Step 2. The first forward pass compiles each operator, and explain() reports that as a dispatch/recompile signal, so run the module once before opening the region to see its steady-state behavior:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
import torch

net = (
    torch.nn.Sequential(
        torch.nn.Linear(64, 128),
        torch.nn.SiLU(),
        torch.nn.Linear(128, 64),
    )
    .to(torch.float16)
    .to("rbln")
    .eval()
)

x = torch.ones(4, 64, dtype=torch.float16, device="rbln")

with torch.inference_mode():
    net(x)  # warm-up: compiles the three operators

with torch.inference_mode(), torch.rbln.explain() as p:
    y = net(x)

print(p.report())

A region with no hidden overhead prints the marker [clean] and lists what it checked, so an empty report is a positive statement. Abbreviated output; the header lines with timings and runtime residency differ per machine:

[clean]  RBLN EXPLAIN   (region wall 801.076us | device mem 264.33 MB peak, reserved)
  ...
  checked: host_bounce, v2v_slow, cpu_fallback, recompile -- none fired
  note: clean = no hidden host overhead, not "fast"

Now introduce a fallback on purpose. An fp32 operand is outside the device dtype policy, and the conversion that produces it is a copy across dtypes, which travels through the host:

1
2
3
4
with torch.rbln.explain(with_stack=True) as p:
    w = x.float() * 2.0

print(p.report())

[overhead: 2 signals]  RBLN EXPLAIN   (region wall 230.989us | device mem 264.33 MB peak, reserved)
  ...
---------------------  -----  -----  -------------------------------------------------------
Signal                 Count  Bytes  Note
---------------------  -----  -----  -------------------------------------------------------
host_bounce/d2d_copy       1  512 B  no DMA (low cost) | try: make contiguous, or v2v engine
dispatch/cpu_fallback      1     --  try: graph mode, or a supported dtype
---------------------  -----  -----  -------------------------------------------------------
  host_bounce/d2d_copy:
    at explain_demo.py:19(<module>)
  dispatch/cpu_fallback:  (sum 67.500us wall)
    aten::mul.out  1  at explain_demo.py:19(<module>)
    why: dtype-not-fp16 1
  >> physical d2h (real transfer): 0 -- everything above served on host, no crossing
  (detail: p.help(signal) | raw: p.dump())

One line of Python produced two signals, which is the point of the tool. x.float() is a device dtype cast that round-tripped the host (host_bounce/d2d_copy, 512 B); the * 2.0 on the resulting float32 tensor ran on the CPU (dispatch/cpu_fallback), and why: names the reason, dtype-not-fp16, the same case as Step 3. For copies, the >> physical d2h line reports whether data actually crossed the device boundary; here x was still host-resident, so nothing did, and the Note column rates the bounce as low cost.

Two more calls cover the common workflow:

1
2
p.verdict()  # {'clean': False, 'reasons': ['hidden host bounce (torch-side): 1x, 512 B', 'cpu_fallback: 1x (ran on CPU)'], ...}  -> usable as a CI gate
p.dump()  # every counter as raw numbers -> programmatic checks

To separate one-time cost (first-call compilation) from recurring overhead, let explain_steady() place two regions around a callable: one around the first call it makes and one around a later call after warmup more:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
# New layer sizes: the compile cache is keyed by operator and shape, so these
# programs do not exist yet even though a module of the same kind has run.
net_wide = (
    torch.nn.Sequential(
        torch.nn.Linear(64, 256),
        torch.nn.SiLU(),
        torch.nn.Linear(256, 64),
    )
    .to(torch.float16)
    .to("rbln")
    .eval()
)

with torch.inference_mode():
    cold, warm = torch.rbln.explain_steady(
        lambda: net_wide(x), warmup=2, return_cold=True
    )

print(cold.verdict()["reasons"])
print(warm.verdict()["clean"])

['recompile/cold-compile: 3x']
True

The first call compiled the three operators; the later call replayed them and is clean. A second module with the same layer sizes as Step 2 would have shown no compilation at all, because the cache is keyed by operator and input shapes rather than by the module object. Signals that disappear between the two regions were one-time; signals that persist are the recurring overhead worth fixing. cold.diff(warm) shows that comparison as a signal-by-signal table (explain_steady(..., as_diff=True) returns the same object directly):

1
print(cold.diff(warm).report())

RBLN EXPLAIN DIFF   (A -> B; you placed both points)
  explain does not know your run structure; this compares ONLY your two regions.
--------------------  -  -  ---------  -------  -----------
Signal                A  B    A bytes  B bytes
--------------------  -  -  ---------  -------  -----------
host_bounce           0  0         --       --
runtime/v2v_slow      0  0    4.00 KB  4.00 KB
runtime/physical_d2h  2  2    4.00 KB  4.00 KB  >> persists
runtime/h2d_push      7  0  499.25 KB       --  gone in B
cpu_fallback          0  0         --       --
recompile             3  0         --       --  gone in B
--------------------  -  -  ---------  -------  -----------

The helper assumes the callable does the same work on every call and was not run before; see APIs for the parameters.

Next steps