RLOO Trainer

Обзор

TRL поддерживает RLOO Trainer для обучения языковых моделей, как описано в статье Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs авторов Arash Ahmadian, Chris Cremer, Matthias Gallé, Marzieh Fadaee, Julia Kreutzer, Ahmet Üstün и Sara Hooker.

Аннотация из статьи следующая:

Выравнивание ИИ в форме обучения с подкреплением на основе обратной связи от человека (RLHF) всё чаще рассматривается как ключевой компонент для высокопроизводительных больших языковых моделей. В недавней литературе метод проксимальной оптимизации политики (PPO) позиционируется как канонический подход для части RL в RLHF. Однако он связан с высокими вычислительными затратами и чувствительной настройкой гиперпараметров. Мы утверждаем, что большинство мотивационных принципов, приведших к разработке PPO, имеют меньшее практическое значение в контексте RLHF, и предлагаем менее вычислительно затратный метод, который сохраняет и даже повышает производительность. Мы пересматриваем формулировку выравнивания на основе человеческих предпочтений в контексте RL. Простота как руководящий принцип позволяет нам показать, что многие компоненты PPO не нужны в контексте RLHF, и что гораздо более простые варианты оптимизации в стиле REINFORCE превосходят как PPO, так и недавно предложенные методы «без RL», такие как DPO и RAFT. Наша работа показывает, что тщательная адаптация к характеристикам выравнивания LLM позволяет извлекать выгоду из онлайн-оптимизации RL при низких затратах.

Этот метод постобучения был предложен Costa Huang и впоследствии переписан Shirin Yamani.

Быстрый старт

В этом примере показано, как обучать модель с использованием метода RLOO. Мы обучаем модель Qwen 0.5B Instruct на промптах из набора данных DeepMath-103K. Данные из набора можно посмотреть здесь:

Ниже приведён скрипт для обучения модели. accuracy_reward требует math-verify, который можно установить с помощью pip install trl[math_verify].

# train_rloo.py
from datasets import load_dataset
from trl import RLOOTrainer
from trl.rewards import accuracy_reward

dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

trainer = RLOOTrainer(
    model="Qwen/Qwen2-0.5B-Instruct",
    reward_funcs=accuracy_reward,
    train_dataset=dataset,
)
trainer.train()

Выполните скрипт с помощью следующей команды:

torchrun --nproc_per_node 8 train_rloo.py

Более подробный обзор метода RLOO

RLOO — это алгоритм обучения в режиме онлайн, то есть он совершенствуется итеративно, используя данные, сгенерированные самим обученным моделью в процессе обучения. Интуиция, лежащая в основе целевой функции RLOO, состоит в том, чтобы максимизировать преимущество сгенерированных завершений, при этом обеспечивая близость модели к эталонной политике. Чтобы понять, как работает RLOO, его можно разбить на четыре основных шага: генерация завершений, вычисление преимущества, оценка расхождения Кульбака–Лейблера и вычисление функции потерь.

Генерация завершений

На каждом шаге обучения мы выбираем батч промптов и генерируем набор $G$ дополнений для каждого промпта (обозначаются как $o_i$).

Вычисление вознаграждения

В RLOO награда состоит из двух компонентов: награды, предоставляемой моделью наград (или функцией наград), и штрафного члена KL, который не позволяет политике отклоняться слишком сильно от фиксированной опорной политики

1. Для каждой из $G$ сгенерированных последовательностей $o_i = (o_{i,1}, \dots, o_{i,T})$ в условиях запроса $q$, мы вычисляем скалярное вознаграждение с помощью модели вознаграждения $R(o_i, q)$.
2. Одновременно мы оцениваем расхождение Кульбака–Лейблера между текущей политикой $\pi_\theta$ и фиксированной опорной политикой $\pi_{\text{ref}}$ на последовательности. Оценка KL для последовательности $o_i$ имеет вид: $$\mathbb{D}_{\mathrm{KL}}\!\left[\pi_\theta\|\pi_{\mathrm{ref}}\right] = \sum_{t=1}^T \log \frac{\pi_\theta(o_{i,t} \mid q, o_{i,<t})}{\pi_{\mathrm{ref}}(o_{i,t} \mid q, o_{i,<t})}.$$

