Skip to content

Porting and validating a transformers model: Llama3.2-1B

Running and debugging used one operator and a three-layer module that you wrote yourself. A real model adds three things that tutorial did not have: hundreds of operators per forward pass, library code that you did not write and cannot easily change, and shapes that change while the model generates. This tutorial ports the Llama3.2-1B model from HuggingFace Transformers to an RBLN NPU and applies the same two tools, torch.rbln.explain() and the fallback log, at that scale: which operators stay on the host and why, where data crosses between host and device, what is compiled when, and whether the result is correct against a float32 CPU reference.

Eager execution is the development path. Nothing is compiled ahead of time, so this is where you observe what a model makes the backend do before deciding on fixes or moving to a compiled or serving path. Throughput is not the goal here; understanding the port is.

Note

Llama 3.2 is licensed under the Llama 3.2 Community License, Copyright (c) Meta Platforms, Inc. All Rights Reserved.

Prerequisites

If you have not installed the RBLN 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

Then install the model dependencies:

$ pip install "transformers>=4.49.0,<4.50.0" "huggingface_hub[cli]"

Note

rebel-compiler requires an RBLN Portal account.

Note

The meta-llama/Llama-3.2-1B model on HuggingFace Hub has restricted access. Once access is granted, log in with the hf command:

$ hf auth login
Enter your token (input will not be visible):

Step 1: Cast, move, and generate

The script llama.py applies the two rules from Running and debugging. It loads the model with torch_dtype=torch.float16, so the model is in a device dtype before it moves, and then model.to(device) and .to(device) on the tokenizer outputs are the only lines that differ from a CUDA script.

llama.py
import re

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "meta-llama/Llama-3.2-1B"
device = "rbln"

tokenizer = AutoTokenizer.from_pretrained(model_name)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map=None,
)
model.to(device)

prompt = "Is Sun bigger than Earth?"
inputs = tokenizer(prompt, return_tensors="pt")

input_ids = inputs["input_ids"].to(device)
attention_mask = inputs["attention_mask"].to(device)

outputs = model.generate(
    input_ids,
    attention_mask=attention_mask,
    pad_token_id=tokenizer.pad_token_id,
    max_new_tokens=4,
    num_return_sequences=1,
    do_sample=False,
    top_p=None,
    temperature=None,
)

prompt_length_tokens = input_ids.shape[1]
generated_text = tokenizer.decode(
    outputs[0][prompt_length_tokens:], skip_special_tokens=True
).strip()
generated_text = re.sub(r"\[duplicate\]\n?", "", generated_text)

print(f"Q: {prompt}")
print(f"A: {generated_text}")
$ python3 llama.py
Q: Is Sun bigger than Earth?
A: The answer is yes

The port works without any model change. The first generation takes noticeably longer than on CUDA: each operator is compiled for each input shape it sees, and the programs live in the process, so a second process compiles again. Step 4 measures what generation compiles and when, and Step 5 compiles the forward pass as one program; Execution Semantics › Operator dispatch path explains the cache behind it.

Step 2: Explain one forward pass

A working port says nothing about what the backend did to make it work. llama_explain.py runs one forward pass over the prompt through torch.rbln.explain_steady(), which places one region around the first call and another around a later call, so the one-time compilation is separated from what happens on every call. The last block of the script is used in Step 3.

llama_explain.py
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "meta-llama/Llama-3.2-1B"
device = "rbln"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map=None,
)
model.to(device).eval()

inputs = tokenizer("Is Sun bigger than Earth?", return_tensors="pt").to(device)


def step():
    # One forward pass over the prompt: the same shapes on every call.
    with torch.inference_mode():
        return model(**inputs).logits


# Two regions around the same call: the first call compiles, a later call replays.
cold, warm = torch.rbln.explain_steady(
    step, warmup=2, return_cold=True, with_stack=True
)

