RLHF(Reinforcement Learning from Human Feedback)_(2) - RM(Reward Model)

이번 포스팅에서는 지난번 포스팅에 이어서 RM(Reward Model)에 대해서 이야기해보겠습니다.

 

RM의 역할을 단순하게 정의하면 '모델의 방향성'으로 이야기할 수 있습니다. 조금 다르게 이야기해보면 거대한 하나의 자연어 모델(LLM)이 한 명의 아이라면 이를 교육 시켜주는 사람이 RM입니다. 그리고 그 교육하는 방법이 지난번 포스팅한 SFT입니다. 이 과정이 반복이 되면 하나의 특정 지식에 특화된 자연어 모델이 완성됩니다. 이를 통해서 완성된 모델이 바로 이루다 모델입니다. 카톡 대화를 통해서 훈련하고 RM을 통해 어떤 대화가 좋은 대화인가에 대해서 지속적인 강화 학습이 된 모델입니다.

 

이번에 알아볼 RM은 ColossalAI의 RM모델입니다. 해당 모델은 LoRA를 부모 클래스로 상속받아 만들어진 리워드 모델입니다. 이번 포스팅에서 LoRA에 대해서 다루지 않지만 간단하게 이야기해보겠습니다. LoRA(Low-Rank Adaptation of Large Language Models)는 2022년에 발표된 논문으로 대용량 자연어 모델(LLM)들이 탄생하였고 이제는 이를 사용자의 목적에 맞게 downstream task하여 용도에 맞게 사용하게 되는데 이때 LLM의 모든 파라미터를 업데이터하면서 모델을 맞춰나가는것은 비효율적이기에 이를 타개하고자 고안된 방법론입니다. 자세한 내용은 위의 LoRA링크와 논문을 읽어보시면 확인하실 수 있습니다.

 

ColossalAI에서는 LoRA Moduel을 상속받아 사용하는데 핵심은 'value_head'라는 변수에 있습니다. 해당 변수는 아래와 같이 정의되어 있음을 확인할 수 있습니다. 해당 코드는 모델의 임베딩 벡터를 1차원으로 바꿔 출력하여 하나의 값을 만드는데 이게 RM이 우리에게 최종적으로 전달해줄 입력 문장에 대한 Score가 됩니다.

self.value_head = nn.Linear(model.config.n_embd, 1)

 

 

RM 역시 하나의 모델이기 때문에 베이스가 되는 모델이 있고 그 모델을 학습 시킴으로서 RM이 다른 모델에서 얻은 결과에 대한 점수 계산을 할 수 있게 됩니다. 즉, RM은 초기에 이 모델이 어떤 방향으로 타 모델의 결과를 평가할 것인지에 대해서 사람이 직접 평가한 데이터를 기반으로 학습한다는 것이죠. 이 부분이 바로 HF(Human Feedback)파트를 담당하게 됩니다. 처음에는 실제로 사람이 피드백한 결과를 바탕으로 RM모델을 훈련시키지만 데이터가 충분히 많았다면 이후에는 RM이 스스로 그 사람이 판단한 것과 같은 스코어를 계산 함으로서 더 좋은 문장이 무엇인지에 대해서 결론을 내려준다는 것이죠. 어떻게 보면 RM이 하나이 정체성이라고 볼 수 있습니다. RM을 훈련시키는 것은 결국 사람의 판단 결과를 바탕으로 학습하기 때문이죠. HF가 편향적이라면 그 편향을 그대로 반영하여 RM은 LLM 강화 학습에 적용하게 됩니다. 이 과정을 단순히 생각해보면 그냥 사람이 하나의 정체성, 나라는 사람의 성격을 형성해 나가는 과정과 유사합니다.

 

그러면 이제 RM모델이 어떻게 훈련되고 입력 문장에 대한 점수화를 하는지 코드를 통해서 자세히 알아보도록 하겠습니다.

LM 실습(KoGPT2)

해당 클래스는 chatgpt.models.base에 있는 RewardModel을 상속 받아 커스텀한 RM 모델입니다. 모델을 전달 받았다면 해당 모델을 불러오고 그렇지 않았따면 config에서 정의된 모델을 가져오도록 되어 있습니다. 그리고 위에서 언급한 핵심적인 부분은  value_head로 모델의 임베딩층을 1차원으로 바꾼 결과를 value_head에 저장하게되고 이게 최종적으로 score가 됩니다.

