콘텐츠로 이동

PyTorch RBLN으로 transformers 모델 실행하기: Llama3.2-1B

개요

HuggingFace Transformers의 Llama3.2-1B 모델을 실행해보는 방법에 대한 설명입니다.

Note

Llama 3는 LLAMA 커뮤니티 라이선스에 따라 라이선스가 부여되며, 저작권은 Meta Platforms, Inc.가 소유하며 모든 권리를 보유합니다.

설정 및 설치

Note

torch-rbln을 사용하려면 RBLN Portal 계정이 필요합니다.

Note

HuggingFace Hub의 meta-llama/Llama-3.2-1B 모델은 접근이 제한되어 있습니다. 접근 권한을 부여받은 후, 아래와 같이 huggingface-cli 명령어를 사용하여 로그인할 수 있습니다:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
$ huggingface-cli login

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

    To login, `huggingface_hub` requires a token generated from [https://huggingface.co/settings/tokens](https://huggingface.co/settings/tokens) .
Token: *****

Llama3.2-1B 예제

다음 llama.py는 HuggingFace Hub에서 다운로드한 Llama3.2-1B 모델을 실행하는 예제입니다. cuda나 cpu 대신 rbln을 쓴다는 점을 제외하면 GPU나 CPU에서 돌려보는 코드와 동일합니다.

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
import torch
import re
from transformers import AutoTokenizer, AutoModelForCausalLM

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으로 실행하던 다음과 같이 결과는 같을 것입니다.

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

Note

위 예제는 HuggingFace Hub에서 Llama3.2-1B model을 받아서 실행하는 예제입니다. device를 "rbln"으로 지정하는 것 외에는 GPU eager mode와 동일한 코드를 사용할 수 있습니다. torch-rbln은 각 op의 구현을 rebel_compiler로 컴파일 및 실행하므로, 동일한 입력 텐서가 들어오는 op이 두 번째로 실행될 때는 컴파일 시간이 제외되어 더 빠르게 실행됩니다. 다만, Llama3.2-1B 모델은 토큰을 하나씩 생성할 때마다 입력 시퀀스 길이가 증가하므로 단계마다 새로운 컴파일이 필요합니다. 이로 인해 실행 시간에 컴파일 시간이 포함되어 다소 느리게 동작하지만 기능적으로는 정상적으로 동작합니다.