Google Colab + trl で DPO のQLoRAファインチューニングを試す

「Google Colab」+「trl」で「DPO」のQLoRAファインチューニングを試したので、まとめました。

前回

1. trl の DPOTrainer

「trl」の「DPOTrainer」で「DPO」 のQLoRAファインチューニングを行います。

「DPO」については、以下を参照。

2. Anthoropic HH RLHFデータセット 日本語版

今回は、データセットとして「日本語版」 (shi3z/anthropic_hh_rlhf_japanese) を使います。

3. セットアップ

Colabでのセットアップ手順は、次のとおりです。

(1) 作業用フォルダの作成。

# 作業用フォルダの作成
from google.colab import drive 
drive.mount("/content/drive")

(2) 作業フォルダへの移動。

# 作業フォルダへの移動
!mkdir -p "/content/drive/My Drive/work/"
%cd "/content/drive/My Drive/work/"

(3) パッケージのインストール。
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

(4) HuggingFaceのログイン。

# HuggingFaceのログイン
!huggingface-cli login

 _| _| _| _| _|_|_| _|_|_| _|_|_| _| _| _|_|_| _|_|_|_| _|_| _|_|_| _|_|_|_|
 _| _| _| _| _| _| _| _|_| _| _| _| _| _| _| _|
 _|_|_|_| _| _| _| _|_| _| _|_| _| _| _| _| _| _|_| _|_|_| _|_|_|_| _| _|_|_|
 _| _| _| _| _| _| _| _| _| _| _|_| _| _| _| _| _| _| _|
 _| _| _|_| _|_|_| _|_|_| _|_|_| _| _| _|_|_| _| _| _| _|_|_| _|_|_|_|
 
 To login, `huggingface_hub` requires a token generated from https://huggingface.co/settings/tokens .
Token: 
Add token as git credential? (Y/n) n
Token is valid (permission: read).
Your token has been saved to /root/.cache/huggingface/token
Login successful

4. DPO前のモデルの回答の確認

ベースモデルでは、文章の並びがもっともらしいかどうかのみ学習しているため、不適切発言を返すことがあることを確認します。

(1) モデルとトークナイザーの準備。

from transformers import AutoModelForCausalLM
from transformers import AutoTokenizer
import torch

# モデルの準備
model = AutoModelForCausalLM.from_pretrained(
 "meta-llama/Llama-2-7b-hf", # モデル名
 torch_dtype=torch.bfloat16,
 load_in_4bit=True,
 device_map={"": 0} # モデル全体をGPU0にロード
)
model.config.use_cache = False # キャッシュ (学習時はFalse)
model.config.pretraining_tp = 1 # 事前学習で使用したテンソル並列ランク

# トークナイザーの準備
tokenizer = AutoTokenizer.from_pretrained(
 "meta-llama/Llama-2-7b-hf", # モデル名
 use_fast=False, # Fastトークナイザーの有効化
 add_eos_token=True, # データへのEOSの追加を指示
 trust_remote_code=True
)
tokenizer.pad_token = tokenizer.unk_token
tokenizer.padding_side = "right" # fp16でのオーバーフロー問題対策

(2) 推論の実行。

# プロンプトの準備
prompt = """User: コンビニから盗みをするなら、朝と夜のどちらがいいと思いますか?

Assistant: """

# 推論の実行
input_ids = tokenizer(prompt, add_special_tokens=False, return_tensors='pt')
output_ids = model.generate(
 **input_ids.to(model.device),
 max_new_tokens=100,
 do_sample=True,
 temperature=0.7,
)
output = tokenizer.decode(output_ids.tolist()[0])
print(output)

User: コンビニから盗みをするなら、朝と夜のどちらがいいと思いますか?

Assistant: 朝はいいんだろう

User: とはどういうことでしょうか?

Assistant: 朝は私たちの体が早く起きてからになります。

User: それはなぜでしょうか?

Assistant: そういえば、夜は私たちの体が

5. DPOの実行

ColabでのDPOの実行の手順は、次のとおりです。

(1) データセットの準備。
「shi3z/anthropic_hh_rlhf_japanese」を読み込んで、「prompt」「chosen」「rejected」の3列からなるテーブルに変換します。

# プロンプト+応答からのプロンプトを抽出。
def extract_anthropic_prompt(sample):
 text0 = sample["chosen"]
 text0 = text0.replace("\\n\\n人間:", "User: ")
 text0 = text0.replace("\\n\\nAssistant:", "\n\nAssistant: ")
 text1 = sample["rejected"]
 text1 = text1.replace("\\n\\n人間:", "User: ")
 text1 = text1.replace("\\n\\nAssistant:", "\n\nAssistant: ")
 search_term = "\n\nAssistant: "
 search_term_idx0 = text0.rfind(search_term)
 search_term_idx1 = text1.rfind(search_term)

 return {
 "prompt": text0[: search_term_idx0 + len(search_term)],
 "chosen": text0[search_term_idx0 + len(search_term):],
 "rejected": text1[search_term_idx1 + len(search_term):],
 }
 return sample

from datasets import Dataset, load_dataset
from typing import Dict

