RLOO Trainer

Visão geral

O TRL oferece suporte ao RLOO Trainer para o treinamento de modelos de linguagem, conforme descrito no artigo Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs de Arash Ahmadian, Chris Cremer, Matthias Gallé, Marzieh Fadaee, Julia Kreutzer, Ahmet Üstün e Sara Hooker.

O resumo do artigo é o seguinte:

O alinhamento de IA na forma de Aprendizado por Reforço a partir de Feedback Humano (RLHF) é cada vez mais tratado como um ingrediente crucial para modelos de linguagem de grande escala de alto desempenho. A Otimização de Política Próxima (PPO) tem sido posicionada pela literatura recente como o método canônico para a parte de RL do RLHF. No entanto, ela envolve tanto alto custo computacional quanto ajuste sensível de hiperparâmetros. Postulamos que a maioria dos princípios motivacionais que levaram ao desenvolvimento da PPO é menos uma preocupação prática no RLHF e defendemos um método menos custoso computacionalmente que preserva e até aumenta o desempenho. Revisitamos a formulação do alinhamento a partir de preferências humanas no contexto de RL. Mantendo a simplicidade como princípio orientador, mostramos que muitos componentes da PPO são desnecessários em um contexto de RLHF e que variantes de otimização no estilo REINFORCE, muito mais simples, superam tanto a PPO quanto os métodos “livres de RL” recém-propostos, como DPO e RAFT. Nosso trabalho sugere que a adaptação cuidadosa às características de alinhamento de LLMs permite se beneficiar da otimização de RL online a baixo custo.

Este método de pós-treinamento foi contribuído por Costa Huang e posteriormente refatorado por Shirin Yamani.

Início rápido

Este exemplo demonstra como treinar um modelo usando o método RLOO. Treinamos um modelo Qwen 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_rloo.py
from datasets import load_dataset
from trl import RLOOTrainer
from trl.rewards import accuracy_reward

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

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

Execute o script usando o seguinte comando:

torchrun --nproc_per_node 8 train_rloo.py

Olhando mais a fundo para o método RLOO

RLOO é 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 RLOO é maximizar a vantagem das conclusões geradas, garantindo que o modelo permaneça próximo à política de referência. Para entender como o RLOO 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 recompensa

No RLOO, a recompensa consiste em dois componentes: a recompensa fornecida pelo modelo de recompensa (ou função de recompensa) e uma penalidade KL que desencoraja a política de se desviar demais de uma política de referência fixa

1. Para cada uma das $G$ sequências geradas $o_i = (o_{i,1}, \dots, o_{i,T})$ condicionadas a uma consulta $q$, calculamos uma recompensa escalar usando um modelo de recompensa $R(o_i, q)$.
2. Concorrentemente, estimamos a divergência de KL entre a política atual $\pi_\theta$ e a política de referência fixa $\pi_{\text{ref}}$ ao longo da sequência. A estimativa de KL para a sequência $o_i$ é: $$\mathbb{D}_{\mathrm{KL}}\!\left[\pi_\theta\|\pi_{\mathrm{ref}}\right] = \sum_{t=1}^T \log \frac{\pi_\theta(o_{i,t} \mid q, o_{i,<t})}{\pi_{\mathrm{ref}}(o_{i,t} \mid q, o_{i,<t})}.$$

A recompensa final atribuída à sequência $o_i$ é a seguinte: $$r_i = R(o_i, q) - \beta \, \mathbb{D}_{\mathrm{KL}}\!\left[\pi_\theta \|\pi_{\mathrm{ref}}\right],$$

onde $\beta > 0$ controla a intensidade da penalidade KL.

Em um cenário puramente online (num_iterations = 1, padrão), os dados são gerados pela política atual. Nesse caso, a penalidade KL é calculada diretamente usando a política atual.

No cenário mais geral (por exemplo, múltiplas etapas de gradiente por lote), os dados são gerados por um snapshot anterior $\pi_{\text{old}}$. Para manter a penalidade consistente com a distribuição de amostragem, o KL é definido em relação a esta política:

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

Equivalentemente, para uma sequência amostrada $o$, a estimativa de Monte Carlo é:

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

Calculando a vantagem

Uma vez que as recompensas para cada conclusão tenham sido calculadas, calculamos uma linha de base como a recompensa média de todos os outros amostras no mesmo lote, excluindo a amostra atual. Essa linha de base é usada para reduzir a variância da estimativa do gradiente da política. A vantagem para cada conclusão é então obtida como a diferença entre a própria recompensa e essa linha de base de exclusão de um elemento.

Formalmente, para um lote de G conclusões, a linha de base para a conclusão é:$$b_i = \frac{1}{G-1} \sum_{j \neq i} r_j$$

e em seguida a vantagem para cada conclusão é calculada como a diferença entre sua recompensa e a linha de base: $$A_i = r_i - b_i$$

Calculando a perda

A perda REINFORCE é simplesmente definida como: $$\mathcal{L}_{\text{RLOO}}(\theta) = - \frac{1}{G} \sum_{i=1}^G \hat{A}_i \, \log \pi_\theta(o_i \mid q)$$

Na prática, executar múltiplas etapas de gradiente no mesmo lote torna as ações efetivamente off-policy em relação aos parâmetros atuais. Para corrigir isso, introduzimos a razão de amostragem de importância. Para evitar atualizações excessivamente grandes quando a política muda entre a amostragem e as etapas de gradiente, aplicamos um limite (clip) a essa razão: $$\mathcal{L}_{\text{RLOO}}(\theta) = - \frac{1}{G} \sum_{i=1}^G \min \left( \frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)} \hat{A}_i, \, \text{clip}\left(\frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)}, 1-\epsilon, 1+\epsilon\right) \hat{A}_i \right)$$

