Treinador GRPO

Visão geral

O TRL oferece suporte ao GRPO Trainer para o treinamento de modelos de linguagem, conforme descrito no artigo DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models, de Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Mingchuan Zhang, Y. K. Li, Y. Wu, Daya Guo.

O resumo do artigo é o seguinte:

O raciocínio matemático representa um desafio significativo para os modelos de linguagem devido à sua natureza complexa e estruturada. Neste artigo, apresentamos o DeepSeekMath 7B, que continua o pré-treinamento do DeepSeek-Coder-Base-v1.5 7B com 120B de tokens relacionados a matemática obtidos do Common Crawl, juntamente com dados de linguagem natural e código. O DeepSeekMath 7B alcançou uma impressionante pontuação de 51,7% no benchmark de nível competitivo MATH, sem recorrer a toolkits externos e técnicas de votação, aproximando-se do nível de desempenho do Gemini-Ultra e do GPT-4. A autoconsistência sobre 64 amostras do DeepSeekMath 7B atinge 60,9% no MATH. A capacidade de raciocínio matemático do DeepSeekMath é atribuída a dois fatores-chave: primeiro, aproveitamos o potencial significativo dos dados web publicamente disponíveis por meio de um pipeline de seleção de dados meticulosamente projetado. Segundo, introduzimos a Group Relative Policy Optimization (GRPO), uma variante da Proximal Policy Optimization (PPO), que aprimora as habilidades de raciocínio matemático enquanto otimiza simultaneamente o uso de memória da PPO.

Este método de pós-treinamento foi contribuído por Quentin Gallouédec.

Início rápido

Este exemplo demonstra como treinar um modelo usando o método GRPO. Treinamos um modelo Qwen2.5 0.5B Instruct com os prompts do conjunto de dados DeepMath-103K. Você pode visualizar os dados do conjunto de dados aqui:

Abaixo está o script para treinar o modelo. accuracy_reward precisa do math-verify, que você pode instalar com 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()

Execute o script usando o seguinte comando:

torchrun --nproc_per_node 8 train_grpo.py

Distribuído em 8 GPUs, o treinamento leva aproximadamente 1 dia.

Observação: As curvas de recompensa acima foram geradas com Qwen/Qwen2-0.5B-Instruct. Os resultados com Qwen/Qwen2.5-0.5B-Instruct serão qualitativamente semelhantes.

Olhando mais a fundo para o método GRPO

GRPO é um algoritmo de aprendizado online, o que significa que ele melhora iterativamente usando os dados gerados pelo próprio modelo treinado durante o treinamento. A intuição por trás do objetivo do GRPO é maximizar a vantagem das conclusões geradas, garantindo que o modelo permaneça próximo da política de referência. Para entender como o GRPO funciona, ele pode ser dividido em quatro etapas principais: gerar conclusões, calcular a vantagem, estimar a divergência de KL e calcular a perda.

Gerando complementos

Em cada etapa de treinamento, amostramos um lote de prompts e geramos um conjunto de $G$ completions para cada prompt (denotado como $o_i$).

Calculando a vantagem

Para cada uma das $G$ sequências, calculamos a recompensa usando um modelo de recompensa ou uma função de recompensa. Para estar alinhado com a natureza comparativa dos modelos de recompensa — geralmente treinados em conjuntos de dados de comparações entre saídas para a mesma pergunta —, a vantagem é calculada para refletir essas comparações relativas. Ela é normalizada da seguinte forma: $$\hat{A}_{i,t} = \frac{r_i - \text{mean}(\mathbf{r})}{\text{std}(\mathbf{r})}$$

Essa abordagem dá nome ao método: Group Relative Policy Optimization (GRPO).

Foi demonstrado no artigo Understanding R1-Zero-Like Training: A Critical Perspective que a escalação por $\text{std}(\mathbf{r})$ pode causar um viés de dificuldade no nível da pergunta. Você pode desativar essa escalação definindo scale_rewards=False em GRPOConfig. Observe que desativar a escalação baseada em std também remove a normalização de variância, de modo que as magnitudes das atualizações dependem diretamente da escala bruta da recompensa e da composição do lote.

Como mostrado em Parte I: Truques ou Armadilhas? Uma Análise Aprofundada do RL para Raciocínio de LLMs (Lite PPO), calcular a média no nível local (grupo) e o desvio padrão no nível global (lote) permite um ajuste de recompensa mais robusto. Você pode usar essa estratégia de escalonamento definindo scale_rewards="batch" em GRPOConfig.

Estimando a divergência de KL

A divergência de KL é estimada usando o aproximador introduzido por Schulman et al. (2020). O aproximador é definido da seguinte forma: $$\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,$$

Calculando a perda

O objetivo é maximizar a vantagem, garantindo que o modelo permaneça próximo da política de referência. Consequentemente, a perda é definida da seguinte forma: $$\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],$$

onde o primeiro termo representa a vantagem escalonada e o segundo termo penaliza as desvios em relação à política de referência por meio da divergência de KL.

Observe que, em comparação com a formulação original em DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models, não escalamos por $\frac{1}{|o_i|}$ porque foi demonstrado no artigo Understanding R1-Zero-Like Training: A Critical Perspective que isso introduz um viés de comprimento no nível da resposta. Mais detalhes em tipos de perda.