class GPTRM_custom(RewardModel):

    def __init__(self,
                 pretrained: Optional[str] = None,
                 config: Optional[GPT2Config] = None,
                 checkpoint: bool = False,
                 lora_rank: int = 0,
                 lora_train_bias: str = 'none',
                 tokenizer=None) -> None:
        if pretrained is not None:
            model = GPT2Model.from_pretrained(pretrained)
            model.resize_token_embeddings(len(tokenizer))
        elif config is not None:
            model = GPT2Model(config)
        else:
            model = GPT2Model(GPT2Config())
        if checkpoint:
            model.gradient_checkpointing_enable()

        value_head = nn.Linear(model.config.n_embd, 1)
        super().__init__(model, value_head, lora_rank, lora_train_bias)

        if pretrained is not None:
            self.model = model
            self.pretrained = pretrained


    def save_pretrained(self, dir):
        if self.pretrained is not None:
            self.model.save_pretrained(dir)

 

사전에 설정한 strategy에 맞춘 tokenizer와 model초기화를 진행합니다.

with strategy.model_init_context():
    # load pretrained gpt2
    if args.model == 'gpt2':
#         tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
        # tokenizer = AutoTokenizer.from_pretrained(args.pretrain)
        tokenizer = AutoTokenizer.from_pretrained(args.pretrain, padding_side="right", model_max_length=512)
        tokenizer.add_special_tokens(
            {
                "eos_token": DEFAULT_EOS_TOKEN,
                "bos_token": DEFAULT_BOS_TOKEN,
                "unk_token": DEFAULT_UNK_TOKEN,
            }
        )
        tokenizer.pad_token = tokenizer.eos_token
        model = GPTRM_custom(pretrained=args.pretrain, lora_rank=args.lora_rank, tokenizer=tokenizer).cuda()
       
       ...

 

RM모델을 학습시키기 위한 데이터를 불러와 그 구조를 확인합니다. 이번 포스팅에서 사용하는 데이터는 KoChatGPT-replica의 RM모델 훈련용 데이터를 사용하였습니다. prompt는 입력 문장이고 이를 통한 답변은 completion_0~2번이 됩니다. 그리고 해당 문장들이 잘 답변한지에 대해서는 0~2로 ranking(순위)로 평가합니다. 이제 이 구조를 RM모델에 훈련할 수 있도록 형태를 변형합니다.

with open(args.data_path_2_RM, "r", encoding='utf-8-sig') as json_file:
    list_data_dict = json.load(json_file)
    if args.verbose:
        print('## data check ##')
        print((list_data_dict[0]))
        
#출력값
## data check ##
{'prompt': '번디는 자신이 탐정잡지, 범죄소설 그리고 성범죄 관련 실제 범죄 다큐멘터리들을 탐독했다고 누구에게 말했나?',
'completion_0': 'Allow me to answer your question. I know that you are curious about me.',
'completion_1': '번디는 다양한 인터뷰자들과 뉴스홍보 담당자들과의 면담 때 밝혔다.',
'completion_2': '라이언에게 말했다.',
'ranking': [2, 1, 0]}

 

위에서 확인한 데이터 구조를 통해서 prompt는 prompt그대로 저장합니다. 다른점은 3개를 동시에 비교하는게 아니라 [0,1], [0,2], [1,2]로 단일 비교를 통해서 모델이 학습할 데이터를 만들어 냅니다. chosen이 더 나은 문장, rejceted가 좋지 못한 문장으로 모델에게 알려줄 문장입니다. 즉, 가장 좋은 문장이 데이터 안에 있지만 그렇지 못한 두 개의 문장으로 통해서도 그 중에서 누가 더 좋은 문장인가에 대해서도 학습하는 데이터를 만들어준다는 것을 확인하시면 됩니다.