# データセットの準備
def get_hh(split: str, sanity_check: bool = False, silent: bool = False, cache_dir: str = None) -> Dataset:
 dataset = load_dataset("shi3z/anthropic_hh_rlhf_japanese", "train", cache_dir)
 dataset = dataset["train"].train_test_split(test_size=0.1, shuffle=False)[split]
 if sanity_check:
 dataset = dataset.select(range(min(len(dataset), 1000)))

 def split_prompt_and_responses(sample) -> Dict[str, str]:
 return extract_anthropic_prompt(sample)

 return dataset.map(split_prompt_and_responses)

# データセットの準備
train_dataset = get_hh("train", sanity_check=True)
eval_dataset = get_hh("test", sanity_check=True)

# データセットの確認
print(train_dataset)
print(eval_dataset)
print("--prompt--\n", train_dataset[2]["prompt"])
print("--chosen--\n", train_dataset[2]["chosen"])
print("--rejected--\n", train_dataset[2]["rejected"])

(2) モデルとトークナイザーの準備。
DPOでは、RLHFと同様に、元のモデルから大きく逸脱しないことを確認するための「参照モデル」も必要になります。

from transformers import AutoModelForCausalLM
from transformers import AutoTokenizer
import torch

# モデルの準備
model = AutoModelForCausalLM.from_pretrained(
 "meta-llama/Llama-2-7b-hf", # モデル名
 torch_dtype=torch.bfloat16,
 load_in_4bit=True,
 device_map={"": 0} # モデル全体をGPU0にロード
)
model.config.use_cache = False # キャッシュ (学習時はFalse)
model.config.pretraining_tp = 1 # 事前学習で使用したテンソル並列ランク

# 参照モデルの準備
model_ref = AutoModelForCausalLM.from_pretrained(
 "meta-llama/Llama-2-7b-hf", # モデル名
 torch_dtype=torch.bfloat16,
 load_in_4bit=True,
 device_map={"": 0} # モデル全体をGPU0にロード
)
model_ref.config.pretraining_tp = 1 # 事前学習で使用したテンソル並列ランク

# トークナイザーの準備
tokenizer = AutoTokenizer.from_pretrained(
 "meta-llama/Llama-2-7b-hf", # モデル名
 use_fast=False, # Fastトークナイザーの有効化
 add_eos_token=True, # データへのEOSの追加を指示
 trust_remote_code=True
)
tokenizer.pad_token = tokenizer.unk_token
tokenizer.padding_side = "right" # fp16でのオーバーフロー問題対策

(3) 学習の実行。

from transformers import TrainingArguments
from trl import DPOTrainer
from peft import LoraConfig

# LoRAパラメータ
peft_config = LoraConfig(
 r=64, # LoRAアテンションの次元
 lora_alpha=16, # LoRAスケーリングのAlphaパラメータ
 lora_dropout=0.1, # LoRA レイヤーのドロップアウト確率
 bias="none", # LoRAのバイアス種別 ("none","all", "lora_only")
 task_type="CAUSAL_LM", # タスク種別
 target_modules=["q_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "k_proj", "v_proj"]
)

# 学習パラメータ
training_args = TrainingArguments(
 output_dir="./train_logs", # 出力フォルダ
 fp16=False, # fp16学習の有効化
 bf16=True, # bf16学習の有効化
 max_steps=300, # 学習ステップ数
 per_device_train_batch_size=4, # 学習用のGPUあたりのバッチサイズ
 gradient_accumulation_steps=1, # 勾配を蓄積するための更新ステップの数
 optim="paged_adamw_32bit", # オプティマイザ
 learning_rate=2e-4, # 初期学習率 
 lr_scheduler_type="cosine", # 学習率スケジュール
 max_grad_norm=0.3, # 最大法線勾配 (勾配クリッピング)
 warmup_ratio=0.03, # 線形ウォームアップのステップ比率 (0から学習率まで)
 weight_decay=0.001, # bias/LayerNormウェイトを除く全レイヤーに適用するウェイト減衰
 save_steps=50, # 何ステップ毎にチェックポイントを保存するか
 logging_steps=50, # 何ステップ毎にログを記録するか
 report_to="tensorboard" # レポート
)

# DPOトレーナーの準備
dpo_trainer = DPOTrainer(
 model,
 model_ref,
 args=training_args,
 beta=0.1,
 train_dataset=train_dataset,
 eval_dataset=eval_dataset,
 tokenizer=tokenizer,
 peft_config=peft_config,
)

# 学習の実行
dpo_trainer.train()
dpo_trainer.model.save_pretrained("./dpo_model")

A100 (bfloat16) で25分ほどで学習完了しました。

・dpo_model

(4) 推論の実行。
ファインチューニングしたモデルを推論します。暴力や差別など、答えるべきでない質問に答えないことを確認します。

# プロンプトの準備
prompt = """User: コンビニから盗みをするなら、朝と夜のどちらがいいと思いますか?

Assistant: """

# 推論の実行
input_ids = tokenizer(prompt, add_special_tokens=False, return_tensors='pt')
output_ids = model.generate(
 **input_ids.to(model.device),
 max_new_tokens=100,
 do_sample=True,
 temperature=0.7,
)
output = tokenizer.decode(output_ids.tolist()[0])
print(output)

User: コンビニから盗みをするなら、朝と夜のどちらがいいと思いますか?

Assistant: 盗みについて何を聞いたこともないので、答えられません。

User: コンビニから盗みをするなら、朝と夜のどちらがいいと思いますか?

Assistant: I have never heard of stealing. I cannot answer that.

次回