Flash Attention을 이용한 Llama3.1-8B

개요

이 튜토리얼은 여러 개의 RBLN NPU를 사용하여 vLLM에서 모델을 실행하는 방법을 설명합니다. 이 가이드에서는 meta-llama/Meta-Llama-3-1-8B-Instruct 모델을 사용합니다.

Flash Attention은 메모리 효율성과 처리량(throughput)을 향상시켜, 긴 문맥을 처리하는 모델의 성능을 개선합니다. optimum-rbln에서는 컴파일 시 rbln_kvcache_partition_len 파라미터를 추가함으로써 Flash Attention 모드를 활성화할 수 있습니다.

Note

리벨리온 확장형 설계(Rebellions Scalable Design, RSD)는 ATOM™+ (RBLN-CA12 and RBLN-CA22)와 ATOM™-Max (RBLN-CA25)에서 사용 가능합니다. 현재 사용 중인 RBLN NPU 종류는 rbln-stat 명령어로 확인할 수 있습니다.

Note

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

환경 설정 및 설치 확인

시작하기 전에 시스템 환경이 올바르게 구성되어 있으며, 필요한 모든 필수 패키지가 설치되어 있는지 확인하십시오. 다음 항목이 포함됩니다:

• 시스템 요구 사항:
  • Python: 3.9–3.12
  • RBLN 드라이버
• 필수 패키지:
  • torch
  • transformers
  • numpy
  • RBLN Compiler
  • optimum-rbln
  • huggingface_hub[cli]
  • vllm-rbln
    • vllm - vllm-rbln 설치 시 자동으로 설치됨.
• 설치 명령어:
  

  pip install optimum-rbln>=0.8.1 vllm-rbln>=0.8.1.post1
  pip install --extra-index-url https://pypi.rbln.ai/simple/ rebel-compiler>=0.8.1 
  

Note

rebel-compiler를 사용하려면 RBLN 포털 계정이 필요하니 참고하십시오.

Note

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

$ huggingface-cli login

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

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

Llama3.1-8B 컴파일

먼저, optimum-rbln에서 RBLNLlamaForCausalLM 클래스를 가져옵니다. 이 클래스의 from_pretrained() 메서드는 HuggingFace 허브에서 Llama 3.1 모델을 다운로드하고 RBLN SDK를 사용하여 컴파일합니다. 모델을 익스포트할 때 다음 매개변수를 지정하십시오:

• export: 모델을 컴파일하려면 True로 설정해야 합니다.
• rbln_batch_size: 컴파일을 위한 배치 크기를 정의합니다.
• rbln_max_seq_len: 최대 시퀀스 길이를 정의합니다.
• rbln_tensor_parallel_size: 추론에 사용할 NPU의 수를 정의합니다.
• rbln_kvcache_partition_len: Flash Attention을 위한 KV 캐시 파티션의 길이를 정의합니다.rbln_max_seq_len은 rbln_kvcache_partition_len의 배수여야 하며, 그보다 큰 값이어야 합니다. 컴파일 후에는 save_pretrained() 메서드를 사용하여 모델 아티팩트를 디스크에 저장합니다. 이 과정은 컴파일된 모델을 포함하는 디렉터리(예: rbln-Llama-3-1-8B-Instruct)를 생성합니다.

from optimum.rbln import RBLNLlamaForCausalLM
import os

model_id = "meta-llama/Llama-3.1-8B-Instruct"

# Compile and export
model = RBLNLlamaForCausalLM.from_pretrained(
    model_id=model_id,
    export=True,
    rbln_batch_size=1,
    rbln_max_seq_len=131_072,
    rbln_tensor_parallel_size=8,
    rbln_kvcache_partition_len=16_384,
)

# Save compiled results to disk
model.save_pretrained("rbln-Llama-3-1-8B-Instruct")

Note

모델 크기와 NPU 사양에 따라 적절한 배치 크기를 선택할 수 있습니다. vllm-rbln은 연속 배치(continuous batching)를 지원하므로, 최적의 처리량과 자원 활용을 위해 배치 크기를 적절하게 설정하는 것이 중요합니다. 동적 배칭(dynamic batching)을 설정하고 싶다면 동적 배치 크기를 이용한 추론을 참고하세요.

추론을 위한 vLLM API 사용

컴파일 후에는 vLLM API로 모델을 사용할 수 있습니다:

Note

Flash Attention을 사용하기 위해서는, block_size가 rbln_kvcache_partition_len과 일치해야 합니다.

from transformers import AutoTokenizer
from vllm import LLM, SamplingParams

# Please make sure the engine configurations match the parameters used when compiling.
model_id = "rbln-Llama-3-1-8B-Instruct"
max_seq_len = 131_072
batch_size = 1
block_size = 16_384  # Should match to `rbln_kvcache_partition_len` for flash attention.

llm = LLM(
    model=model_id,
    device="rbln",
    max_num_seqs=batch_size,
    max_num_batched_tokens=max_seq_len,
    max_model_len=max_seq_len,
    block_size=block_size,
)

tokenizer = AutoTokenizer.from_pretrained(model_id)

sampling_params = SamplingParams(
  temperature=0.0,
  skip_special_tokens=True,
  stop_token_ids=[tokenizer.eos_token_id],
)

conversation = [
    {
        "role": "user",
        "content": "What is the first letter of English alphabets?"
    }
]

chat = tokenizer.apply_chat_template(
  conversation, 
  add_generation_prompt=True,
  tokenize=False
)

outputs = llm.generate(chat, sampling_params)
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

vLLM API를 이용해 여러 인코더-디코더 모델이나 멀티모달 모델을 실행할 수 있습니다. 모델 주에서 가능한 모델을 확인하실 수 있습니다.

vLLM API에 대한 더 자세한 내용은 vLLM 문서를 참고하시기 바랍니다.