콘텐츠로 이동

Transformers 모델 포팅 및 검증: Llama3.2-1B

실행 및 디버깅에서는 연산 하나와 직접 작성한 3계층 모듈을 다뤘습니다. 실제 모델에는 그 튜토리얼에 없던 세 가지 요소가 추가됩니다. 순전파 한 번에 수백 개의 연산을 실행합니다. 라이브러리 코드는 직접 작성하지 않았으므로 쉽게 수정하기 어렵고 생성 중에는 shape가 바뀝니다. 이 튜토리얼에서는 HuggingFace Transformers의 Llama3.2-1B 모델을 RBLN NPU로 포팅하고 같은 두 도구인 torch.rbln.explain()과 fallback 로그를 모델 규모로 적용합니다. 어떤 연산이 왜 호스트에 남는지, 호스트와 디바이스 사이에서 데이터가 어디로 이동하는지, 무엇을 언제 컴파일하는지, 결과가 float32 CPU 기준값과 비교해 올바른지 확인합니다.

Eager 실행은 개발 단계에서 사용하는 경로입니다. 사전 컴파일이 없으므로 수정 방향을 정하거나 컴파일·서빙 경로로 옮기기 전에 모델이 백엔드에서 어떤 작업을 일으키는지 관찰할 수 있습니다. 여기서는 처리량보다 포팅 동작을 이해하는 데 초점을 맞춥니다.

Note

Llama 3.2는 Llama 3.2 Community License에 따라 사용이 허가됩니다. Copyright(c) Meta Platforms, Inc. All Rights Reserved.

사전 준비

RBLN 패키지를 아직 설치하지 않았다면 다음 명령을 실행합니다.

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

이어서 모델 의존성을 설치합니다.

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

Note

rebel-compiler를 사용하려면 RBLN Portal 계정이 필요합니다.

Note

HuggingFace Hub의 meta-llama/Llama-3.2-1B 모델은 접근이 제한되어 있습니다. 접근 권한을 받은 뒤 hf 명령으로 로그인해야 합니다.

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

1단계: dtype 변환, 이동 및 생성

llama.py 스크립트는 실행 및 디버깅의 두 규칙을 적용합니다. 모델을 torch_dtype=torch.float16으로 불러와 이동 전에 디바이스 dtype으로 맞춥니다. CUDA 스크립트와 다른 부분은 model.to(device)와 토크나이저 출력의 .to(device)뿐입니다.

llama.py
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
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

모델을 수정하지 않아도 포팅은 동작합니다. 다만 첫 생성은 CUDA보다 눈에 띄게 오래 걸립니다. 각 연산은 입력 shape별로 컴파일되고 프로그램은 프로세스 안에 유지되므로 다른 프로세스에서는 다시 컴파일합니다. 4단계에서는 생성 중 무엇을 언제 컴파일하는지 측정하고 5단계에서는 순전파 전체를 하나의 프로그램으로 컴파일합니다. 이때 사용하는 캐시는 실행 방식 › 연산자 디스패치 경로에서 설명합니다.

2단계: 순전파 한 번 분석

포팅이 동작한다는 사실만으로는 백엔드가 어떤 작업을 했는지 알 수 없습니다. llama_explain.py는 torch.rbln.explain_steady()로 프롬프트의 순전파를 실행합니다. 첫 호출과 이후 호출에 각각 region을 설정해 일회성 컴파일과 매번 발생하는 작업을 구분합니다. 스크립트의 마지막 블록은 3단계에서 사용합니다.

llama_explain.py
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
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())

첫 두 print 호출의 출력을 줄여서 표시했습니다. 횟수와 시간은 머신마다 다르며 헤더 줄은 생략했습니다.

$ 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)

