LLM Trend Note2 (1) Base model and Dataset for RLHF

안녕하세요 여러분:) LLM Trend Note2에 오신 걸 환영합니다!

LLM Trend Note에서 우리는 최신 LLM의 흐름을 살펴보면서
Foundation model의 조건과 Emergent Abilities의 특징에 대해 알아보았습니다.
그기저엔 대규모 분산 컴퓨팅 기술과
보다 혁신적인 모델 아키텍쳐 및 학습기법에 대한 고민이
단단한 기초를 이루고 있다는 사실도 함께 말이죠.

이번 노트에서는 RLHF를 간접적으로 구현한 언어모델에 한국어 말뭉치를 학습시켜보도록 하겠습니다.
이른바 KoChatGPT라고 할 수 있겠죠?
KoChatGPT를 구현한 레퍼런스 코드는 고우영님의 깃헙에서 참고했습니다.

KoChatGPT 학습을 위해 필요한 환경과 라이브러리는 아래와 같습니다.

wget https://developer.download.nvidia.com/compute/cuda/11.6.0/local_installers/cuda_11.6.0_510.39.01_linux.run
sudo sh cuda_11.6.0_510.39.01_linux.run

pip uninstall torch -y
pip install torch==1.13.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116

pip install transformers==4.28.0
pip install colossalai==0.2.7

git clone https://github.com/airobotlab/KoChatGPT

cd KoChatGPT/colossalai_ChatGPT_230319/
pip install .

pip install openai
pip install langchain==0.0.113
pip install pandas>=1.4.1

pip install --upgrade accelerate
pip install bitsandbytes  
pip install loralib

LLM Trend Note에서 보았던
FLAN의 Instruction Tuning이나 PaLM의 Prompt Engineering 이 효과를 보기 위해선
언어모델의 입력을 단순한 query 형태로 주기보단
정교한 입력 시퀀스를 설계해야 한다고 배웠습니다.

예컨대, 작업의 지시사항이 담긴 instruction과
실제 모델이 작업 내용이 담긴 input,
그리고 CoT(Chain of thought) 형태의 예시답안 등을 prompt로 주는 식으로 말이죠.

그런데 이렇게 긴 prompt를 입력할 수 있으려면 그만한 모델 capacity가 뒷받침 되어야 합니다.
일단 수백에서 수천개의 token을 입력 벡터로 받아낼 수 있어야 하고,
각 토큰에 대한 충분한 셀프어텐션 연산이 가능한 트랜스포머 모듈이 쌓여야 하죠.
우리가 backbone 모델로 사용할 KoGPT-2의 성능을 잠시 확인해볼까요?

허깅페이스의 transformers를 사용하면 토크나이저와 모델을 간단히 불러올 수 있습니다.

import torch  
from transformers import AutoTokenizer, AutoModelForCausalLM  
import pandas as pd  
import numpy

device = "cuda" if torch.cuda.is\_available() else "cpu"  
model\_name = "skt/kogpt2-base-v2"  
tokenizer = AutoTokenizer.from\_pretrained(model\_name)  
model = AutoModelForCausalLM.from\_pretrained(model\_name).to(device)

우리가 사용할 모델의 토크나이저가 입력받아 처리할 수 있는 최대 토큰 수를 확인해봅시다.

tokenizer.max_model_input_sizes  

kogpt-2는 어떻게 토크나이징을 하는지 잠시 확인해 볼까요?

input_txt = "바람도 없는 공중에 수직의 파문을 내이며 고요히 떨어지는 오동잎은 누구의 발자취 입니까."  

tokens = tokenizer(input_txt).tokens()  
input_ids = tokenizer(input_txt, return_tensors="pt")["input_ids"].numpy()  

pd.options.display.max_columns = 40  
pd.options.display.max_rows = 60  
df = pd.DataFrame([tokens, input_ids[0]], index=["kogpt-2_tokens", "Input_IDs"])  
df  

내친 김에 디코딩 성능도 확인해보겠습니다.

max_length = 128
input_ids = tokenizer(input_txt, return_tensors="pt")["input_ids"].to(device)  
output_greedy = model.generate(input_ids, max_length=max_length, do_sample=False)  
print(tokenizer.decode(output_greedy[0]))  

시퀀스가 반복되어 출력되는군요.
그리디 서치 디코딩시 발견되는 전형적인 현상입니다.

이번엔 빔 서치 디코딩을 사용하고 n-gram 패널티까지 부과해보겠습니다.

input_ids = tokenizer(input_txt, return_tensors="pt")["input_ids"].to(device)
output_beam = model.generate(input_ids, max_length=max_length, num_beams=10, no_repeat_ngram_size=2,
                             do_sample=False)
print(tokenizer.decode(output_beam[0]))  

입력 시퀀스와 별 상관 없어 보이는 긴 문단이 생성됩니다.
그럼에도 생성된 문단은 제법 맥락을 갖춘 듯 보입니다.
하지만 문장 간의 정합성이나 일관성은 다소 떨어지는 부분도 관찰됩니다.
이번엔 샘플링 기법까지 추가해 보겠습니다.

output_beam = model.generate(input_ids, max_length=max_length, num_beams=7, no_repeat_ngram_size=2,
                             do_sample=True, temperature=2.0, top_k=50)
print(tokenizer.decode(output_beam[0]))  

generate함수의 인자로 사용한 temperature, top_k 값은 어떤 효과를 주는 옵션인가요?

top_p 샘플링 기법도 사용해보겠습니다.

output_beam = model.generate(input_ids, max_length=max_length, num_beams=7, no_repeat_ngram_size=2,
                             do_sample=True, top_p=0.90)
print(tokenizer.decode(output_beam[0]))  

top_p 인자의 기능은 무엇인가요?

최선의 디코딩 방법을 찾기 위해선
빔사이즈와 n-gram 패널티, temperature와 샘플링 인자로
조합할 수 있는 최선의 값을 찾아보는 실험이 필요합니다.
다양한 입력 시퀀스로 실험을 해보면 더 좋겠죠?

베이스라인 모델로 사용한 kogpt-2의 일반적인 성능을 확인해봤으니
구체적인 instruction과 prompting을 사용해 어떻게 디코딩을 해내는지도 확인해보세요.
RLHF를 적용하기 전의 실험값을 정리해보면
KoChatGPT의 성능 개선 여부를 확인하는데 도움이 될 것입니다.

우리가 사용하는 kogpt-2는 오리지널 GPT2의 가장 작은 버전입니다.

emergent abilities를 기대할 수 있는 foundation model로는 많이 부족한 감이 있죠.

위에서 살펴본 것처럼
단순한 Causal LM에 불과한 kogpt-2는 생성해 낼 문장의 품질을
디코딩 단계에서 인위적으로 조절해주는 방법 밖에 쓸 수 없습니다.
하지만 RLHF를 kogpt-2에 적용한다면
더 좋은 문장을 생성해내는 방법을 모델이 스스로 학습해낼 수 있으리라 기대해 볼 수 있습니다.
고도의 prompting은 어렵겠지만
현재 상태에서 특정 task에 fine-tuning했을 때보다는 성능이 한층 더 개선될 수 있지 않을까요?

kogpt-2에 RLHF를 적용하기기 위해선
새로운 데이터셋으로 일련의 재학습을 해줘야 합니다.
그럼 각 단계별 모델 구현에 앞서 우리가 사용할 데이터셋을 확인해보도록 하겠습니다.

먼저 SFT를 시도할 initial 모델에 쓸 데이터셋을 살펴보겠습니다.

데이터셋 확인

import json 
data_path_1_SFT = './data_kochatgpt/kochatgpt_1_SFT.jsonl' 
with open(data_path_1_SFT, "r", encoding='utf-8-sig') as json_file:
    list_data_dict = json.load(json_file)

print(len(list_data_dict))
list_data_dict[:3]  

다음으로 RM에 사용할 데이터셋을 살펴보겠습니다.

data_path_2_RM = './data_kochatgpt/kochatgpt_2_RM.jsonl'
with open(data_path_2_RM, "r", encoding='utf-8-sig') as json_file:
    list_data_dict = json.load(json_file)

print(len(list_data_dict))
list_data_dict[:3]  

마지막으로 PPO 학습에 쓰일 데이터를 살펴보겠습니다.

data_path_3_PPO = './data_kochatgpt/kochatgpt_3_PPO.jsonl'
with open(data_path_3_PPO, "r", encoding='utf-8-sig') as json_file:
    list_data_dict = json.load(json_file)

print(len(list_data_dict))
list_data_dict[:3]

베이스라인 모델과 데이터셋을 확인했으니
다음 노트에서 본격적으로 진행해보도록 하겠습니다.