はじめに

DeepSeekモデルでGRPO (Group Relative Policy Optimization) トレーニングします。

GRPOとは

GRPO (Group Relative Policy Optimization) は、強化学習の一種で、複数の生成結果をグループ化し、モデルの複数の異なる出力を比較し、そのグループ内での相対的な良さに基づいて学習を進めます。これにより、より安定した学習と、多様な出力の生成が期待できます。

DeepSeek R1でのGRPOトレーニング

DeepSeek R1 - Distilled Qwen 1.5B モデルをGRPOトレーニングでファインチューニングし、数学の問題解決能力を向上させます。
NVIDIAのGPUは、A100 80GBか、RTX 6000 48GBあたりが必要です。L4 24GBだとGPUのメモリオーバーにならないように学習の設定を下げる必要がありました。

1. 必要なライブラリのインストール

trl (Transformer Reinforcement Learningライブラリ)、levenshtein (文字列の類似度計算用)、bitsandbytes (量子化用) をインストールします。

!pip install trl
!pip install levenshtein
!pip install bitsandbytes
!pip install kaggle
!pip install kagglehub

2. 環境変数の設定

GPUの使用設定とトークナイザーの並列処理設定を行います。

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3" # 使用するGPUを指定
os.environ["TOKENIZERS_PARALLELISM"] = "false" # トークナイザーの並列処理をオフ

3. 必要なモジュールのインポート

必要なPythonモジュールをインポートします。

from datasets import load_dataset, Dataset
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, PeftModel
from trl import GRPOConfig, GRPOTrainer
import datetime
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    HfArgumentParser,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    GenerationConfig,
    PrinterCallback,
)
from tqdm import tqdm
import torch
import time
import transformers
import pandas as pd
import numpy as np
from Levenshtein import ratio as levenshtein_ratio
transformers.set_seed(42)

4. 設定クラスの定義

設定を管理するためのクラス CFG を定義します。

class CFG:
    MAX_TRAIN = 100 # トレーニングデータセットの最大サイズ
    MAX_TOKENS = 2048 # 最大トークン数
    NUM_GENERATIONS = 4 # 生成数
    USE_PEFT = True # PEFT (Parameter-Efficient Fine-Tuning) の使用フラグ
    BATCH_SIZE = 1 # バッチサイズ
    MAX_STEPS = 80 # 最大ステップ数
    BETA = 0.04 # GRPOのベータパラメータ
    LR = 1.e-5 # 学習率
    model_name = '/root/.cache/kagglehub/models/deepseek-ai/deepseek-r1/transformers/deepseek-r1-distill-qwen-1.5b/1' # モデル名
    splitter = '< | Assistant | >' # テキスト分割子
    step_count = 10 # ログと保存のステップ間隔
    DEBUG = False # デバッグモードフラグ

NVIDIA L4 24GB GPUでは、以下の設定にするとメモリ使用量が16GB程度になり動きました。

MAX_TRAIN = 20
MAX_TOKENS = 1024
NUM_GENERATIONS = 3

モデルは、Kaggleから取得します。
https://www.kaggle.com/models/deepseek-ai/deepseek-r1/Transformers/deepseek-r1-distill-qwen-1.5b/1

import kagglehub
# Download latest version
path = kagglehub.model_download("deepseek-ai/deepseek-r1/transformers/deepseek-r1-distill-qwen-1.5b")
print("Path to model files:", path)

5. ボックステキスト抽出関数の定義

テキストから \boxed{} で囲まれた部分を抽出する関数 extract_boxed_text を定義します。

import re
def extract_boxed_text(text):
    pattern = r'\\boxed{(.*?)}'
    matches = re.findall(pattern, text)
    if not matches:
        return ""
    for match in matches[::-1]:
        if match != "":
            return match
    return ""

6. 回答検証関数の定義

抽出された回答が有効な数値かどうかを検証する関数 is_valid_answer を定義します。

def is_valid_answer(s):
    try:
        if float(s) == int(s):
            i = int(s)
            return 0<=i<1000
        else:
            return False
    except ValueError:
        return False

7. math_problems データセットの準備

