GRPO Trainer

Обзор

TRL поддерживает GRPO Trainer для обучения языковых моделей, как описано в статье DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models авторов Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Mingchuan Zhang, Y. K. Li, Y. Wu, Daya Guo.

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

Математическое рассуждение представляет собой значительный вызов для языковых моделей из-за своей сложной и структурированной природы. В данной статье мы представляем DeepSeekMath 7B, который продолжает предварительное обучение DeepSeek-Coder-Base-v1.5 7B на 120 млрд токенов, связанных с математикой, полученных из Common Crawl, а также на данных естественного языка и кода. DeepSeekMath 7B достиг впечатляющего результата 51,7% на соревновательном бенчмарке MATH без использования внешних библиотек и методов голосования, приближаясь к уровню производительности Gemini-Ultra и GPT-4. Метод self-consistency по 64 выборкам из DeepSeekMath 7B достигает 60,9% на MATH. Способность к математическому рассуждению в DeepSeekMath обусловлена двумя ключевыми факторами: во-первых, мы используем значительный потенциал общедоступных веб-данных с помощью тщательно спроектированного конвейера отбора данных. Во-вторых, мы вводим Group Relative Policy Optimization (GRPO), вариант Proximal Policy Optimization (PPO), который повышает способности к математическому рассуждению, одновременно оптимизируя использование памяти в PPO.

Метод постобучения был предложен Quentin Gallouédec.

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

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

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

# train_grpo.py
from datasets import load_dataset
from trl import GRPOTrainer
from trl.rewards import accuracy_reward

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

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

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

accelerate launch train_grpo.py

Обучение распределено по 8 GPU и занимает примерно 1 день.

Примечание: Кривые вознаграждения выше были получены с использованием Qwen/Qwen2-0.5B-Instruct. Результаты с Qwen/Qwen2.5-0.5B-Instruct будут качественно схожи.

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

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

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

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

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

Для каждой из $G$ последовательностей мы вычисляем вознаграждение с помощью модели или функции вознаграждения. Чтобы согласовать с сравнительной природой моделей вознаграждения — как правило, обученных на наборах данных со сравнениями выходов для одного и того же вопроса — преимущество рассчитывается так, чтобы отражать эти относительные сравнения. Оно нормализуется следующим образом: $$\hat{A}_{i,t} = \frac{r_i - \text{mean}(\mathbf{r})}{\text{std}(\mathbf{r})}$$

Этот подход и дал методу его название: Group Relative Policy Optimization (GRPO).

В статье Understanding R1-Zero-Like Training: A Critical Perspective показано, что масштабирование на $\text{std}(\mathbf{r})$ может привести к смещению по уровню сложности вопросов. Вы можете отключить это масштабирование, установив scale_rewards=False в GRPOConfig. Обратите внимание, что отключение масштабирования на std также отключает нормализацию дисперсии, поэтому величина обновлений напрямую зависит от исходного масштаба вознаграждения и состава батча.

Как показано в Части I: Хитрости или ловушки? Глубокий анализ RL для рассуждений LLM (Lite PPO), вычисление среднего на локальном (групповом) уровне и стандартного отклонения на глобальном (батчевом) уровне обеспечивает более устойчивое формирование вознаграждения. Вы можете использовать эту стратегию масштабирования, установив scale_rewards="batch" в GRPOConfig.

Оценка расхождения Кульбака–Лейблера

Расхождение KL оценивается с помощью аппроксиматора, предложенного Schulman et al. (2020). Аппроксиматор определяется следующим образом: $$\mathbb{D}_{\text{KL}}\left[\pi_\theta \|\pi_{\text{ref}}\right] = \frac{\pi_{\text{ref}}(o_{i,t} \mid q, o_{i,<t})}{\pi_\theta(o_{i,t} \mid q, o_{i,<t})} - \log \frac{\pi_{\text{ref}}(o_{i,t} \mid q, o_{i,<t})}{\pi_\theta(o_{i,t} \mid q, o_{i,<t})} - 1,$$

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

Цель состоит в том, чтобы максимизировать преимущество, при этом обеспечивая близость модели к эталонной политике. Следовательно, функция потерь определяется следующим образом: $$\mathcal{L}_{\text{GRPO}}(\theta) = -\frac{1}{\sum_{i=1}^G |o_i|} \sum_{i=1}^G \sum_{t=1}^{|o_i|} \left[ \frac{\pi_\theta(o_{i,t} \mid q, o_{i,< t})}{\left[\pi_\theta(o_{i,t} \mid q, o_{i,< t})\right]_{\text{no grad}}} \hat{A}_{i,t} - \beta \mathbb{D}_{\text{KL}}\left[\pi_\theta \| \pi_{\text{ref}}\right] \right],$$

где первый член представляет собой масштабированное преимущество, а второй член штрафует за отклонения от эталонной политики через дивергенцию Кульбака-Лейблера.

Обратите внимание, что в отличие от исходной формулировки в DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models, мы не масштабируем на $\frac{1}{|o_i|}$ , поскольку в статье Understanding R1-Zero-Like Training: A Critical Perspective показано, что это вносит смещение по длине ответа. Подробнее в разделе типы потерь.

