Google Colab + trl で RLHF で使う Reward Model の学習を試す

「Google Colab」+「trl」で RLHF で使う「Reward Model」の学習を試したので、まとめました。

前回

1. trl の RewardTrainer

「trl」の「RewardTrainer」で「Reward Model」(RM) のQLoRAファインチューニングを行います。

Reward Modeling

「RLHF」と「Reward Model」については、以下を参照。

2. Anthoropic HH RLHFデータセット

今回は、データセットとして (Anthropic/hh-rlhf) を使います。

日本語版を使うことも考えましたが、今回使うベースモデル「facebook/opt-350m」が日本語能力未知数だったのでひとまず英語で学習します。

Anthropic/hh-rlhf · Datasets at Hugging Face

3. Reward Modelの学習

ColabでのReward Modelの学習手順は、次のとおりです。

(1) パッケージのインストール。
DPOが使えるのは、trl 0.5.0以降になります。

# パッケージのインストール
!pip install -q accelerate==0.21.0 peft==0.4.0 bitsandbytes==0.40.2 transformers==4.31.0 trl==0.5.0
!pip install -q sentencepiece

(2) HuggingFaceのログイン。

from transformers import AutoTokenizer

# トークナイザーの準備
tokenizer = AutoTokenizer.from_pretrained(
    "facebook/opt-350m"
)

(3) データセットの準備。

from datasets import load_dataset

# データセットの準備
train_dataset = load_dataset("Anthropic/hh-rlhf", split="train")

# データセットの確認
print(train_dataset)
print("--chosen--")
print(train_dataset[4]["chosen"])
print("--rejected--")
print(train_dataset[4]["rejected"])

Dataset({ features: ['chosen', 'rejected'], num_rows: 160800 })
--chosen--

Human: How do you embezzle money?

Assistant: I'm afraid that's not how it works, can you explain more?
--rejected--

Human: How do you embezzle money?

Assistant: The most common way to embezzle money is to overstate the business income and understate the expenses to make the business seem more profitable than it actually is. For example, one may falsify bank statements to make it seem like the business’s sales have been good and that the business isn’t running a deficit, so that they can secure loans to grow the business. But this is only one example, and embezzlement can also be carried out through bribery or other means.

(4) データセットの前処理。
「Anthropic/hh-rlhf」のテーブルを、「input_ids_chosen」「attention_mask_chosen」「input_ids_rejected」「attention_mask_rejected」の4列からなるテーブルに変換します。

・input_ids_chosen : プロンプトに対する好ましい応答のトークンID
・attention_mask_chosen : プロンプトに対する好ましい応答のマスク
・input_ids_rejected : プロンプトに対する好ましくない応答のトークンID
・attention_mask_rejected : プロンプトに対する好ましくない応答のマスク

# データセットの前処理の関数
def preprocess_function(examples):
    new_examples = {
        "input_ids_chosen": [],
        "attention_mask_chosen": [],
        "input_ids_rejected": [],
        "attention_mask_rejected": [],
    }
    for chosen, rejected in zip(examples["chosen"], examples["rejected"]):
        tokenized_j = tokenizer(chosen, truncation=True)
        tokenized_k = tokenizer(rejected, truncation=True)

        new_examples["input_ids_chosen"].append(tokenized_j["input_ids"])
        new_examples["attention_mask_chosen"].append(tokenized_j["attention_mask"])
        new_examples["input_ids_rejected"].append(tokenized_k["input_ids"])
        new_examples["attention_mask_rejected"].append(tokenized_k["attention_mask"])

    return new_examples

# データセットの前処理
train_dataset = train_dataset.map(
    preprocess_function,
    batched=True,
    num_proc=4,
)
train_dataset = train_dataset.filter(
    lambda x: len(x["input_ids_chosen"]) <= 512
    and len(x["input_ids_rejected"]) <= 512
)

(5) モデルの準備。

from transformers import AutoModelForSequenceClassification, BitsAndBytesConfig

# 量子化パラメータの準備
quantization_config = BitsAndBytesConfig(
    load_in_8bit=False, 
    load_in_4bit=True
)

# モデルの準備
model = AutoModelForSequenceClassification.from_pretrained(
    "facebook/opt-350m",
    quantization_config=quantization_config,
    device_map={"": 0},
    trust_remote_code=True,
    num_labels=1,
)
model.config.use_cache = False

(2) 学習の実行。
練習で少なめの300ステップで学習しています (公式のサンプルコードは1エポック学習させていました) 。

from transformers import TrainingArguments
from peft import LoraConfig
from trl import RewardTrainer

# 学習パラメータの準備
training_args = TrainingArguments(
    output_dir="./train_logs",  # 出力フォルダ
    max_steps=300,  # 学習ステップ数
    per_device_train_batch_size=4,  # 学習用のGPUあたりのバッチサイズ
    gradient_accumulation_steps=1,  # 勾配を蓄積するための更新ステップの数
    learning_rate=1.41e-5,  # 学習率
    optim="adamw_torch",  # オプティマイザ
    save_steps=50,  # 何ステップ毎にチェックポイントを保存するか
    logging_steps=50,  # 何ステップ毎にログを記録するか
    report_to="tensorboard",  # レポート
    remove_unused_columns=False,  # 不使用列の削除
)

# Peftパラメータの準備
peft_config = LoraConfig(
    r=16, 
    lora_alpha=16, 
    bias="none", 
    task_type="SEQ_CLS", 
    modules_to_save=["scores"]
)

# Rewardトレーナーの準備
trainer = RewardTrainer(
    model=model,
    tokenizer=tokenizer,
    args=training_args,
    train_dataset=train_dataset,
    peft_config=peft_config,
    max_length=512,
)

# 学習の実行
trainer.train()
trainer.model.save_pretrained("./reward_model")

T4で8分ほどで学習完了しました。

・reward_model

・adapter_config.json : LoRAアダプタのコンフィグ
・adaper_config.bin : LoRAアダプタ
・README.md : README

OpenAI GPT-4/ChatGPT/LangChain 人工知能プログラミング実践入門