Google Colab で trlX による大規模言語モデルのファインチューニングを試す

「Google Colab」で「trlX」による大規模言語モデルのファインチューニングを試したので、まとめました。

1. trlX

「trlX」 (Transformer Reinforcement Learning X)は、「報酬関数」または「報酬ラベル付きデータセット」のいずれかを使用して、強化学習で大規模言語モデル (LLM) をファインチューニングするために分散学習フレームワークです。

「facebook/opt-6.7b」「EleutherAI/gpt-neox-20b」など、最大200億のパラメータの「causal」および「T5」ベースの言語モデルをファインチューニングできます。

現在、次の強化学習アルゴリズムが実装されています。

・PPO (Proximal Policy Optimization)
・ILQL (Implicit Language Q-Learning)

Welcome to trlX’s documentation! — trlX documentation

2. 学習方法

「報酬関数」または「報酬ラベル付きデータセット」でLLMをファインチューニングできます。

◎ 報酬関数の場合

trainer = trlx.train('gpt2', reward_fn=lambda samples, **kwargs: [sample.count('cats') for sample in samples])

◎ 報酬ラベル付きデータセットの場合

trainer = trlx.train('EleutherAI/gpt-j-6B', samples=['dolphins', 'geese'], rewards=[1.0, 100.0])

「trainer」は、ベースとなるモデルのラッパーを提供します。generate()で推論できます。

trainer.generate(**tokenizer('Q: Who rules the world? A:', return_tensors='pt'), do_sample=True)

save_pretrained()で事前学習済み言語モデルを保存できます。

trainer.save_pretrained('/path/to/output/folder/')

3. Colabでの実行

映画レビューのコーパス「IMDB」を使って、肯定的な映画レビューを生成するよう「GPT-2」をファインチューニングします。

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

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

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

# パッケージのインストール
!git clone https://github.com/CarperAI/trlx.git
!git config --global --add safe.directory /content/trlx && cd /content/trlx && pip install -e .
!pip uninstall -y scikit_learn jax
%cd trlx

(3) コンフィグの読み込み。
「ILQL」によるファインチューニングの設定を読み込みます。モデルやトークナイザーや学習などの設定があります。今回は、テスト用にtrackerなし、バッチサイズ16、エポック数10に変更してます。

import yaml
from datasets import load_dataset
from transformers import pipeline
import pathlib
from typing import Dict, List
import trlx
from trlx.data.configs import TRLConfig

# コンフィグの読み込み
with open('configs/ilql_config.yml') as f:
    default_config = yaml.safe_load(f)
default_config['train']['tracker'] = None
default_config['train']['batch_size'] = 16
default_config['train']['epochs'] = 10
config = TRLConfig.update(default_config, {})
print(config)

{
    "method": {
        "name": "ilqlconfig",
        "tau": 0.7,
        "gamma": 0.99,
        "cql_scale": 0.1,
        "awac_scale": 1,
        "alpha": 0.001,
        "beta": 0,
        "steps_for_target_q_sync": 5,
        "two_qs": true,
        "gen_kwargs": {
            "max_new_tokens": 56,
            "top_k": 20,
            "beta": 4,
            "temperature": 1.0
        }
    },
    "model": {
        "model_path": "gpt2",
        "model_arch_type": "causal",
        "num_layers_unfrozen": -1,
        "delta_kwargs": null
    },
    "optimizer": {
        "name": "adamw",
        "kwargs": {
            "lr": 5e-05,
            "betas": [
                0.9,
                0.95
            ],
            "eps": 1e-08,
            "weight_decay": 1e-06
        }
    },
    "scheduler": {
        "name": "cosine_annealing",
        "kwargs": {
            "T_max": 1000,
            "eta_min": 5e-05
        }
    },
    "tokenizer": {
        "tokenizer_path": "gpt2",
        "padding_side": "left",
        "truncation_side": "right"
    },
    "train": {
        "total_steps": 1000,
        "seq_length": 64,
        "epochs": 10,
        "batch_size": 16,
        "checkpoint_interval": 1000,
        "eval_interval": 100,
        "pipeline": "PromptPipeline",
        "trainer": "AccelerateILQLTrainer",
        "trainer_kwargs": {},
        "project_name": "trlx",
        "entity_name": null,
        "group_name": null,
        "checkpoint_dir": "ckpts",
        "rollout_logging_dir": null,
        "save_best": true,
        "tracker": null,
        "logging_dir": null,
        "seed": 1000
    }
}

(4) 報酬関数の準備。
報酬関数は、samplesでテキスト群を受け取り、戻り値でPOSITIVEスコア群を返してます。

# POSITIVEスコアの取得
def get_positive_score(scores):
    "Extract value associated with a positive sentiment from pipeline's output"
    return dict(map(lambda x: tuple(x.values()), scores))["POSITIVE"]

# 感情分析の推論
sentiment_fn = pipeline(
    "sentiment-analysis",
    "lvwerra/distilbert-imdb",
    top_k=2,
    truncation=True,
    batch_size=256,
    device=0,
)

# 報酬関数
def metric_fn(samples: List[str], **kwargs) -> Dict[str, List[float]]:
    sentiments = list(map(get_positive_score, sentiment_fn(samples)))
    return {"sentiments": sentiments}

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

# IMDBデータセットの読み込み
imdb = load_dataset("imdb", split="train+test")

(6) 学習の実行。

# 学習の実行
trainer = trlx.train(
    samples=imdb["text"], 
    rewards=imdb["label"],
    eval_prompts=[
        "I don't know much about Hungarian underground",
        "What made this movie so distinctly",
        "Like the sandwich I just bought at the grocery store,",
        "I cannot believe how much this movie made me want to"
    ] * 20,
    metric_fn=metric_fn,
    config=config,

(7) 動作確認。
動作確認で「One thing you should know about this Movie is that」に続くテキストを生成します。

# 動作確認
input_str = 'One thing you should know about this Movie is that'
trainer_output = trainer.generate_eval(
    **trainer.tokenizer(input_str, return_tensors='pt'))[0]
print(trainer.tokenizer.decode(trainer_output))

One thing you should know about this Movie is that it was set in a real world that is very low budget but still has an atmosphere that is not very realistic. You have your character who gets into the car and starts getting into trouble. The scene where someone (Beth) has her hair turned into a dog is

このMovieについて知っておいていただきたいのは、非常に低予算でありながら、現実の世界を舞台にした雰囲気を持っていることです。車に乗り込んだ主人公がトラブルを起こし始めるんですね。誰か（ベス）が髪を犬に変身させるシーンは

肯定的な映画レビューが生成されました。

関連

Illustrating Reinforcement Learning from Human Feedback (RLHF)

Implementing RLHF: Learning to Summarize with trlX

Tuning Large Language Models with Reinforcement Learning on a Single GPU | Hippocampus's Garden

Google Colaboratory

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