print(
    cold.verdict()["reasons"]
)  # one-time cost: expect a 'recompile/cold-compile: Nx' entry
print(warm.report())  # what remains on every call

# Narrow the region to one module to see its operator chain exactly.
hidden = torch.randn(
    1,
    inputs["input_ids"].shape[1],
    model.config.hidden_size,
    dtype=torch.float16,
    device=device,
)
with torch.inference_mode():
    model.model.norm(hidden)  # warm-up: compile what the norm runs on the device
    with torch.rbln.explain(with_stack=True) as norm:
        model.model.norm(hidden)

print(norm.report())

Abbreviated output of the first two print calls; counts and timings differ per machine, and the header lines are omitted:

$ python3 llama_explain.py
['hidden host bounce (torch-side): 101x, 2.98 MB', ..., 'recompile/cold-compile: 137x', 'cpu_fallback: 169x (ran on CPU)']
[overhead: 4 signals]  RBLN EXPLAIN   (region wall 615.279ms | device mem 4.90 GB peak, reserved)
  ...
-----------------------------  ------  ---------  ------------------------------------------------------------
Signal                          Count      Bytes  Note
-----------------------------  ------  ---------  ------------------------------------------------------------
host_bounce/d2d_copy               69    2.71 MB  see physical d2h below | try: make contiguous, or v2v engine
host_bounce/h2d_noncontig_dst      32  280.00 KB  see physical d2h below | try: contiguous staging, then h2d
runtime/v2v_slow               12,999    1.01 MB  see physical d2h below
dispatch/cpu_fallback             169         --  try: graph mode, or a supported dtype
dispatch/recompile                128         --  try: stabilize shapes, or graph mode
-----------------------------  ------  ---------  ------------------------------------------------------------

  host_bounce/d2d_copy:
    at modeling_llama.py:127(forward) <- _contextlib.py:124(decorate_context)

  host_bounce/h2d_noncontig_dst:
    at ops_utils.py:2008(cpu_fallback_path) <- ops_utils.py:1764(compile_and_run_view_aware)

  runtime/v2v_slow:
    state:  src_not_on_device 8,519  src_device_only_real_d2h 4,480
    reject: internal_fallback 12,999 (runtime-internal, not user-fixable)

  dispatch/cpu_fallback:  (sum 1.216ms wall)
    aten::mul.out                35  at modeling_llama.py:138(forward) <- _contextlib.py:124(decorate_context)
    aten::pow.Tensor_Scalar_out  33  at modeling_llama.py:72(forward) <- module.py:1790(_call_impl)
    aten::rsqrt.out              33  at modeling_llama.py:73(forward) <- module.py:1790(_call_impl)
    aten::add.out                33  at modeling_llama.py:73(forward) <- module.py:1790(_call_impl)
    aten::mean.out               33  at modeling_llama.py:72(forward) <- module.py:1790(_call_impl)
    aten::eq.Scalar_out           1  at modeling_attn_mask_utils.py:288(_ignore_causal_mask_sdpa) <- modeling_llama.py:646(_update_causal_mask)
    aten::bmm.out                 1  at modeling_llama.py:132(forward) <- _contextlib.py:124(decorate_context)
    why: dtype-not-fp16 169
    candidates (no fast-path handler): add.out, eq.Scalar_out, bmm.out

  dispatch/recompile:
    aten::mul.out  64  at modeling_llama.py:138(forward) <- _contextlib.py:124(decorate_context)
    aten::neg.out  32  at modeling_llama.py:148(rotate_half) <- modeling_llama.py:173(apply_rotary_pos_emb)
    aten::add.out  32  at modeling_llama.py:73(forward) <- module.py:1790(_call_impl)

  >> physical d2h (real transfer): 432 copies, 15.09 MB -- real device crossing (see runtime/v2v_slow)