Итоговое вознаграждение, назначенное последовательности $o_i$ затем равно: $$r_i = R(o_i, q) - \beta \, \mathbb{D}_{\mathrm{KL}}\!\left[\pi_\theta \|\pi_{\mathrm{ref}}\right],$$

где $\beta > 0$ управляет силой штрафного члена KL.

В чисто онлайн-настройке (num_iterations = 1, значение по умолчанию) данные генерируются текущей политикой. В этом случае штраф за KL вычисляется непосредственно с использованием текущей политики.

В более общем случае (например, несколько шагов градиента на батч) данные генерируются более ранним снимком $\pi_{\text{old}}$. Чтобы штраф оставался согласованным с распределением выборки, KL определяется относительно этой политики:

$\mathbb{D}_{\mathrm{KL}}\!\left[\pi_{\text{old}} \,\|\, \pi_{\text{ref}}\right].$

Эквивалентно, для выборочной последовательности $o$, оценка Монте-Карло равна:

$\mathbb{D}_{\mathrm{KL}}\!\left[\pi_{\text{old}} \|\pi_{\mathrm{ref}}\right] = \sum_{t=1}^T \log \frac{\pi_{\text{old}}(o_{i,t} \mid q, o_{i,<t})}{\pi_{\mathrm{ref}}(o_{i,t} \mid q, o_{i,<t})}.$

Вычисление преимущества

После того как вознаграждения для каждого завершения вычислены, мы рассчитываем базовую линию (baseline) как среднее вознаграждение всех остальных образцов в том же батче, исключая текущий образец. Эта базовая линия используется для снижения дисперсии оценки градиента политики. Затем преимущество (advantage) для каждого завершения получается как разница между его собственным вознаграждением и этой базовой линией, исключающей текущий образец (leave-one-out).

Формально, для батча из G завершений базовая линия для завершения i:$$b_i = \frac{1}{G-1} \sum_{j \neq i} r_j$$

а затем преимущество для каждого завершения вычисляется как разница между его наградой и базовой линией: $$A_i = r_i - b_i$$

Вычисление функции потерь

Потеря REINFORCE определяется следующим образом: $$\mathcal{L}_{\text{RLOO}}(\theta) = - \frac{1}{G} \sum_{i=1}^G \hat{A}_i \, \log \pi_\theta(o_i \mid q)$$

На практике выполнение нескольких шагов градиента на одной и той же партии делает действия фактически off-policy относительно текущих параметров. Для исправления этого мы вводим коэффициент важности выборки (importance sampling ratio). Чтобы предотвратить чрезмерно большие обновления, когда политика меняется между выборкой и шагами градиента, мы ограничиваем (clip) этот коэффициент: $$\mathcal{L}_{\text{RLOO}}(\theta) = - \frac{1}{G} \sum_{i=1}^G \min \left( \frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)} \hat{A}_i, \, \text{clip}\left(\frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)}, 1-\epsilon, 1+\epsilon\right) \hat{A}_i \right)$$

В полностью онлайн-настройке с одним шагом (по умолчанию), $\frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)} = 1$ и это сводится к стандартному REINFORCE.

В отличие от GRPO, RLOO не выполняет обратное распространение градиентов через член KL. Здесь KL используется исключительно для формирования вознаграждения: он вычисляется в режиме no_grad и включается в скалярное вознаграждение $r_i$, которое становится отсоединённым преимуществом $\hat{A}_i$.

Иными словами, единственный путь, по которому градиент поступает в политику, — это член отношения важности $\frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)}$, а член KL не распространяется на политику при обратном распространении ошибки.

В отличие от этого, GRPO добавляет в свою целевую функцию явный дифференцируемый член KL, вычисляемый на лету с использованием текущей политики $\pi_\theta$, поэтому её градиент проходит и через член отношения, и через штраф KL.

Записываемые метрики

