본문 바로가기
python/딥러닝

vLLM 사용방법 with LoRA

by _avocado_ 2024. 6. 18.

vLLM은 LLM 모델의 간단한 GPU 병렬처리등을 제공하며 허깅페이스 대비 최대 24배 빠른 추론 속도를 높일 수 있는 라이브러리다.


1. 설치방법

pip install vllm

2. 모델 로드

로컬 모델 및 허깅페이스 모델 모두 사용가능

사용할 GPU의 개수를 입력하여 병렬처리 가능

from vllm import LLM
import torch

model_path = 'mymodel/mymodel'
llm = LLM(
	model=model_path, # 로드할 모델 path
    dtype=torch.float16, # dtype
    tensor_parallel_size=2, # 병렬처리 GPU 개수
    enforce_engaer=True # CUDA 그래프 사용여부
    )

3. quantization

vllm에서 아직 bitsandbytes qunatization 모델은 서빙을 제공하지 않는다. awq 방식의 quantization을 제공하고 있으며 AWQ 라이브러리를 통해 쉽게 quantization이 가능하다.

 

3-1. AWQ install

pip install autoawq

 

3-2. awq quantization

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_path = 'mymodel/mymodel'
quant_path = 'mymodel/mymodel_qauntization_awq'
quant_config = { "zero_point": True, "q_group_size": 128, "w_bit": 4, "version": "GEMM" }

# Load model
model = AutoAWQForCausalLM.from_pretrained(model_path, **{"low_cpu_mem_usage": True})
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

# Quantize
model.quantize(tokenizer, quant_config=quant_config)

# Save quantized model
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)

 

3-3. quantization model 로드하기

from vllm import LLM
import torch

llm = LLM(
    model_path,
    dtype=torch.float16,
    quantization="AWQ", # quantization 종류
    tensor_parallel_size=2,
    enforce_eager=True
    )

4. 결과 생성

SamplingParams를 이용하여 generation config를 입력하고 결과를 생성할 수 있다.

from vllm import SamplingParams

# generation config
sampling_params = SamplingParams( 
    temperature=0,
    max_tokens=512,
    stop=["[/assistant]"],
) 

output = llm.generate(
            prompt,
            sampling_params
        )

5. LoRA adapter 적용

LoRA adapter를 적용하는 방법은 다른 라이브러리에선 모델 로드 시에 적용하는 경우가 많다. 하지만 vllm에서는 모델 로드시에 LoRA를 적용할 수 있는 parameter들을 넣어주고 답변생성 시에 adapter를 입력하여 적용한다. 

from vllm import LLM, SamplingParams
from vllm.lora.request import LoRARequest
import torch

model_path = "mymodel/mymodel"
adapter_path = "mymodel/adapter"

llm = LLM(
    model=model_path,
    dtype=torch.float16,
    quantization="AWQ",
    tensor_parallel_size=2,
    enable_lora=True, # LoRA 적용 여부
    max_lora_rank=64, # LoRA 최대 r 값
    enforce_eager=True,
    )
    
    output = llm.generate(
            prompt,
            sampling_params,
            lora_request=LoRARequest("adapter", 1, adaper_path), # LoRA 적용
        )

댓글