How to read it:

  • The cold region compiled, the warm region did not. The first list is the cold region's verdict(); its recompile/cold-compile: 137x entry is the first pass compiling every operator for the prompt's shapes. That entry is absent from the warm region, whose report is printed next: the second pass had the same shapes and replayed the programs.
  • dispatch/cpu_fallback is the RMSNorm. why: dtype-not-fp16 169 says that every fallback is a dtype decision, not a missing kernel. HuggingFace's LlamaRMSNorm upcasts the hidden states to float32 before pow, mean, the epsilon add, rsqrt, and the scaling mul, and the count of 33 per operator is the number of norms in the model: two per decoder layer for 16 layers, plus the final norm. The remaining eq, bmm, and two mul calls are the float32 rotary-frequency and attention-mask setup. The candidates line names the fallbacks that have no host fast path; the norm's pow, mean, and rsqrt have one, which is why the block's total wall time stays at a millisecond.
  • dispatch/recompile in a warm region counts warm-cache misses, not compilations. Nothing was compiled in the second pass, yet 128 calls missed the eager warm cache and re-entered Python dispatch. rotate_half negates the second half of each 64-wide head, a 32-wide slice below the 64-element alignment that device kernels need, so the Python layer hands that neg to the CPU on every call. q * cos and k * cos multiply by a broadcast view of the rotary tables, and a view input is replayed in Python rather than served from the C++ cache. Each such call costs Python dispatch time, and the neg costs a CPU round trip as well.
  • host_bounce rows are copies that went through the host. d2d_copy is a copy that changes dtype; its count covers every such conversion in the pass, mostly the .to(torch.float32) and .to(input_dtype) pairs inside the norms, while the location shown is the first one seen. h2d_noncontig_dst is the CPU fallback path writing results back into strided device tensors. The >> physical d2h line at the bottom gives the volume that really crossed the device boundary inside the region.
  • runtime/v2v_slow is runtime-internal. Its reject: internal_fallback ... (runtime-internal, not user-fixable) line says so; read it as information and move on.
  • with_stack=True attributes one call site per operator name, the first one seen in the region. That is why aten::mul.out points at the rotary scaling on line 138 although most of its 35 fallbacks are the norm's scaling on line 73. To attribute precisely, narrow the region, which is what Step 3 does.

warm.verdict() returns the same facts as a dictionary and warm.dump() returns raw counters, so the check can run in CI. See APIs › Explain and profiling for the signals and their notes.

Step 3: Narrow the region to one module

A region around the whole model shows totals. To see one operator chain exactly, place the region around one module. The last block of llama_explain.py feeds a float16 hidden-state tensor of the prompt's shape through the model's final LlamaRMSNorm, once as a warm-up and once inside a region:

[overhead: 2 signals]  RBLN EXPLAIN   (region wall 334.689us | device mem 4.90 GB peak, reserved)
  ...
---------------------  -----  --------  -------------------------------------------------------
Signal                 Count     Bytes  Note
---------------------  -----  --------  -------------------------------------------------------
host_bounce/d2d_copy       2  84.00 KB  no DMA (low cost) | try: make contiguous, or v2v engine
dispatch/cpu_fallback      5        --  try: graph mode, or a supported dtype
---------------------  -----  --------  -------------------------------------------------------

  host_bounce/d2d_copy:
    at modeling_llama.py:71(forward) <- module.py:1790(_call_impl)

  dispatch/cpu_fallback:  (sum 31.220us wall)
    aten::pow.Tensor_Scalar_out  1  at modeling_llama.py:72(forward) <- module.py:1790(_call_impl)
    aten::rsqrt.out              1  at modeling_llama.py:73(forward) <- module.py:1790(_call_impl)
    aten::mul.out                1  at modeling_llama.py:73(forward) <- module.py:1790(_call_impl)
    aten::add.out                1  at modeling_llama.py:73(forward) <- module.py:1790(_call_impl)
    aten::mean.out               1  at modeling_llama.py:72(forward) <- module.py:1790(_call_impl)
    why: dtype-not-fp16 5
    candidates (no fast-path handler): add.out

  >> physical d2h (real transfer): 0 -- everything above served on host, no crossing