Во время обучения и оценки мы фиксируем следующие метрики:

• num_tokens: общее количество обработанных на данный момент токенов, включая как промпты, так и завершения.
• step_time: среднее время (в секундах), затрачиваемое на один шаг обучения (включая генерацию).
• completions/mean_length: средняя длина сгенерированных завершений.
• completions/min_length: минимальная длина генерируемых завершений.
• completions/max_length: максимальная длина генерируемых дополнений.
• completions/mean_terminated_length: средняя длина сгенерированных завершений, которые заканчиваются токеном EOS.
• completions/min_terminated_length: минимальная длина сгенерированных завершений, которые заканчиваются токеном EOS.
• completions/max_terminated_length: максимальная длина сгенерированных завершений, которые заканчиваются токеном EOS.
• completions/clipped_ratio: Доля обрезанных (усечённых) завершений.
• rewards/{reward_func_name}/mean: среднее значение вознаграждения от конкретной функции вознаграждения.
• rewards/{reward_func_name}/std: стандартное отклонение вознаграждения от конкретной функции вознаграждения.
• reward: Общая средняя награда после суммирования наград по функциям (с весами reward_weights).
• reward_std: стандартное отклонение суммарных наград по функциям (с учётом весов reward_weights), вычисленное для всей выборки.
• frac_reward_zero_std: доля образцов в батче генерации, у которых стандартное отклонение награды равно нулю, что указывает на малое разнообразие для данного промпта (все ответы верные или неверные).
• entropy: Средняя энтропия предсказаний токенов по всем сгенерированным завершениям. (Если mask_truncated_completions=True, токены замаскированных последовательностей исключаются.)
• aux_loss: вспомогательная функция потерь для балансировки нагрузки в модели Mixture-of-Experts, до её масштабирования коэффициентом router_aux_loss_coef и добавления к общей функции потерь. Записывается в журнал только в том случае, если модель является MoE и router_aux_loss_coef не равен нулю.
• kl: средняя дивергенция Кульбака–Лейблера между моделью и эталонной моделью, рассчитанная по сгенерированным завершениям. Записывается в журнал только если beta не равно нулю.
• clip_ratio/region_mean: Соотношение вероятностей последовательностей, для которых целевая функция RLOO обрезается, чтобы оставаться в области доверия: $\text{clip}\left( r_{i}(\theta), 1 - \epsilon_\mathrm{low}, 1 + \epsilon_\mathrm{high} \right)\,, \quad r_{i}(\theta) = \frac{\pi_\theta(o_{i} \mid q)}{\pi_{\theta_{\text{old}}}(o_{i} \mid q)}$. Более высокое значение означает, что обрезается больше выборок, что ограничивает степень изменения политики $\pi_\theta$ .
• clip_ratio/low_mean: Среднее отношение вероятностей последовательностей, обрезанных по нижней границе области доверия: $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/low_min: Наименьший индикатор обрезки на нижнюю границу области доверия для каждого завершения: $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/high_mean: Среднее отношение вероятностей последовательностей, обрезанных по верхней границе области доверия: $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.
• clip_ratio/high_max: Наибольший индикатор обрезки на верхнюю границу области доверия для каждого завершения: $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.

Настройка

Ускорение обучения с генерацией на базе vLLM

Генерация часто является основным узким местом при обучении с использованием онлайн-методов. Для ускорения генерации можно использовать vLLM — высокопроизводительный движок инференса с низкой задержкой для LLM. Для его включения сначала установите пакет с

pip install trl[vllm]

Мы поддерживаем два способа использования vLLM во время обучения: режим сервера и режим colocation.

Вариант 1: Режим Colocate

В этом режиме vLLM работает внутри процесса тренера и разделяет память GPU с моделью обучения. Это позволяет избежать запуска отдельного сервера и может повысить использование GPU, но может привести к конкуренции за память на GPU, используемых для обучения. Это режим по умолчанию.

from trl import RLOOConfig

training_args = RLOOConfig(
    ...,
    use_vllm=True,  # vllm_mode="colocate" by default
)

Вариант 2: серверный режим