total_data_ranking2chosen = []
for tmp in list_data_dict:
    one_data_ranking2chosen = []

    # data 1) 0 VS 1
    data = {}
    data['prompt'] = tmp['prompt']
    if tmp['ranking'][0] < tmp['ranking'][1]:
        data['chosen'] = tmp['completion_0']
        data['rejected'] = tmp['completion_1']
    else:
        data['chosen'] = tmp['completion_1']
        data['rejected'] = tmp['completion_0']
    one_data_ranking2chosen.append(data)


    # data 2) 0 VS 2
    data = {}
    data['prompt'] = tmp['prompt']
    if tmp['ranking'][0] < tmp['ranking'][2]:
        data['chosen'] = tmp['completion_0']
        data['rejected'] = tmp['completion_2']
    else:
        data['chosen'] = tmp['completion_2']
        data['rejected'] = tmp['completion_0']
    one_data_ranking2chosen.append(data)

    # data 1) 1 VS 2
    data = {}
    data['prompt'] = tmp['prompt']
    if tmp['ranking'][1] < tmp['ranking'][2]:
        data['chosen'] = tmp['completion_1']
        data['rejected'] = tmp['completion_2']
    else:
        data['chosen'] = tmp['completion_2']
        data['rejected'] = tmp['completion_1']
    one_data_ranking2chosen.append(data)



    total_data_ranking2chosen.extend(one_data_ranking2chosen)

 

 

위에서 RM 훈련을 위한 데이터 형태로 변환을 시켰으니 이제 모델에 넣을 데이터 형태로 만들기 위해서 RewardDataset을 사용하여 변환해주고 검증 데이터로 일부 데이터를 떼어내어 만들었습니다. 이후에는 RewardModelTrainer를 통해 훈련을 진행시키면 RM모델에 대한 학습이 끝나게 됩니다.

 

import random
random.seed(230319)
random.shuffle(total_data_ranking2chosen)


train_data = total_data_ranking2chosen[:-1000]  # 29000 학습
eval_data = total_data_ranking2chosen[-1000:0]  # 1000개만 평가

train_dataset = RewardDataset(train_data, tokenizer, args.max_len)
eval_dataset = RewardDataset(eval_data, tokenizer, args.max_len)

trainer = RewardModelTrainer(model=model,
                             strategy=strategy,
                             optim=optim,
                             train_dataset=train_dataset,
                             eval_dataset=eval_dataset,
                             batch_size=args.batch_size,
                             max_epochs=args.max_epochs)
                             
trainer.fit(use_lora=args.lora_rank)

 

훈련을 통해 나온 RM에 입력 문장에 대해서 인코딩후 모델에 넣어주면 문장에 대한 Score를 얻을 수 있습니다.  '인공지능은 무엇인가?'에 대한 답변으로 두 개의 예시 문장을 PLM이 내놓았다고 가정해보았고 이를 RM모델에 적용시키면 어떻게 되는지 확인해보았습니다. RM은 사전 학습된 결과로 두 번째 문장이 좀 더 나은 문장이라는 것을 Score로 보여주었습니다. 이를 강화 학습을 진행하여 RM이 제시해 준 지표를 바탕으로 좀 더 강한 편향을 가지는 모델로 학습하게 됩니다.

# 보상모델 체크
def inference_RM(input_text):
    input_ids = tokenizer.encode(input_text, return_tensors='pt').to(
        torch.cuda.current_device())
    output = model(input_ids)
    output_reward = output.cpu().detach().numpy()[0]

    print('input: %s\nreward score: %.1f'%(input_text, output_reward))

    return output_reward



input_text1 = '인공지능은 인공지능 입니다'
input_text2 = '인공지능(AI)은 다양한 분야에서 활용되고 있습니다. 음성, 이미지, 자연어처리 등 생성AI가 가장 큰 주목을 받고 있습니다.'

output_reward1 = inference_RM(input_text=input_text1)
output_reward2 = inference_RM(input_text=input_text2)

#출력값
input: 인공지능은 인공지능 입니다
reward score: -0.3

input: 인공지능(AI)은 다양한 분야에서 활용되고 있습니다. 음성, 이미지, 자연어처리 등 생성AI가 가장 큰 주목을 받고 있습니다.
reward score: 0.9

 

---

LM 코드 전문

import argparse

import os
import json
from typing import Optional
import torch
import torch.nn as nn
from torch.optim import Adam
from chatgpt.dataset import RewardDataset
from chatgpt.models.base import RewardModel
from chatgpt.trainer import RewardModelTrainer
from chatgpt.trainer.strategies import NaiveStrategy
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModel, AutoConfig
from transformers.models.gpt2.configuration_gpt2 import GPT2Config
from transformers.models.gpt2.modeling_gpt2 import GPT2Model
import loralib as lora

import os
import json

