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

AI・機械学習

この記事では、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 とは

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点です。

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

重要な用語の解説

  • バッチサイズ: 1ステップで学習するサンプル数。バッチサイズが大きいほど学習が安定しますが、メモリ使用量が増えます。
  • 勾配蓄積: バックプロパゲーションを複数ステップ分まとめて行う手法。バッチサイズを実効的に大きくできます。
  • エポック: データセット全体を1周する単位。エポック数が多いほど学習が進みますが、時間がかかります。
  • ステップ: バックプロパゲーションを1回行う単位。ステップ数が多いほど学習が進みますが、時間がかかります。

以上のように、SFTTrainerとTrainingArgumentsを活用し、バッチサイズとステップ数を調整することで、限られたリソースでも効率的に言語モデルをファインチューニングできます。ぜひ試してみてください!

SFTTrainer の重要な引数

SFTTrainer の初期化時には、以下のような重要な引数があります。

  • model: 学習対象のモデル。
  • train_dataset: 学習用のデータセット。
  • dataset_text_field: データセット内のテキストフィールド名。
  • tokenizer: トークナイザー。
  • args: TrainingArguments のインスタンス。

また、以下のような便利な引数もあります。

  • eval_dataset: 評価用のデータセット。
  • peft_config: PeftConfig のインスタンス。モデルを PeftModel に変換する際に使用。
  • packing: バッチ内のシーケンス長を統一するかどうか。メモリ効率が向上。
  • neftune_noise_alpha: NEFTune ノイズを追加する際のハイパーパラメータ。
  • formatting_func: データセットの各サンプルをフォーマットするための関数。

TrainingArguments の重要な引数

TrainingArguments では、以下のような重要な引数があります。

  • output_dir: チェックポイントの保存先ディレクトリ。
  • per_device_train_batch_size: 1デバイスあたりのバッチサイズ。
  • gradient_accumulation_steps: 勾配を蓄積するステップ数。
  • max_steps: 学習の最大ステップ数。
  • num_train_epochs: 学習のエポック数。
  • learning_rate: 学習率。
  • weight_decay: 重み減衰率。
  • warmup_steps: ウォームアップのステップ数。
  • logging_steps: ロギングの間隔(ステップ数)。
  • logging_strategy: ロギングの戦略("steps", "epoch", "no")。
  • evaluation_strategy: 評価の戦略("steps", "epoch", "no")。
  • eval_steps: 評価の間隔(ステップ数)。
  • save_strategy: チェックポイントの保存戦略("steps", "epoch", "no")。
  • save_steps: チェックポイントの保存間隔(ステップ数)。

その他にも多くの引数がありますが、上記の引数は特に重要です。
以上のように、SFTTrainer と TrainingArguments を適切に組み合わせることで、様々な学習設定を柔軟に行えます。

ノートブック

Google Colaboratory

参考サイト

Trainer
We’re on a journey to advance and democratize artificial intelligence through open source and open science.
Trainer
We’re on a journey to advance and democratize artificial intelligence through open source and open science.

コメント

タイトルとURLをコピーしました