В этом режиме vLLM работает в отдельном процессе (и на отдельных GPU) и взаимодействует с тренером через HTTP. Это идеальный вариант, если у вас есть выделенные GPU для инференса.

1. Запустите сервер vLLM:

   Скопировано

   VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
       --weight-transfer-config '{"backend": "nccl"}' \
       --logprobs-mode processed_logprobs \
       --max-logprobs -1

2. Включите серверный режим в вашем скрипте обучения:

   Скопировано

   from trl import RLOOConfig
   
   training_args = RLOOConfig(
       ...,
       use_vllm=True,
       vllm_mode="server",
   )

Убедитесь, что сервер использует другие GPU, чем тренажёр, иначе могут возникнуть ошибки NCCL. Вы можете указать используемые GPU с помощью переменной окружения CUDA_VISIBLE_DEVICES.

В зависимости от размера модели и общих требований к памяти GPU для обучения может потребоваться отрегулировать параметр vllm_gpu_memory_utilization в RLOOConfig, чтобы избежать недоиспользования ресурсов или ошибок нехватки памяти.

Мы предоставляем HF Space, который помогает оценить рекомендуемое использование памяти GPU на основе конфигурации вашей модели и настроек эксперимента. Просто используйте его следующим образом, чтобы получить рекомендацию по vllm_gpu_memory_utilization:

Если рекомендуемое значение не работает в вашей среде, мы рекомендуем добавить небольшой запас (например, +0.05 или +0.1) к рекомендуемому значению для обеспечения стабильности.

Если вы по-прежнему получаете ошибки нехватки памяти, установите vllm_enable_sleep_mode в True, и параметры vllm и кэш будут выгружены во время шага оптимизации. Подробнее см. в разделе Снижение потребления памяти с помощью режима сна vLLM.

По умолчанию RLOO использует MASTER_ADDR=localhost и MASTER_PORT=12345 для vLLM, но вы можете переопределить эти значения, установив соответствующие переменные окружения.

Для получения дополнительной информации см. Ускорение обучения с помощью vLLM.

RLOO в масштабе: обучение модели 70B+ на нескольких узлах

При обучении крупных моделей, таких как Qwen2.5-72B, для обеспечения эффективности и масштабируемости обучения на нескольких GPU и узлах необходимы несколько ключевых оптимизаций. К ним относятся:

• DeepSpeed ZeRO Stage 3: ZeRO использует параллелизм данных для распределения состояний модели (весов, градиентов, состояний оптимизатора) между несколькими GPU и CPU, что снижает требования к памяти и вычислительным ресурсам на каждом устройстве. Поскольку большие модели не помещаются на одном GPU, для обучения таких моделей необходимо использовать ZeRO Stage 3. Подробнее см. в разделе DeepSpeed Integration.
• torchrun: запускатель PyTorch запускает по одному процессу на каждый GPU на каждом узле и связывает их через общий rendezvous. Подробнее см. в разделе Распределение обучения.
• vLLM: см. предыдущий раздел о том, как использовать vLLM для ускорения генерации.

Ниже приведён пример SLURM-скрипта для обучения модели 70B с использованием RLOO на нескольких узлах. Этот скрипт обучает модель на 4 узлах и использует 5-й узел для генерации на базе vLLM.

#!/bin/bash
#SBATCH --nodes=5
#SBATCH --gres=gpu:8

# Get the list of allocated nodes
NODELIST=($(scontrol show hostnames $SLURM_JOB_NODELIST))

# Assign the first 4 nodes for training and the 5th node for vLLM
TRAIN_NODES="${NODELIST[@]:0:4}"  # Nodes 0, 1, 2, 3 for training
VLLM_NODE="${NODELIST[4]}"  # Node 4 for vLLM

# Run training on the first 4 nodes (Group 1)
srun --nodes=4 --ntasks=4 --nodelist="${NODELIST[@]:0:4}" torchrun \
     --nnodes 4 \
     --nproc_per_node 8 \
     --node_rank $SLURM_PROCID \
     --rdzv_backend c10d \
     --rdzv_endpoint ${NODELIST[0]}:29500 \
     train_rloo.py \
     --server_ip $VLLM_NODE &