Обратите внимание, что по сравнению с исходной формулировкой в DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models, мы используем $\beta = 0.0$ по умолчанию, что означает, что член KL-расхождения не используется. Этот выбор обусловлен рядом недавних исследований (например, Open-Reasoner-Zero: An Open Source Approach to Scaling Up Reinforcement Learning on the Base Model) ), которые показали, что член KL-расхождения не является обязательным для обучения с GRPO. В результате стало общепринятой практикой исключать его (например, Understanding R1-Zero-Like Training: A Critical Perspective, DAPO: An Open-Source LLM Reinforcement Learning System at Scale). Если вы хотите включить член KL-расхождения, вы можете установить beta в GRPOConfig в ненулевое значение.

В исходной статье эта формулировка обобщена с учётом нескольких обновлений после каждой генерации (обозначается $\mu$, и может быть задано с помощью num_iterations в GRPOConfig) за счёт использования обрезанной целевой функции суррогата: $$\mathcal{L}_{\text{GRPO}}(\theta) = - \frac{1}{\sum_{i=1}^G |o_i|} \sum_{i=1}^G \sum_{t=1}^{|o_i|} \left[ \min \left( \frac{\pi_\theta(o_{i,t} \mid q, o_{i,< t})}{\pi_{\theta_{\text{old}}}(o_{i,t} \mid q, o_{i,< t})} \hat{A}_{i,t}, \, \text{clip}\left( \frac{\pi_\theta(o_{i,t} \mid q, o_{i,< t})}{\pi_{\theta_{\text{old}}}(o_{i,t} \mid q, o_{i,< t})}, 1 - \epsilon, 1 + \epsilon \right) \hat{A}_{i,t} \right) - \beta \mathbb{D}_{\text{KL}}\left[\pi_\theta \| \pi_{\text{ref}}\right] \right],$$

где $\text{clip}(\cdot, 1 - \epsilon, 1 + \epsilon)$ обеспечивает, чтобы обновления не отклонялись чрезмерно от эталонной политики, ограничивая отношение политик значениями от $1 - \epsilon$ до $1 + \epsilon$. Если $\mu = 1$ (по умолчанию в TRL), обрезанная суррогатная целевая функция упрощается до исходной целевой функции.

Типы потерь

В литературе предложено несколько формулировок целевой функции. Изначально целевая функция GRPO определялась следующим образом: $$\mathcal{L}_{\text{GRPO}}(\theta) = - \frac{1}{G} \sum_{i=1}^G \frac{1}{|o_i|} \sum_{t=1}^{|o_i|} l_{i,t},$$

где $$l_{i,t} = \frac{\pi_\theta(o_{i,t} \mid q, o_{i,< t})}{\left[\pi_\theta(o_{i,t} \mid q, o_{i,< t})\right]_{\text{no grad}}} \hat{A}_{i,t} - \beta \mathbb{D}_{\text{KL}}\left[\pi_\theta \| \pi_{\text{ref}}\right].$$

В статье DAPO подчеркиваются ограничения функции потерь на уровне образцов в алгоритме GRPO в сценариях с длинными цепочками рассуждений (long-CoT), где более длинные ответы получают недостаточное наказание, что приводит к ухудшению качества вывода. Предложенное решение — нормализация на уровне токенов, которая лучше справляется с длинными последовательностями, распределяя более сбалансированные награды между отдельными токенами независимо от длины ответа: $$\mathcal{L}_{\text{DAPO}}(\theta) = - \frac{1}{\sum_{i=1}^G |o_i|} \sum_{i=1}^G \sum_{t=1}^{|o_i|} l_{i,t},$$

Чтобы использовать эту формулировку, установите loss_type="dapo" в GRPOConfig.

Кроме того, в статье Understanding R1-Zero-Like Training: A Critical Perspective было показано, что исходная формулировка GRPO вносит смещение в длину ответа. Авторы демонстрируют, что хотя формулировка DAPO снижает это смещение, она не устраняет его полностью. Чтобы полностью избавиться от этого смещения, они предлагают делить на константу вместо длины последовательности, что приводит к следующей формулировке: $$\mathcal{L}_{\text{Dr. GRPO}}(\theta) = - \frac{1}{LG} \sum_{i=1}^G \sum_{t=1}^{|o_i|} l_{i,t},$$

Рекомендуется устанавливать это значение равным максимальной длине завершения. Чтобы использовать эту формулировку, задайте loss_type="dr_grpo" в GRPOConfig.

В качестве альтернативы в статье SAPO команда Qwen предлагает заменить «жёдкий» механизм обрезки в GRPO на плавный механизм мягкого гейтинга с управлением температурой. Если GRPO обнуляет градиенты, когда политика слишком сильно отклоняется от эталонной, то SAPO использует мягкую область доверия, которая плавно снижает вес градиента. Это позволяет модели сохранять полезные сигналы обучения от токенов, близких к текущей политике, подавляя при этом шум от экстремальных отклонений.

