Google Colab で trl によるTransformerモデルの強化学習を試す

「Google Colab」で「trl」によるTransformerモデルの強化学習を試したので、まとめました。

【注意】「trl」を動作させるには、「Google Colab Pro/Pro+」のプレミアム (A100 40GB) が必要です。

1. trl

「trl」(Transformer Reanforcement Learning)は、強化学習でTransformerモデルを強化学習するためのパッケージです。

PPOによるTransformerモデルの強化学習は、次の3つのステップで構成されます。

(1) ロールアウト：言語モデルは、文頭のクエリに基づいて応答や継続を生成。
(2) 評価 : クエリとレスポンスは、関数、モデル、人間のフィードバック、またはそれらの組み合わせで評価。
(3) 最適化 : クエリとレスポンスのペアを使用して、シーケンス内のトークンの対数確率を計算後、PPOで学習。

2. Colabでの実行

公式サンプル「gpt2-sentiment.ipynb」を試します。

IMDBデータセットを使って、肯定的な映画レビューを生成するように「GPT-2」をファインチューニングします。肯定的な文章に報酬を与えるため、BERT分類器による感情分析を、PPOの報酬シグナルとして使用します。
(wandbとHuggingFace Hubへのpushは無効化してます)

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

(1) 新規のColabのノートブックを開き、メニュー「編集 → ノートブックの設定」で「GPU」の「プレミアム」を選択。

(2) 自動リロードの設定。

# 自動リロード
%load_ext autoreload
%autoreload 2

(3) パッケージのインストール。

# パッケージのインストール
%pip install transformers trl

(4) パッケージのインポート。

# パッケージのインポート
import torch
from tqdm import tqdm
import pandas as pd
tqdm.pandas()
from transformers import pipeline, AutoTokenizer
from datasets import load_dataset
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead
from trl.core import LengthSampler

(5) コンフィグの準備

# コンフィグの準備
config = PPOConfig(
    model_name="lvwerra/gpt2-imdb",
    learning_rate=1.41e-5,
)

sent_kwargs = {
    "return_all_scores": True,
    "function_to_apply": "none",
    "batch_size": 16
}

「gpt2_imdb」という「GPT2」を読み込もうとしてることがわかります。このモデルは、1エポックの IMDB データセットでファインチューニングされています。。他のパラメータは、主に元の論文"Fine-Tuning Language Models from Human Preferences"から取得しています。

(6) IMDBデータセットの読み込み。

# IMDBデータセットの読み込み
def build_dataset(config, dataset_name="imdb", input_min_text_length=2, input_max_text_length=8):
    tokenizer = AutoTokenizer.from_pretrained(config.model_name)
    tokenizer.pad_token = tokenizer.eos_token

    ds = load_dataset(dataset_name, split='train')
    ds = ds.rename_columns({'text': 'review'})
    ds = ds.filter(lambda x: len(x["review"])>200, batched=False)

    input_size = LengthSampler(input_min_text_length, input_max_text_length)

    def tokenize(sample):
        sample["input_ids"] = tokenizer.encode(sample["review"])[:input_size()]
        sample["query"] = tokenizer.decode(sample["input_ids"])
        return sample

    ds = ds.map(tokenize, batched=False)
    ds.set_format(type='torch')
    return ds


dataset = build_dataset(config)

def collator(data):
    return dict((key, [d[key] for d in data]) for key in data[0])

IMDBデータセットには、感情を示す POSITIVE / NEGATIVE でラベルが付けられた 50,000 の映画レビューが含まれています。IMDB データセットをDataFrameに読み込み、200文字以上のコメントをフィルタで処理します。次に、各テキストをトークン化し、LengthSamplerでランダムなサイズにカットしています。

(7) モデルの読み込み。

# モデルの読み込み
model = AutoModelForCausalLMWithValueHead.from_pretrained(config.model_name)
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained(config.model_name)
tokenizer = AutoTokenizer.from_pretrained(config.model_name)

tokenizer.pad_token = tokenizer.eos_token

GPT2モデルのバリューヘッドとトークナイザーを読み込みます。 モデルは2回読み込みます。最初のモデルは最適化、2 番目のモデルは開始点からの KL 発散を計算するための参照として機能します。これは、最適化モデルが元の言語モデルから大きく逸脱しないようにするためのPPOの追加の報酬シグナルとして機能します。

(8) PPOトレーナーの準備。

# PPOトレーナーの準備
ppo_trainer = PPOTrainer(
    config, 
    model, 
    ref_model, 
    tokenizer, 
    dataset=dataset, 
    data_collator=collator)

(9) BERT分類器の準備。

device = ppo_trainer.accelerator.device
if ppo_trainer.accelerator.num_processes == 1:
    device = 0 if torch.cuda.is_available() else "cpu" # to avoid a `pipeline` bug
sentiment_pipe = pipeline("sentiment-analysis", model="lvwerra/distilbert-imdb", device=device)

モデル出力は、POSITIVE / NEGATIVEの2値分類で、POSITIVEを言語モデルの報酬シグナルとして使用します。

・POSITIVEの動作確認

text = 'this movie was really bad!!'
sentiment_pipe(text, **sent_kwargs)

[[{'label': 'NEGATIVE', 'score': 2.3350484371185303},
  {'label': 'POSITIVE', 'score': -2.726576328277588}]]

・NEGATIVEの動作確認

text = 'this movie was really good!!'
sentiment_pipe(text, **sent_kwargs)

[[{'label': 'NEGATIVE', 'score': -2.294790267944336},
  {'label': 'POSITIVE', 'score': 2.557040214538574}]]

(10) 生成設定の準備。
応答の生成には、サンプリングを使用するだけで、top_kとnucleus samplingとmin_lengthがオフになっていることを確認します。

gen_kwargs = {
    "min_length":-1,
    "top_k": 0.0,
    "top_p": 1.0,
    "do_sample": True,
    "pad_token_id": tokenizer.eos_token_id
}

(11) モデルの最適化
A100で2時間30分ほどかかりました。

output_min_length = 4
output_max_length = 16
output_length_sampler = LengthSampler(output_min_length, output_max_length)


generation_kwargs = {
    "min_length":-1,
    "top_k": 0.0,
    "top_p": 1.0,
    "do_sample": True,
    "pad_token_id": tokenizer.eos_token_id
}


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

    #### GPT-2からの回答を得る
    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]

    #### 感情分析のスコアを計算
    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]

    #### PPOのステップを実行 
    stats = ppo_trainer.step(query_tensors, response_tensors, rewards)
    ppo_trainer.log_stats(stats, batch, rewards)

モデルの最適化の学習ループは、次のステップで構成されます。

(1) ポリシー ネットワークからクエリ応答を取得 (GPT-2)
(2) BERT からのクエリ/応答に対する感情分析を取得
(3) トリプレット(クエリ、応答、報酬) を使用してPPOでポリシーを最適化

(12) モデルの動作確認
IMDBデータセットの例をいくつかつかって確認します。最適化前より最適化後の方がスコアが上がっていることがわかります。

#### データセットからバッチを取得
bs = 16
game_data = dict()
dataset.set_format("pandas")
df_batch = dataset[:].sample(bs)
game_data['query'] = df_batch['query'].tolist()
query_tensors = df_batch['input_ids'].tolist()

response_tensors_ref, response_tensors = [], []

#### gpt2およびgpt2_refから応答を取得
for i in range(bs):
    gen_len = output_length_sampler()
    output = ref_model.generate(torch.tensor(query_tensors[i]).unsqueeze(dim=0).to(device),
        max_new_tokens=gen_len, **gen_kwargs).squeeze()[-gen_len:]
    response_tensors_ref.append(output)
    output = model.generate(torch.tensor(query_tensors[i]).unsqueeze(dim=0).to(device),
        max_new_tokens=gen_len, **gen_kwargs).squeeze()[-gen_len:]
    response_tensors.append(output)

#### レスポンスのデコード
game_data['response (before)'] = [tokenizer.decode(response_tensors_ref[i]) for i in range(bs)]
game_data['response (after)'] = [tokenizer.decode(response_tensors[i]) for i in range(bs)]

#### クエリ/レスポンスペアの前後における感情分析
texts = [q + r for q,r in zip(game_data['query'], game_data['response (before)'])]
game_data['rewards (before)'] = [output[1]["score"] for output in sentiment_pipe(texts, **sent_kwargs)]

texts = [q + r for q,r in zip(game_data['query'], game_data['response (after)'])]
game_data['rewards (after)'] = [output[1]["score"] for output in sentiment_pipe(texts, **sent_kwargs)]

# 結果をデータフレームに格納
df_results = pd.DataFrame(game_data)
df_results

報酬の平均/中央値を見ると、大きな違いが見られます。

print('mean:')
display(df_results[["rewards (before)", "rewards (after)"]].mean())
print()
print('median:')
display(df_results[["rewards (before)", "rewards (after)"]].median())

mean:
rewards (before)    0.090049
rewards (after)     1.672922
dtype: float64

median:
rewards (before)    0.193053
rewards (after)     1.915701
dtype: float64

(13) モデルの保存。
HuggingFace HubにはPushしない設定にしてます。

# モデルの保存
model.save_pretrained('gpt2-imdb-pos-v2', push_to_hub=False)
tokenizer.save_pretrained('gpt2-imdb-pos-v2', push_to_hub=False)

('gpt2-imdb-pos-v2/tokenizer_config.json',
 'gpt2-imdb-pos-v2/special_tokens_map.json',
 'gpt2-imdb-pos-v2/vocab.json',
 'gpt2-imdb-pos-v2/merges.txt',
 'gpt2-imdb-pos-v2/added_tokens.json',
 'gpt2-imdb-pos-v2/tokenizer.json')

関連

Fine-tuning 20B LLMs with RLHF on a 24GB consumer GPU

BERT/GPT-3/DALL-E 自然言語処理・画像処理・音声処理 人工知能プログラミング実践入門