# Run vLLM server on the 5th node (Group 2)
srun --nodes=1 --ntasks=1 --nodelist="${NODELIST[4]}" env VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen2.5-72B --tensor-parallel-size 8 \
    --weight-transfer-config '{"backend": "nccl"}' \
    --logprobs-mode processed_logprobs \
    --max-logprobs -1 &

wait

import argparse

from datasets import load_dataset
from trl import RLOOTrainer, RLOOConfig

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--vllm_server_host", type=str, default="", help="The server IP")
    args = parser.parse_args()

    # Example dataset from TLDR
    dataset = load_dataset("trl-lib/tldr", split="train")

    # Dummy reward function: count the number of unique characters in the completions
    def reward_num_unique_chars(completions, **kwargs):
        return [len(set(c)) for c in completions]

    training_args = RLOOConfig(
        output_dir="Qwen2.5-72B-RLOO",
        per_device_train_batch_size=4,
        bf16=True,
        deepspeed="examples/deepspeed_configs/zero3.json",
        use_vllm=True,
        vllm_mode="server",
        vllm_server_host=args.vllm_server_host.replace("ip-", "").replace("-", "."),  # from ip-X-X-X-X to X.X.X.X
    )

    trainer = RLOOTrainer(model="Qwen/Qwen2.5-72B", args=training_args, reward_funcs=reward_num_unique_chars, train_dataset=dataset)
    trainer.train()

if __name__=="__main__":
    main()

Использование пользовательской функции вознаграждения

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

Функции вознаграждения могут быть как синхронными вызываемыми объектами Python, так и асинхронными корутинами async def. Если вы указываете несколько асинхронных функций вознаграждения, они ожидаются одновременно (выполняются параллельно через asyncio.gather), что позволяет перекрыть их задержки.

1. Аргументы ввода:

   • Функция должна принимать следующие аргументы по имени:

     • prompts (содержит промпты),
     • completions (содержит сгенерированные завершения),
     • completion_ids (содержит токенизированные завершения),
     • trainer_state (TrainerState): текущее состояние тренера. Это можно использовать для реализации динамических функций вознаграждения, например, для обучения по программе (curriculum learning), где вознаграждение корректируется в зависимости от прогресса обучения.
     • log_extra: вызываемый объект log_extra(column: str, values: list) для добавления дополнительных столбцов в таблицу с результатами. См. Пример 6.
     • log_metric: вызываемый объект log_metric(name: str, value: float) для записи скалярных метрик в виде графиков наряду с kl, entropy и т. д. См. Пример 6.
     • Все имена столбцов (кроме prompt), которые могут присутствовать в наборе данных. Например, если набор данных содержит столбец с именем ground_truth, функция будет вызвана с ground_truth в качестве именованного аргумента.

     Самый простой способ выполнить это требование — использовать **kwargs в сигнатуре функции.

   • В зависимости от формата набора данных входные данные будут различаться:

     • Для стандартного формата prompts и completions будут списками строк.
     • Для диалогового формата prompts и completions будут списками словарей сообщений.

2. Возвращаемое значение: функция должна возвращать список чисел с плавающей точкой. Каждое число представляет награду, соответствующую одному завершению.

Пример 1: Поощрение более длинных завершений

Ниже приведен пример функции вознаграждения для стандартного формата, которая поощряет более длинные завершения:

def reward_func(completion_ids, **kwargs):
    """Reward function that assigns higher scores to longer completions (in terms of token count)."""
    return [float(len(ids)) for ids in completion_ids]

Вы можете протестировать это следующим образом:

>>> prompts = ["The sky is", "The sun is"]  # not used in the reward function, but the trainer will pass it
>>> completions = [" blue.", " in the sky."]  # not used in the reward function, but the trainer will pass it
>>> completion_ids = [[6303, 13], [304, 279, 12884, 13]]
>>> reward_func(prompts=prompts, completions=completions, completion_ids=completion_ids)
[2.0, 4.0]