# config
IGNORE_INDEX = -100
DEFAULT_PAD_TOKEN = "[PAD]"
DEFAULT_EOS_TOKEN = "</s>"
DEFAULT_BOS_TOKEN = "</s>"
DEFAULT_UNK_TOKEN = "</s>"
PROMPT_DICT = {
    "prompt_input": (
        "Below is an instruction that describes a task, paired with an input that provides further context.\n"
        "아래는 작업을 설명하는 명령어와 추가적 맥락을 제공하는 입력이 짝을 이루는 예제입니다.\n\n"
        "Write a response that appropriately completes the request.\n요청을 적절히 완료하는 응답을 작성하세요.\n\n"
        "### Instruction(명령어):\n{prompt}\n\n### Input(입력):\n{input}\n\n### Response(응답):"
    ),
    "prompt_no_input": (
        "Below is an instruction that describes a task.\n"
        "아래는 작업을 설명하는 명령어입니다.\n\n"
        "Write a response that appropriately completes the request.\n명령어에 따른 요청을 적절히 완료하는 응답을 작성하세요.\n\n"
        "### Instruction(명령어):\n{prompt}\n\n### Response(응답):"
    ),
}

#Argment
parser = argparse.ArgumentParser()
parser.add_argument('--output_dir', type=str, default='./output_2_RM')
parser.add_argument('--data_path_2_RM', type=str, default='./data_kochatgpt/kochatgpt_2_RM.jsonl', help='https://huggingface.co/datasets/fka/awesome-chatgpt-prompts/blob/main/prompts.csv')
parser.add_argument('--strategy',
                    choices=['naive', 'ddp', 'colossalai_gemini', 'colossalai_zero2'],
                    default='naive')
parser.add_argument('--model', type=str, default='gpt2', choices=['gpt2', 'bloom', 'opt'])
parser.add_argument('--pretrain', type=str, default=None)
parser.add_argument('--dataset', type=str, default='Dahoas/rm-static')
parser.add_argument('--save_path', type=str, default='rm_ckpt.pth')
parser.add_argument('--max_epochs', type=int, default=10)
parser.add_argument('--batch_size', type=int, default=4)
parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank")
parser.add_argument('--max_len', type=int, default=512)  # wygo 추가

args = parser.parse_args(args=[])


args.pretrain = 'skt/kogpt2-base-v2'  # pretrained 모델 가져오기
args.verbose = True

print(args)
if not os.path.exists(args.output_dir):
    os.makedirs(args.output_dir)
    
    
 # strategy
if args.strategy == 'naive':
    strategy = NaiveStrategy()
elif args.strategy == 'ddp':
    strategy = DDPStrategy()
elif args.strategy == 'colossalai_gemini':
    strategy = ColossalAIStrategy(stage=3, placement_policy='cuda')
elif args.strategy == 'colossalai_zero2':
    strategy = ColossalAIStrategy(stage=2, placement_policy='cuda')
else:
    raise ValueError(f'Unsupported strategy "{args.strategy}"')
    
    
from typing import Optional

import torch.nn as nn
from transformers.models.gpt2.configuration_gpt2 import GPT2Config
from transformers.models.gpt2.modeling_gpt2 import GPT2Model

#RewardModel
from chatgpt.models.base import RewardModel

class GPTRM_custom(RewardModel):

    def __init__(self,
                 pretrained: Optional[str] = None,
                 config: Optional[GPT2Config] = None,
                 checkpoint: bool = False,
                 lora_rank: int = 0,
                 lora_train_bias: str = 'none',
                 tokenizer=None) -> None:
        if pretrained is not None:
            model = GPT2Model.from_pretrained(pretrained)
            model.resize_token_embeddings(len(tokenizer))
        elif config is not None:
            model = GPT2Model(config)
        else:
            model = GPT2Model(GPT2Config())
        if checkpoint:
            model.gradient_checkpointing_enable()

        value_head = nn.Linear(model.config.n_embd, 1)
        super().__init__(model, value_head, lora_rank, lora_train_bias)

        if pretrained is not None:
            self.model = model
            self.pretrained = pretrained


    def save_pretrained(self, dir):
        if self.pretrained is not None:
            self.model.save_pretrained(dir)
            

# configure model, tokenizer
with strategy.model_init_context():
    # load pretrained gpt2
    if args.model == 'gpt2':
#         tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
        # tokenizer = AutoTokenizer.from_pretrained(args.pretrain)
        tokenizer = AutoTokenizer.from_pretrained(args.pretrain, padding_side="right", model_max_length=512)
        tokenizer.add_special_tokens(
            {
                "eos_token": DEFAULT_EOS_TOKEN,
                "bos_token": DEFAULT_BOS_TOKEN,
                "unk_token": DEFAULT_UNK_TOKEN,
            }
        )
        tokenizer.pad_token = tokenizer.eos_token
        model = GPTRM_custom(pretrained=args.pretrain, lora_rank=args.lora_rank, tokenizer=tokenizer).cuda()

    elif args.model == 'bloom':
        model = BLOOMRM(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda()
        tokenizer = BloomTokenizerFast.from_pretrained(args.pretrain)

    elif args.model == 'opt':
        model = OPTRM(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda()
        tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")

    else:
        raise ValueError(f'Unsupported model "{args.model}"')


    # model.resize_token_embeddings(len(tokenizer))
    

# make ranking data to chosen, rejetced data
with open(args.data_path_2_RM, "r", encoding='utf-8-sig') as json_file:
    list_data_dict = json.load(json_file)
    if args.verbose:
        print('## data check ##')
        print((list_data_dict[0]))

total_data_ranking2chosen = []
for tmp in list_data_dict:
    one_data_ranking2chosen = []

    # data 1) 0 VS 1
    data = {}
    data['prompt'] = tmp['prompt']
    if tmp['ranking'][0] < tmp['ranking'][1]:
        data['chosen'] = tmp['completion_0']
        data['rejected'] = tmp['completion_1']
    else:
        data['chosen'] = tmp['completion_1']
        data['rejected'] = tmp['completion_0']
    one_data_ranking2chosen.append(data)


    # data 2) 0 VS 2
    data = {}
    data['prompt'] = tmp['prompt']
    if tmp['ranking'][0] < tmp['ranking'][2]:
        data['chosen'] = tmp['completion_0']
        data['rejected'] = tmp['completion_2']
    else:
        data['chosen'] = tmp['completion_2']
        data['rejected'] = tmp['completion_0']
    one_data_ranking2chosen.append(data)

    # data 1) 1 VS 2
    data = {}
    data['prompt'] = tmp['prompt']
    if tmp['ranking'][1] < tmp['ranking'][2]:
        data['chosen'] = tmp['completion_1']
        data['rejected'] = tmp['completion_2']
    else:
        data['chosen'] = tmp['completion_2']
        data['rejected'] = tmp['completion_1']
    one_data_ranking2chosen.append(data)



    total_data_ranking2chosen.extend(one_data_ranking2chosen)


import random
random.seed(230319)
# list_tmp = list(range(10))
random.shuffle(total_data_ranking2chosen)
print(total_data_ranking2chosen[45])

# train_data = total_data_ranking2chosen[:-1000]  # 29000 학습
# eval_data = total_data_ranking2chosen[-1000:0]  # 1000개만 평가

train_data = total_data_ranking2chosen[:100]  # 29000 학습
eval_data = total_data_ranking2chosen[100:130]  # 1000개만 평가


train_dataset = RewardDataset(train_data, tokenizer, args.max_len)
eval_dataset = RewardDataset(eval_data, tokenizer, args.max_len)


# configure optimizer
if args.strategy.startswith('colossalai'):
    optim = HybridAdam(model.parameters(), lr=5e-5)
else:
    optim = Adam(model.parameters(), lr=5e-5)
    
    
trainer = RewardModelTrainer(model=model,
                             strategy=strategy,
                             optim=optim,
                             train_dataset=train_dataset,
                             eval_dataset=eval_dataset,
                             batch_size=args.batch_size,
                             max_epochs=args.max_epochs)
                             
trainer.fit(use_lora=args.lora_rank)



strategy.save_model(model, os.path.join(args.output_dir, 'RM.pt'), only_rank0=True)
strategy.save_optimizer(optim,
                        os.path.join(args.output_dir, 'RM_optim_checkpoint_%d.pt' % (torch.cuda.current_device())),
                        only_rank0=False)

model.save_pretrained(args.output_dir)  # config.json 생성

 

 

---

Reference

https://tech.scatterlab.co.kr/luda-rlhf/

https://github.com/SKT-AI/KoGPT2?tab=readme-ov-file#classification-or-regression

https://ai.shop2world.net/%EC%8A%A4%ED%83%A0%ED%8F%AC%EB%93%9C-%EC%95%8C%ED%8C%8C%EC%B9%B4alpaca-%EC%BD%94%EB%93%9C%EB%B6%84%EC%84%9D-%EB%88%84%EA%B5%AC%EB%82%98-%EC%B1%97gpt-3-5%EC%84%B1%EB%8A%A5%EC%9D%98-%EB%AA%A8/