Observe que, em comparação com a formulação original em DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models, usamos $\beta = 0.0$ por padrão, o que significa que o termo de divergência KL não é utilizado. Essa escolha é motivada por vários estudos recentes (por exemplo, Open-Reasoner-Zero: An Open Source Approach to Scaling Up Reinforcement Learning on the Base Model) ) que demonstraram que o termo de divergência KL não é essencial para o treinamento com GRPO. Como resultado, tornou-se prática comum excluí-lo (por exemplo, Understanding R1-Zero-Like Training: A Critical Perspective, DAPO: An Open-Source LLM Reinforcement Learning System at Scale). Se você deseja incluir o termo de divergência KL, pode definir beta em GRPOConfig para um valor diferente de zero.

No artigo original, essa formulação é generalizada para considerar múltiplas atualizações após cada geração (denotadas por $\mu$, pode ser definido com num_iterations em GRPOConfig) aproveitando o objetivo substituto com corte: $$\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],$$

onde $\text{clip}(\cdot, 1 - \epsilon, 1 + \epsilon)$ garante que as atualizações não se desviem excessivamente da política de referência, limitando a razão de política entre $1 - \epsilon$ e $1 + \epsilon$. Quando $\mu = 1$ (padrão no TRL), o objetivo substituto com recorte simplifica-se para o objetivo original.

Tipos de Perda

Várias formulações do objetivo foram propostas na literatura. Inicialmente, o objetivo do GRPO foi definido da seguinte forma: $$\mathcal{L}_{\text{GRPO}}(\theta) = - \frac{1}{G} \sum_{i=1}^G \frac{1}{|o_i|} \sum_{t=1}^{|o_i|} l_{i,t},$$

onde $$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].$$

O artigo DAPO destaca as limitações da perda em nível de amostra do algoritmo GRPO em cenários de CoT longo, onde respostas mais longas são penalizadas de forma insuficiente, resultando em saídas de menor qualidade. A solução proposta é uma normalização em nível de token, que lida melhor com sequências mais longas ao atribuir recompensas mais equilibradas aos tokens individuais, independentemente do comprimento da resposta: $$\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},$$

Para usar essa formulação, defina loss_type="dapo" em GRPOConfig.

Além disso, foi demonstrado no artigo Understanding R1-Zero-Like Training: A Critical Perspective que a formulação inicial do GRPO introduz um viés de comprimento da resposta. Eles mostram que, embora a formulação do DAPO reduza esse viés, ela não o elimina completamente. Para remover totalmente esse viés, eles propõem dividir por uma constante em vez do comprimento da sequência, resultando na seguinte formulação: $$\mathcal{L}_{\text{Dr. GRPO}}(\theta) = - \frac{1}{LG} \sum_{i=1}^G \sum_{t=1}^{|o_i|} l_{i,t},$$

Recomenda-se que esta constante seja o comprimento máximo de conclusão. Para usar essa formulação, defina loss_type="dr_grpo" no GRPOConfig.

Alternativamente, no artigo SAPO, a equipe do Qwen propõe substituir o mecanismo de “hard” clipping do GRPO por um mecanismo de soft gating suave e controlado por temperatura. Enquanto o GRPO zera os gradientes quando a política se desvia demais da referência, o SAPO usa uma região de confiança suave que decai o peso do gradiente de forma gradual. Isso permite que o modelo retenha sinais de aprendizado úteis de tokens “near-on-policy”, enquanto suprime o ruído proveniente de desvios extremos.

A função de perda é definida como: $$\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}$$

A função de soft-gating $f_{i,t}$ é definida usando a função sigmoide $\sigma$ como: $$f_{i,t}(x) = \sigma \left( \tau_{i,t} (x - 1) \right) \cdot \frac{4}{\tau_{i,t}}$$

A temperatura $\tau_{i,t}$ é escolhida com base no sinal da vantagem $\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}$$

Eles recomendam o uso de temperaturas assimétricas, $\tau_{\text{neg}} > \tau_{\text{pos}}$ (os valores padrão são $\tau_{\text{pos}}=1.0, \tau_{\text{neg}}=1.05$ ). Isso garante que o modelo seja penalizado de forma mais rigorosa por ações “ruins” para evitar instabilidade, enquanto é mais permissivo com ações “boas”.

Para usar essa formulação, defina loss_type="sapo" no GRPOConfig.

Métricas registradas

Durante o treinamento e a avaliação, registramos as seguintes métricas:

• num_tokens: O número total de tokens processados até o momento, incluindo tanto prompts quanto completions. Ao usar ferramentas, apenas tokens que não são de ferramentas são contabilizados.
• step_time: O tempo médio (em segundos) por etapa de treinamento (incluindo a geração).
• completions/mean_length: O comprimento médio das conclusões geradas. Ao usar ferramentas, apenas os tokens que não são de ferramentas são contados.
• completions/min_length: O comprimento mínimo das conclusões geradas. Ao usar ferramentas, apenas os tokens que não são de ferramentas são contados.
• completions/max_length: O comprimento máximo das conclusões geradas. Ao usar ferramentas, apenas os tokens que não são de ferramentas são contados.
• completions/mean_terminated_length: O comprimento médio das conclusões geradas que terminam com EOS. Ao usar ferramentas, apenas os tokens que não são de ferramentas são contados.
• completions/min_terminated_length: O comprimento mínimo das conclusões geradas que terminam com EOS. Ao usar ferramentas, apenas os tokens que não são de ferramentas são contados.
• completions/max_terminated_length: O comprimento máximo das completions geradas que terminam com EOS. Ao usar ferramentas, apenas os tokens que não são de ferramentas são contados.
• completions/clipped_ratio: A proporção de conclusões truncadas (clipped).
• tools/call_frequency: O número médio de chamadas de ferramentas por conclusão no lote de geração. Registrado apenas quando tools são fornecidas.
• tools/failure_frequency: A fração de chamadas de ferramenta que falharam (a ferramenta não foi encontrada, lançou uma exceção ou o tipo de chamada não é suportado). É 0.0 quando nenhuma ferramenta foi chamada. Registrado apenas quando tools são fornecidas.
• rewards/{reward_func_name}/mean: A recompensa média de uma função de recompensa específica. Quando um ambiente possui a recompensa por meio de get_reward, {reward_func_name} é o nome da classe do ambiente.
• rewards/{reward_func_name}/std: O desvio padrão da recompensa de uma função de recompensa específica.
• reward: A recompensa média geral após somar as recompensas entre as funções (ponderada por reward_weights).
• reward_std: O desvio padrão dos recompensas somadas entre as funções (ponderado por reward_weights), calculado sobre o lote completo.
• frac_reward_zero_std: A fração de amostras no lote de geração com desvio padrão de recompensa igual a zero, o que implica que há pouca diversidade para aquele prompt (todas as respostas estão corretas ou incorretas).
• sampling/sampling_logp_difference/mean: A diferença absoluta média entre as probabilidades logarítmicas retornadas pelo amostrador (vLLM) e as recalculadas pelo modelo de treinamento, nos tokens de conclusão. Um valor crescente indica um aumento na divergência entre treinamento e inferência. Registrado apenas quando use_vllm=True e vllm_importance_sampling_correction=True.
• sampling/sampling_logp_difference/max: A maior diferença absoluta desse tipo. Registrado apenas quando use_vllm=True e vllm_importance_sampling_correction=True.
• sampling/importance_sampling_ratio/min: a menor razão de amostragem de importância usada para corrigir a divergência entre treinamento e inferência, após a aplicação da restrição selecionada por vllm_importance_sampling_mode (limitada a [C_min, C_max] para os modos *_truncate, definida como zero para os modos *_mask). Calculada sobre os tokens de conclusão, ou sobre as sequências para os modos sequence_*. Registrada apenas quando use_vllm=True e vllm_importance_sampling_correction=True.
• sampling/importance_sampling_ratio/mean: A média da razão de amostragem de importância restrita. Registrado apenas quando use_vllm=True e vllm_importance_sampling_correction=True.
• sampling/importance_sampling_ratio/max: A maior razão de amostragem por importância restrita. Registrado apenas quando use_vllm=True e vllm_importance_sampling_correction=True.
• policy_loss: O valor da perda do gradiente de política (antes de qualquer bônus de entropia). Registrado quando entropy_coef é diferente de zero ou use_adaptive_entropy=True.
• entropy: Entropia média das previsões de tokens entre as conclusões geradas. (Se mask_truncated_completions=True, os tokens das sequências mascaradas são excluídos.)
• entropy_coef: O coeficiente atual de regularização de entropia. Registrado quando entropy_coef é diferente de zero ou use_adaptive_entropy=True. Atualizado uma vez por passo do otimizador quando use_adaptive_entropy=True.
• aux_loss: A perda auxiliar de balanceamento de carga de um modelo Mixture-of-Experts, antes de ser escalada por router_aux_loss_coef e adicionada à perda. Registrada apenas quando o modelo é um modelo MoE e router_aux_loss_coef é diferente de zero.
• kl: A divergência KL média entre o modelo e o modelo de referência, calculada sobre as conclusões geradas. Registrado apenas se beta for diferente de zero.
• clip_ratio/region_mean: A proporção de tokens (ou sequências, se importance_sampling_level="sequence") probabilidades em que o objetivo do GRPO é limitado para permanecer dentro da região de confiança: $\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})}$. Um valor mais alto significa que mais tokens são limitados, o que restringe o quanto a política $\pi_\theta$ pode mudar.
• clip_ratio/low_mean: A média da razão de probabilidades de tokens (ou sequências, se importance_sampling_level="sequence") foram cortadas no limite inferior da região de confiança: $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/low_min: A menor fração por conclusão de tokens (ou da própria sequência, se importance_sampling_level="sequence") for truncada no limite inferior da região de confiança: $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/high_mean: A proporção média de probabilidades de tokens (ou sequências, se importance_sampling_level="sequence") foram cortadas no limite superior da região de confiança: $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.
• clip_ratio/high_max: A maior fração de tokens por conclusão (ou a própria sequência, se importance_sampling_level="sequence") for truncada no limite superior da região de confiança: $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.
• cispo_clip_ratio: A proporção de pesos de amostragem de importância de tokens (ou sequências, se importance_sampling_level="sequence") ) que foram truncados em epsilon_high enquanto tinham uma vantagem positiva: $r_{i,t}(\theta) > \epsilon_\mathrm{high}$. Registrado apenas quando loss_type="cispo".
• vespo/phi_seq_mean: O valor médio da ponderação Gamma do VESPO $\varphi(w)$ aplicada aos pesos de importância no nível da sequência. Valores abaixo de 1.0 significam que as sequências estão sendo desvalorizadas. Registrado apenas quando loss_type="vespo".

Personalização

Acelere o treinamento com geração impulsionada por vLLM