출력은 다음과 같이 해석합니다.

  • Cold region에서는 컴파일했고 warm region에서는 컴파일하지 않았습니다. 첫 목록은 cold region의 verdict()입니다. recompile/cold-compile: 137x 항목은 첫 순전파에서 프롬프트 shape에 맞춰 각 연산을 컴파일한 결과입니다. 이어서 출력되는 warm region 보고서에는 이 항목이 없습니다. 두 번째 순전파는 shape가 같아 기존 프로그램을 재실행했습니다.
  • dispatch/cpu_fallback은 RMSNorm에서 발생합니다. why: dtype-not-fp16 169는 모든 fallback이 커널 누락이 아닌 dtype 판단에 따른 것임을 나타냅니다. HuggingFace의 LlamaRMSNorm은 pow, mean, epsilon을 더하는 add, rsqrt, 스케일링 mul 전에 hidden state를 float32로 올립니다. 연산별 횟수 33은 모델의 norm 개수로, 16개 decoder 계층에 각각 두 개씩 있고 마지막 norm이 하나 더 있습니다. 나머지 eq, bmm, 두 번의 mul은 float32 rotary frequency와 attention mask를 준비하는 작업입니다. candidates 줄에는 호스트 fast path가 없는 fallback이 표시됩니다. Norm의 pow, mean, rsqrt에는 fast path가 있어 블록 전체 wall time이 밀리초 수준에 머뭅니다.
  • Warm region의 dispatch/recompile은 컴파일 횟수가 아니라 warm cache miss 횟수입니다. 두 번째 순전파에서는 아무것도 컴파일하지 않았지만 128번의 호출이 eager warm cache를 찾지 못해 Python 디스패치로 다시 진입했습니다. rotate_half는 너비 64인 각 head의 뒤쪽 절반을 음수로 바꿉니다. 이 slice의 너비는 32로, 디바이스 커널에 필요한 64원소 정렬을 만족하지 않습니다. Python 계층은 매번 neg를 CPU로 보냅니다. q * cos와 k * cos는 rotary table의 broadcast view를 곱합니다. View 입력은 C++ 캐시에서 처리하지 않고 Python에서 재실행합니다. 이러한 호출마다 Python 디스패치 시간이 들고 neg에는 CPU 왕복 비용도 발생합니다.
  • host_bounce 행은 호스트를 거친 복사입니다. d2d_copy는 dtype을 바꾸는 복사입니다. 횟수에는 해당 순전파의 모든 dtype 변환이 포함되며 대부분 norm 내부의 .to(torch.float32)와 .to(input_dtype) 쌍입니다. 표시된 위치는 처음 관측한 호출입니다. h2d_noncontig_dst는 CPU fallback 경로가 결과를 stride가 있는 디바이스 텐서에 쓰는 작업입니다. 마지막 >> physical d2h 줄은 region 안에서 실제로 디바이스 경계를 넘은 데이터량을 나타냅니다.
  • runtime/v2v_slow는 런타임 내부 작업입니다. reject: internal_fallback ... (runtime-internal, not user-fixable) 줄에도 명시되어 있습니다. 참고 정보로 확인하면 됩니다.
  • with_stack=True는 연산 이름별로 호출 위치 하나를 표시합니다. Region 안에서 처음 관측한 위치입니다. aten::mul.out의 35번 fallback 대부분이 73행의 norm 스케일링에서 발생해도 표시 위치는 138행의 rotary 스케일링입니다. 정확한 위치를 파악하려면 3단계처럼 region을 좁힙니다.

warm.verdict()는 같은 정보를 딕셔너리로 반환하고 warm.dump()는 원시 카운터를 반환하므로 CI에서 검사할 수 있습니다. Signal과 설명은 API 및 지원 연산 › Explain과 프로파일링에서 확인할 수 있습니다.

3단계: 모듈 하나로 region 범위 축소

모델 전체를 감싼 region은 합계를 보여 줍니다. 연산 체인 하나를 정확히 보려면 모듈 하나에 region을 설정합니다. llama_explain.py의 마지막 블록은 프롬프트 shape의 float16 hidden-state 텐서를 모델의 마지막 LlamaRMSNorm에 입력합니다. 한 번은 warm-up으로, 한 번은 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

이 출력은 LlamaRMSNorm.forward의 각 행에 대응합니다. 71행의 hidden_states.to(torch.float32)가 첫 d2d_copy이고 끝에서 float16으로 되돌리는 변환이 두 번째입니다. 72행과 73행의 다섯 float32 연산은 각각 CPU에서 한 번 실행됩니다. >> physical d2h는 0입니다. 입력 텐서를 호스트에서 만들었고 아직 디바이스 연산이 사용하지 않았으므로 CPU 커널이 다운로드 없이 읽었습니다. 모델 내부에서는 각 norm의 입력이 이전 계층이 디바이스에 출력한 텐서입니다. 같은 연산 체인도 모델 안에서는 2단계의 >> physical d2h에 집계되는 전송 비용이 발생합니다.

