> ## Documentation Index
> Fetch the complete documentation index at: https://mintlify.com/huggingface/trl/llms.txt
> Use this file to discover all available pages before exploring further.

# Callbacks

> Trainer callbacks for TRL, including BEMA, reference model sync, completion logging, rich progress display, and Weave integration.

TRL provides a set of `TrainerCallback` subclasses that extend Hugging Face `Trainer` with reinforcement-learning-specific features such as exponential moving average weight tracking, reference model synchronization, completion logging, and third-party observability integrations.

Import all callbacks from the top-level `trl` package:

```python theme={null}
from trl import (
    BEMACallback,
    LogCompletionsCallback,
    RichProgressCallback,
    SyncRefModelCallback,
    WeaveCallback,
)
```

***

## BEMACallback

`BEMACallback` implements **Bias-Corrected Exponential Moving Average (BEMA)**, introduced in [Block & Zhang (2025)](https://huggingface.co/papers/2508.00180). It maintains a running shadow model whose weights track the training model via a bias-corrected EMA scheme:

$$
\theta_t' = \alpha_t \cdot (\theta_t - \theta_0) + \text{EMA}_t
$$

where $\alpha_t = (\rho + \gamma \cdot t)^{-\eta}$ decays with the step count. The EMA itself is updated as:

$$
\text{EMA}_t = (1 - \beta_t) \cdot \text{EMA}_{t-1} + \beta_t \cdot \theta_t, \quad \beta_t = (\rho + \gamma \cdot t)^{-\kappa}
$$

At the end of training the shadow model is saved to `{output_dir}/bema/`.

<Note>
  The BEMA buffers live on a separate device (default `"cpu"`) to avoid out-of-memory errors on the training accelerator.
</Note>

### Signature

```python theme={null}
class BEMACallback(TrainerCallback):
    def __init__(
        self,
        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

<ParamField path="update_freq" type="int" default="400">
  Update the BEMA shadow model every this many steps. Denoted $\phi$ in the paper.
</ParamField>

<ParamField path="ema_power" type="float" default="0.5">
  Exponent $\kappa$ controlling the EMA decay factor $\beta_t$. Set to `0.0` to disable EMA.
</ParamField>

<ParamField path="bias_power" type="float" default="0.2">
  Exponent $\eta$ controlling the BEMA scaling factor $\alpha_t$. Set to `0.0` to disable bias correction.
</ParamField>

<ParamField path="lag" type="int" default="10">
  Initial offset $\rho$ in the weight decay schedule. Controls smoothness in early training by acting as a virtual starting age.
</ParamField>

<ParamField path="update_after" type="int" default="0">
  Burn-in steps $\tau$ before BEMA updates begin. The snapshot $\theta_0$ is taken at this step.
</ParamField>

<ParamField path="multiplier" type="float" default="1.0">
  Step multiplier $\gamma$ applied to the step count inside the decay schedule.
</ParamField>

<ParamField path="min_ema_multiplier" type="float" default="0.0">
  Floor value for the EMA decay factor $\beta_t$.
</ParamField>

<ParamField path="device" type="str" default="'cpu'">
  Device for BEMA buffers. Should differ from the training device to avoid OOM errors.
</ParamField>

### Example

```python theme={null}
from trl import BEMACallback
from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(output_dir="./output", max_steps=2000)
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    callbacks=[BEMACallback(update_freq=400, device="cpu")],
)
trainer.train()
# Shadow model saved to ./output/bema/
```

***

## SyncRefModelCallback

`SyncRefModelCallback` periodically synchronizes a reference model toward the current training model using an exponential moving average blend controlled by `ref_model_mixup_alpha`. It is used by trainers such as `DPOTrainer` when a soft-update reference policy is desired.

The sync is triggered at every step where `global_step % args.ref_model_sync_steps == 0`.

<Note>
  DeepSpeed ZeRO Stage 3 is handled automatically: parameters are gathered across ranks before the blend is applied.
</Note>

### Signature

```python theme={null}
class SyncRefModelCallback(TrainerCallback):
    def __init__(
        self,
        ref_model: PreTrainedModel | torch.nn.Module,
        accelerator: Accelerator | None,
    )
```

### Parameters

<ParamField path="ref_model" type="PreTrainedModel | torch.nn.Module">
  The reference model to keep synchronized with the training model.
</ParamField>

<ParamField path="accelerator" type="Accelerator | None">
  Accelerate `Accelerator` instance used to unwrap the model before syncing. Pass `None` if not using Accelerate.
</ParamField>

<Note>
  The sync frequency (`ref_model_sync_steps`) and blend coefficient (`ref_model_mixup_alpha`) are read from `TrainingArguments` at runtime, not from the callback constructor.
</Note>

### Example

```python theme={null}
from trl import SyncRefModelCallback
from accelerate import Accelerator

