Callbacks

SyncRefModelCallback

class trl.SyncRefModelCallback

( ref_model: typing.Union[transformers.modeling_utils.PreTrainedModel, torch.nn.modules.module.Module] accelerator: typing.Optional[accelerate.accelerator.Accelerator] )

Callback to synchronize the model with a reference model.

RichProgressCallback

class trl.RichProgressCallback

< source >

( )

A TrainerCallback that displays the progress of training or evaluation using Rich.

WinRateCallback

class trl.WinRateCallback

< source >

( judge: BasePairwiseJudge trainer: Trainer generation_config: typing.Optional[transformers.generation.configuration_utils.GenerationConfig] = None num_prompts: typing.Optional[int] = None shuffle_order: bool = True use_soft_judge: bool = False )

Parameters

judge (BasePairwiseJudge) — The judge to use for comparing completions.
trainer (Trainer) — Trainer to which the callback will be attached. The trainer’s evaluation dataset must include a "prompt" column containing the prompts for generating completions. If the Trainer has a reference model (via the ref_model attribute), it will use this reference model for generating the reference completions; otherwise, it defaults to using the initial model.
generation_config (GenerationConfig, optional) — The generation config to use for generating completions.
num_prompts (int, optional) — The number of prompts to generate completions for. If not provided, defaults to the number of examples in the evaluation dataset.
shuffle_order (bool, optional, defaults to True) — Whether to shuffle the order of the completions before judging.
use_soft_judge (bool, optional, defaults to False) — Whether to use a soft judge that returns a win probability between 0 and 1 for the first completion vs the second.

A TrainerCallback that computes the win rate of a model based on a reference.

It generates completions using prompts from the evaluation dataset and compares the trained model’s outputs against a reference. The reference is either the initial version of the model (before training) or the reference model, if available in the trainer. During each evaluation step, a judge determines how often the trained model’s completions win against the reference using a judge. The win rate is then logged in the trainer’s logs under the key "eval_win_rate".

Usage:

trainer = DPOTrainer(...)
judge = PairRMJudge()
win_rate_callback = WinRateCallback(judge=judge, trainer=trainer)
trainer.add_callback(win_rate_callback)

LogCompletionsCallback

class trl.LogCompletionsCallback

< source >

( trainer: Trainer generation_config: typing.Optional[transformers.generation.configuration_utils.GenerationConfig] = None num_prompts: typing.Optional[int] = None freq: typing.Optional[int] = None )

Parameters

trainer (Trainer) — Trainer to which the callback will be attached. The trainer’s evaluation dataset must include a "prompt" column containing the prompts for generating completions.
generation_config (GenerationConfig, optional) — The generation config to use for generating completions.
num_prompts (int, optional) — The number of prompts to generate completions for. If not provided, defaults to the number of examples in the evaluation dataset.
freq (int, optional) — The frequency at which to log completions. If not provided, defaults to the trainer’s eval_steps.

A TrainerCallback that logs completions to Weights & Biases and/or Comet.

Usage:

trainer = DPOTrainer(...)
completions_callback = LogCompletionsCallback(trainer=trainer)
trainer.add_callback(completions_callback)

MergeModelCallback

class trl.MergeModelCallback

< source >

( merge_config: typing.Optional[ForwardRef('MergeConfig')] = None merge_at_every_checkpoint: bool = False push_to_hub: bool = False )

Parameters

merge_config (MergeConfig, optional) — Configuration used for the merging process. If not provided, the default MergeConfig is used.
merge_at_every_checkpoint (bool, optional, defaults to False) — Whether to merge the model at every checkpoint.
push_to_hub (bool, optional, defaults to False) — Whether to push the merged model to the Hub after merging.

A TrainerCallback that merges the policy model (the model being trained) with another model based on a merge configuration.

Example:

from trl.mergekit_utils import MergeConfig
from trl import MergeModelCallback

config = MergeConfig()
merge_callback = MergeModelCallback(config)
trainer = DPOTrainer(..., callbacks=[merge_callback])

BEMACallback

class trl.BEMACallback

< source >

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

Parameters

update_freq (int, optional, defaults to 400) — Update the BEMA weights every X steps. Denoted this as {@html " $\\phi$ "} in the paper.
ema_power (float, optional, defaults to 0.5) — Power for the EMA decay factor. Denoted {@html " $\\kappa$ "} in the paper. To disable EMA, set this to 0.0.
bias_power (float, optional, defaults to 0.2) — Power for the BEMA scaling factor. Denoted {@html " $\\eta$ "} in the paper. To disable BEMA, set this to 0.0.
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$ "} in the paper.
update_after (int, optional, defaults to 0) — Burn-in time before starting to update the BEMA weights. Denoted {@html " $\\tau$ "} in the paper.
multiplier (float, optional, defaults to 1.0) — Initial value for the EMA decay factor. Denoted as {@html " $\\gamma$ "} in the paper.
min_ema_multiplier (float, optional, defaults to 0.0) — Minimum value for the EMA decay factor.
device (str, optional, defaults to "cpu") — Device to use for the BEMA buffers, e.g. "cpu" or "cuda". Note that in most cases, this device SHOULD BE DIFFERENT from the device used for training in order to avoid OOM.