math_problems.parquet データセットはKaggleから取得する。

import kagglehub
# Download latest version
path = kagglehub.dataset_download("artemgoncarov/math-problems-imo")
print("Path to dataset files:", path)

前処理を行います。

df = pd.read_parquet('/root/.cache/kagglehub/datasets/artemgoncarov/math-problems-imo/versions/1')
df = df.reset_index().rename({'index': 'id'}, axis=1)
df['answer'] = df['solution'].map(extract_boxed_text) # 解答を抽出
mask = df['answer'].map(is_valid_answer) # 有効な回答を持つデータのみ抽出
df = df[mask]
df = df.iloc[:CFG.MAX_TRAIN] # データセットサイズを制限
dataset = Dataset.from_pandas(df)
dataset = dataset.train_test_split(test_size=0.1) # データセットをtrainとtestに分割
dataset

8. プロンプト作成関数の定義

トークナイザーをロードします。

tokenizer = AutoTokenizer.from_pretrained(CFG.model_name, trust_remote_code=True, padding_side="left")

モデルへの入力プロンプトを生成する関数 create_prompt を定義します。

def create_prompt(sample):
    question = sample['problem']
    chat = [{"role": "system", "content": "A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within <think> </think> and <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think> <answer> answer here </answer>"},
            {"role": "user", "content": question + ' Return final answer within \\boxed{}, after taking modulo 1000.'},]
    sample['prompt'] = tokenizer.apply_chat_template(
        conversation=chat,
        tokenize=False,
        add_generation_prompt=True
    )
    return sample

9. 報酬関数の定義

GRPOトレーニングで使用する報酬関数を定義します。ここでは3つの報酬関数を定義します。

format_reward_func: 生成されたテキストが特定のフォーマット <think>...</think>...\boxed{} に従っているかを評価します。
accuracy_reward_func: 生成された回答の精度を評価します。正解と一致すれば1.0、そうでなければ0.0の報酬を与えます。
levenshtein_reward_func: 生成されたソリューションの質をLevenshtein距離に基づいて評価します。

def format_reward_func(completions, **kwargs):
    pattern = r"^<think>.*?</think>.*?\\boxed{(.*?)}.*?$"
    matches = [re.match(pattern, content, re.DOTALL) for content in completions]
    return [1.0 if match else 0.0 for match in matches]

def accuracy_reward_func(completions, answer, **kwargs):
    contents = [extract_boxed_text(completion) for completion in completions]
    return [1.0 if c == str(gt) else 0.0 for c, gt in zip(contents, answer)]

def levenshtein_reward_func(completions, solution, **kwargs):
    res = []
    for completion, sol in zip(completions, solution):
        if '</think>' in completion:
            t = completion.split('</think>')[-1]
            res.append(levenshtein_ratio(t, sol))
        else:
            res.append(0.0)
    return res

reward_functions = {'formatting': format_reward_func, 'accuracy': accuracy_reward_func, 'solution_quality': levenshtein_reward_func}

10. モデルのロード

DeepSeek R1モデルをロードします。PEFTを使用するかどうかに応じてモデルのロード方法が異なります。

device_map = 'auto'
if CFG.USE_PEFT:
    compute_dtype = getattr(torch, "float16")
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type='nf4',
        bnb_4bit_compute_dtype=compute_dtype,
        bnb_4bit_use_double_quant=False,
    )
    original_model = AutoModelForCausalLM.from_pretrained(CFG.model_name,
                                                device_map=device_map,
                                                quantization_config=bnb_config,
                                                trust_remote_code=True)
else:
    original_model = AutoModelForCausalLM.from_pretrained(CFG.model_name,
                                                device_map=device_map,
                                                trust_remote_code=True)

11. データセットへのプロンプトの適用

create_prompt 関数をデータセットに適用し、プロンプトを生成します。

dataset = dataset.map(create_prompt)

12. 評価関数の定義と初期評価

モデルの性能を評価する関数 evaluate_rewards を定義し、学習前のモデルで評価を行います。

def gen(model, text, max_tokens):
    model_input = tokenizer(text, return_tensors='pt').to(model.device)
    model.eval()
    with torch.no_grad():
        tok = model.generate(**model_input, max_new_tokens=max_tokens, pad_token_id=tokenizer.pad_token_type_id)
        outputs = []
        for i in range(len(tok)):
            res = tokenizer.decode(tok[i], skip_special_tokens=True)
            output = res.split(CFG.splitter)[-1]
            outputs.append(output)
        return outputs[0] if len(outputs) == 1 else outputs

def evaluate_rewards(model, dataset, reward_functions: dict[str, callable], max_tokens: int, num_generations: int):
    completions = []
    other_info = []
    for example in tqdm(dataset):
        txt = example['prompt']
        kw = {k: v for k, v in example.items() if k not in {'prompt', 'completion'}}
        for _ in range(num_generations):
            other_info.append(kw)
        completion = gen(model, [txt] * num_generations, max_tokens)
        if isinstance(completion, str):
            completions.append(completion)
        else:
            completions += completion

    kwargs = {k: [d[k] for d in other_info] for k in other_info[0].keys()}
    res = {}
    for nm, reward_func in reward_functions.items():
        v = reward_func(completions=completions, **kwargs)
        print(nm, np.mean(v))
        res[nm] = np.mean(v)
    return res

if not CFG.DEBUG:
    original_rewards = evaluate_rewards(model=original_model, dataset=dataset['test'], reward_functions=reward_functions, max_tokens=CFG.MAX_TOKENS, num_generations=CFG.NUM_GENERATIONS)

13. GRPO Trainer の設定

GRPOConfig を使用して、GRPO Trainer の設定を行います。

dtstr = datetime.datetime.now().strftime('%Y%m%d%H%M%S')
output_directory = f"./DEEPSEEK-GRPO-{dtstr}"
training_args = GRPOConfig(
    output_dir=output_directory,
    learning_rate=CFG.LR,
    per_device_train_batch_size=CFG.BATCH_SIZE,
    gradient_accumulation_steps=1,
    max_steps=CFG.MAX_STEPS,
    max_completion_length=CFG.MAX_TOKENS, #8192
    num_generations=CFG.NUM_GENERATIONS,
    beta=CFG.BETA,
    logging_steps=CFG.step_count,
    logging_dir="./logs",
    save_strategy="steps",
    save_steps=CFG.step_count,
    report_to="none",
    overwrite_output_dir='True',
)

14. PEFT設定

PEFTを使用する場合、LoRA (Low-Rank Adaptation) の設定を行います。

if CFG.USE_PEFT:
    peft_config = LoraConfig(
        r=32,
        lora_alpha=32,
        target_modules=[
            'q_proj',
            'k_proj',
            'v_proj',
            'dense'
        ],
        bias="none",
        lora_dropout=0.05,
        task_type="CAUSAL_LM",
    )

15. GRPO Trainer の初期化とトレーニングの実行

GRPOTrainer を初期化し、トレーニングを実行します。

if CFG.USE_PEFT:
    trainer = GRPOTrainer(
        model=original_model,
        reward_funcs=list(reward_functions.values()),
        args=training_args,
        train_dataset=dataset['train'],
        peft_config=peft_config,
        callbacks=[PrinterCallback()]
    )
else:
    trainer = GRPOTrainer(
        model=original_model,
        reward_funcs=list(reward_functions.values()),
        args=training_args,
        train_dataset=dataset['train'],
        callbacks=[PrinterCallback()]
    )

trainer.train()

16. 学習済みモデルのロードと評価

トレーニング後、学習済みモデルをロードし、再度評価を行います。

if CFG.USE_PEFT:
    print('Loading trained model')
    CHKPT = CFG.MAX_STEPS
    adapter_model_name = f'{output_directory}/checkpoint-{CHKPT}/'
    new_model = PeftModel.from_pretrained(original_model, adapter_model_name)
else:
    new_model = original_model

rewards = evaluate_rewards(model=new_model, dataset=dataset['test'], reward_functions=reward_functions, max_tokens=CFG.MAX_TOKENS, num_generations=CFG.NUM_GENERATIONS)
rewards