LLM Trend Note2 (4) PPO(Proximal Policy Optimization)

PPO

드디어 RLHF의 마지막 세번째 단계인 PPO를 실습해볼 차례입니다.
사용할 라이브러리들을 불러오도록 하겠습니다.

from copy import deepcopy

import torch
from torch.optim import Adam
from chatgpt.models.base import RewardModel
from chatgpt.models.gpt import GPTActor, GPTCritic
from chatgpt.trainer import PPOTrainer
from chatgpt.trainer.strategies import NaiveStrategy
from transformers import AutoTokenizer

여기에서 소개하는 KoChatGPT의 경우
PPO에 사용할 actor모델은 1단계 SFT 모델을,
critic모델은 2단계 RM 모델을 사용합니다.

그리고 actor 모델이 critic 모델로부터 피드백을 받아 파라미터를 업데이트 할 때
적절한 페널티를 줄 수 있도록 하는 initial model은
SFT모델을 그대로 freezing 하여 사용합니다.

토크나이저는 pretrain 모델인 kogpt-2의 토크나이저를 그대로 사용해야겠죠?

with NaiveStrategy().model_init_context():
    actor = GPTActor(pretrained='./output_1_SFT', lora_rank=0).to(torch.cuda.current_device())
    critic = GPTCritic(pretrained='./output_2_RM', lora_rank=0).to(torch.cuda.current_device())

    tokenizer = AutoTokenizer.from_pretrained(
        'skt/kogpt2-base-v2', bos_token='</s>', eos_token='</s>', unk_token='</s>', pad_token='</s>',
        padding_side="right", 
        model_max_length=512
    )

    initial_model = deepcopy(actor)
    reward_model = RewardModel(deepcopy(critic.model), deepcopy(critic.value_head)).to(torch.cuda.current_device())

모델학습에 사용할 옵티마이저와 모델을 준비합니다.

actor_optim = Adam(actor.parameters(), lr=5e-6)
critic_optim = Adam(critic.parameters(), lr=5e-6)

(actor, actor_optim), (critic, critic_optim), reward_model, initial_model = NaiveStrategy().prepare(
    (actor, actor_optim), (critic, critic_optim), reward_model, initial_model)

PPO 학습에 쓸 데이터를 불러와 토크나이징 해줍니다.

with open('./data_kochatgpt/kochatgpt_3_PPO.jsonl', "r", encoding='utf-8-sig') as json_file:
    list_data_dict = json.load(json_file)
    list_prompt = [tmp['prompt'] for tmp in list_data_dict]

def tokenize_fn(texts):
    batch = tokenizer(texts, return_tensors='pt', max_length=96, padding=True, truncation=True)
    return {k: v.cuda() for k, v in batch.items()}

print(tokenize_fn(“It takes something more than intelligence to act intelligently.”))

len(list_prompt)

PPO는 별도의 PPOTrainer 클래스를 설계하여 학습시켜줘야 합니다.
빠르게 실습해보기 위해 1epoch만 돌려보겠습니다.

trainer = PPOTrainer(NaiveStrategy(),
                     actor,
                     critic,
                     reward_model,
                     initial_model,
                     actor_optim,
                     critic_optim,
                     max_epochs=1,  
                     train_batch_size=8, 
                     tokenizer=tokenize_fn,
                     max_length=128,
                     do_sample=True,
                     temperature=1.0,
                     top_k=50,
                     pad_token_id=tokenizer.pad_token_id,
                     eos_token_id=tokenizer.eos_token_id)

위 코드블럭의 원본 코드는 chatgpt/trainer 폴더 내의 ppo.py 모듈에서 확인할 수 있습니다.
PPO는 SFT, RM 보다 훨씬 복잡한 단계로 설계되는 강화학습 알고리즘입니다.
PPO의 loss function은 chatgpt/models 폴더 내의 loss.py 모듈에서
PolicyLoss와 ValueLoss 클래스에 정의되어 있습니다.

이 링크의 3 - Fine-tuning the LM with RL 부분을 읽어보고
PolicyLoss클래스의 forward메서드에 있는 코드와 비교해보세요.

PPO에 대해 좀 더 자세히 알고 싶으신 분들은 허깅페이스에서 제공하는 Deep RL Course를 참고하세요.