보고서에서 가능한 변경과 각각의 비용도 확인할 수 있습니다.

  • Norm을 그대로 둡니다. Fallback은 HuggingFace의 float32 수치 동작을 유지하고 호스트 fast path는 비용을 낮게 유지합니다. 2단계에서는 순전파당 밀리초 수준입니다.
  • Norm을 float16으로 계산합니다. Fallback과 두 번의 변환을 없애지만 수치 동작이 달라지므로 6단계의 비교를 다시 수행해야 합니다.
  • 블록을 torch.compile(backend="rbln")으로 실행합니다. 컴파일러가 그래프 안에서 float32 구간을 처리합니다. 자세한 내용은 실행 방식 › 그래프 실행과 입력 shape에서 확인할 수 있습니다.

Region 대신 실행 전체를 보려면 소스 행이 포함된 fallback 로그를 사용합니다. TORCH_RBLN_LOG_LEVEL=INFO로 순전파를 한 번 실행하면 rotate_half의 neg fallback과 attention fallback이 기록됩니다. Fused 디바이스 커널은 query 길이가 32의 배수가 아니면 attention을 거부합니다.

[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

로그는 CPU에서 실행한 연산과 호출 위치를 표시하지만 이유는 표시하지 않습니다. explain()은 norm의 dtype fallback처럼 디스패치 전 검사에서 결정한 fallback의 why: 이유와 region별 횟수를 추가합니다. 디바이스 커널이 없는 연산은 집계하지 않으므로 두 도구를 함께 사용합니다.

4단계: 생성 중 컴파일 관찰

2단계에서는 shape가 바뀌지 않는 순전파를 분석했습니다. 생성에서는 프롬프트를 해당 길이로 처리한 뒤 decode 단계마다 토큰 하나를 처리합니다. llama_shapes.py는 generate() 호출 세 번에 region을 설정합니다. 두 번은 같은 프롬프트를 사용하고 한 번은 토큰 수가 다른 프롬프트를 사용합니다. Region의 wall time과 그 안에서 컴파일한 aten::linear 프로그램 수를 함께 출력합니다. linear는 항상 디바이스에서 실행되므로 이 횟수는 컴파일만 반영합니다.

llama_shapes.py
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
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

출력은 다음과 같이 해석합니다.

  • 프로그램은 입력 shape를 키로 사용하고 계층 간에 공유됩니다. 첫 호출은 16개 계층에 projection 7개씩을 곱한 수가 아닌 10개의 linear 프로그램을 컴파일했습니다. 모델에는 서로 다른 projection shape가 다섯 개 있으며 각각 프롬프트 길이와 단일 토큰 decode shape마다 한 번씩 컴파일했습니다. Shape가 같은 계층은 모두 같은 프로그램을 재실행합니다.
  • 같은 프롬프트를 다시 처리할 때는 컴파일하지 않습니다. 두 번째 호출은 sequence 길이가 같아 훨씬 짧은 시간에 실행됩니다. 남아 있는 linear 항목 하나는 컴파일이 아니라 2단계처럼 매번 Python 디스패치로 다시 진입하는 호출입니다.
  • 토큰 수가 바뀌면 프롬프트 shape를 다시 컴파일합니다. 세 번째 호출은 12토큰 프롬프트입니다. Prefill shape용 새 프로그램 다섯 개를 컴파일하면서 wall time도 늘어납니다. Decode 단계는 항상 토큰 하나를 처리하므로 첫 호출의 프로그램을 재사용합니다.

서비스에서는 요청마다 프롬프트 길이가 달라지므로 eager 실행은 새로운 길이마다 컴파일합니다. 해결 방법은 dispatch/recompile 설명과 같습니다. Padding이나 bucketing으로 프롬프트를 소수의 길이에 맞추거나, 고정 shape의 전체 그래프를 컴파일하는 경로로 옮깁니다. 서빙에는 RBLN vLLM을, 사전 컴파일한 모델 클래스에는 RBLN Optimum을 사용합니다. Shape가 계속 바뀌면 결국 Dynamo 캐시 한도에 도달합니다. 이때의 대처 방법은 문제 해결 › 재컴파일 한도 초과에 정리되어 있습니다.

5단계: 순전파 컴파일

지금까지는 각 연산이 별도 프로그램이고 RMSNorm의 float32 체인이 매번 호스트에서 실행되는 eager 경로를 관찰했습니다. llama_compile.py는 순전파 전체를 하나의 프로그램으로 컴파일합니다. 프롬프트를 고정 길이 64토큰으로 padding하거나 잘라 하나의 프로그램이 모든 프롬프트를 처리하도록 합니다. 실패가 CPU 실행으로 숨지 않도록 compile_error fallback을 오류로 바꾸고 결과를 eager 실행과 비교합니다.

llama_compile.py
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
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

출력은 다음과 같이 해석합니다.

  • 한 번 컴파일한 뒤 재실행합니다. 첫 순전파에서는 64토큰 shape를 컴파일하는 데 13초가 걸립니다. 이후 순전파는 5밀리초이며 region에 오버헤드 signal이 없습니다. 같은 호스트에서 2단계의 eager 순전파는 약 600밀리초가 걸렸고 fallback이 169번 발생했습니다. 그래프 안에서는 컴파일러가 float32 norm, rotary view, attention을 직접 처리하므로 호스트로 넘어가지 않습니다.
  • Shape가 실행 조건입니다. 토크나이저가 모든 프롬프트를 64토큰으로 padding하거나 잘라 같은 프로그램을 재사용합니다. truncation=True가 없으면 65토큰 프롬프트는 새 프로그램을 컴파일합니다. 토큰을 버릴 수 없는 서비스라면 더 긴 프롬프트를 거부하거나 더 큰 고정 shape로 보내야 합니다. 생성은 단계마다 sequence 길이가 달라지므로(4단계), 현재 코드의 generate()를 컴파일해도 이점이 작습니다. 서빙 스택은 고정 shape bucket과 static KV cache로 sequence 길이 변화를 처리합니다. RBLN vLLM과 RBLN Optimum이 이 방식을 사용합니다.
  • 경고 두 개가 한 번씩 출력됩니다. 런타임이 분류하지 못하는 attention view 두 개에 대해 view-on-device fallback (host materialize via .contiguous())가 표시됩니다. 이는 프로세스당 한 번 발생하는 비용이며 호출마다 발생하지 않습니다.
  • Compiled와 eager 결과가 일치합니다. Argmax가 같고 최대 logit 차이는 약 0.13입니다. 6단계에서 설명하는 디바이스 형식의 반올림 차이를 float32 대신 eager 결과와 비교한 값입니다.

6단계: CPU 기준값과 비교 검증

포팅이 실행되는지와 결과가 올바른지는 따로 확인해야 합니다. llama_validate.py는 같은 프롬프트를 CPU의 float32와 rbln의 float16으로 각각 실행하고 greedy 토큰과 첫 단계 logits를 비교합니다.

llama_validate.py
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
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

출력은 다음과 같이 해석합니다.

  • 기준값은 CPU의 float32이며 CPU의 float16이 아닙니다. 디바이스의 16비트 연산은 런타임 고유 연산 형식을 사용합니다. 이 형식은 IEEE float16보다 가수 비트가 적습니다. CPU의 float16 결과는 기준값이 아니라 또 다른 수치 결과입니다. 자세한 내용은 실행 방식 › 정밀도와 dtype에서 확인할 수 있습니다.
  • max |diff|가 0이 아닌 것은 예상되는 결과입니다. 각 원소에 작은 반올림 차이가 있습니다. Logit 크기에 비해 차이가 작고 생성이 진행되어도 커지지 않는지 확인합니다. 신뢰할 수 있는 실행의 값을 기록하고 torch.testing.assert_close(dev_logits, ref_logits, atol=..., rtol=...)로 회귀 검사를 구성합니다.
  • greedy tokens match는 전체 생성 결과를 확인하는 검사입니다. 기준값의 top-2 gap이 max |diff|보다 명백히 클 때만 토큰 불일치를 결함으로 판단합니다. 두 후보의 값이 거의 같으면 반올림 차이만으로 선택이 바뀔 수 있으며 어느 토큰이든 올바른 결과입니다.
  • 2단계에서 fallback한 연산은 호스트에서 float32로 계산하므로 eager 실행에서는 디바이스 오차를 추가하지 않습니다. 관측한 차이는 디바이스에서 실행한 연산에서 발생하며 이것이 측정하려는 값입니다.

Top-2 gap이 큰데 토큰이 다르거나 sequence 길이가 늘수록 차이가 커지면 TORCH_RBLN_LOG_LEVEL=INFO를 설정하고 TORCH_RBLN_DEPLOY는 해제한 상태로 실행합니다. 이렇게 하면 NaN이나 Inf가 있는 입력도 CPU로 보내져 로그에 기록됩니다. 이후 확인 절차는 문제 해결 › CPU 실행과 다른 결과를 따릅니다.

다음 단계