Em um ambiente totalmente online, de etapa única (padrão), $\frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)} = 1$ e isso se reduz ao REINFORCE padrão.

Diferente do GRPO, o RLOO não faz backpropagation de gradientes através do termo KL. Aqui, o KL é usado puramente para modelagem de recompensa: ele é calculado sob no_grad e incorporado à recompensa escalar $r_i$, que se torna a vantagem destacada $\hat{A}_i$.

Em outras palavras, o único caminho pelo qual o gradiente entra na política é o termo da razão de importância $\frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)}$, e o termo KL não faz backpropagation para a política.

Em contraste, GRPO adiciona um termo KL explícito e diferenciável ao seu objetivo, calculado em tempo real com a política atual $\pi_\theta$, de modo que seu gradiente flui tanto pelo termo de razão quanto pela penalidade KL.

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 os prompts quanto as conclusões.
• 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.
• completions/min_length: O comprimento mínimo das conclusões geradas.
• completions/max_length: O comprimento máximo das conclusões geradas.
• completions/mean_terminated_length: O comprimento médio das conclusões geradas que terminam com EOS.
• completions/min_terminated_length: O comprimento mínimo das conclusões geradas que terminam com EOS.
• completions/max_terminated_length: O comprimento máximo das conclusões geradas que terminam com EOS.
• completions/clipped_ratio: A proporção de conclusões truncadas (clipped).
• rewards/{reward_func_name}/mean: A recompensa média de uma função de recompensa específica.
• 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).
• 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.)
• 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 probabilidades de sequência em que o objetivo RLOO é limitado para permanecer dentro da região de confiança: $\text{clip}\left( r_{i}(\theta), 1 - \epsilon_\mathrm{low}, 1 + \epsilon_\mathrm{high} \right)\,, \quad r_{i}(\theta) = \frac{\pi_\theta(o_{i} \mid q)}{\pi_{\theta_{\text{old}}}(o_{i} \mid q)}$. Um valor mais alto significa que mais amostras são limitadas, o que restringe o quanto a política $\pi_\theta$ pode mudar.
• clip_ratio/low_mean: A proporção média das probabilidades de sequência que foram cortadas no limite inferior da região de confiança: $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/low_min: O menor indicador de recorte por conclusão no limite inferior da região de confiança: $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/high_mean: A razão média das probabilidades de sequência que foram cortadas no limite superior da região de confiança: $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.
• clip_ratio/high_max: O maior indicador de recorte por conclusão no limite superior da região de confiança: $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.

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.

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 RLOOConfig

training_args = RLOOConfig(
    ...,
    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 RLOOConfig
   
   training_args = RLOOConfig(
       ...,
       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 RLOOConfig 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 RLOO 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.

RLOO em escala: treinar 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 RLOO 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_rloo.py \
     --server_ip $VLLM_NODE &

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

wait

import argparse

from datasets import load_dataset
from trl import RLOOTrainer, RLOOConfig

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

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

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

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

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

if __name__=="__main__":
    main()

Usando uma função de recompensa personalizada

O RLOOTrainer 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.
     • 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 RLOOTrainer. 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 RLOOTrainer

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

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

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

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

trainer.train()

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 RLOOTrainer continuará com as funções e tarefas válidas. Isso permite que o RLOOTrainer lide com múltiplas funções de recompensa com diferentes aplicabilidades.

Observe que o RLOOTrainer 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 RLOOTrainer 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 RLOOTrainer da seguinte forma:

from trl import RLOOTrainer

trainer = RLOOTrainer(
    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 RLOOTrainer

trainer = RLOOTrainer(
    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 RLOOTrainer suporta múltiplas funções de recompensa de diferentes tipos. Consulte a documentação dos parâmetros para mais detalhes.

Treinando Modelos de Visão e Linguagem

O RLOO 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 rloo_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/rloo_visual_math/rloo_visual_math.py \
  --deepspeed examples/deepspeed_configs/zero3.json \
  --model_name_or_path Qwen/Qwen2.5-VL-3B-Instruct \
  --output_dir rloo-Qwen2.5-VL-3B-Instruct \
  --learning_rate 1e-5 \
  --dtype bfloat16 \
  --max_completion_length 1024 \
  --use_vllm \
  --vllm_mode colocate \
  --use_peft \
  --lora_target_modules "q_proj", "v_proj" \
  --log_completions

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.

RLOOTrainer

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

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

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

RLOOConfig

Referências

1. Artigo RLOO
2. Artigo Back to Basics: Revisitando a Otimização no Estilo REINFORCE para Aprendizado a partir de Feedback Humano em LLMs
3. Artigo - REINFORCE++: Uma abordagem simples e eficiente para alinhar grandes modelos de linguagem
4. Post no blog - Colocando RL de volta no RLHF
5. Post no blog - Desvendando o RLHF e suas variantes: avanços e insights práticos de engenharia
6. Youtube - RLOO: Uma otimização de baixo custo para aprendizado com feedback humano em LLMs

Guia de migração da implementação antiga (0.21 e anteriores)

Com o lançamento da versão 0.22.0, reformulamos o RLOOTrainer para que ele esteja mais alinhado com outros treinadores online da biblioteca, como o GRPOTrainer. Esta nova implementação introduz várias alterações nos parâmetros de configuração e na estrutura geral do treinador. Abaixo está um resumo das principais alterações para o RLOOConfig:

Abaixo está um resumo das principais alterações para o RLOOTrainer: