PPO 구현을 위한 TRL패키지 살펴보기

이번 포스팅에는 NLP분야에서 PPO 알고리즘을 구현하기 위한 TRL패키지에 대한 이야기와 이전에 포스팅했던 ColossalAI에서 제공하는 패키지의 다른점들들 몇가지에 대해서 기록합니다.

 

현재 던전앤파이터의 세계관에 대해서 이야기할 수 있는 챗봇을 만드는 작업을 진행하고 있습니다. 지난번 RLHF에 대한 이야기를 하면서 작성했던 코드들은 모두 ColossalAI에서 제공하는 패키지로 KoGPT모델을 통한 훈련을 진행하였습니다. 그런데 해당 실습을 하면서 문제점은 작은 모델이여서 생각만큼 원하는 결과가 나오지 않는다는 것이었고, 두번째로 사용하는 패키지가 최근 버전과는 맞지 않아서 강제로 다운그레이드를 해서 진행을 해야하는 번거로움이 존재했습니다. 그래서 던파 챗봇 프로젝트에서는 해당 패키지를 사용하지 않고 Huggingface에서 제공하는 TRL패키지를 사용하여 진행하고 있습니다.

 

TRL패키지를 선택한 이유는 위에서 언급한 두 가지를 모두 해결할 수 있었기 때문이었습니다. KoGPT보다 큰 모델을 불러와서 베이스 모델로 사용하는 것이 목표였습니다. 그래서 선택한 모델은 'EleutherAI/polyglot' 모델입니다. 해당 모델의 소형버전부터 원본 버전까지 존재하는데 저는 SFT모델에는 12.8B를 선택하였고 RM은 1.2B모델을 선택하였습니다. 이부분에 대해서는 해당 프로젝트 포스팅때 자세히 다루겠습니다. 여하튼 이러한 큰 모델을 불러와 작업하기에는 저의 컴퓨팅 자원이 그렇게 녹록치 않기 때문에 지난번 포스팅한 QLoRA를 통해 모델을 불러와 적은 자원에서 모델을 훈련하여 이 문제를 해결하였습니다. 다음 문제가 이 QLoRA와 이어지는 문제인데 큰 모델을 불러오기 위해서는 bitbyteconfg 패키지를 사용해 모델을 불러올때 나눠서 불러오는게 필요했습니다. 그런데 이 부분을 사용하기에는 ColossalAI의 패키지와 버전이 맞지 않는 상황이 발생했습니다. 그래서 다른 방안을 찾아보니 TRL패키지에서도 PPO를 위한 모듈을 제공해주고 있음을 알고 이에 맞는 코드로 수정하는 작업을 진행하고 있습니다. 현재는 PPO 훈련만 남겨둔 상황입니다. 여튼 그래서 이러한 이유로 지난번에 학습과정에서 사용한 모듈이 아닌 TRL패키지에 대한 이야기를 해보겠습니다.

 

TRL패키지에 대한 코드 내용들은 모두 아래 링크를 통해서 확인이 가능합니다. 이번 포스팅은 아래 링크와 huggingface 문서를 토대로 작성하였습니다.

GitHub - huggingface/trl: Train transformer language models with reinforcement learning.

 

SFTTrainer

SFT를 위한 Trainer를 제공합니다. 해당 Trainer는 기존 huggingface에서 사용하던 trainer와 크게 다르지 않습니다. 아래는 가장 기초적인 호출 방법으로 기존의 Trainer에서 쉽게 보던 모습입니다. 모델을 직접 불러와 Trainer에게 넘겨주어도 되고 아래처럼 huggingface에 업로드된 경로 혹은 로컬의 path를 넘겨주면 해당 경로에 있는 모델과 토크나이저를 자동으로 모델이 불러올 수 있게 됩니다.

from datasets import load_dataset
from trl import SFTTrainer

dataset = load_dataset("imdb", split="train")

trainer = SFTTrainer(
    "facebook/opt-350m",
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=512,
)
trainer.train()

 

 

TRL에 좋은점은 PEFT를 사용할 수 있게 설계되어 있다는데 있습니다. 이 부분이 Finetuning에서 크게 시간을 단축할 수 있는 LoRA를 쉽게 사용할 수 있도록 되어 있습니다. 아래와 같이 peft를 불러와 LoraConfig만 작성하여 SFTTrainer에 넘겨주면 됩니다.

from datasets import load_dataset
from trl import SFTTrainer
from peft import LoraConfig

dataset = load_dataset("imdb", split="train")

peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

trainer = SFTTrainer(
    "EleutherAI/gpt-neo-125m",
    train_dataset=dataset,
    dataset_text_field="text",
    peft_config=peft_config
)

trainer.train()

 

 

여기에 더해 제가 봉착했던 대형 자연어 모델을 불러오기 위한 방법도 적용이 가능합니다. 이 방법에서는 Trainer밖에서 모델을 불러오고 Trainer에게 넘겨주어야만 합니다. load_in_8bit or load_in_4bit옵션을 통해 불러오면 되겠습니다.

peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

model = AutoModelForCausalLM.from_pretrained(
    "EleutherAI/gpt-neo-125m",
    load_in_8bit=True,
    device_map="auto",
)

trainer = SFTTrainer(
    model,
    train_dataset=dataset,
    dataset_text_field="text",
    peft_config=peft_config,
)

trainer.train()

 

SFTTrainer는 훈련을 위한 데이터는 기존에 사용하던 datasets를 통해서 사용할 수 있으니 따로 언급하지는 않겠습니다.

 

RewardTrainer

RewardTrainer에서도 크게 다를게 없습니다. 주의할점은 모델을 불러올때와 데이터의 형태입니다. 모델을 불러올때 AutoModelForSequenceClassification을 사용하여 모델을 불러오는 것을 추천하고 있습니다. 옵션을 설정하지 않고 불러올 경우 아래서 볼 데이터셋인 비교군에 대한 점수를 얻을 수 있게 됩니다.

from peft import LoraConfig, TaskType
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from trl import RewardTrainer, RewardConfig

model = AutoModelForSequenceClassification.from_pretrained("gpt2")
peft_config = LoraConfig(
    task_type=TaskType.SEQ_CLS,
    inference_mode=False,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
)

...

trainer = RewardTrainer(
    model=model,
    args=training_args,
    tokenizer=tokenizer,
    train_dataset=dataset,
    peft_config=peft_config,
)

trainer.train()

 

RewardTrainer에 넘겨줄 dataset은 아래 4가지에 대한 값을 갖고 있어야만 정상적으로 작동합니다. 따라서 이에 대한 전처리를 토크나이저를 통해 만들어주시면 되겠습니다.

• input_ids_chosen
• attention_mask_chosen
• input_ids_rejected
• attention_mask_rejected

PPOTrainer

PPOTrainer의 사용법은 간단합니다. 이전에 훈련했던 모델들을 불러와 지정해주면 됩니다. PPO를 위해서 필요한 모델은 총 3개입니다. Trainable Model, Freeze Model, Reward Model입니다. 초기화시 Trainable, Freeze 모델 두 가지는 같은 모델입니다. 그래서 불러오는 작업을 똑같이 두 번 진행해도 되지만 trl 패키지 내에서 모델을 deepcopy하는 기능도 제공하고 있습니다. 또한, freeze model을 제공하지 않을 경우 훈련 모델을 copy하도록 내부적으로 설계되어 있습니다.

 

이후에 진행될 작업은 PPO과정을 그대로 구현해주면 됩니다. trainer.step에서는 입력 tensor, 생성된 tensor, 보상값 이렇게 전달하여 step을 진행하면됩니다. 이 과정을 통해서 trainalbe model의 가중치가 업데이트가 이뤄집니다.

# trl 패키지 내에 있는 freeze model create function
#from trl import create_reference_model
#ref_model = create_reference_model(model)

ppo_trainer = PPOTrainer(config, model, ref_model, tokenizer, dataset=dataset)

for epoch, batch in tqdm(enumerate(ppo_trainer.dataloader)):
    query_tensors = batch["input_ids"]

    #### Get response from gpt2
    response_tensors = []
    for query in query_tensors:
        gen_len = output_length_sampler()
        generation_kwargs["max_new_tokens"] = gen_len
        response = ppo_trainer.generate(query, **generation_kwargs)
        response_tensors.append(response.squeeze()[-gen_len:])
    batch["response"] = [tokenizer.decode(r.squeeze()) for r in response_tensors]

    #### Compute sentiment score
    texts = [q + r for q, r in zip(batch["query"], batch["response"])]
    pipe_outputs = sentiment_pipe(texts, **sent_kwargs)
    rewards = [torch.tensor(output[1]["score"]) for output in pipe_outputs]

    #### Run PPO step
    stats = ppo_trainer.step(query_tensors, response_tensors, rewards)
    ppo_trainer.log_stats(stats, batch, rewards)