Функция потерь определяется следующим образом: $$\mathcal{L}_{\text{SAPO}}(\theta) = - \frac{1}{G} \sum_{i=1}^G \frac{1}{|o_i|} \sum_{t=1}^{|o_i|} f_{i,t} \left( \frac{\pi_\theta(o_{i,t} | q, o_{i,<t})}{\pi_{\theta_{old}}(o_{i,t} | q, o_{i,<t})} \right) \hat{A}_{i,t}$$

Функция мягкого гейтинга $f_{i,t}$ определяется с помощью сигмоидной функции $\sigma$ следующим образом: $$f_{i,t}(x) = \sigma \left( \tau_{i,t} (x - 1) \right) \cdot \frac{4}{\tau_{i,t}}$$

Температура $\tau_{i,t}$ выбирается на основе знака преимущества $\hat{A}_{i,t}$: $$\tau_{i,t} = \begin{cases} \tau_{\text{pos}}, & \text{if } \hat{A}_{i,t} > 0 \\ \tau_{\text{neg}}, & \text{otherwise} \end{cases}$$

Они рекомендуют использовать асимметричные температуры, $\tau_{\text{neg}} > \tau_{\text{pos}}$ (значения по умолчанию: $\tau_{\text{pos}}=1.0, \tau_{\text{neg}}=1.05$ ). Это обеспечивает более строгое наказание модели за «плохие» действия для предотвращения нестабильности, при этом допуская более свободное поведение в отношении «хороших» действий.

Чтобы использовать эту формулировку, установите loss_type="sapo" в GRPOConfig.

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

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

• 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: Доля обрезанных (усечённых) завершений.
• tools/call_frequency: среднее количество вызовов инструментов на одно завершение в батче генерации. Записывается в журнал только при наличии tools.
• tools/failure_frequency: доля вызовов инструментов, завершившихся ошибкой (инструмент не найден, возникло исключение или тип вызова не поддерживается). Значение равно 0.0, если ни один инструмент не был вызван. Записывается в журнал только при наличии tools.
• rewards/{reward_func_name}/mean: среднее значение вознаграждения от конкретной функции вознаграждения. Если среда владеет вознаграждением через get_reward, то {reward_func_name} — это имя класса среды.
• rewards/{reward_func_name}/std: стандартное отклонение вознаграждения от конкретной функции вознаграждения.
• reward: Общая средняя награда после суммирования наград по функциям (с весами reward_weights).
• reward_std: стандартное отклонение суммарных наград по функциям (с учётом весов reward_weights), вычисленное для всей выборки.
• frac_reward_zero_std: доля образцов в батче генерации, у которых стандартное отклонение награды равно нулю, что указывает на малое разнообразие для данного промпта (все ответы верные или неверные).
• sampling/sampling_logp_difference/mean: средняя абсолютная разница между лог-вероятностями, возвращаемыми сэмплером (vLLM), и пересчитанными моделью обучения, по токенам завершения. Рост значения указывает на увеличение рассогласования между обучением и инференсом. Записывается в лог только при use_vllm=True и vllm_importance_sampling_correction=True.
• sampling/sampling_logp_difference/max: максимальное такое абсолютное значение. Записывается в журнал только при use_vllm=True и vllm_importance_sampling_correction=True.
• sampling/importance_sampling_ratio/min: минимальное отношение важности (importance sampling ratio), используемое для коррекции рассогласования между обучением и инференсом, после применения ограничения, выбранного параметром vllm_importance_sampling_mode (обрезается до [C_min, C_max] для режимов *_truncate, устанавливается в ноль для режимов *_mask). Вычисляется по токенам завершения (completion tokens) или по последовательностям для режимов sequence_*. Записывается в лог только при use_vllm=True и vllm_importance_sampling_correction=True.
• sampling/importance_sampling_ratio/mean: среднее значение ограниченного отношения важности (importance sampling ratio). Записывается в журнал только при use_vllm=True и vllm_importance_sampling_correction=True.
• sampling/importance_sampling_ratio/max: максимальное ограниченное отношение важности выборки. Записывается только при use_vllm=True и vllm_importance_sampling_correction=True.
• policy_loss: значение потери градиента политики (до добавления бонуса за энтропию). Записывается в журнал, если entropy_coef не равно нулю или use_adaptive_entropy=True.
• entropy: Средняя энтропия предсказаний токенов по всем сгенерированным завершениям. (Если mask_truncated_completions=True, токены замаскированных последовательностей исключаются.)
• entropy_coef: Текущий коэффициент регуляризации энтропии. Записывается в журнал, если entropy_coef не равен нулю или use_adaptive_entropy=True. Обновляется один раз на каждый шаг оптимизатора, если use_adaptive_entropy=True.
• aux_loss: вспомогательная функция потерь для балансировки нагрузки в модели Mixture-of-Experts, до её масштабирования коэффициентом router_aux_loss_coef и добавления к общей функции потерь. Записывается в журнал только в том случае, если модель является MoE и router_aux_loss_coef не равен нулю.
• kl: средняя дивергенция Кульбака–Лейблера между моделью и эталонной моделью, рассчитанная по сгенерированным завершениям. Записывается в журнал только если beta не равно нулю.
• clip_ratio/region_mean: Соотношение токенов (или последовательностей, если importance_sampling_level="sequence") вероятностей, для которых цель GRPO обрезается, чтобы оставаться в области доверия: $\text{clip}\left( r_{i,t}(\theta), 1 - \epsilon_\mathrm{low}, 1 + \epsilon_\mathrm{high} \right)\,, \quad r_{i,t}(\theta) = \frac{\pi_\theta(o_{i,t} \mid q, o_{i,< t})}{\pi_{\theta_{\text{old}}}(o_{i,t} \mid q, o_{i,< t})}$. Более высокое значение означает, что обрезается больше токенов, что ограничивает степень изменения политики $\pi_\theta$ .
• clip_ratio/low_mean: Среднее отношение вероятностей токенов (или последовательностей, если importance_sampling_level="sequence") было обрезано по нижней границе области доверия: $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/low_min: Минимальная доля токенов на завершение (или самой последовательности, если importance_sampling_level="sequence") обрезается по нижней границе области доверия: $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/high_mean: Среднее отношение вероятностей токенов (или последовательностей, если importance_sampling_level="sequence") было обрезано по верхней границе области доверия: $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.
• clip_ratio/high_max: Максимальная доля токенов на завершение (или последовательность целиком, если importance_sampling_level="sequence") обрезается по верхней границе области доверия: $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.
• cispo_clip_ratio: Доля токенов (или последовательностей, если importance_sampling_level="sequence") весов importance sampling, которые были обрезаны на уровне epsilon_high при положительном преимуществе: $r_{i,t}(\theta) > \epsilon_\mathrm{high}$. Записывается в журнал только при loss_type="cispo".
• vespo/phi_seq_mean: Среднее значение весов VESPO Gamma $\varphi(w)$ применяется к весам важности на уровне последовательности. Значения ниже 1.0 означают, что последовательности получают меньший вес. Записывается в журнал только при loss_type="vespo".