A geração é frequentemente o principal gargalo ao treinar com métodos online. Para acelerar a geração, você pode usar vLLM, um motor de inferência de alto rendimento e baixa latência para LLMs. Para ativá-lo, primeiro instale o pacote com

pip install trl[vllm]

Suportamos duas formas de usar o vLLM durante o treinamento: modo servidor e modo colocalizado.

Por padrão, a Amostragem de Importância com Máscara em nível de sequência (vllm_importance_sampling_mode="sequence_mask") é ativada para a geração do vLLM para lidar com a incompatibilidade entre geração e treinamento que ocorre ao usar diferentes frameworks. Isso pode ser desativado definindo vllm_importance_sampling_correction=False. Para mais informações, consulte Masked Importance Sampling e Sequence-level Importance Sampling

Opção 1: Modo Colocate

Neste modo, o vLLM é executado dentro do processo do trainer e compartilha a memória da GPU com o modelo de treinamento. Isso evita a inicialização de um servidor separado e pode melhorar a utilização da GPU, mas pode levar a disputas de memória nas GPUs de treinamento. Este é o modo padrão.

from trl import GRPOConfig

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

Opção 2: Modo servidor

Neste modo, o vLLM é executado em um processo separado (e usando GPUs separadas) e se comunica com o treinador via HTTP. Isso é ideal se você tiver GPUs dedicadas para inferência.

1. Inicie o servidor vLLM:

   Copiado

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

2. Ative o modo de servidor no seu script de treinamento:

   Copiado

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

Certifique-se de que o servidor está usando GPUs diferentes das do trainer; caso contrário, você pode encontrar erros do NCCL. Você pode especificar as GPUs a serem usadas com a variável de ambiente CUDA_VISIBLE_DEVICES.

Dependendo do tamanho do modelo e dos requisitos gerais de memória de GPU para o treinamento, você pode precisar ajustar o parâmetro vllm_gpu_memory_utilization em GRPOConfig para evitar subutilização ou erros de memória insuficiente.

Fornecemos um HF Space para ajudar a estimar a utilização recomendada de memória de GPU com base na configuração do seu modelo e nas configurações do experimento. Basta usá-lo da seguinte forma para obter a recomendação de vllm_gpu_memory_utilization:

Se o valor recomendado não funcionar no seu ambiente, sugerimos adicionar uma pequena margem (por exemplo, +0.05 ou +0.1) ao valor recomendado para garantir estabilidade.

Se você ainda estiver recebendo erros de memória insuficiente, defina vllm_enable_sleep_mode como True e os parâmetros e o cache do vllm serão descarregados durante a etapa de otimização. Para mais informações, consulte Reduzindo o uso de memória com o vLLM Sleep Mode.

Por padrão, o GRPO usa MASTER_ADDR=localhost e MASTER_PORT=12345 para o vLLM, mas você pode sobrescrever esses valores definindo as variáveis de ambiente de acordo.

Para mais informações, consulte Aceleração do treinamento com vLLM.

Lidando com a Divergência entre Treinamento e Inferência

Embora o vLLM acelere significativamente a inferência, ele também desacopla o motor de inferência do motor de treinamento. Em teoria, esses motores são matematicamente idênticos, mas, na prática, podem produzir saídas diferentes devido a efeitos de precisão e otimizações específicas de hardware. Essa divergência reflete os diferentes objetivos de otimização dos dois sistemas. Os motores de inferência visam maximizar o throughput de amostragem, tipicamente medido em tokens por segundo, mantendo uma fidelidade de amostragem aceitável. Os frameworks de treinamento, por outro lado, focam na estabilidade numérica e na precisão para o cálculo de gradientes, frequentemente usando formatos de maior precisão, como FP32, para pesos mestres e estados do otimizador. Essas prioridades e restrições diferentes introduzem um descompasso inevitável, embora sutil, entre o treinamento e a inferência.

Esse descompasso leva a uma atualização de gradiente enviesada, que foi observada para desestabilizar o treinamento ([1][2][3][4][5]). Para simplificar, considere o gradiente de política 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]$$

Aqui $x$ denota prompts amostrados de alguma distribuição de dados, e $\pi^\text{train}$ é a política implementada pelo mecanismo de treinamento. Com o vLLM no fluxo, obtemos uma política de inferência separada $\pi^\text{inference}$, para que o gradiente de política efetivo se torne $$\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].$$

Isso transforma um problema de RL que, de outra forma, seria on policy em um problema off policy.

A maneira padrão de corrigir esse deslocamento de distribuição é a amostragem de importância (IS). Oferecemos duas variantes de IS: Amostragem de Importância Truncada (TIS) e Amostragem de Importância com Máscara (MIS). Ambas as variantes podem ser aplicadas no nível de token ou no nível de sequência. Seja $\rho$ o peso de importância, por exemplo $\rho_t$ por token ou $\rho_{\text{seq}}$ por sequência. Sob o TIS, as proporções fora da faixa [vllm_importance_sampling_clip_min, vllm_importance_sampling_clip_max] são truncadas, $$\rho \leftarrow \text{clip}(\rho, C_{\min}, C_{\max}).$$

O original artigo TIS propôs um único mecanismo de recorte de limite superior, ou seja, $\min(\rho, C_{\max})$. A implementação no TRL generalizou isso introduzindo também um limite inferior, resultando na formulação de dois lados mostrada acima, inspirada no IcePop método.
Observe que no IcePop, os limites são rotulados como $\alpha$ e $\beta$ enquanto no TRL usamos $C_{\min}$ e $C_{\max}$.

