Trainerメソッドのサブクラス化

Trainer のメソッドをサブクラス化することで、ループ全体を書き直すことなく学習の動作を変更できます。サブクラス化により、training loop（例えばフォワードパスや損失の計算）が変更されます。

サブクラス化を検討する前に、Trainerが何を計算するか、あるいはいつおよびどの条件で動作するかを変更する必要があるかを検討してください。タイミングや条件付きロジックを制御する場合は、代わりにCallbackを使用してください。Callbackはイベントの発生タイミング（ログ出力、評価、早期停止など）を制御し、サブクラス化は実行される処理の内容（損失の計算、データ読み込み、最適化など）を変更します。

サブクラス化できるメソッドの完全なリストについては、Trainer の API ドキュメントを参照してください。_ で始まるプライベートメソッド（_save_checkpoint や _evaluate など）もオーバーライドできますが、これらは予告なく変更される場合があります。

get_train_dataloader

標準的なget_train_dataloader()メソッドは、1つのバッチを読み込み、それに対して学習を行い、破棄して、次のバッチを読み込みます。

def get_train_dataloader(self):
    return self._get_dataloader(
        batch_size=self._train_batch_size,
        ...
)

GRPO は、学習前に補完を生成するオンライン強化学習アルゴリズムです。補完の生成は自己回帰的であるため、各ステップで実行するとコストが高くなります。512トークンの補完には、学習ステップの1回のフォワードパスに対して約512回の逐次的なフォワードパスが必要です。GRPOTrainer は get_train_dataloader() をサブクラス化し、複数ステップにわたって生成をバッチ処理します。

trl.GRPOTrainer.get_train_dataloader は、バッチサイズに steps_per_generation 引数を乗算することで、複数の学習ステップ分の バッチ を一度に読み込みます。train_batch_size=4 かつ steps_per_generation=8 の場合、dataloader は 32 のバッチを生成し、生成コストを 8 分の 1 に削減します。

def get_train_dataloader(self):
    dataloader_params = {
        "batch_size": self._train_batch_size * self.args.steps_per_generation, # this is the only change
        ...
    }

損失を計算

compute_loss() は、モデルによって計算されたクロスエントロピー損失を返します。

def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
    ...
    outputs = model(**inputs)
    ...
    loss = outputs["loss"] # get loss from model

    return (loss, outputs) if return_outputs else loss

DPOは、参照モデルに対する相対的な観点から、ポリシーモデルが拒否された応答よりも選択された応答をどの程度強く好むかを測定します。DPOTrainerは、損失の計算が標準的なクロスエントロピーと複数の点で異なるため、compute_loss()をサブクラス化しています。

モデルはラベルを一切参照せず、DPOが対数確率を計算するためのlogitsのみを返します
選択された応答と拒否された応答が連結されます
参照モデルが自身のlog-probsを計算します
損失はπ_chosen、π_rejected、π_ref_chosen、π_ref_rejectedの関数です

上記のいずれも、標準的なTrainer.compute_loss()メソッドには当てはまりません。

def compute_loss(
    self,
    model: PreTrainedModel | nn.Module,
    inputs: dict[str, torch.Tensor | Any],
    return_outputs=False,
    num_items_in_batch=None,
) -> torch.Tensor | tuple[torch.Tensor, dict[str, float]]:
    ...
    outputs = model(**inputs)
    logits = outputs.logits
    logps = get_logps(logits, inputs)
    chosen_logps, rejected_logps = logps.chunk(2, dim=0)  # batch is [chosen, rejected]
    ref_logits = self.ref_model(**inputs).logits
    ref_logps = get_logps(ref_logits, inputs)
    ref_chosen_logps, ref_rejected_logps = ref_logps.chunk(2, dim=0)  # batch is [chosen, rejected]
    chosen_scores = chosen_logps - ref_chosen_logps
    rejected_scores = rejected_logps - ref_rejected_logps
    per_sequence_loss = -F.logsigmoid(self.beta * chosen_scores - rejected_scores)
    loss = per_sequence_loss.mean()
    return (loss, outputs) if return_outputs else loss

次のステップ

より実用的な例として、TRL における GRPOTrainer と DPOTrainer が Trainer をどのように拡張しているか、または Axolotl がその上にカスタムトレーナーをどのように構築しているかを参照してください。
トレーニングステップの終了時にメトリクスを記録するなど、トレーニングイベント中に発生する処理のみをカスタマイズする必要がある場合は、Callbacksガイドを確認してください。

Update on GitHub

Transformers

Trainerメソッドのサブクラス化

get_train_dataloader

損失を計算

次のステップ