DeepSeek R1モデルで数学問題をGRPO (Group Relative Policy Optimization) トレーニング

# はじめに
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` (量子化用) をインストールします。

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


## **2. 環境変数の設定**

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

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

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

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

```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` を定義します。

```python
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程度になり動きました。

```python
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

```python
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` を定義します。

```python
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` を定義します。

```python
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から取得する。

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

前処理を行います。

```python
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. プロンプト作成関数の定義**

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

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

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

```python
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距離に基づいて評価します。

```python
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を使用するかどうかに応じてモデルのロード方法が異なります。

```python
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` 関数をデータセットに適用し、プロンプトを生成します。

```python
dataset = dataset.map(create_prompt)
```

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

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

```python
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 の設定を行います。

```python
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) の設定を行います。

```python
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` を初期化し、トレーニングを実行します。

```python
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. 学習済みモデルのロードと評価**

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

```python
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
```