Sob o MIS, as proporções fora dessa faixa são definidas como zero, de modo que essas amostras não contribuem para o gradiente. Em outras palavras, amostras atípicas têm seu peso reduzido sob o TIS e são descartadas sob o MIS. A flag de configuração vllm_importance_sampling_mode seleciona tanto a variante de IS (mascaramento ou truncamento) quanto a granularidade (nível de token ou nível de sequência).

A amostragem de importância é a resposta algorítmica fundamentada para a incompatibilidade entre treinamento e inferência. No entanto, também existem abordagens mais diretas que tentam reduzir a incompatibilidade entre os próprios dois motores. A maioria delas são soluções de engenharia. Por exemplo, o MiniMax M1 usa uma cabeça de modelo de linguagem em FP32 no motor de inferência. A Thinking Machines explorou kernels de inferência determinísticos, embora isso venha com um custo significativo de eficiência. O vLLM demonstrou políticas consistentes em nível de bit ao se basear nos kernels determinísticos invariantes em lote da Thinking Machines, mas, até novembro de 2025, ainda há uma penalidade substancial de throughput em relação à inferência padrão do vLLM.

Acelere o treinamento com o loteamento contínuo do transformers

Como alternativa ao vLLM, você pode usar o mecanismo de lote contínuo integrado do transformers para uma geração mais rápida. O lote contínuo remove as sequências finalizadas do lote imediatamente, em vez de esperar a mais lenta terminar. Para tarefas com comprimentos de conclusão variáveis (por exemplo, raciocínio matemático), isso resulta em uma geração mais rápida e menor uso de VRAM em comparação com o generate() padrão em tamanhos de lote grandes (N≥32).

O loteamento contínuo é uma atualização direta, sem necessidade de configuração de servidor ou sincronização de pesos. Ele é executado no processo e é adequado para treinamento em uma única GPU ou ambientes com restrições de memória. Para obter o máximo de throughput de geração em escala, use o 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
    },
)

O TRL define max_memory_percent como 0.5 (em vez de 0.9 do Transformers) para deixar VRAM suficiente para a passagem de retropropagação do treinamento. Ajuste para 0.3–0.4 em lotes de geração grandes (N≥32) ou se você vir erros de memória insuficiente.

Para um exemplo completo de treinamento, consulte examples/grpo_continuous_batching/grpo_continuous_batching.py.

GRPO em escala: treine um modelo de 70B+ em vários nós

Ao treinar modelos grandes como Qwen2.5-72B, você precisa de várias otimizações-chave para tornar o treinamento eficiente e escalável em várias GPUs e nós. Essas incluem:

• DeepSpeed ZeRO Stage 3: o ZeRO utiliza paralelismo de dados para distribuir os estados do modelo (pesos, gradientes, estados do otimizador) entre várias GPUs e CPUs, reduzindo os requisitos de memória e computação em cada dispositivo. Como modelos grandes não cabem em uma única GPU, o uso do ZeRO Stage 3 é obrigatório para treinar esse tipo de modelo. Para mais detalhes, consulte DeepSpeed Integration.
• torchrun: o lançador do PyTorch inicia um processo por GPU em cada nó e os conecta por meio de um ponto de encontro comum. Para mais detalhes, consulte Distributing Training.
• vLLM: Consulte a seção anterior para saber como usar o vLLM para acelerar a geração.

Abaixo está um exemplo de script SLURM para treinar um modelo de 70B com GRPO em vários nós. Este script treina um modelo em 4 nós e usa o 5º nó para geração com 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_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,
        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 = GRPOTrainer(
        model="Qwen/Qwen2.5-72B",
        args=training_args,
        reward_funcs=accuracy_reward,
        train_dataset=dataset
    )
    trainer.train()

if __name__=="__main__":
    main()

Usando uma função de recompensa personalizada

O GRPOTrainer suporta o uso de funções de recompensa personalizadas em vez de modelos de recompensa densos. Para garantir compatibilidade, sua função de recompensa deve atender aos seguintes requisitos:

As funções de recompensa podem ser chamadas síncronas em Python ou corrotinas assíncronas async def. Quando você fornece várias funções de recompensa assíncronas, elas são aguardadas de forma concorrente (executadas em paralelo via asyncio.gather), de modo que suas latências se sobrepõem.

1. Argumentos de entrada:

   • A função deve aceitar os seguintes argumentos como palavras-chave:

     • prompts (contém os prompts),
     • completions (contém as completions geradas),
     • completion_ids (contém as conclusões tokenizadas),
     • trainer_state (TrainerState): O estado atual do trainer. Isso pode ser usado para implementar funções de recompensa dinâmicas, como aprendizado por currículo, em que a recompensa é ajustada com base no progresso do treinamento.
     • log_extra: uma função log_extra(column: str, values: list) para adicionar colunas extras à tabela de completions. Veja o Exemplo 6.
     • log_metric: uma função chamável log_metric(name: str, value: float) para registrar métricas escalares como gráficos ao lado de kl, entropy, etc. Veja o Exemplo 6.
     • environments: uma lista de instâncias de ambiente, uma por conclusão. Presente apenas quando environment_factory é fornecido. Use isso para ler o estado acumulado durante o episódio (por exemplo, env.counter). Para ambientes com estado, prefira deixar que o ambiente possua a recompensa por meio de um método get_reward — veja Recompensas.
     • Todos os nomes de colunas (exceto prompt) que o conjunto de dados pode ter. Por exemplo, se o conjunto de dados contiver uma coluna chamada ground_truth, a função será chamada com ground_truth como argumento de palavra-chave.

     A maneira mais fácil de atender a esse requisito é usar **kwargs na assinatura da função.

   • Dependendo do formato do conjunto de dados, a entrada variará:

     • Para o formato padrão, prompts e completions serão listas de strings.
     • Para o formato conversacional, prompts e completions serão listas de dicionários de mensagens.

