# BEMA для эталонной модели

Эта функция реализует алгоритм BEMA для обновления эталонной модели во время обучения DPO.

## Использование

```python
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]]

#### trl.DPOTrainer[[trl.DPOTrainer]]

```python
trl.DPOTrainer(*args, **kwargs)
```

[Источник](https://github.com/huggingface/trl/blob/main/trl/experimental/bema_for_ref_model/dpo_trainer.py#L19)

#### train[[trl.DPOTrainer.train]]

```python
train(resume_from_checkpoint: str | bool | None = None, trial: optuna.Trial | dict[str, Any] | None = None, ignore_keys_for_eval: list[str] | None = None)
```

[Источник](https://github.com/huggingface/trl/blob/main/transformers/trainer.py#L1450)

**Параметры:**

resume_from_checkpoint (`str` или `bool`, *необязательный*) : Если передано значение типа `str`, это локальный путь к сохранённой контрольной точке, созданной предыдущим экземпляром `Trainer`. Если передано значение типа `bool` и оно равно `True`, загружается последняя контрольная точка из *args.output_dir*, сохранённая предыдущим экземпляром `Trainer`. Если параметр указан, обучение возобновляется из состояний модели/оптимизатора/планировщика, загруженных отсюда.

trial (`optuna.Trial` или `dict[str, Any]`, *необязательный*) : Запуск испытания или словарь гиперпараметров для поиска гиперпараметров.

ignore_keys_for_eval (`list[str]`, *необязательный*) : Список ключей в выводе вашей модели (если он является словарём), которые следует игнорировать при сборе предсказаний для оценки во время обучения.

**Возвращает:** `~trainer_utils.TrainOutput`

Объект, содержащий количество глобальных шагов, потерю обучения и метрики.

Основная точка входа для обучения.

#### save_model[[trl.DPOTrainer.save_model]]

```python
save_model(output_dir: str | None = None, _internal_call: bool = False)
```

[Источник](https://github.com/huggingface/trl/blob/main/transformers/trainer.py#L4077)

Сохранит модель, чтобы вы могли загрузить её снова с помощью `from_pretrained()`.

Сохранение будет выполняться только из основного процесса.

#### push_to_hub[[trl.DPOTrainer.push_to_hub]]

```python
push_to_hub(commit_message: str | None = 'End of training', blocking: bool = True, token: str | None = None, revision: str | None = None, private: bool | None = None, **kwargs)
```

[Источник](https://github.com/huggingface/trl/blob/main/transformers/trainer.py#L4337)

**Параметры:**

commit_message (`str`, *необязательный*, по умолчанию `"End of training"`) : Сообщение для коммита при отправке.

blocking (`bool`, *необязательный*, по умолчанию `True`) : Указывает, должна ли функция возвращаться только после завершения `git push`.

token (`str`, *необязательный*, по умолчанию `None`) : Токен с правом записи для перезаписи исходных аргументов Trainer.

revision (`str`, *необязательный*) : git-ревизия, от которой выполняется коммит. По умолчанию используется head ветки "main".

private (`bool`, *необязательный*) : Указывает, следует ли сделать репозиторий приватным при его создании. По умолчанию используется значение `args.hub_private_repo`. Не имеет эффекта, если репозиторий уже существует.

kwargs (`dict[str, Any]`, *необязательный*) : Дополнительные именованные аргументы, передаваемые в `~Trainer.create_model_card`.

**Возвращает:**

URL репозитория, в который была загружена модель, если `blocking=False`, или объект `Future`, отслеживающий ход коммита, если `blocking=True`.

Загрузите `self.model` и `self.processing_class` в модельный хаб 🤗 в репозиторий `self.args.hub_model_id`.

## BEMACallback[[trl.BEMACallback]]

#### trl.BEMACallback[[trl.BEMACallback]]

```python
trl.BEMACallback(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', update_ref_model: bool = False, ref_model_update_freq: int = 400, ref_model_update_after: int = 0)
```

[Источник](https://github.com/huggingface/trl/blob/main/trl/experimental/bema_for_ref_model/callback.py#L59)

**Параметры:**

update_freq (`int`, *optional*, по умолчанию `400`) : обновлять веса BEMA каждые X шагов. В статье это обозначено как \\( \phi \\).

ema_power (`float`, *optional*, по умолчанию `0.5`) : Степень для коэффициента затухания EMA. Обозначается как \\( \kappa \\) в статье. Чтобы отключить EMA, установите это значение в `0.0`.

bias_power (`float`, *optional*, по умолчанию `0.2`) : Степень для коэффициента масштабирования BEMA. Обозначается как \\( \eta \\) в статье. Большое значение (например, `8.0`) приводит к тому, что \\( \alpha_t \\) стремится к `0`, что приближает отключённую коррекцию смещения; `0.0` вместо этого фиксирует \\( \alpha_t \\) на уровне `1` на каждом шаге (максимальная, не затухающая коррекция).

lag (`int`, *optional*, по умолчанию `10`) : Начальное смещение в расписании весового затухания, которое контролирует плавность на ранних этапах, выступая в роли виртуального начального возраста для обновлений. Обозначается как \\( \rho \\) в статье.

update_after (`int`, *optional*, по умолчанию `0`) : Время прогрева перед началом обновления весов BEMA. Обозначается как \\( \tau \\) в статье.

multiplier (`float`, *optional*, по умолчанию `1.0`) : Начальное значение для коэффициента затухания EMA. Обозначается как  \\( \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`) : Количество шагов ожидания перед началом обновления эталонной модели.

Класс [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback), реализующий [BEMA](https://huggingface.co/papers/2508.00180) (Bias-Corrected Exponential Moving Average) от [Adam Block](https://huggingface.co/abblock) и [Cyril Zhang](https://huggingface.co/cyrilzhang). Код взят с https://github.com/abblock/bema под лицензией MIT.

BEMA вычисляет веса модели, которые масштабируются следующим образом:

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

где \\( \theta_t \\) — текущие веса модели, \\( \theta_0 \\) — снимок весов модели на первом шаге `update_after`, \\( \text{EMA}_t  \\) — экспоненциальное скользящее среднее весов модели, а \\( \alpha_t \\) — коэффициент масштабирования, который затухает с числом шагов \\( t \\) следующим образом

$$ \alpha_t = (\rho + \gamma \cdot t)^{-\eta}. $$

EMA вычисляется следующим образом:

$$ \text{EMA}_t = (1 - \beta_t) \cdot \text{EMA}_{t-1} + \beta_t \cdot \theta_t $$

где \\( \beta_t \\) — коэффициент затухания, который уменьшается с числом шагов \\( t \\) следующим образом

$$ \beta_t = (\rho + \gamma \cdot t)^{-\kappa}. $$

Пример:

```python
>>> from trl import BEMACallback

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