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

此功能实现了 BEMA 算法，用于在 DPO 训练过程中更新参考模型。

用法

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 )

训练

< >

( 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

包含全局步数、训练损失和指标的 Object。

主要训练入口。

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) — 具有写入权限的 Token，用于覆盖 Trainer 的原始参数。
  • revision (str, 可选) — 要从中提交内容的 git 版本。默认为 “main” 分支的头部。
  • kwargs (dict[str, Any]，可选) — 传递给 ~Trainer.create_model_card 的额外关键字参数。

将 self.model 和 self.processing_class 上传到 🤗 模型库中的仓库 self.args.hub_model_id。

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) — 每隔 X 步更新一次 BEMA 权重。在论文中记为 {@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 τ"}。multiplier (float, 可选, 默认为 1.0) — EMA 衰减因子的初始值。论文中记为 {@html "γ \\gamma γ "}τ \\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) — 在开始更新参考模型之前需要等待的步数。

一个实现了 BEMA（偏差校正指数移动平均）的 TrainerCallback，由 Adam Block 和 Cyril Zhang 提出。代码来自 https://github.com/abblock/bema，采用 MIT 许可证。

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→