accelerator = Accelerator()
callback = SyncRefModelCallback(ref_model=ref_model, accelerator=accelerator)
trainer.add_callback(callback)
```

***

## LogCompletionsCallback

`LogCompletionsCallback` generates model completions for prompts from the evaluation dataset at regular intervals and logs them as a table to **Weights & Biases** and/or **Comet ML**. This makes it easy to track qualitative output quality throughout training.

<Warning>
  The trainer must have an evaluation dataset with a `"prompt"` column. A `ValueError` is raised at construction time if the dataset is absent.
</Warning>

### Signature

```python theme={null}
class LogCompletionsCallback(TrainerCallback):
    def __init__(
        self,
        trainer: Trainer,
        generation_config: GenerationConfig | None = None,
        num_prompts: int | None = None,
        freq: int | None = None,
    )
```

### Parameters

<ParamField path="trainer" type="Trainer">
  The trainer instance to attach the callback to. Used to access the model, tokenizer, accelerator, and evaluation dataset.
</ParamField>

<ParamField path="generation_config" type="GenerationConfig" optional>
  Generation configuration used when producing completions. If not provided the model's default config is used.
</ParamField>

<ParamField path="num_prompts" type="int" optional>
  Number of prompts sampled from the evaluation dataset. Defaults to the full evaluation dataset.
</ParamField>

<ParamField path="freq" type="int" optional>
  Logging frequency in steps. Defaults to `trainer.args.eval_steps`.
</ParamField>

### Example

```python theme={null}
from trl import DPOTrainer, LogCompletionsCallback

trainer = DPOTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,  # must have a "prompt" column
)
completions_callback = LogCompletionsCallback(trainer=trainer, num_prompts=32)
trainer.add_callback(completions_callback)
trainer.train()
```

***

## RichProgressCallback

`RichProgressCallback` replaces the default tqdm-based progress display with a [Rich](https://github.com/Textualize/rich) layout that shows training and evaluation progress bars alongside a live metrics table grouped by prefix.

<Note>
  This callback requires the `rich` package: `pip install rich`.
</Note>

### Signature

```python theme={null}
class RichProgressCallback(TrainerCallback):
    def __init__(self)
```

No constructor arguments are required.

### Example

```python theme={null}
from trl import RichProgressCallback
from transformers import Trainer

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    callbacks=[RichProgressCallback()],
)
trainer.train()
```

***

## WeaveCallback

`WeaveCallback` logs completions and optional scorer evaluations to [Weights & Biases Weave](https://weave-docs.wandb.ai/) during evaluation steps. It supports two modes:

* **Tracing mode** (`scorers=None`): logs predictions for data exploration.
* **Evaluation mode** (`scorers` provided): logs predictions with per-scorer scores and summary statistics.

Both modes use Weave's `EvaluationLogger` for structured logging.

<Warning>
  The trainer must have an evaluation dataset with a `"prompt"` column. A `ValueError` is raised at construction time if absent.
</Warning>

### Signature

```python theme={null}
class WeaveCallback(TrainerCallback):
    def __init__(
        self,
        trainer: Trainer,
        project_name: str | None = None,
        scorers: dict[str, Callable] | None = None,
        generation_config: GenerationConfig | None = None,
        num_prompts: int | None = None,
        dataset_name: str = "eval_dataset",
        model_name: str | None = None,
    )
```

### Parameters

<ParamField path="trainer" type="Trainer">
  Trainer instance to attach the callback to.
</ParamField>

<ParamField path="project_name" type="str" optional>
  Weave project name for logging. If not provided, the callback tries the existing Weave client, then the active wandb run. Raises a `ValueError` if none is available.
</ParamField>

<ParamField path="scorers" type="dict[str, Callable]" optional>
  Mapping of scorer names to scorer functions with signature `scorer(prompt: str, completion: str) -> float | int`. When provided, enables evaluation mode.
</ParamField>

<ParamField path="generation_config" type="GenerationConfig" optional>
  Generation configuration for producing completions.
</ParamField>

<ParamField path="num_prompts" type="int" optional>
  Number of evaluation prompts to use. Defaults to the full evaluation dataset.
</ParamField>

<ParamField path="dataset_name" type="str" default="'eval_dataset'">
  Name label for the dataset metadata in Weave.
</ParamField>

<ParamField path="model_name" type="str" optional>
  Name label for the model metadata in Weave. Extracted automatically from `model.config._name_or_path` if not provided.
</ParamField>

### Example

<CodeGroup>
  ```python tracing_mode.py theme={null}
  from trl import DPOTrainer, WeaveCallback

  trainer = DPOTrainer(
      model=model,
      args=training_args,
      train_dataset=train_dataset,
      eval_dataset=eval_dataset,
  )

  # Tracing mode — log predictions only
  weave_callback = WeaveCallback(trainer=trainer, project_name="my-llm-training")
  trainer.add_callback(weave_callback)
  trainer.train()
  ```

  ```python evaluation_mode.py theme={null}
  from trl import DPOTrainer, WeaveCallback

  trainer = DPOTrainer(
      model=model,
      args=training_args,
      train_dataset=train_dataset,
      eval_dataset=eval_dataset,
  )

  def accuracy_scorer(prompt: str, completion: str) -> float:
      # custom scoring logic
      return 1.0 if "correct" in completion else 0.0

  # Evaluation mode — log predictions + scores + summary
  weave_callback = WeaveCallback(
      trainer=trainer,
      project_name="my-llm-training",
      scorers={"accuracy": accuracy_scorer},
  )
  trainer.add_callback(weave_callback)
  trainer.train()
  ```
</CodeGroup>