SFTTrainer を使って、簡単に CausalLM をファインチューニングをしよう

# SFTTrainer を使って、簡単に CausalLM をファインチューニングをしよう

https://huggingface.co/docs/trl/sft_trainer

## はじめに

はじめまして。現在、私は大学院生（修士課程）です。

この記事では、[TRL](https://huggingface.co/docs/trl/index) というライブラリを使って、簡単に Causal Language Model (CausalLM) をファインチューニングしようと思います。

この記事が「CausalLM の勉強を始めました」や「コードを書くのが楽になった」などの貢献ができるように書かせていただきます。

---

https://qiita.com/Natu_ja/items/b1fc52e20b0bdd8d5b14

https://qiita.com/Natu_ja/items/8031fd0f07d2da86005b

類似記事もありますので、是非こちらもご覧になってください。

## 環境構築

### Python バージョン

- Python >= 3.0.0

https://www.python.org/

### ライブラリ

```text:requirements.txt
datasets
peft
torch>=2.0.0
transformers
trl
```

https://huggingface.co/docs/datasets

https://huggingface.co/docs/peft

https://pytorch.org/

https://huggingface.co/docs/transformers

https://huggingface.co/docs/trl

## 使用するモデルとデータセット

### モデル

今回使用するモデルは、[rinna GPT](https://rinna.co.jp/news/2021/08/20210825.html) を使用します。

https://huggingface.co/rinna/japanese-gpt2-xsmall

### データセット

今回使用するデータセットは、SciQ [^SciQ] を日本語に翻訳したデータセットを使用します。

https://huggingface.co/datasets/izumi-lab/sciq-ja-mbartm2m

## コード

### ライブラリのインポート

```python
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForLanguageModeling
from datasets import load_dataset
from peft import LoraConfig
from trl import DataCollatorForCompletionOnlyLM, SFTConfig, SFTTrainer
```

### トークナイザとモデル

#### ロード

`cuda` が使えれば、`cuda` の方にモデルをロードします。

```python
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained("rinna/japanese-gpt2-xsmall")
model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt2-xsmall").to(device)
```

モデルの構造を確認します。

```python
model
```

#### LoRA の設定

![LoRA architecture](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3836663/5ea0d877-b508-9833-f53d-ff101a0560af.png)

LoRA [^LoRA] を適応したモデルの入出力は、以下の式で表せます。ここで、$W_0 \in \mathbb{R}^{d \times k}$ は凍結された事前学習済み重み行列、$B \in \mathbb{R}^{d \times r}$、$A \in \mathbb{R}^{r \times k}$とし、$r \ll \min(d, k)$です。初期化方法は、$B$ を零行列に従い初期化し、$A$ を正規分布に従い初期化します。

```math
h = W_0 x + \Delta W x = W_0 x + B A x
```

LoRA を適応したとき、学習可能なパラメータ数の割合は、以下の式のように表せます。

```math
\frac{d \times r + r \times k}{d \times k + d \times r + r \times k} = \frac{r (d + k)}{d k + r (d + k)}
```

LoRA を使うことで、ファインチューニングの時のコストを大幅に削減することができます。

この記事では、モデルサイズが $43.7 \mathrm{M}$ のモデルを使っているので、あまり恩恵は感じられないかもしれません。
数 $\mathrm{B}$ 程度のモデルでは、LoRA の恩恵を受けられると思います。

今回は、全ての線形層に LoRA を適応したいため、`target_modules="all-linear"` としました。

```python
peft_config = LoraConfig(
    peft_type="LORA",
    task_type="CAUSAL_LM",
    r=8,
    target_modules="all-linear",
    lora_alpha=8,
    lora_dropout=0.0
)
```

### データセット

#### ロード

```python
train_dataset = load_dataset("izumi-lab/sciq-ja-mbartm2m", split="train")
eval_dataset = load_dataset("izumi-lab/sciq-ja-mbartm2m", split="validation")
```

#### 学習に使用するフォーマットの作成とデータコレーター

`formatting_func` に、学習の時に使うプロンプトを書きます。

```python
def formatting_func(example):
    output_texts = [
        f"# 問題: {example['question'][i]}\n#ヒント: {example['support'][i]}\n# 答え: {example['correct_answer'][i]}" for i in range(len(example))
    ]
    return output_texts
```

```python
data_collator = DataCollatorForCompletionOnlyLM(
    response_template=tokenizer.encode("# 答え: ", add_special_tokens=False),
    instruction_template=tokenizer.encode("# 問題: ", add_special_tokens=False),
    tokenizer=tokenizer
)
```

:::note warn
LLaMA 2 のようないくつかのトークナイザでは、他とは異なるトークン化戦略のため、上記のような書き方の `data_collator` では学習ができません。
:::

https://huggingface.co/docs/trl/sft_trainer#using-tokenids-directly-for-responsetemplate

:::note info
LLaMA 2 のようないくつかのトークナイザの場合は、以下のようなコードに変更して下さい。
:::

```python
data_collator = DataCollatorForCompletionOnlyLM(
    response_template=tokenizer.encode("\n# 答え: ", add_special_tokens=False)[2:],
    instruction_template=tokenizer.encode("# 問題: ", add_special_tokens=False),
    tokenizer=tokenizer
)
```


あるいは、[Transformers](https://huggingface.co/docs/transformers/index) の [`DataCollatorForLanguageModeling`](https://huggingface.co/docs/transformers/main_classes/data_collator#transformers.DataCollatorForLanguageModeling) を使って、以下のようなコードに変更して下さい。

```python
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,
    return_tensors="pt"
)
```

### 学習

#### 設定

[`SFTConfig`](https://huggingface.co/docs/trl/sft_trainer#trl.SFTConfig) は、[Transformers](https://huggingface.co/docs/transformers/index) の [`TrainingArguments`](https://huggingface.co/docs/transformers/main/main_classes/trainer#transformers.TrainingArguments) を継承しています。
今回は、基本的なパラメータのみをここに書いていますが、より細かく設定したい人は、調べてみて下さい。

最適化関数は、AdamW [^AdamW] がデフォルトで設定されています。

- `output_dir`：チェックポイントなどの保存先のパス
- `evaluation_strategy`：評価戦略
- `per_device_train_batch_size`：訓練時のバッチサイズ
- `per_device_eval_batch_size`：評価時のバッチサイズ
- `learning_rate`：初期の学習率
- `num_train_epochs`：エポック数
- `lr_scheduler_type`：スケジューラーのタイプ
- `warmup_ratio`：学習率を $0$ から `learning_rate` で指定した値までの線形ウォームアップの割合
- `logging_strategy`：ロギング戦略
- `save_strategy`：保存戦略
- `report_to`：結果やログを記録する外部サービス

[TensorBoard](https://www.tensorflow.org/tensorboard) や [WandB](https://www.wandb.jp/) などの設定ができていれば、損失などの記録を確認することができます。

```python
args = SFTConfig(
    output_dir="./outputs",
    evaluation_strategy="epoch",
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    learning_rate=5e-5,
    num_train_epochs=3.0,
    lr_scheduler_type="linear",
    warmup_ratio=0.0,
    logging_strategy="epoch",
    save_strategy="epoch",
    report_to="all"
)
```

#### Trainer の用意

```python
trainer = SFTTrainer(
    model=model,
    args=args,
    data_collator=data_collator,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=tokenizer,
    peft_config=peft_config,
    formatting_func=formatting_func
)
```

#### Trainer の確認

この時点で、LoRA [^LoRA] を適応済みです。

```python
trainer.model
```

```python
trainer.model.print_trainable_parameters()
```

学習の設定は、`SFTConfig` で指定した値以外はデフォルトに設定されています。

```python
trainer.args
```

データセットは、[`formatting_func`](#学習に使用するフォーマットの作成とデータコレーター) で指定したフォーマットに従い、トークン化されています。

```python
trainer.tokenizer.decode(trainer.train_dataset[0]["input_ids"])
```

#### 学習開始

```python
trainer.train()
```

### 推論

#### 学習済みのモデルをロード

`outputs/checkpoint-xxx` の `xxx` は、`outputs` フォルダ内の存在する整数に書き換えて下さい。

```python
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained("outputs/checkpoint-xxx")
model = AutoModelForCausalLM.from_pretrained("outputs/checkpoint-xxx").to(device)
```

#### 評価用のデータセット

```python
test_dataset = load_dataset("izumi-lab/sciq-ja-mbartm2m", split="test")
```

#### 生成

```python
with torch.no_grad():

    id = 0
    input_text = f"# 問題: {test_dataset['question'][id]}\n#ヒント: {test_dataset['support'][id]}\n# 答え: "
    
    tokenized_input_text = tokenizer(input_text, return_tensors="pt").to(model.device)
    
    tokenized_output_text_list = model.generate(**tokenized_input_text, max_new_tokens=128)
    
    output_text_list = [
        tokenizer.decode(tokenized_output_text, skip_special_tokens=True) for tokenized_output_text in tokenized_output_text_list
    ]
    
    print(output_text_list)
```

## おわりに

この記事では、ファインチューニングを簡単にできるようにしました。また、学習したモデルを使って、推論をするため方法について書きました。

この記事が CausalLM を学習・研究する人に貢献できたら、幸いです。

## 参考文献

[^AdamW]: Ilya Loshchilov and Frank Hutter. Decoupled Weight Decay Regularization. In *7th International Conference on Learning Representations*, 2019.

[^LoRA]: Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang and Weizhu Chen. LoRA: Low-Rank Adaptation of Large Language Models. In *The Tenth International Conference on Learning Representations*, 2022.

[^SciQ]: Johannes Welbl, Nelson F. Liu and Matt Gardner. Crowdsourcing Multiple Choice Science Questions. In *Proceedings of the 3rd Workshop on Noisy User-generated Text*, 2017.