Пример 1.1: Поощрение более длинных завершений (на основе количества символов)

Аналогично предыдущему примеру, но на этот раз функция вознаграждения основана на количестве символов, а не токенов.

def reward_func(completions, **kwargs):
    """Reward function that assigns higher scores to longer completions (in terms of character count)."""
    return [float(len(completion)) for completion in completions]

Вы можете протестировать это следующим образом:

>>> prompts = ["The sky is", "The sun is"]
>>> completions = [" blue.", " in the sky."]
>>> completion_ids = [[6303, 13], [304, 279, 12884, 13]]  # not used in the reward function, but the trainer will pass it
>>> reward_func(prompts=prompts, completions=completions, completion_ids=completion_ids)
[6.0, 12.0]

Пример 2: Награждайте завершения с определенным форматом

Ниже приведен пример функции вознаграждения, которая проверяет, соответствует ли ответ определенному формату. Этот пример вдохновлен функцией format reward, используемой в статье DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. Она предназначена для диалогового формата, в котором промпты и ответы состоят из структурированных сообщений.

import re

def format_reward_func(completions, **kwargs):
    """Reward function that checks if the completion has a specific format."""
    pattern = r"^<think>.*?</think><answer>.*?</answer>$"
    completion_contents = [completion[0]["content"] for completion in completions]
    matches = [re.match(pattern, content) for content in completion_contents]
    return [1.0 if match else 0.0 for match in matches]

Вы можете протестировать эту функцию следующим образом:

>>> prompts = [
...     [{"role": "user", "content": "What is the result of (1 + 2) * 4?"}],
...     [{"role": "user", "content": "What is the result of (3 + 1) * 2?"}],
... ]
>>> completions = [
...     [{"role": "assistant", "content": "<think>The sum of 1 and 2 is 3, which we multiply by 4 to get 12.</think><answer>(1 + 2) * 4 = 12</answer>"}],
...     [{"role": "assistant", "content": "The sum of 3 and 1 is 4, which we multiply by 2 to get 8. So (3 + 1) * 2 = 8."}],
... ]
>>> format_reward_func(prompts=prompts, completions=completions)
[1.0, 0.0]

Пример 3: Награды за завершения на основе эталонного

Ниже приведен пример функции вознаграждения, которая проверяет правильность завершения. Этот пример вдохновлен функцией accuracy reward, используемой в статье DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. Этот пример предназначен для стандартного формата, в котором набор данных содержит столбец с именем ground_truth.

import re

def reward_func(completions, ground_truth, **kwargs):
    # Regular expression to capture content inside \boxed{}
    matches = [re.search(r"\\boxed\{(.*?)\}", completion) for completion in completions]
    contents = [match.group(1) if match else "" for match in matches]
    # Reward 1 if the content is the same as the ground truth, 0 otherwise
    return [1.0 if c == gt else 0.0 for c, gt in zip(contents, ground_truth)]

Вы можете протестировать эту функцию следующим образом:

>>> prompts = ["Problem: Solve the equation $2x + 3 = 7$. Solution:", "Problem: Solve the equation $3x - 5 = 10$."]
>>> completions = [r" The solution is \boxed{2}.", r" The solution is \boxed{6}."]
>>> ground_truth = ["2", "5"]
>>> reward_func(prompts=prompts, completions=completions, ground_truth=ground_truth)
[1.0, 0.0]

Пример 4: Многозадачные функции вознаграждения

Ниже приведен пример использования нескольких функций вознаграждения в RLOOTrainer. В этом примере мы определяем две функции вознаграждения, специфичные для задач: math_reward_func и coding_reward_func. Функция math_reward_func оценивает математические задачи на основе их правильности, а coding_reward_func — задачи по программированию на основе того, работает ли решение.

from datasets import Dataset
from trl import RLOOTrainer

# Define a dataset that contains both math and coding problems
dataset = Dataset.from_list(
    [
        {"prompt": "What is 2+2?", "task": "math"},
        {"prompt": "Write a function that returns the sum of two numbers.", "task": "code"},
        {"prompt": "What is 3*4?", "task": "math"},
        {"prompt": "Write a function that returns the product of two numbers.", "task": "code"},
    ]
)