2. Valor de retorno: A função deve retornar uma lista de floats. Cada float representa a recompensa correspondente a uma única conclusão.

Exemplo 1: Recompensar conclusões mais longas

Abaixo está um exemplo de uma função de recompensa para um formato padrão que recompensa conclusões mais longas:

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]

Você pode testá-lo da seguinte forma:

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

Exemplo 1.1: Recompensar conclusões mais longas (com base no número de caracteres)

Mesmo que no exemplo anterior, mas desta vez a função de recompensa é baseada no número de caracteres em vez de tokens.

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]

Você pode testá-lo da seguinte forma:

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

Exemplo 2: Recompensar conclusões com um formato específico

Abaixo está um exemplo de uma função de recompensa que verifica se a conclusão possui um formato específico. Este exemplo é inspirado na função format reward usada no artigo DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. Ele foi projetado para um formato conversacional, em que os prompts e as conclusões consistem em mensagens estruturadas.

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]

Você pode testar esta função da seguinte forma:

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

Exemplo 3: Recompensar conclusões com base em uma referência

Abaixo está um exemplo de uma função de recompensa que verifica se a conclusão está correta. Este exemplo é inspirado na função de recompensa de precisão usada no artigo DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. Este exemplo é projetado para o formato padrão, onde o conjunto de dados contém uma coluna chamada 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)]

Você pode testar esta função da seguinte forma:

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

Exemplo 4: Funções de recompensa multitarefa

Abaixo está um exemplo de uso de múltiplas funções de recompensa no GRPOTrainer. Neste exemplo, definimos duas funções de recompensa específicas para tarefas: math_reward_func e coding_reward_func. A math_reward_func recompensa problemas de matemática com base na sua correção, enquanto a coding_reward_func recompensa problemas de programação com base em se a solução funciona.

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()

Neste exemplo, o math_reward_func e o coding_reward_func são projetados para funcionar com um conjunto de dados misto que contém problemas de matemática e de programação. A coluna task no conjunto de dados é usada para determinar qual função de recompensa aplicar a cada problema. Se não houver uma função de recompensa relevante para uma amostra no conjunto de dados, a função de recompensa retornará None, e o GRPOTrainer continuará com as funções e tarefas válidas. Isso permite que o GRPOTrainer lide com múltiplas funções de recompensa com diferentes alcances de aplicação.

Observe que o GRPOTrainer ignorará as recompensas None retornadas pelas funções de recompensa e considerará apenas as recompensas retornadas pelas funções relevantes. Isso garante que o modelo seja treinado nas tarefas relevantes e ignore as tarefas para as quais não há uma função de recompensa relevante.

Exemplo 5: Funções de recompensa assíncronas

Funções de recompensa personalizadas também podem ser definidas como corrotinas async def. Isso é útil se a sua recompensa depender de I/O lento (por exemplo, chamar um serviço remoto). Quando você passa várias funções de recompensa assíncronas, o GRPOTrainer as executa de forma concorrente, de modo que a latência delas se sobrepõe.

Abaixo está um exemplo mínimo de uma função de recompensa assíncrona que simula uma operação limitada por E/S:

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]

Exemplo 6: Registro de colunas e métricas extras

Abaixo está um exemplo de uma função de recompensa que registra colunas extras na tabela de conclusões e métricas escalares como gráficos.

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

Passando a função de recompensa para o treinador

Para usar sua função de recompensa personalizada, passe-a para o GRPOTrainer da seguinte forma:

from trl import GRPOTrainer

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

Você pode passar várias funções de recompensa como uma lista; essa lista pode incluir tanto funções síncronas quanto assíncronas:

from trl import GRPOTrainer

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

e a recompensa será calculada como a soma das recompensas de cada função, ou como a soma ponderada, se reward_weights for fornecido na configuração.

Observe que o GRPOTrainer suporta múltiplas funções de recompensa de diferentes tipos. Consulte a documentação dos parâmetros para mais detalhes.

Regularização de entropia

Para incentivar a exploração e evitar que a política colapse para saídas quase determinísticas, você pode adicionar um bônus de entropia ao objetivo de treinamento. A regularização de entropia aumenta a perda do GRPO da seguinte forma: $$\mathcal{L}(\theta) = \mathcal{L}_{\text{GRPO}}(\theta) - \alpha \cdot \mathcal{H}(\pi_\theta),$$

em que$\mathcal{H}(\pi_\theta)$ é a entropia média por token da política e$\alpha$ é o coeficiente de entropia. O bônus é sempre a entropia média por token, independentemente de loss_type; ele não é reescalado para corresponder à normalização da política de um tipo de perda (por exemplo, o denominador batch_size * max_completion_length do Dr. GRPO), portanto entropy_coef tem o mesmo significado para todos os tipos de perda.

Entropia estática — um coeficiente fixo ao longo do treinamento:

from trl import GRPOConfig, GRPOTrainer

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

Entropia adaptativa — o coeficiente é atualizado a cada passo do otimizador com base em uma entropia alvo, conforme introduzido em Skywork-OR1. Quando a entropia atual cai para ou abaixo de entropy_target, o coeficiente é incrementado em entropy_coef_delta; caso contrário, é decrementado. O coeficiente é aplicado apenas (ou seja, é não nulo) enquanto a entropia estiver no nível ou abaixo do alvo:

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,
    ...
)

Modelos de linguagem típicos apresentam entropias por token de 2 a 10 nats, portanto o valor padrão entropy_target=0.2 quase nunca aciona a regularização — o bônus só entra em ação quando a entropia é igual ou inferior ao alvo, ou seja, em colapso quase completo. Defina-o com um valor significativo para o seu modelo, por exemplo, próximo da entropia observada no início do treinamento (registrada como a métrica entropy). Ao usar top_entropy_quantile < 1.0, entropy_target se aplica ao subconjunto de tokens de alta entropia — a entropia desse subconjunto será maior que a entropy registrada para todos os tokens, portanto, calibre de acordo.

Quando use_adaptive_entropy=True, o coeficiente de entropia atual entropy_coef é salvo junto com cada checkpoint e restaurado ao retomar, de modo que o treinamento é totalmente retomável.

Experimentação Rápida para GRPO

RapidFire AI é um motor de experimentação de código aberto que se integra ao TRL e permite iniciar múltiplas configurações de GRPO simultaneamente, mesmo em uma única GPU. Em vez de testar configurações sequencialmente, o RapidFire permite ver todas as suas curvas de aprendizado mais cedo, interromper execuções com baixo desempenho e clonar as promissoras com novas configurações em andamento sem reiniciar. Para mais informações, consulte RapidFire AI Integration.

Treinamento de Agentes

O GRPO oferece suporte a treinamento de agentes: o modelo chama ferramentas durante a geração e aprende com o resultado.

• Uma ferramenta é uma função Python simples (síncrona ou assíncrona) exposta ao modelo. Use tools para chamadas sem estado (uma calculadora, uma busca na web).
• Um environment é a forma mais geral: um objeto com estado, criado do zero para cada rollout, cujos métodos públicos são expostos como ferramentas, além de um hook de ciclo de vida reset e um get_reward opcional que permite que ele possua a recompensa. Use environment_factory quando precisar de estado por rollout, de um hook de reset ou de recompensa gerenciada pelo environment.

Elas se compõem — você pode passar tools independentes junto com um environment_factory.

Ferramentas

O argumento tools espera uma lista de funções Python (síncronas ou assíncronas) que definem as ferramentas disponíveis para o agente:

from trl import GRPOTrainer

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

Cada ferramenta deve ser uma função Python padrão com argumentos e tipos de retorno com type hints, além de uma docstring no estilo Google descrevendo seu propósito, argumentos e valor de retorno. Para mais detalhes, consulte o Guia de passagem de ferramentas.

O loop de chamada de ferramentas do GRPO exige que o modelo de chat seja preservador de prefixo (a adição de uma mensagem de ferramenta não deve alterar a forma como as mensagens anteriores são renderizadas). Para famílias de modelos conhecidas (por exemplo, Qwen3, DeepSeek-V3), o TRL substitui automaticamente o modelo de treinamento por um modelo corrigido quando as ferramentas estão habilitadas. Consulte Modelos de Chat para a lista completa.

Use max_tool_calling_iterations no GRPOConfig para limitar o número de turnos de chamada de ferramentas. Por padrão, não há limite, e a geração é interrompida quando o modelo produz um turno de resposta sem chamadas de ferramentas, quando a conclusão (todos os turnos e resultados das ferramentas) atinge max_completion_length tokens, ou quando a sequência completa atinge o comprimento máximo do modelo (vllm_max_model_length com vllm_mode="colocate", ou o max_position_embeddings do modelo, caso contrário).

Exemplo:

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],
    ...,
)

Ambientes

Você também pode fornecer ferramentas por meio de environment_factory. Neste modo, o GRPOTrainer cria uma instância de ambiente por rollout e expõe os métodos públicos do ambiente como ferramentas.

environment_factory requer transformers>=5.2.0.

Ao usar um environment_factory, o ambiente possui os dados: não há um train_dataset externo. Em cada rollout, reset() produz a tarefa — amostragem própria de um corpus mantido pelo ambiente, ou geração de um estado proceduralmente — e retorna o prompt. max_steps define o comprimento do treinamento. O seguinte é um exemplo mínimo de um ambiente que amostra um alvo e expõe um método increment (exposto como uma ferramenta).

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()

Uma classe de ambiente possui dois métodos reservados: reset e get_reward. Diferentemente de qualquer outro método público, esses não são expostos ao modelo como ferramentas.

• reset (obrigatório) é chamado quando o rollout começa. Ele pode retornar None ou uma string. No GRPO, quando retorna uma string, essa string é o prompt do usuário.
• get_reward (opcional, síncrona ou assíncrona) não recebe argumentos e retorna um float: o ambiente pontua o episódio que acabou de executar com base no seu próprio estado interno (o jogo terminou em uma vitória? a palavra foi adivinhada?). É chamada uma vez por rollout concluído e atua como fonte de recompensa.

