Hugging Face's logo

TRL documentation

参照モデルのBEMA

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.17.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

参照モデルのBEMA

この機能は、DPOトレーニング中に参照モデルを更新するためのBEMAアルゴリズムを実装しています。

使用方法

from trl.experimental.bema_for_ref_model import BEMACallback, DPOTrainer
from datasets import load_dataset

dataset = load_dataset("trl-internal-testing/zen", "standard_preference", split="train")

bema_callback = BEMACallback(update_ref_model=True)

trainer = DPOTrainer(
    model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
    train_dataset=dataset,
    callbacks=[bema_callback],
)
trainer.train()

DPOTrainer

クラス trl.DPOTrainer

< >

( *args**kwargs )

train

< >

( resume_from_checkpoint: str | bool | None = Nonetrial: optuna.Trial | dict[str, Any] | None = Noneignore_keys_for_eval: list[str] | None = None ) → ~trainer_utils.TrainOutput

パラメータ

  • resume_from_checkpoint (str または bool、省略可) — str の場合、以前の Trainer インスタンスによって保存されたチェックポイントのローカルパスを指定します。bool で True の場合、以前の Trainer インスタンスによって args.output_dir に保存された最後のチェックポイントを読み込みます。指定がある場合、ここで読み込まれたモデル/オプティマイザ/スケジューラの状態からトレーニングを再開します。
  • trial (optuna.Trial または dict[str, Any]、省略可) — ハイパーパラメータ探索のための試行またはハイパーパラメータの辞書。
  • ignore_keys_for_eval (list[str]、省略可) — モデルの出力が辞書である場合、学習中の評価で予測を収集する際に無視すべきキーのリストです。

戻り値

~trainer_utils.TrainOutput

グローバルステップ数、トレーニング損失、およびメトリクスを含むオブジェクト。

メインの学習エントリポイント。

save_model

< >

( output_dir: str | None = None_internal_call: bool = False )

モデルを保存するため、from_pretrained()を使用して再読み込みできます。

メインプロセスからのみ保存します。

push_to_hub

< >

( commit_message: str | None = 'End of training'blocking: bool = Truetoken: str | None = Nonerevision: str | None = None**kwargs )

パラメータ

  • commit_message (str、任意、デフォルトは"End of training") — プッシュ時にコミットするメッセージ。
  • blocking (bool、省略可能、デフォルトはTrue) — git pushが完了した後にのみ関数が返すかどうか。
  • token (str、任意、デフォルトはNone) — Trainer の元の引数を上書きするための書き込み権限を持つトークン。
  • revision (str, 省略可) — コミット元の git リビジョンです。“main” ブランチのヘッドがデフォルトです。
  • kwargs (dict[str, Any]、省略可) — ~Trainer.create_model_card に渡される追加のキーワード引数。

self.args.hub_model_id のリポジトリに、self.model と self.processing_class を 🤗 model hub にアップロードします。

BEMACallback

クラス trl.BEMACallback

< >

( update_freq: int = 400ema_power: float = 0.5bias_power: float = 0.2lag: int = 10update_after: int = 0multiplier: float = 1.0min_ema_multiplier: float = 0.0device: str = 'cpu'update_ref_model: bool = Falseref_model_update_freq: int = 400ref_model_update_after: int = 0 )

パラメータ

  • update_freq (int, 省略可能, デフォルト値は 400) — BEMA の重みを X ステップごとに更新します。論文では {@html "ϕ \\phi ϕ"} in the paper.
  • ema_power (float, optional, defaults to 0.5) — Power for the EMA decay factor. Denoted {@html "κ \\kappa κ "} と表記されています。EMA を無効にするには、この値を 0.0.
  • bias_power (float, 省略可, デフォルト値は 0.2) — BEMA スケーリング係数のべき乗です。論文では {@html "η \\eta η"} で表されます。大きな値 (例: 8.0) を設定すると {@html "αt \\alpha_t αt​"} が 0 に減衰し、バイアス補正が無効化された状態を近似します。一方、0.0 を設定すると、すべてのステップで {@html "αt \\alpha_t αt​"} が 1 に固定されます (最大かつ減衰しない補正)。lag (int, 省略可, デフォルト値は 10) — ウェイト減衰スケジュールの初期オフセットです。更新の仮想の開始年齢として機能し、初期段階の滑らかさを制御します。論文では {@html "ρ \\rho ρ "}η \\eta η"} in the paper. A large value (e.g. 8.0) makes {@html "αt \\alpha_t αt​"} decay to 0, approximating disabled bias-correction; 0.0 instead pins {@html "αt \\alpha_t αt​"} at 1 for every step (maximum, undecayed correction).
  • lag (int, optional, defaults to 10) — Initial offset in the weight decay schedule that controls early-stage smoothness by acting as a virtual starting age for the updates. Denoted as {@html "ρ \\rho ρ"} で表されます。
  • update_after (int, 省略可能, デフォルト値は 0) — BEMA 重みの更新を開始するまでのバーンイン時間。論文では {@html "τ \\tau τ"} で表されます。τ \\tau τ"} in the paper.
  • multiplier (float, optional, defaults to 1.0) — Initial value for the EMA decay factor. Denoted as {@html "γ \\gamma γ"} 論文では γ で表されます。
  • min_ema_multiplier (float、任意、デフォルトは0.0) — EMA減衰係数の最小値。
  • device (str, 任意、デフォルトは "cpu") — BEMA バッファに使用するデバイス。例: "cpu" または "cuda"。ほとんどの場合、OOM を回避するために、このデバイスは学習に使用するデバイスとは異なるべきです。
  • update_ref_model (bool、省略可、デフォルトはFalse) — BEMA 重みで参照モデルを更新するかどうか。これにより、メインモデルの遅延・平滑化されたバージョンが参照モデルとして作成されます。
  • ref_model_update_freq (int、任意、デフォルトは400) — BEMA重みを使用して、このステップ数ごとに参照モデルを更新します。
  • ref_model_update_after (int、任意、デフォルトは0) — 参照モデルの更新を開始する前に待機するステップ数。

Adam Block と Cyril Zhang による BEMA（Bias-Corrected Exponential Moving Average）を実装した TrainerCallback です。コードは MIT ライセンスの下で https://github.com/abblock/bema から提供されています。

BEMA は、次のようにスケーリングされるモデルの重みを計算します: θt′=αt⋅(θt−θ0)+EMAt \theta_t' = \alpha_t \cdot (\theta_t - \theta_0) + \text{EMA}_t

ここで θt \theta_t は現在のモデルの重み、 θ0 \theta_0 は最初の update_after ステップ時点のモデルの重みのスナップショット、 EMAt \text{EMA}_t はステップ数αt \alpha_t とともに減衰するスケーリング係数であり、 t t は αt=(ρ+γ⋅t)−η. \alpha_t = (\rho + \gamma \cdot t)^{-\eta}.

EMA は次のように計算されます: EMAt=(1−βt)⋅EMAt−1+βt⋅θt \text{EMA}_t = (1 - \beta_t) \cdot \text{EMA}_{t-1} + \beta_t \cdot \theta_t

ここで βt \beta_t はステップ数 t t に対して次のように減衰する減衰係数です βt=(ρ+γ⋅t)−κ. \beta_t = (\rho + \gamma \cdot t)^{-\kappa}.

例:

>>> from trl import BEMACallback

>>> trainer = Trainer(..., callbacks=[BEMACallback()])
Update on GitHub

←非同期GRPO CPO→