# Math-specific reward function
def math_reward_func(prompts, completions, task, **kwargs):
    rewards = []
    for prompt, completion, t in zip(prompts, completions, task):
        if t == "math":
            # Calculate math-specific reward
            correct = check_math_solution(prompt, completion)
            reward = 1.0 if correct else -1.0
            rewards.append(reward)
        else:
            # Return None for non-math tasks
            rewards.append(None)
    return rewards

# Coding-specific reward function
def coding_reward_func(prompts, completions, task, **kwargs):
    rewards = []
    for prompt, completion, t in zip(prompts, completions, task):
        if t == "code":
            # Calculate coding-specific reward
            works = test_code_solution(prompt, completion)
            reward = 1.0 if works else -1.0
            rewards.append(reward)
        else:
            # Return None for non-coding tasks
            rewards.append(None)
    return rewards

# Use both task-specific reward functions
trainer = RLOOTrainer(
    model="Qwen/Qwen2-0.5B-Instruct",
    reward_funcs=[math_reward_func, coding_reward_func],
    train_dataset=dataset,
)

trainer.train()

В этом примере math_reward_func и coding_reward_func разработаны для работы со смешанным набором данных, содержащим как математические, так и программные задачи. Колонка task в наборе данных используется для определения, какую функцию вознаграждения применить к каждой задаче. Если для образца в наборе данных нет соответствующей функции вознаграждения, функция вернёт None, и RLOOTrainer продолжит работу с валидными функциями и задачами. Это позволяет RLOOTrainer обрабатывать несколько функций вознаграждения с разной областью применения.

Обратите внимание, что RLOOTrainer игнорирует вознаграждения None, возвращаемые функциями вознаграждения, и учитывает только вознаграждения, возвращаемые соответствующими функциями. Это гарантирует, что модель обучается на релевантных задачах и игнорирует те, для которых нет соответствующей функции вознаграждения.

Пример 5: Асинхронные функции вознаграждения

Пользовательские функции вознаграждения также можно определять как корутины async def. Это полезно, если ваше вознаграждение зависит от медленного ввода-вывода (например, вызова удалённого сервиса). При передаче нескольких асинхронных функций вознаграждения RLOOTrainer выполняет их параллельно, чтобы их задержки перекрывались.

Ниже приведен минимальный пример асинхронной функции вознаграждения, имитирующей операцию, ограниченную вводом-выводом:

import asyncio

async def async_reward_func(prompts, completions, **kwargs):
    # Simulate an I/O-bound call (e.g., HTTP request, database lookup)
    await asyncio.sleep(0.01)
    # Simple toy reward: 1.0 if the completion is non-empty, else 0.0
    return [1.0 if completion else 0.0 for completion in completions]

Пример 6: Логирование дополнительных столбцов и метрик

Ниже приведен пример функции вознаграждения, которая записывает дополнительные столбцы в таблицу завершений и скалярные метрики в виде графиков.

import re

def reward_func(completions, ground_truth, log_extra=None, log_metric=None, **kwargs):
    extracted = [re.search(r"\\boxed\{(.*?)\}", c) for c in completions]
    extracted = [m.group(1) if m else None for m in extracted]
    rewards = [1.0 if e == gt else 0.0 for e, gt in zip(extracted, ground_truth)]

    if log_extra:
        log_extra("golden_answer", list(ground_truth))
        log_extra("extracted_answer", [e or "[none]" for e in extracted])

    if log_metric:
        log_metric("accuracy", sum(rewards) / len(rewards))

    return rewards

Передача функции вознаграждения в тренер

Чтобы использовать собственную функцию вознаграждения, передайте её в RLOOTrainer следующим образом:

from trl import RLOOTrainer

trainer = RLOOTrainer(
    reward_funcs=reward_func,
    ...,
)

Можно передать несколько функций вознаграждения в виде списка; этот список может включать как синхронные, так и асинхронные функции:

from trl import RLOOTrainer

trainer = RLOOTrainer(
    reward_funcs=[reward_func, async_reward_func1, async_reward_func2],
    ...,
)

а награда будет вычисляться как сумма наград от каждой функции или как взвешенная сумма, если в конфигурации указано reward_weights.

Обратите внимание, что RLOOTrainer поддерживает несколько функций вознаграждения разных типов. Подробности см. в документации параметров.

Обучение моделей «зрение-язык»

RLOO поддерживает обучение моделей «зрение-язык» (VLM) на мультимодальных наборах данных, содержащих как текст, так и изображения.

Поддерживаемые модели

Проверено с:

• Gemma3 — например, google/gemma-3-4b-it
• LLaVA-NeXT — например, llava-hf/llava-v1.6-mistral-7b-hf
• Qwen2-VL — например, Qwen/Qwen2-VL-2B-Instruct
• Qwen2.5-VL — например, Qwen/Qwen2.5-VL-3B-Instruct
• SmolVLM2 — например, HuggingFaceTB/SmolVLM2-2.2B-Instruct

Совместимость со всеми VLM не гарантируется. Если вы считаете, что модель должна поддерживаться, не стесняйтесь открыть issue на GitHub — или, что ещё лучше, отправить pull request с необходимыми изменениями.

Быстрый старт

Используйте rloo_vlm.py для дообучения VLM. Пример команды для обучения на lmms-lab/multimodal-open-r1-8k-verified:

torchrun --nproc_per_node 8 examples/rloo_visual_math/rloo_visual_math.py \
  --deepspeed examples/deepspeed_configs/zero3.json \
  --model_name_or_path Qwen/Qwen2.5-VL-3B-Instruct \
  --output_dir rloo-Qwen2.5-VL-3B-Instruct \
  --learning_rate 1e-5 \
  --dtype bfloat16 \
  --max_completion_length 1024 \
  --use_vllm \
  --vllm_mode colocate \
  --use_peft \
  --lora_target_modules "q_proj", "v_proj" \
  --log_completions

Советы по конфигурации

• Использование LoRA на проекционных слоях зрительно-языковой модели
• Включите 4-битное квантование для снижения потребления памяти
• VLMs требуют много памяти — начните с меньших размеров батчей
• Большинство моделей совместимы с vLLM (режимы server и colocate)

Формат набора данных

Каждая обучающая выборка должна включать:

• prompt: Текст, отформатированный через шаблон чата процессора
• image/images: изображение PIL или список изображений PIL

Тренер автоматически обрабатывает преобразование изображений в тензоры с помощью процессора изображений модели.

RLOOTrainer

>>> from trl import RLOOTrainer
>>> from trl.rewards import accuracy_reward
>>> from datasets import load_dataset

>>> dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

>>> trainer = RLOOTrainer(
...     model="Qwen/Qwen2.5-0.5B-Instruct",
...     reward_funcs=accuracy_reward,
...     train_dataset=dataset,
... )
>>> trainer.train()

RLOOConfig

Ссылки

1. Статья RLOO
2. Статья Back to Basics: Пересмотр оптимизации в стиле REINFORCE для обучения по обратной связи от человека в LLM
3. Статья - REINFORCE++: простой и эффективный подход к выравниванию больших языковых моделей
4. Статья в блоге — Возвращаем RL в RLHF
5. Статья в блоге — Разбираемся с RLHF и его вариантами: прогресс и практические инженерные наблюдения
6. YouTube - RLOO: экономичная оптимизация для обучения с обратной связью от человека в LLM

Руководство по миграции со старой реализации (0.21 и ниже)

С выходом версии 0.22.0 мы переработали RLOOTrainer, чтобы он лучше соответствовал другим онлайн-тренажерам в библиотеке, таким как GRPOTrainer. Эта новая реализация вносит несколько изменений в параметры конфигурации и общую структуру тренажера. Ниже приведено краткое описание ключевых изменений для RLOOConfig:

Ниже приведено краткое описание ключевых изменений для RLOOTrainer: