PyTorch RBLN을 사용한 Transformers 모델 실행: Llama3.2-1B¶

개요¶

본 문서는 HuggingFace Transformers의 Llama3.2-1B 모델을 PyTorch RBLN 환경에서 실행하는 절차를 설명합니다.

Note

설정 및 설치¶

시스템 요구 사항
- Python: 3.10–3.13
- RBLN Driver
주요 패키지 요구 사항
- PyTorch
- transformers: >=4.49.0,<4.50.0
- RBLN Compiler
- torch-rbln
- HuggingFace Hub CLI (huggingface_hub[cli])
설치

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

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

Note

rebel-compiler 설치에는 RBLN Portal 계정이 필요합니다.

Note

HuggingFace Hub에 게시된 meta-llama/Llama-3.2-1B 모델은 접근 권한이 제한되어 있습니다. 사용 권한이 부여된 경우, 아래와 같이 hf(HuggingFace CLI) 명령으로 인증할 수 있습니다.

$ hf auth login

    _|    _|  _|    _|    _|_|_|    _|_|_|  _|_|_|  _|      _|    _|_|_|      _|_|_|_|    _|_|      _|_|_|  _|_|_|_|
    _|    _|  _|    _|  _|        _|          _|    _|_|    _|  _|            _|        _|    _|  _|        _|
    _|_|_|_|  _|    _|  _|  _|_|  _|  _|_|    _|    _|  _|  _|  _|  _|_|      _|_|_|    _|_|_|_|  _|        _|_|_|
    _|    _|  _|    _|  _|    _|  _|    _|    _|    _|    _|_|  _|    _|      _|        _|    _|  _|        _|
    _|    _|    _|_|      _|_|_|    _|_|_|  _|_|_|  _|      _|    _|_|_|      _|        _|    _|    _|_|_|  _|_|_|_|

    To login, `huggingface_hub` requires a token generated from [https://huggingface.co/settings/tokens](https://huggingface.co/settings/tokens) .
Enter your token (input will not be visible):

Llama3.2-1B 예제¶

다음 llama.py 스크립트는 HuggingFace Hub에서 내려받은 Llama3.2-1B 모델을 실행합니다. 연산 장치 지정을 cuda 또는 cpu 대신 rbln으로 변경하는 부분을 제외하면, GPU 또는 CPU 환경에서 사용하는 코드와 동일한 구성입니다.

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}")

동일한 스크립트를 cuda 또는 rbln 장치로 실행할 경우, 아래와 유사한 결과가 출력됩니다.

$ python3 llama.py
Q: Is Sun bigger than Earth?
A: The answer is yes.

Note

RBLN NPU에서의 eager mode: 일반적인 GPU eager mode 스크립트처럼 작성하되, 장치 인자만 "rbln"으로 지정합니다.

Shape 기반 컴파일 캐시: PyTorch 연산이 rbln 장치에서 수행될 때마다 torch-rbln은 rebel-compiler를 호출하여 해당 연산을 RBLN NPU용 실행 파일로 컴파일합니다. 동일한 연산을 동일한 입력 텐서 shape로 다시 실행하는 경우, 전체를 재컴파일하지 않고 기존에 생성된 컴파일 산출물을 재사용합니다.

본 예제에서의 디코딩 동작: 자기회귀(autoregressive) 방식의 생성 과정에서는 토큰 단위로 시퀀스 길이가 증가하므로, 단계마다 텐서 shape가 달라지는 경우가 빈번합니다. 이에 따라 캐시 적중률은 낮아지고 재컴파일 빈도는 높아질 수 있으며, 응답 지연 시간이 증가할 수 있습니다. 다만 모델의 출력 정확도에는 문제가 없습니다. 시퀀스 길이가 고정된 입력 배치를 사용하는 경우에는 컴파일 비용이 상대적으로 잘 분산됩니다.