Настройка

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

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

pip install trl[vllm]

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

По умолчанию для генерации в vLLM включается выборка с маскированием важности на уровне последовательности (vllm_importance_sampling_mode="sequence_mask") для устранения рассогласования между генерацией и обучением, возникающего при использовании разных фреймворков. Это можно отключить, установив vllm_importance_sampling_correction=False. Подробнее см. Masked Importance Sampling и Sequence-level Importance Sampling

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

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

from trl import GRPOConfig

training_args = GRPOConfig(
    ...,
    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 GRPOConfig
   
   training_args = GRPOConfig(
       ...,
       use_vllm=True,
       vllm_mode="server",
   )

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

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

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

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

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

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

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

Устранение рассогласования между обучением и инференсом

Хотя vLLM значительно ускоряет инференс, он также разделяет движок инференса и движок обучения. Теоретически эти движки математически идентичны, однако на практике они могут выдавать разные результаты из-за эффектов точности и аппаратно-специфичных оптимизаций. Это расхождение отражает различные цели оптимизации двух систем. Движки инференса стремятся максимизировать пропускную способность сэмплирования, обычно измеряемую в токенах в секунду, при сохранении приемлемой точности сэмплирования. Фреймворки обучения, напротив, фокусируются на численной стабильности и точности для вычисления градиентов, часто используя форматы с более высокой точностью, такие как FP32, для мастер-весов и состояний оптимизатора. Эти различающиеся приоритеты и ограничения вносят неизбежное, хотя и незначительное, несоответствие между обучением и инференсом.

Это несоответствие приводит к смещённым обновлениям градиента, что, как было замечено, дестабилизирует обучение ([1][2][3][4][5]). Для простоты рассмотрим градиент политики REINFORCE: $$\nabla_\theta \mathcal{J}(x,\theta) = \mathbb{E}_{y \sim \pi^\text{train}(\cdot \mid x,\theta)} \left[ \nabla_\theta \log \pi^\text{train}(y \mid x,\theta) \cdot R(x,y) \right]$$

Здесь $x$ обозначает промпты, выбранные из некоторого распределения данных, а $\pi^\text{train}$ — это политика, реализованная движком обучения. При использовании vLLM в цикле мы получаем отдельную политику инференса $\pi^\text{inference}$, так что эффективный градиент политики становится $$\nabla_\theta \mathcal{J}_{\text{biased}}(x,\theta) = \mathbb{E}_{y \sim \pi^\text{inference}(\cdot \mid x,\theta)} \left[ \nabla_\theta \log \pi^\text{train}(y \mid x,\theta) \cdot R(x,y) \right].$$

Это превращает задачу RL, которая иначе была бы on-policy, в off-policy.

Стандартный способ коррекции этого смещения распределения — importance sampling (IS). Мы предлагаем два варианта IS: Truncated Importance Sampling (TIS) и Masked Importance Sampling (MIS). Оба варианта можно применять как на уровне токенов, так и на уровне последовательностей. Пусть $\rho$ обозначает вес importance, например $\rho_t$ на токен или $\rho_{\text{seq}}$ на уровне последовательности. В TIS значения вне диапазона [vllm_importance_sampling_clip_min, vllm_importance_sampling_clip_max] обрезаются: $$\rho \leftarrow \text{clip}(\rho, C_{\min}, C_{\max}).$$

В исходной статье TIS был предложен механизм обрезки по верхней границе, а именно $\min(\rho, C_{\max})$. Реализация в TRL обобщила этот подход, введя также нижнюю границу, что привело к двусторонней формулировке, показанной выше, по мотивам метода IcePop .
Обратите внимание, что в IcePop границы обозначаются как $\alpha$ и $\beta$ тогда как в TRL мы используем $C_{\min}$ и $C_{\max}$.

В режиме MIS отношения, выходящие за пределы этого диапазона, обнуляются, поэтому такие сэмплы не вносят вклад в градиент. Иными словами, выбросы понизятся по весу в режиме TIS и отбрасываются в режиме MIS. Флаг конфигурации vllm_importance_sampling_mode определяет как вариант IS (маскирование или усечение), так и гранулярность (на уровне токенов или последовательностей).

Импортанс-сэмплинг — это принципиальный алгоритмический ответ на рассогласование между обучением и инференсом. Однако существуют и более прямые подходы, направленные на сокращение рассогласования между самыми движками. Большинство из них представляют собой инженерные решения. Например, MiniMax M1 использует языковую модель с головой в формате FP32 в движке инференса. Thinking Machines исследовала детерминированные ядра инференса, хотя это сопряжено со значительными потерями в эффективности. vLLM продемонстрировала битово согласованные политики, опираясь на детерминированные ядра с инвариантностью по батчу от Thinking Machines, но по состоянию на ноябрь 2025 года сохраняется существенное снижение пропускной способности по сравнению со стандартным инференсом в vLLM.

Ускорение обучения с помощью непрерывного батчинга в transformers

В качестве альтернативы vLLM можно использовать встроенный в transformers движок непрерывного батчинга (continuous batching) для более быстрой генерации. Непрерывный батчинг сразу удаляет завершённые последовательности из батча, а не ждёт завершения самой медленной. Для задач с переменной длиной завершения (например, математическое рассуждение) это обеспечивает более быструю генерацию и меньшее потребление VRAM по сравнению с функцией generate() по умолчанию при больших размерах батча (N≥32).

Continuous batching — это решение, которое можно внедрить без настройки сервера или синхронизации весов. Оно работает в рамках процесса и хорошо подходит для обучения на одном GPU или в средах с ограниченной памятью. Для максимальной пропускной способности генерации в масштабе используйте vLLM.

from trl import GRPOConfig

training_args = GRPOConfig(
    ...,
    use_transformers_continuous_batching=True,
    transformers_continuous_batching_config={
        "use_cuda_graph": False,
        "max_memory_percent": 0.4,  # lower values leave more VRAM for the training backward pass
    },
)

TRL по умолчанию устанавливает max_memory_percent в 0.5 (вместо 0.9 в transformers), чтобы оставить достаточно VRAM для обратного прохода при обучении. Уменьшите это значение до 0.3–0.4 для больших батчей генерации (N≥32) или если вы видите ошибки нехватки памяти.

Полный пример обучения см. в examples/grpo_continuous_batching/grpo_continuous_batching.py.

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

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

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

Ниже приведён пример SLURM-скрипта для обучения модели 70B с GRPO на нескольких узлах. Этот скрипт обучает модель на 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}" accelerate launch \
     --config_file examples/accelerate_configs/deepspeed_zero3.yaml \
     --num_processes 32 \
     --num_machines 4 \
     --main_process_ip ${NODELIST[0]} \
     --machine_rank $SLURM_PROCID \
     --rdzv_backend c10d \
     train_grpo.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 GRPOTrainer, GRPOConfig