This is LlamaRMSNorm.forward line by line. Line 71, hidden_states.to(torch.float32), is the first d2d_copy; the second is the cast back to float16 at the end. Lines 72 and 73 are the five float32 operators, each executed once on the CPU. The >> physical d2h line reports zero: the input tensor was created on the host and had not been consumed by a device operator yet, so the CPU kernels read it without a download. Inside the model the input of every norm is the device-resident output of the previous layer, so the same chain there costs the kind of transfer the >> physical d2h line in Step 2 counts.

The report also tells you which changes are available, and what each costs:

  • Leave the norm as it is. The fallback keeps HuggingFace's float32 numerics, and the host fast paths keep it cheap: a millisecond per forward pass in Step 2.
  • Compute the norm in float16. This removes the fallback and the two conversions, but changes the numerics, so the comparison in Step 6 has to be repeated.
  • Run the block under torch.compile(backend="rbln"). The compiler handles the float32 segment inside the graph; see Execution Semantics › Graph execution.

For a whole run rather than a region, the fallback log gives a stream with source lines. With TORCH_RBLN_LOG_LEVEL=INFO, one forward pass logs the neg fallbacks of rotate_half and the attention fallback, which the fused device kernel refuses for a query length that is not a multiple of 32:

[2026-09-14 19:49:00.145][I] `aten::neg` op ran on CPU instead of RBLN
[2026-09-14 19:49:00.145][W] TRACE
/transformers/models/llama/modeling_llama.py:148: UserWarning: TRACE (Triggered internally at c10/rbln/RBLNLogging.cpp:266.)
  return torch.cat((-x2, x1), dim=-1)
[2026-09-14 19:49:00.152][I] `sdpa_overrideable (query sequence length 7 not aligned to 32)` op ran on CPU instead of RBLN

The log names each operator that ran on the CPU and where it was called from, but not why. explain() adds the why: reason and a per-region count for the fallbacks a pre-dispatch check decides, such as the norm's dtype fallbacks, and does not count operators that have no device kernel, which is why the two tools are used together.

Step 4: Observe what generation compiles

Step 2 explained a forward pass whose shapes never change. Generation is different: the prompt is processed at its own length, then each decode step processes one token. llama_shapes.py wraps three generate() calls in regions, two with the same prompt and one with a prompt of a different token count, and prints the region wall time together with the number of aten::linear programs compiled inside the region. linear is used because it always runs on the device, so its count reflects compilation alone.

llama_shapes.py
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "meta-llama/Llama-3.2-1B"
device = "rbln"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map=None,
)
model.to(device).eval()

prompts = [
    "Is Sun bigger than Earth?",  # first call: nothing compiled yet
    "Is Sun bigger than Earth?",  # same token count: same shapes
    "Is the Moon bigger than the Earth or the Sun?",  # a new token count: new shapes
]

for prompt in prompts:
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    with torch.inference_mode(), torch.rbln.explain() as p:
        model.generate(
            **inputs,
            max_new_tokens=4,
            do_sample=False,
            top_p=None,
            temperature=None,
            pad_token_id=tokenizer.eos_token_id,
        )
    d = p.dump()
    print(
        f"{inputs['input_ids'].shape[1]:2d} prompt tokens: "
        f"region wall {d['wall_ns'] / 1e9:5.2f} s, "
        f"new `linear` programs {d['recompile_by_op'].get('aten::linear', 0)}"
    )
$ python3 llama_shapes.py
 7 prompt tokens: region wall 17.82 s, new `linear` programs 10
 7 prompt tokens: region wall  2.43 s, new `linear` programs 1
12 prompt tokens: region wall  8.20 s, new `linear` programs 5

How to read it:

  • Programs are keyed by input shape and shared across layers. The first call compiled 10 linear programs, not 16 layers times 7 projections: the model has five distinct projection shapes, and each was compiled once for the prompt's length and once for the single-token decode steps. Every layer with the same shapes replays the same program.
  • The same prompt again compiles nothing. The second call sees the same sequence lengths and runs in a fraction of the time. The single linear entry that remains is a call that re-enters Python dispatch on every call, as in Step 2, not a compilation.
  • A new token count compiles the prompt shapes again. The third call is a 12-token prompt: five new programs for its prefill shapes, and the wall time grows with them. The decode steps reuse the programs from the first call, because a decode step always processes one token.

In a service, prompt lengths vary on every request, so eager execution compiles on every new length. The remedies are the ones the dispatch/recompile note names: fix the shapes by padding or bucketing prompts to a small set of lengths, or move the model to a path that compiles the whole graph for fixed shapes, RBLN vLLM for serving or RBLN Optimum for ahead-of-time compiled model classes. See Troubleshoot › Recompile limit exceeded for the Dynamo cache limit that shape churn eventually hits.

Step 5: Compile the forward pass

Every observation so far came from eager execution, where each operator is its own program and the RMSNorm's float32 chain runs on the host on every call. llama_compile.py compiles the whole forward pass as one program instead. It pads or truncates the prompt to a fixed length of 64 tokens so that one program serves every prompt, turns the compile_error fallback into an error so a failure cannot hide as a CPU run, and compares the result with eager.

llama_compile.py
import os
import time

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Make a compilation failure an error instead of a silent run on the CPU.
os.environ["TORCH_RBLN_DISABLE_FALLBACK"] = "compile_error"

model_name = "meta-llama/Llama-3.2-1B"
device = "rbln"

tokenizer = AutoTokenizer.from_pretrained(model_name)
# The base tokenizer ships without a pad token. Pad on the left so that the last
# position of every padded prompt is its final token.
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map=None,
)
model.to(device).eval()

# A fixed shape: pad or truncate every prompt to 64 tokens so one program serves them all.
inputs = tokenizer(
    "Is Sun bigger than Earth?",
    return_tensors="pt",
    padding="max_length",
    truncation=True,
    max_length=64,
).to(device)

compiled = torch.compile(model, backend="rbln", dynamic=False)

with torch.inference_mode():
    start = time.perf_counter()
    logits = compiled(**inputs).logits.cpu()
    print(
        f"first compiled forward: {time.perf_counter() - start:6.1f} s  {tuple(logits.shape)}"
    )

    with torch.rbln.explain() as steady:
        start = time.perf_counter()
        logits = compiled(**inputs).logits.cpu()
        print(f"steady compiled forward: {time.perf_counter() - start:6.3f} s")
    print("steady region:", steady.verdict()["reasons"])

    eager_logits = model(**inputs).logits.cpu()

next_token = tokenizer.decode(logits[0, -1].argmax())
print("next token from the compiled forward:", repr(next_token))
print(
    "eager and compiled argmax match:",
    torch.equal(eager_logits[0, -1].argmax(), logits[0, -1].argmax()),
)
print("max |eager - compiled|:", (eager_logits - logits).abs().max().item())
$ python3 llama_compile.py
first compiled forward:   13.5 s  (1, 64, 128256)
steady compiled forward:  0.005 s
steady region: ['no hidden host bounce, no CPU fallback, no recompile']
next token from the compiled forward: ' The'
eager and compiled argmax match: True
max |eager - compiled|: 0.13037109375

How to read it:

  • One compilation, then replay. The first forward pays 13 seconds of compilation for the 64-token shape; the steady forward takes 5 milliseconds and the region is clean. The eager forward in Step 2 took about 600 milliseconds on the same host and carried 169 fallbacks; inside the graph the compiler handles the float32 norm, the rotary views, and the attention itself, so none of that reaches the host.
  • The shape is the contract. The tokenizer pads or truncates every prompt to 64 tokens, so every prompt reuses this program. Without truncation=True, a 65-token prompt would compile a new one; a service that cannot drop tokens rejects longer prompts or routes them to a larger fixed shape instead. Generation still changes the sequence length on every step (Step 4), so compiling generate() as written gains little. Serving stacks solve this with fixed-shape buckets and a static KV cache; that is what RBLN vLLM and RBLN Optimum do.
  • Two warnings print once. view-on-device fallback (host materialize via .contiguous()) appears for two attention views the runtime cannot classify; it is a one-time cost per process, not a per-call one.
  • Compiled and eager agree. The argmax matches and the largest logit difference is about 0.13, the device-format rounding from Step 6 measured against eager rather than against float32.

Step 6: Validate against a CPU reference

A port that runs and a port that is correct are different things. llama_validate.py runs the same prompt twice, in float32 on CPU and in float16 on rbln, and compares the greedy tokens and the first-step logits.

llama_validate.py
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "meta-llama/Llama-3.2-1B"
prompt = "Is Sun bigger than Earth?"
generate_kwargs = dict(max_new_tokens=8, do_sample=False, top_p=None, temperature=None)

tokenizer = AutoTokenizer.from_pretrained(model_name)
inputs = tokenizer(prompt, return_tensors="pt")

# Reference: the same model in float32 on CPU.
ref_model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32)
ref_model.eval()
with torch.inference_mode():
    ref_logits = ref_model(**inputs).logits[0, -1]
    ref_tokens = ref_model.generate(
        **inputs, pad_token_id=tokenizer.eos_token_id, **generate_kwargs
    )[0]

# Port: float16 on rbln.
dev_model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
dev_model.to("rbln").eval()
dev_inputs = inputs.to("rbln")
with torch.inference_mode():
    dev_logits = dev_model(**dev_inputs).logits[0, -1].float().cpu()
    dev_tokens = dev_model.generate(
        **dev_inputs, pad_token_id=tokenizer.eos_token_id, **generate_kwargs
    )[0].cpu()

print("greedy tokens match:      ", torch.equal(ref_tokens, dev_tokens))
print(
    "first-step argmax match:  ",
    ref_logits.argmax().item() == dev_logits.argmax().item(),
)
print("first-step max |diff|:    ", (ref_logits - dev_logits).abs().max().item())
print(
    "first-step top-2 gap (ref):",
    (ref_logits.topk(2).values[0] - ref_logits.topk(2).values[1]).item(),
)
$ python3 llama_validate.py
greedy tokens match:       True
first-step argmax match:   True
first-step max |diff|:     0.044589996337890625
first-step top-2 gap (ref): 0.2722015380859375

How to read it:

  • The reference is float32 on CPU, not float16 on CPU. Device 16-bit arithmetic uses the runtime's own compute format, which carries fewer mantissa bits than IEEE float16, so a float16 CPU run is a third set of numbers rather than the truth. See Execution Semantics › Precision and dtypes.
  • A non-zero max |diff| is expected; every element carries a small rounding difference. What matters is that it stays small relative to the logit scale and does not grow as generation proceeds. Record the value from a run you trust and encode it with torch.testing.assert_close(dev_logits, ref_logits, atol=..., rtol=...) as a regression check.
  • greedy tokens match is the end-to-end check. A mismatch is a defect only when the reference's top-2 gap is clearly larger than max |diff|; when two candidates are nearly tied, the rounding difference alone can flip the choice and either token is a correct answer.
  • The operators that fell back in Step 2 were computed on the host in float32, so they contribute no device error under eager execution. The difference you see comes from the operators that ran on the device, which is what you want to measure.

If the tokens differ at a clear gap, or the difference grows with sequence length, run with TORCH_RBLN_LOG_LEVEL=INFO and TORCH_RBLN_DEPLOY unset so that an input with NaN or Inf still routes to CPU and shows up in the log, and follow Troubleshoot › Results differ from a CPU run.

Next steps