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¶
- An RBLN NPU host with the RBLN Driver installed
- Python: 3.10–3.14
torch,torch-rbln, andrebel-compilerinstalled as described in Installation- transformers
>=4.49.0,<4.50.0and the HuggingFace Hub CLI (huggingface_hub[cli])
If you have not installed the RBLN packages yet:
Then install the model dependencies:
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:
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.
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.
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(); itsrecompile/cold-compile: 137xentry 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_fallbackis the RMSNorm.why: dtype-not-fp16 169says that every fallback is a dtype decision, not a missing kernel. HuggingFace'sLlamaRMSNormupcasts the hidden states tofloat32beforepow,mean, the epsilonadd,rsqrt, and the scalingmul, 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 remainingeq,bmm, and twomulcalls are thefloat32rotary-frequency and attention-mask setup. Thecandidatesline names the fallbacks that have no host fast path; the norm'spow,mean, andrsqrthave one, which is why the block's total wall time stays at a millisecond.dispatch/recompilein 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_halfnegates 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 thatnegto the CPU on every call.q * cosandk * cosmultiply 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 thenegcosts a CPU round trip as well.host_bouncerows are copies that went through the host.d2d_copyis 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_dstis the CPU fallback path writing results back into strided device tensors. The>> physical d2hline at the bottom gives the volume that really crossed the device boundary inside the region.runtime/v2v_slowis runtime-internal. Itsreject: internal_fallback ... (runtime-internal, not user-fixable)line says so; read it as information and move on.with_stack=Trueattributes one call site per operator name, the first one seen in the region. That is whyaten::mul.outpoints 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
float32numerics, 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 thefloat32segment 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.
$ 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
linearprograms, 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
linearentry 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.
$ 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
float32norm, 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 compilinggenerate()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.
$ 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
float32on CPU, notfloat16on CPU. Device 16-bit arithmetic uses the runtime's own compute format, which carries fewer mantissa bits than IEEEfloat16, so afloat16CPU 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 withtorch.testing.assert_close(dev_logits, ref_logits, atol=..., rtol=...)as a regression check. greedy tokens matchis the end-to-end check. A mismatch is a defect only when the reference's top-2 gap is clearly larger thanmax |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¶
- Execution Semantics - Shape stability, the precision model behind Step 6, and where
torch.compilefits - Execution Semantics › Graph execution and input shapes - The rules that keep a compiled step on the device, and where the
torch.compilebackend itself is documented - APIs and Supported Operations - The
explain()signals and which operators a transformer block runs on the device in eager execution - RBLN vLLM - Serving LLMs with continuous batching instead of eager generation
- RBLN Optimum - Ahead-of-time compiled HuggingFace model classes