from trl.rewards import accuracy_reward

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

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

    training_args = GRPOConfig(
        per_device_train_batch_size=4,
        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 = GRPOTrainer(
        model="Qwen/Qwen2.5-72B",
        args=training_args,
        reward_funcs=accuracy_reward,
        train_dataset=dataset
    )
    trainer.train()

if __name__=="__main__":
    main()

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

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

Функции вознаграждения могут быть как синхронными вызываемыми объектами 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.
     • environments: список экземпляров сред, по одному на каждое завершение. Присутствует только при указании environment_factory. Используйте этот параметр для чтения состояния, накопленного в ходе эпизода (например, env.counter). Для сред с состоянием предпочтительнее передавать управление наградой самой среде через метод get_reward — см. Награды.
     • Все имена столбцов (кроме 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: Многозадачные функции вознаграждения

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

from datasets import Dataset
from trl import GRPOTrainer

# 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 = GRPOTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    reward_funcs=[math_reward_func, coding_reward_func],
    train_dataset=dataset,
)

trainer.train()

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

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

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

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

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

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

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

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

from trl import GRPOTrainer

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

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

from trl import GRPOTrainer

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

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

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

Регуляризация энтропии

Чтобы стимулировать исследование и предотвратить схлопывание политики к почти детерминированным выводам, можно добавить бонус за энтропию в целевую функцию обучения. Регуляризация энтропии дополняет потерю GRPO следующим образом: $$\mathcal{L}(\theta) = \mathcal{L}_{\text{GRPO}}(\theta) - \alpha \cdot \mathcal{H}(\pi_\theta),$$

где$\mathcal{H}(\pi_\theta)$ — средняя энтропия на токен для текущей политики, а$\alpha$ — коэффициент энтропии. Бонус всегда равен средней энтропии на токен, независимо от loss_type; он не масштабируется под нормализацию политики, соответствующую типу потерь (например, знаменатель batch_size * max_completion_length в Dr. GRPO), поэтому entropy_coef имеет одинаковый смысл для всех типов потерь.

Статическая энтропия — фиксированный коэффициент на протяжении всего обучения:

from trl import GRPOConfig, GRPOTrainer

training_args = GRPOConfig(entropy_coef=0.05, ...)

Адаптивная энтропия — коэффициент обновляется на каждом шаге оптимизатора на основе целевой энтропии, как предложено в Skywork-OR1. Если текущая энтропия равна entropy_target или меньше, коэффициент увеличивается на entropy_coef_delta; в противном случае он уменьшается. Коэффициент применяется (т. е. не равен нулю) только тогда, когда энтропия равна целевому значению или меньше:

training_args = GRPOConfig(
    entropy_coef=0.01,          # initial coefficient
    use_adaptive_entropy=True,
    entropy_target=5.0,         # target mean per-token entropy (nats); tune for your model
    entropy_coef_delta=0.005,   # step size per optimizer step
    entropy_coef_min=0.0,
    entropy_coef_max=1.0,
    ...
)

Типичные языковые модели имеют энтропию на токен в диапазоне 2–10 натов, поэтому значение по умолчанию entropy_target=0.2 почти никогда не активирует регуляризацию — бонус срабатывает только тогда, когда энтропия равна целевому значению или ниже, то есть при практически полном коллапсе. Установите значение, имеющее смысл для вашей модели, например, близкое к энтропии, наблюдаемой на ранних этапах обучения (логгируется как метрика entropy). При использовании top_entropy_quantile < 1.0 параметр entropy_target применяется к подмножеству токенов с высокой энтропией — энтропия этого подмножества будет выше, чем логгируемая энтропия по всем токенам (entropy), поэтому калибруйте значение соответственно.

При use_adaptive_entropy=True текущий коэффициент энтропии entropy_coef сохраняется вместе с каждой контрольной точкой и восстанавливается при возобновлении, что делает обучение полностью возобновляемым.

Быстрая экспериментальная работа с GRPO

RapidFire AI — это экспериментальный движок с открытым исходным кодом, работающий поверх TRL и позволяющий запускать несколько конфигураций GRPO одновременно, даже на одном GPU. Вместо последовательного перебора конфигураций RapidFire позволяет ранее увидеть все их кривые обучения, остановить низкоэффективные запуски и клонировать перспективные с новыми настройками на лету без перезапуска. Подробнее см. RapidFire AI Integration.

Обучение агента

GRPO поддерживает обучение агентов: модель вызывает инструменты в процессе генерации и учится на полученных результатах.

• Инструмент — это обычная функция Python (синхронная или асинхронная), доступная модели. Используйте tools для вызовов без состояния (калькулятор, веб-поиск).
• Environment — более общая форма: объектом с состоянием, создаваемым заново для каждого rollout, публичные методы которого доступны как инструменты, плюс хук жизненного цикла reset и необязательный get_reward, позволяющий среде самостоятельно определять награду. Используйте environment_factory, если вам нужно состояние для каждого rollout, хук сброса или награда, управляемая средой.

Они комбинируются — вы можете передать автономные tools вместе с environment_factory.

Инструменты

Аргумент tools принимает список функций Python (синхронных или асинхронных), которые определяют инструменты, доступные агенту:

from trl import GRPOTrainer

trainer = GRPOTrainer(
    tools=[tool1, tool2],
    ...,
)

Каждый инструмент должен быть стандартной функцией Python с аргументами и типом возвращаемого значения, указанными через type hints, а также документирующей строкой в стиле Google, описывающей его назначение, аргументы и возвращаемое значение. Подробнее см. в руководстве по передаче инструментов.

Цикл вызова инструментов GRPO требует, чтобы шаблон чата был сохраняющим префикс (добавление сообщения инструмента не должно изменять отображение предыдущих сообщений). Для известных семейств моделей (например, Qwen3, DeepSeek-V3) TRL автоматически подставляет исправленный шаблон обучения, когда инструменты включены. Полный список см. в разделе Chat Templates.

Используйте max_tool_calling_iterations в GRPOConfig, чтобы ограничить количество ходов вызова инструментов. По умолчанию ограничение отсутствует, и генерация останавливается, когда модель выдаёт ответный ход без вызовов инструментов.

Пример:

from trl import GRPOTrainer

def multiply(a: int, b: int) -> int:
    """
    Multiplies two integers.

    Args:
        a: The first integer.
        b: The second integer.

    Returns:
        The product of the two integers.
    """
    return a * b

async def async_add(a: int, b: int) -> int:
    """
    Asynchronously adds two integers.

    Args:
        a: The first integer.
        b: The second integer.

    Returns:
        The sum of the two integers.
    """
    return a + b

trainer = GRPOTrainer(
    tools=[multiply, async_add],
    ...,
)

Окружения

Вы также можете предоставить инструменты через environment_factory. В этом режиме GRPOTrainer создаёт один экземпляр среды на каждый rollout и предоставляет публичные методы среды в качестве инструментов.

environment_factory требует transformers>=5.2.0.

При использовании environment_factory среда владеет данными: внешнего train_dataset нет. На каждом rollout reset() формирует задачу — самообразец из корпуса, которым владеет среда, или процедурная генерация состояния — и возвращает промпт. max_steps задаёт длительность обучения. Ниже приведён минимальный пример среды, которая самообразует целевое значение и предоставляет метод increment (представленный как инструмент).

import random

from trl import GRPOConfig, GRPOTrainer

class IncrementEnv:
    # Reserved methods
    def reset(self, **kwargs) -> str | None:  # required; called at the start of each rollout
        self.counter = 0
        self.target = random.randint(1, 6)  # self-sample the task
        return f"Increment the counter by {self.target}."  # returned string becomes the prompt

    def get_reward(self) -> float:  # optional: the environment scores itself from its own state
        return float(self.counter == self.target)

    # Public methods (exposed as tools)
    def increment(self, step: int) -> int:
        """
        Increment the internal counter.

        Args:
            step: Value to add to the counter.

        Returns:
            The updated counter value.
        """
        self.counter += step
        return self.counter

trainer = GRPOTrainer(
    model="Qwen/Qwen3-0.6B",
    args=GRPOConfig(max_steps=1000, chat_template_kwargs={"enable_thinking": False}),
    environment_factory=IncrementEnv,
)
trainer.train()

Класс окружения имеет два зарезервированных метода: reset и get_reward. В отличие от любых других публичных методов, они не предоставляются модели в качестве инструментов.

• reset (обязательный) вызывается при начале rollout. Он может возвращать либо None, либо строку. В GRPO, если возвращается строка, она является пользовательским запросом.
• get_reward (необязательный, синхронный или асинхронный) не принимает аргументов и возвращает float: среда оценивает только что завершённый эпизод на основе своего внутреннего состояния (завершилась ли игра победой? было ли угадано слово?). Вызывается один раз за каждый завершённый rollout и служит источником вознаграждения.

Поскольку среда выполняет самоотборку при каждом вызове reset(), G членов группы GRPO могут не разделять одно и то же начальное состояние, что делает базовую линию группы немного более шумной. Для идентичных состояний внутри группы убедитесь, что reset() детерминированно выводит состояние из общего ключа.

Предоставление внешнего набора данных (необязательно)

Вместо того чтобы позволять среде самостоятельно генерировать примеры, вы можете передать внешний train_dataset. Например, если ваши задачи уже хранятся в датасете, или для маршрутизации между несколькими средами. Для каждого rollout тренер выбирает одну строку, использует её столбец "prompt" как промпт и передаёт остальные столбцы этой строки в reset() в виде именованных аргументов (таким образом, среда считывает задачу из датасета).

from datasets import Dataset
from trl import GRPOTrainer

# Each row carries the prompt and the task data (here, the counter target to reach).
dataset = Dataset.from_dict(
    {
        "prompt": [[{"role": "user", "content": f"Increment the counter to {i}."}] for i in range(1, 7)],
        "target": list(range(1, 7)),
    }
)

class IncrementEnv:
    def reset(self, target, **kwargs) -> None:  # the row's `target` column arrives as a keyword argument
        self.counter = 0
        self.target = target

    def get_reward(self) -> float:
        return float(self.counter == self.target)

    def increment(self, step: int) -> int:  # exposed as a tool (see the full docstring above)
        self.counter += step
        return self.counter

trainer = GRPOTrainer(
    model="Qwen/Qwen3-0.6B",
    train_dataset=dataset,  # the dataset owns the task; the "target" column reaches reset(target=...)
    environment_factory=IncrementEnv,
)
trainer.train()

Награды

Награды, управляемые средой, и награды, управляемые тренером, не являются взаимоисключающими. Передавайте reward_funcs для наград, которые естественно вычисляются по завершению (например, проверка формата текста) и позволяйте среде возвращать свою зависящую от состояния награду через get_reward. Все источники суммируются; reward_weights применяется только к reward_funcs (среда сама управляет своим масштабом).

def format_reward(completions, **kwargs):  # trainer-owned: scores the completion text
    ...

trainer = GRPOTrainer(
    model="Qwen/Qwen3-0.6B",
    args=GRPOConfig(max_steps=1000),
    reward_funcs=format_reward,          # trainer-owned: scores the completion
    environment_factory=IncrementEnv,    # env-owned: scores its internal state via `get_reward`
)

Функция reward_func также может напрямую читать состояние среды: если задан environment_factory, тренер передаёт аргумент environments (по одному экземпляру на каждое завершение). Это бывает полезно, но для наград, зависящих от состояния, предпочтительнее использовать get_reward.

Это работает идентично для AsyncGRPOTrainer.

Несколько сред

Чтобы обучать модель на смеси задач в одном прогоне (например, на задаче программирования и игре), передайте словарь, сопоставляющий имена сред с фабриками, вместо единой вызываемой функции. Каждый rollout выбирает свою среду через поле environment в наборе данных, и доступны только инструменты этой среды. Это предотвращает утечку нерелевантных инструментов (инструмент move из игры в пример с кодом и наоборот). Каждая среда по-прежнему владеет своими данными и возвращает промпт из reset(), поэтому в наборе данных достаточно столбца environment для маршрутизации между средами.

import random

from datasets import Dataset
from trl import GRPOConfig, GRPOTrainer


class CodingEnv:
    def reset(self, **kwargs) -> str:
        n = random.randint(1, 20)
        return f"Compute the {n}th Fibonacci number."

    def run_code(self, code: str) -> str:  # exposed as a tool only for "coding" examples
        """Run the given Python code and return its stdout.

        Args:
            code: The Python source to execute.

        Returns:
            The captured standard output.
        """
        ...

class GameEnv:
    def reset(self, **kwargs) -> str:
        return "Reach the goal tile."

    def move(self, action: str) -> str:  # exposed as a tool only for "game" examples
        """Apply an action to the game and return the new observation.

        Args:
            action: One of "up", "down", "left", "right".

        Returns:
            The observation after the move.
        """
        ...

# The dataset only routes: each row picks the environment for that rollout
dataset = Dataset.from_dict({"environment": ["coding", "game"] * 500})

trainer = GRPOTrainer(
    model="Qwen/Qwen3-0.6B",
    args=GRPOConfig(chat_template_kwargs={"enable_thinking": False}),
    train_dataset=dataset,
    reward_funcs=reward_func,
    environment_factory={"coding": CodingEnv, "game": GameEnv},
)
trainer.train()

Функция вознаграждения по-прежнему получает окружение каждого примера через reward_kwargs["environments"], поэтому она может ветвиться в зависимости от типа окружения (например, isinstance(env, CodingEnv)). Экземпляры окружений переиспользуются между шагами (сбрасываются между эпизодами), поэтому дорогостоящий __init__ выполняется один раз, а не на каждом шаге. Поле environment используется тренером и не передаётся в reset.

Ответы мультимодальных инструментов

Инструменты могут возвращать изображения вместе с текстом, возвращая список блоков контента. Это полезно для обучения агентов VLM, где инструмент предоставляет визуальную обратную связь (например, скриншоты, графики, снимки с камеры).

from PIL import Image

def take_screenshot() -> list:
    """
    Takes a screenshot of the current screen.

    Returns:
        The screenshot image with a description.
    """
    img = Image.open("screenshot.png")
    return [{"type": "image", "image": img}, {"type": "text", "text": "Here is the screenshot."}]

Возвращаемые изображения автоматически добавляются в диалог и передаются в VLM для последующих шагов генерации.

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

Проверено с:

• Beam — например, reflection/Beam-1.0-501B-A23B
• DeepSeek-R1 — например, deepseek-ai/DeepSeek-R1
• Gemma4 — например, google/gemma-4-E2B-it
• GLM-4-MoE (4.5, 4.6 или 4.7) — например, zai-org/GLM-4.7
• GPT-OSS — например, openai/gpt-oss-20b
• Llama 3.1 — например, meta-llama/Llama-3.1-8B-Instruct
• Llama 3.2 — например, meta-llama/Llama-3.2-3B-Instruct
• Qwen2.5 — например, Qwen/Qwen2.5-0.5B-Instruct
• Qwen3 — например, Qwen/Qwen3-0.6B
• Qwen3-VL — например, Qwen/Qwen3-VL-2B-Instruct
• Qwen3.5 — например, Qwen/Qwen3.5-2B
• Qwen3.6 — например, Qwen/Qwen3.6-35B-A3B

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

Интеграции с окружениями

Все среды подключаются к одному и тому же слоту environment_factory, поэтому на уровне TRL они взаимозаменяемы — выберите ту, чья экосистема лучше подходит для вашей задачи:

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

GRPO поддерживает обучение моделей Vision-Language (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 с необходимыми изменениями.

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

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

accelerate launch \
  --config_file=examples/accelerate_configs/deepspeed_zero3.yaml \
  examples/grpo_visual_math/grpo_visual_math.py \
  --model_name_or_path Qwen/Qwen2.5-VL-3B-Instruct \
  --output_dir grpo-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

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

GRPOTrainer

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

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

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

GRPOConfig