이제 PPO 학습을 진행하도록 하겠습니다.

trainer.fit(list_prompt,  
            num_episodes=10,  
            max_timesteps=3,
            update_timesteps=3)

model.save_pretrained('./output_3_PPO')

드디어 SFT, RM 그리고 PPO 학습이 모두 완료되었습니다.
RLHF가 적용된 koGPT-2의 생성능력을 확인해볼까요?

def generation(input_text):
    input_ids = tokenizer.encode(input_text, return_tensors='pt').to(
        torch.cuda.current_device())
    outputs = actor.generate(input_ids,
                             max_length=250,
                             do_sample=True,
                             top_k=50,
                             top_p=0.95,
                             num_return_sequences=1)
    output = tokenizer.batch_decode(outputs[0], skip_special_tokens=True)[0]
    print()
    print(output)
    return output

PROMPT_DICT = {
    "prompt_input": (
        "### Instruction(명령어):\n{prompt}\n\n### Response(응답):"
    )
}

list_prompt = [
    '불고기용 고기 한우에요?', 
    '리처드 닉슨이 43대 부통령직을 수행한 년도는?', 
    '시카고 오헤어 국제공항은 어디에 있어',
    '오늘 미세먼지 어때?']

list_prompt = [PROMPT_DICT['prompt_input'].format_map({'prompt': tmp}) for tmp in list_prompt]

for input_text in list_prompt:
    output = generation(input_text)

최종 모델 학습 결과는 어떠셨나요?
중복 토큰 생성 문제를 비롯해 맥락에서 벗어난 한문이나 영문이 출력되기도 합니다.
kogpt-2에 SFT만 적용했을 때와 비교해보면 큰 차이를 느끼지 못할 수도 있습니다.
각 단계에서 사용되는 데이터셋을 충분히 정제하고, 훈련 사이클을 늘려 정교하게 디코딩한다면
훨씬 나은 성능을 기대해볼 수 있습니다.

RLHF의 진가는 고도로 정제된 instruction dataset와 정교하게 설계된 보상체계로 학습되는 Reward model,
그리고 PPO 학습이 안정적으로 이뤄질 수 있도록 하는 충분한 크기의 foundation model이 뒷받침 되었을 때 발휘될 수 있습니다.

그러나 진짜 인간의 피드백이 반영된 데이터셋을 구축하기란 아주 어려운 일입니다.
첫 단계부터 큰 난관이라 할 수 있죠.
하지만 SFT만으로도 굉장히 훌륭한 언어모델을 만들 수 있습니다.
이준범님이 공개한koalpaca 모델이 대표적인 사례라 할 수 있습니다.

충분한 컴퓨팅 파워가 마련된 여건이라면 koalpaca모델로 RLHF까지 진행해 볼 수도 있을 것 같습니다.

그 전에 KoChatGPT 소스코드를 바탕으로 다양한 모델 개선 전략을 선택해 KoChatGPT를 업그레이드해 볼 수도 있지 않을까요?

아래 제시된 옵션 중 하나를 선택하거나 여러 개를 조합하여
자기만의 custom ChatGPT를 개발해 볼 수 있지 않을까 조심스레 제안을 해봅니다.

 1.  우리가 살펴본 KoChatGPT 모델에 사용한 데이터셋은 아직 완벽히 정제되지 않았습니다.
 2.  Hunman Feedback이 반영된 데이터셋을 대체하기 위해
     SFT와 RM 모델에 사용할 다양한 benchmark 데이터셋도 검토해볼 수 있습니다.
 3.  언어모델의 생성능력을 좌우하는 최선의 디코딩을 위한 하이퍼파라미터 서치가 필요합니다.
 4.  생성된 답변에 대한 주관적인 평가를 보완할 수 있는 정량적인 메트릭은 도입하지 않았었습니다.
 5.  LLM Trend Note1에서 살펴본 다양한 Instruction Tuning 및 Prompting 기법들도 적용해볼만 합니다.
 6.  무엇보다 foundation model로 사용한 KoGPT-2는 Emergent abilities를 기대하기엔 다소 작은 사이즈의 모델입니다.
     더 큰 파라미터 스케일을 가진 모델을 사용해보거나,
 7.  더 효율적인 연산을 수행할 수 있는 LoRA의 적용 또는
     새로운 Instruction Tuning 및 reward ranking 알고리즘을 도입해볼 수도 있습니다.

몇 가지 예시를 구체적으로 적어보겠습니다.

기존 데이터셋 추가 정제

data_kochatgpt 폴더에는 세 파일이 있습니다.
ㄱ. kochatgpt_1_SFT.jsonl : SFT를 위한 prompt와 completion 문장셋
ㄴ. kochatgpt_1_RM.jsonl : RM 학습을 위한 prompt와 세 가지 ranking 문장셋
ㄷ. kochatgpt_1_PPO.jsonl : promt 문장

각 말뭉치를 EDA하여 도메인과 문체, 길이분포, 문장의 완성도 등을 분석합니다.
언어모델의 문장생성능력은 말뭉치의 전처리 수준에 큰 영향을 받습니다.
말뭉치의 분석결과를 토대로 데이터를 정제하여 모델을 재학습시켜봅니다.
(정제후 데이터셋 크기가 줄어들지 않도록, 다양한 augmentation 기법을 활용하여 크기를 유지 내지 증량합니다.)
추가 전처리 후, 기존 인퍼런스 결과와 성능을 비교해봅니다.
(주관적인 평가와 BLEU, ROUGE 등을 활용한 정량적인 평가 결과를 비교 분석하여 제시합니다.)

새로운 데이터셋 추가

KoChatGPT는 human feedback이 반영된 데이터를 직접 사용하는 대신
ChatGPT API를 사용하는 대안을 선택했습니다.
LLM Trend Note1 에서 살펴보았듯이
Anthropic의 RLHF는 StackExchange 같은 온라인 상의 댓글정보를 활용하여
ranking dataset을 구축해 구현되었습니다.
우리도 비슷한 로직을 적용해볼 수 있습니다.

하나의 prompt에 대한 다양한 수준의 품질로 댓글이 달린 한국어로 된 웹사이트를 찾아봅시다.
웹크롤링 기법을 사용해 reward 점수를 차등적으로 적용해볼 수 있는
instruction dataset과 ranking dataset을 구축해봅니다.

KorQuAD 2.0 같은 한국어 이해 benchmark를 활용해 고품질의 데이터셋을 확보하고,
KoGPT-2를 사용해 빠르게 저품질 데이터셋을 페어링해볼 수도 있습니다.
다양한 데이터 증량전략을 구사하여 기존 데이터셋에 새로 구축한 데이터셋을 추가해
모델을 재학습시키고 추론 결과를 비교하 분석하여 제시해보세요.

foundation model 교체

현재 제공되는 LMS GPU 사양으로는 수십 billion 단위 이상의 LLM을 튜닝하기 어렵습니다.
그러나 허깅페이스에서 제공하는 큰 규모의 모델을 적은 컴퓨팅 자원으로도 사용할 수 있게 해주는
경량화, 최적화 라이브러리를 사용하면
속도는 느리지만 우리의 LMS에서도 학습 및 추론이 가능해질 수 있습니다.
(힌트 : LLM Trend Note (6)을 참고해보세요)

허깅페이스에서 제공되는 1.2B 사이즈의 한국어 GPT pretrain model로 skt/ko-gpt-trinity-1.2B-v0.5 가 있습니다.
해당 모델로 foundation model을 교체해보세요.
(단 OOM 문제를 해소하기 위해 허깅페이스에서 제공하는
다양한 training argument들을 조합하여 최상의 하이퍼파라미터를 찾아내야 합니다.)
데이터셋을 아예 바꿔 모델 선택의 폭을 늘려보는 것도 좋은 선택지입니다.

foundation model 교체에 성공했다면, generator 함수를 수정하여 모델 인퍼런스 결과를 제시해보세요.

참고
LLM Trend Note2에서 살펴본 KoChatGPT 소스코드는
빠르게 baseline모델을 설계해 실습해보기 위해 오리지널 코드를 일부 수정한 버전입니다.
프로젝트 진행을 위해 모델을 커스터마이징할 때, 필요시 "colossalai_ChatGPT_230319" 폴더 내의 원본 스크립트들을 참고하세요.