FP16으로 Transformers Pipeline에 모델 로드하기

Huggingface🤗의 Transformers 라이브러리 중 가장 강력한 기능 하나가 바로 pipeline 모듈이라고 생각한다.


from transformers import pipeline

pipe = pipeline('text-generation', model='beomi/kcgpt2', device=0)

pipe("안녕") ## "안녕하세요."

위와 같이 두세줄의 코드 만으로, 복잡한 모델 불러오기 과정 없이 tokenizer encode, pytorch tensor 변환, 모델 Generate(Forward), 그리고 tokenizer decode 과정을 모두 거친 결과를 보여준다.

한편, 아직까지(transformers==4.21.0)는 pipeline에 모델을 불러오는 경우, 기본적으로 CPU로 로드한 뒤 → GPU로 올리는 과정을 거친다. 또한, 이 과정에서 fp16으로 모델을 불러오지 않아서 CPU에 fp32로 로드 후 fp32로 GPU에 올리게 된다.

이때, 카카오브레인의 KoGPT 6B 모델의 경우 fp32로 로드하면 약 30GB 수준의 VRAM과 메모리를 사용한다.

kakaobrain/kogpt · Hugging Face

KakaoBrain's Pre-Trained Language Models. KoGPT (Korean Generative Pre-trained Transformer) The following is the recommended minimum GPU hardware guidance for a handful of example KoGPT. 32GB GPU RAM in the required minimum memory size The following is the recommended minimum GPU hardware guidance for a handful of example KoGPT.

https://huggingface.co/kakaobrain/kogpt

이를 방지하기 위해 KoGPT 6B는 KoGPT6B-ryan1.5b-float16 라는 fp16 버전의 리비전 모델 weight 파일을 제공하고 있다.

이를 로드하기 위해서는 아래와 같이 torch_dtype='auto', low_cpu_mem_usage=True 두 가지 옵션을 활성화해야 한다.


import torch
from transformers import AutoTokenizer, AutoModelForCausalLM 

tokenizer = AutoTokenizer.from_pretrained(
	'kakaobrain/kogpt', revision='KoGPT6B-ryan1.5b-float16',  # or float32 version: revision=KoGPT6B-ryan1.5b
	bos_token='[BOS]', eos_token='[EOS]', unk_token='[UNK]', pad_token='[PAD]', mask_token='[MASK]'
)
model = AutoModelForCausalLM.from_pretrained(
  'kakaobrain/kogpt', revision='KoGPT6B-ryan1.5b-float16',  # or float32 version: revision=KoGPT6B-ryan1.5b
  pad_token_id=tokenizer.eos_token_id,
  torch_dtype='auto', low_cpu_mem_usage=True
).to(device='cuda', non_blocking=True)
_ = model.eval()

low_cpu_mem_usage=True

이 옵션을 사용할 경우 Huggingface의 Accelerator가 필요하다.
해당 옵션을 켜면, CPU Ram에 올리고 → GPU로 복사하는 대신, 모델의 weight를 곧바로 GPU로 업로드한다.
즉, CPU 메모리 사용량이 극도로 줄어든다!

torch_dtype='auto'

float16 리비전의 경우 모델 weight 역시 fp16 으로 저장되어 있기 때문에, 기본값인 fp32 대신 fp16으로 자동으로 로드할 수 있게 만들어줘야 한다.

위 코드를 통해 모델을 읽어오고 나서는 아래와 같이 토크나이저, 모델의 forward/generate를 거쳐 최종 디코딩된 결과물을 생성할 수 있다.


prompt = '인간처럼 생각하고, 행동하는 \'지능\'을 통해 인류가 이제까지 풀지 못했던'
with torch.no_grad():
  tokens = tokenizer.encode(prompt, return_tensors='pt').to(device='cuda', non_blocking=True)
  gen_tokens = model.generate(tokens, do_sample=True, temperature=0.8, max_length=64)
  generated = tokenizer.batch_decode(gen_tokens)[0]
  
print(generated)  # print: 인간처럼 생각하고, 행동하는 '지능'을 통해 인류가 이제까지 풀지 못했던 문제의 해답을 찾을 수 있을 것이다. 과학기술이 고도로 발달한 21세기를 살아갈 우리 아이들에게 가장 필요한 것은 사고력 훈련이다. 사고력 훈련을 통해, 세상

축약 가능한 버전은 다음과 같다.

위 코드를 아래와 같이 바꾸어주자.


from transformers import pipeline

pipe = pipeline(
	'text-generation', 
	model=model, 
	tokenizer=tokenizer, 
	device=0, 
	framework='pt',
)
prompt = '인간처럼 생각하고, 행동하는 \'지능\'을 통해 인류가 이제까지 풀지 못했던'
pipe(prompt, do_sample=True, temperature=0.8, max_length=64)

이렇게 하면 위와 같이 문장 생성 결과가 나오게 된다.