A TrainerCallback that implements BEMA (Bias-Corrected Exponential Moving Average) by Adam Block and Cyril Zhang. Code from https://github.com/abblock/bema under MIT license.

BEMA computes model weights that scale like: $\theta_t' = \alpha_t \cdot (\theta_t - \theta_0) + \text{EMA}_t$

where $\theta_t$ is the current model weights, $\theta_0$ is a snapshot of the model weights at the first update_after step, $\text{EMA}_t$ is the exponential moving average of the model weights, and $\alpha_t$ is a scaling factor that decays with the number of steps $t$ as $\alpha_t = (\rho + \gamma \cdot t)^{-\eta}.$

The EMA is computed as: $\text{EMA}_t = (1 - \beta_t) \cdot \text{EMA}_{t-1} + \beta_t \cdot \theta_t$

where $\beta_t$ is a decay factor that decays with the number of steps $t$ as $\beta_t = (\rho + \gamma \cdot t)^{-\kappa}.$

Example:

from trl import BEMACallback

trainer = Trainer(..., callbacks=[BEMACallback()])

WeaveCallback

class trl.WeaveCallback

< source >

( trainer: Trainer project_name: typing.Optional[str] = None scorers: typing.Optional[dict[str, callable]] = None generation_config: typing.Optional[transformers.generation.configuration_utils.GenerationConfig] = None num_prompts: typing.Optional[int] = None dataset_name: str = 'eval_dataset' model_name: typing.Optional[str] = None )

Parameters

trainer (Trainer) — Trainer to which the callback will be attached. The trainer’s evaluation dataset must include a "prompt" column containing the prompts for generating completions.
project_name (str, optional) — Name of the Weave project where data will be logged. If not provided, will try to use existing weave client or fall back to the active wandb run’s project name. Raises an error if none of these are available.
scorers (dict[str, Callable], optional) — Dictionary mapping scorer names to scorer functions. If None, operates in tracing mode (predictions only). If provided, operates in evaluation mode (predictions + scores + summary). Scorer functions should have signature: scorer(prompt: str, completion: str) -> Union[float, int]
generation_config (GenerationConfig, optional) — Generation config to use for generating completions.
num_prompts (int or None, optional) — Number of prompts to generate completions for. If not provided, defaults to the number of examples in the evaluation dataset.
dataset_name (str, optional, defaults to "eval_dataset") — Name for the dataset metadata in Weave.
model_name (str, optional) — Name for the model metadata in Weave. If not provided, attempts to extract from model config.

A TrainerCallback that logs traces and evaluations to W&B Weave. The callback uses https://weave-docs.wandb.ai/guides/evaluation/evaluation_logger/ to log traces and evaluations at each evaluation step.

Supports two modes based on the scorers parameter:

Tracing Mode (when scorers=None): Logs predictions for data exploration and analysis
Evaluation Mode (when scorers provided): Logs predictions with scoring and summary metrics

Both modes use Weave’s EvaluationLogger for structured, consistent data logging.

The callback logs data during evaluation phases (on_evaluate) rather than training steps, making it more efficient and semantically correct. It gracefully handles missing weave installation by logging warnings and skipping weave-specific functionality. It also checks for existing weave clients before initializing new ones.

Usage:

# Tracing mode (just log predictions)
trainer = DPOTrainer(...)
weave_callback = WeaveTraceCallback(trainer=trainer)  # project_name optional
trainer.add_callback(weave_callback)

# Or specify a project name
weave_callback = WeaveTraceCallback(trainer=trainer, project_name="my-llm-training")
trainer.add_callback(weave_callback)


# Evaluation mode (log predictions + scores + summary)
def accuracy_scorer(prompt: str, completion: str) -> float:
    # Your scoring logic here (metadata available via eval_attributes)
    return score


weave_callback = WeaveTraceCallback(
    trainer=trainer,
    project_name="my-llm-training",  # optional and needed only if weave client is not initialized
    scorers={"accuracy": accuracy_scorer},
)
trainer.add_callback(weave_callback)

on_train_begin

< source >

( args state control **kwargs )

Initialize Weave when training begins.

< > Update on GitHub

TRL

Callbacks

SyncRefModelCallback

class trl.SyncRefModelCallback

RichProgressCallback

class trl.RichProgressCallback

WinRateCallback

class trl.WinRateCallback

LogCompletionsCallback

class trl.LogCompletionsCallback

MergeModelCallback

class trl.MergeModelCallback

BEMACallback

class trl.BEMACallback

WeaveCallback

class trl.WeaveCallback

on_train_begin