Como o ambiente realiza a autoamostragem a cada reset(), os G membros de um grupo GRPO podem não compartilhar o mesmo estado inicial, tornando a linha de base do grupo ligeiramente mais ruidosa. Para estados idênticos dentro de um grupo, faça com que reset() derive o estado deterministicamente a partir de uma chave compartilhada.

Fornecendo um conjunto de dados externo (opcional)

Em vez de deixar que o ambiente faça a amostragem por conta própria, você pode fornecer um train_dataset externo. Por exemplo, quando suas tarefas já estão em um dataset, ou para rotear entre múltiplos ambientes. A cada rollout, o trainer amostra uma linha, usa sua coluna "prompt" como prompt e passa as outras colunas da linha para reset() como argumentos de palavra-chave (para que o ambiente leia a tarefa do dataset).

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()

Recompensas

Recompensas pertencentes ao ambiente e recompensas pertencentes ao treinador não são exclusivas. Passe reward_funcs para recompensas que são naturais de calcular a partir da conclusão (por exemplo, uma verificação de formato no texto) e deixe o ambiente retornar sua recompensa dependente do estado por meio de get_reward. Todas as fontes são somadas; reward_weights se aplica apenas a reward_funcs (o ambiente possui sua própria escala).

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`
)

A reward_func também pode ler o estado do ambiente diretamente: quando environment_factory está definido, o trainer passa um kwarg environments (uma instância por completion). Isso é ocasionalmente útil, mas para recompensas dependentes de estado, prefira get_reward.

Isso funciona de forma idêntica para AsyncGRPOTrainer.

Vários ambientes

Para treinar em uma mistura de tarefas em uma única execução (por exemplo, uma tarefa de codificação e um jogo), passe um dicionário que mapeia nomes de ambientes para fábricas, em vez de um único chamável. Cada rollout seleciona seu ambiente por meio de um campo environment no conjunto de dados, e somente as ferramentas daquele ambiente são expostas. Isso evita vazar ferramentas irrelevantes (a ferramenta move do jogo para um exemplo de codificação, e vice-versa). Cada ambiente ainda possui seus próprios dados e retorna o prompt de reset(), de modo que o conjunto de dados precisa apenas da coluna environment para rotear entre eles.

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()

A função de recompensa ainda recebe o ambiente de cada exemplo através de reward_kwargs["environments"], permitindo que ela faça um branch com base no tipo de ambiente (por exemplo, isinstance(env, CodingEnv)). As instâncias de ambiente são reutilizadas entre os passos (com reset entre episódios), de modo que um __init__ caro é pago uma única vez, em vez de a cada passo. O campo environment é consumido pelo trainer e não é encaminhado para reset.

Respostas de Ferramentas Multimodais

Ferramentas podem retornar imagens junto com texto ao retornar uma lista de blocos de conteúdo. Isso é útil para o treinamento de agentes VLM em que a ferramenta fornece feedback visual (por exemplo, capturas de tela, gráficos, capturas de câmera).

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."}]

As imagens retornadas são injetadas automaticamente na conversa e passadas ao VLM para as subsequentes rodadas de geração.

Modelos Suportados

Testado com:

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

A compatibilidade com todos os LLMs não é garantida. Se você acredita que um modelo deve ser suportado, sinta-se à vontade para abrir uma issue no GitHub — ou, melhor ainda, enviar um pull request com as alterações necessárias.

Integrações de Ambiente

Todos os ambientes se conectam ao mesmo slot environment_factory, portanto, são intercambiáveis no nível do TRL — escolha aquele cujo ecossistema se adequa à sua tarefa:

Treinando Modelos de Visão e Linguagem

O GRPO suporta o treinamento de Vision-Language Models (VLMs) em conjuntos de dados multimodais que contêm tanto texto quanto imagens.

Modelos Suportados

Testado com:

• Gemma3 — por exemplo, google/gemma-3-4b-it
• LLaVA-NeXT — por exemplo, llava-hf/llava-v1.6-mistral-7b-hf
• Qwen2-VL — por exemplo, Qwen/Qwen2-VL-2B-Instruct
• Qwen2.5-VL — por exemplo, Qwen/Qwen2.5-VL-3B-Instruct
• SmolVLM2 — por exemplo, HuggingFaceTB/SmolVLM2-2.2B-Instruct

A compatibilidade com todos os VLMs não é garantida. Se você acredita que um modelo deve ser suportado, sinta-se à vontade para abrir uma issue no GitHub — ou, melhor ainda, enviar um pull request com as alterações necessárias.

Início rápido

Use grpo_vlm.py para fazer o fine-tuning de um VLM. Comando de exemplo para treinamento em lmms-lab/multimodal-open-r1-8k-verified:

torchrun --nproc_per_node 8 examples/grpo_visual_math/grpo_visual_math.py \
  --deepspeed examples/deepspeed_configs/zero3.json \
  --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

Dicas de configuração

• Use LoRA nas camadas de projeção de visão-linguagem
• Ative a quantização de 4 bits para reduzir o uso de memória
• VLMs exigem muita memória — comece com tamanhos de lote menores
• A maioria dos modelos é compatível com o vLLM (modos server e colocate)

Formato do Conjunto de Dados

Cada amostra de treinamento deve incluir:

• prompt: Texto formatado por meio do modelo de chat do processador
• image/images: Imagem PIL ou lista de Imagens PIL

O trainer trata automaticamente a conversão de imagem para tensor por meio do processador de imagens do modelo.

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