SFTTrainer と TrainingArguments を使ってバッチ数を小さくしてステップを少なくした学習コードの解説

この記事では、Hugging Face の Transformers ライブラリの SFTTrainer と TrainingArguments を使って、限られたリソースでも効率的に言語モデルをファインチューニングする方法を解説します。特に、バッチサイズを小さくしてステップ数を減らすことで、メモリ使用量を抑えつつ学習を行う方法に焦点を当てます。

下記の記事もおすすめ

Supervised Fine-tuning Trainer (SFT) 入門

Supervised Fine-tuning Trainer (SFT) 入門Supervised Fine-tuning (SFT) は、Reinforcement Learning from Human Feedback (RLHF) ...

SFTTrainer とは
TrainingArguments とは
バッチサイズとステップ数を調整して学習する方法
重要な用語の解説
SFTTrainer の重要な引数
TrainingArguments の重要な引数
ノートブック
参考サイト
1. 関連

SFTTrainer とは

SFTTrainer は、Supervised Fine-tuning (SFT) を行うための Trainer クラスです。SFT は、教師あり学習によってプレトレーン済みモデルをファインチューニングする手法で、Reinforcement Learning from Human Feedback (RLHF) における重要なステップの1つです。

SFTTrainer は、データセットの前処理から学習までを簡単に行えるようにしてくれます。また、プレトレーン済みモデルをロードし、必要に応じて PeftModel に変換することもできます。

TrainingArguments とは

TrainingArguments は、学習ループに関連する引数を管理するクラスです。バッチサイズ、学習率、エポック数などの各種ハイパーパラメータを設定できます。また、ロギング、チェックポイント、評価などの動作も制御できます。

バッチサイズとステップ数を調整して学習する方法

限られたリソース（特にGPUメモリ）で効率的に学習するには、バッチサイズとステップ数を適切に調整することが重要です。以下のようなコードで実現できます。

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from trl import SFTTrainer

# データセットのロード
dataset = load_dataset("imdb", split="train")

# モデルとトークナイザーのロード
model_name = "facebook/opt-350m"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# TrainingArguments の設定
training_args = TrainingArguments(
    output_dir="output",
    per_device_train_batch_size=4,  # バッチサイズを小さく設定
    gradient_accumulation_steps=8,  # 勾配を蓄積するステップ数を大きく設定
    max_steps=1000,  # 最大ステップ数を設定
    logging_steps=100,
    save_steps=500,
    learning_rate=1e-5,
    weight_decay=0.01,
)

# SFTTrainer の初期化
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    dataset_text_field="text",
    tokenizer=tokenizer,
    args=training_args
)

# 学習の実行
trainer.train()

ポイントは以下の3点です。

per_device_train_batch_size を小さく設定する（例: 4）
- バッチサイズを小さくすることで、メモリ使用量を抑えられます。
gradient_accumulation_steps を大きく設定する（例: 8）
- 勾配を蓄積するステップ数を増やすことで、実効的なバッチサイズを大きくできます。
- これにより、学習の安定性を保ちつつ、メモリ使用量を抑えられます。
max_steps を設定する（例: 1000）
- 最大ステップ数を設定することで、エポック数ではなくステップ数でどの程度学習するかを制御できます。
- これにより、データセットが大きい場合でも、少ないステップ数で学習を終えられます。