RLOO Trainer

Aperçu

TRL prend en charge le RLOO Trainer pour l’entraînement de modèles de langage, tel que décrit dans l’article 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 et Sara Hooker.

L'abstract de l'article est le suivant :

L’alignement de l’IA, sous la forme de l’apprentissage par renforcement à partir des retours humains (RLHF), est de plus en plus considéré comme un ingrédient crucial pour les grands modèles de langage à haute performance. L’optimisation de politique proximale (PPO) a été positionnée par la littérature récente comme la méthode canonique pour la partie RL du RLHF. Cependant, elle implique à la fois un coût de calcul élevé et un réglage sensible des hyperparamètres. Nous posons que la plupart des principes motivationnels qui ont conduit au développement de PPO sont moins une préoccupation pratique dans le contexte du RLHF, et nous plaidons pour une méthode moins coûteuse en calcul qui préserve, voire augmente, les performances. Nous réexaminons la formulation de l’alignement à partir des préférences humaines dans le contexte de l’apprentissage par renforcement. En gardant la simplicité comme principe directeur, nous montrons que de nombreux composants de PPO sont inutiles dans un contexte de RLHF et que des variantes d’optimisation de type REINFORCE, bien plus simples, surpassent à la fois PPO et les méthodes « sans RL » récemment proposées, telles que DPO et RAFT. Notre travail suggère qu’une adaptation soignée aux caractéristiques d’alignement des LLM permet de tirer parti de l’optimisation RL en ligne à faible coût.

Cette méthode de post-entraînement a été proposée par Costa Huang puis refactorisée par Shirin Yamani.

Démarrage rapide

Cet exemple montre comment entraîner un modèle à l'aide de la méthode RLOO. Nous entraînons un modèle Qwen 0.5B Instruct avec les prompts du jeu de données DeepMath-103K. Vous pouvez consulter les données du jeu de données ici :

Voici le script pour entraîner le modèle. accuracy_reward nécessite math-verify, que vous pouvez installer avec 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()

Exécutez le script à l'aide de la commande suivante :

accelerate launch train_rloo.py

Explorer plus en profondeur la méthode RLOO

RLOO est un algorithme d'apprentissage en ligne, ce qui signifie qu'il s'améliore de manière itérative en utilisant les données générées par le modèle entraîné lui-même pendant l'entraînement. L'intuition derrière l'objectif de RLOO est de maximiser l'avantage des complétions générées, tout en veillant à ce que le modèle reste proche de la politique de référence. Pour comprendre le fonctionnement de RLOO, il peut être décomposé en quatre étapes principales : génération des complétions, calcul de l'avantage, estimation de la divergence KL et calcul de la perte.

Générer des complétions

À chaque étape d'entraînement, nous échantillonnons un lot de prompts et générons un ensemble de $G$ complétions pour chaque prompt (notées $o_i$).

Calculer la récompense

Dans RLOO, la récompense se compose de deux éléments : la récompense fournie par le modèle de récompense (ou la fonction de récompense) et une pénalité KL qui décourage la politique de s'écarter trop d'une politique de référence fixe

1. Pour chacune des $G$ séquences générées $o_i = (o_{i,1}, \dots, o_{i,T})$ conditionnées par une requête $q$, nous calculons une récompense scalaire à l'aide d'un modèle de récompense $R(o_i, q)$.
2. Simultanément, nous estimons la divergence de KL entre la politique actuelle $\pi_\theta$ et la politique de référence fixe $\pi_{\text{ref}}$ sur la séquence. L'estimation de KL pour la séquence $o_i$ est : $$\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})}.$$

La récompense finale attribuée à la séquence $o_i$ est alors la suivante : $$r_i = R(o_i, q) - \beta \, \mathbb{D}_{\mathrm{KL}}\!\left[\pi_\theta \|\pi_{\mathrm{ref}}\right],$$

où $\beta > 0$ contrôle l'intensité de la pénalité KL.

Dans un contexte purement en ligne (num_iterations = 1, par défaut), les données sont générées par la politique actuelle. Dans ce cas, la pénalité KL est calculée directement à l'aide de la politique actuelle.

Dans un contexte plus général (par exemple, plusieurs étapes de gradient par lot), les données sont plutôt générées par un instantané antérieur $\pi_{\text{old}}$. Pour maintenir la pénalité cohérente avec la distribution d'échantillonnage, la divergence KL est définie par rapport à cette politique :

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

De manière équivalente, pour une séquence échantillonnée $o$, l'estimation de Monte Carlo est :

$\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})}.$

Calcul de l'avantage

Une fois les récompenses de chaque complétion calculées, nous déterminons une ligne de base en prenant la récompense moyenne de tous les autres échantillons du même lot, en excluant l'échantillon courant. Cette ligne de base est utilisée pour réduire la variance de l'estimation du gradient de politique. L'avantage de chaque complétion est ensuite obtenu en soustrayant cette ligne de base leave-one-out de sa propre récompense.

Formellement, pour un lot de G complétions, la référence pour la complétion est :$$b_i = \frac{1}{G-1} \sum_{j \neq i} r_j$$

puis l'avantage pour chaque complétion est calculé comme la différence entre sa récompense et la ligne de base : $$A_i = r_i - b_i$$

Calcul de la perte

La perte REINFORCE est simplement définie comme : $$\mathcal{L}_{\text{RLOO}}(\theta) = - \frac{1}{G} \sum_{i=1}^G \hat{A}_i \, \log \pi_\theta(o_i \mid q)$$

En pratique, effectuer plusieurs pas de gradient sur le même lot rend les actions effectivement hors politique par rapport aux paramètres actuels. Pour corriger cela, nous introduisons le ratio d'échantillonnage d'importance. Pour éviter des mises à jour excessivement importantes lorsque la politique change entre l'échantillonnage et les pas de gradient, nous tronquons ce ratio : $$\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)$$

Dans un contexte entièrement en ligne, en une seule étape (par défaut), $\frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)} = 1$ et cela se réduit à REINFORCE standard.

Contrairement à GRPO, RLOO ne rétropropage pas les gradients à travers le terme KL. Ici, le KL est utilisé uniquement pour la mise en forme de la récompense : il est calculé sous no_grad et intégré à la récompense scalaire $r_i$, qui devient l'avantage détaché $\hat{A}_i$.

En d'autres termes, le seul chemin par lequel le gradient pénètre dans la politique est le terme du ratio d'importance $\frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)}$, et le terme KL ne rétropropage pas dans la politique.

En revanche, GRPO ajoute un terme KL explicite et différentiable à son objectif, calculé à la volée avec la politique actuelle $\pi_\theta$, de sorte que son gradient traverse à la fois le terme de ratio et la pénalité KL.

Métriques journalisées

Lors de l'entraînement et de l'évaluation, nous enregistrons les métriques suivantes :

• num_tokens : le nombre total de jetons traités jusqu'à présent, prompts et complétions inclus.
• step_time : le temps moyen (en secondes) par étape d'entraînement (génération incluse).
• completions/mean_length : La longueur moyenne des complétions générées.
• completions/min_length : La longueur minimale des complétions générées.
• completions/max_length : La longueur maximale des complétions générées.
• completions/mean_terminated_length : La longueur moyenne des complétions générées qui se terminent par EOS.
• completions/min_terminated_length : Longueur minimale des complétions générées qui se terminent par EOS.
• completions/max_terminated_length : La longueur maximale des complétions générées qui se terminent par EOS.
• completions/clipped_ratio : Le ratio des complétions tronquées (clippées).
• rewards/{reward_func_name}/mean : la récompense moyenne issue d'une fonction de récompense spécifique.
• rewards/{reward_func_name}/std : l'écart-type de la récompense issue d'une fonction de récompense spécifique.
• reward : La récompense moyenne globale après avoir additionné les récompenses sur l'ensemble des fonctions (pondérées par reward_weights).
• reward_std : l'écart-type des récompenses cumulées sur l'ensemble des fonctions (pondérées par reward_weights), calculé sur le lot complet.
• frac_reward_zero_std : La proportion d'échantillons du lot de génération dont l'écart-type de la récompense est nul, ce qui implique qu'il y a peu de diversité pour ce prompt (toutes les réponses sont correctes ou incorrectes).
• entropy : Entropie moyenne des prédictions de jetons sur l’ensemble des complétions générées. (Si mask_truncated_completions=True, les jetons des séquences masquées sont exclus.)
• aux_loss : perte auxiliaire d'équilibrage de charge d'un modèle Mixture-of-Experts, avant d'être mise à l'échelle par router_aux_loss_coef et ajoutée à la perte. Journalisée uniquement lorsque le modèle est un modèle MoE et que router_aux_loss_coef est non nul.
• kl : divergence KL moyenne entre le modèle et le modèle de référence, calculée sur les complétions générées. Journalisé uniquement si beta est non nul.
• clip_ratio/region_mean: Le ratio des probabilités de séquences pour lesquelles l'objectif RLOO est tronqué afin de rester dans la région de confiance : $\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)}$. Une valeur plus élevée signifie que plus d'échantillons sont tronqués, ce qui limite l'ampleur des modifications de la politique $\pi_\theta$ peut subir.
• clip_ratio/low_mean: Le ratio moyen des probabilités de séquences qui ont été rognées sur la borne inférieure de la région de confiance : $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/low_min: L'indicateur de clip minimal par complétion sur la borne inférieure de la région de confiance : $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/high_mean: Le ratio moyen des probabilités de séquence qui ont été rognées sur la borne supérieure de la région de confiance : $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.
• clip_ratio/high_max: L'indicateur de clip maximal par complétion sur la borne supérieure de la région de confiance : $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.

Personnalisation

Accélérer l'entraînement avec la génération propulsée par vLLM

La génération est souvent le principal goulot d'étranglement lors de l'entraînement avec des méthodes en ligne. Pour accélérer la génération, vous pouvez utiliser vLLM, un moteur d'inférence à haut débit et faible latence pour les LLM. Pour l'activer, installez d'abord le package avec

pip install trl[vllm]

Nous prenons en charge deux façons d'utiliser vLLM pendant l'entraînement : le mode serveur et le mode colocalisé.

Option 1 : Mode Colocate

Dans ce mode, vLLM s’exécute dans le processus d’entraînement et partage la mémoire GPU avec le modèle d’entraînement. Cela évite de démarrer un serveur séparé et peut améliorer l’utilisation des GPU, mais peut entraîner une concurrence mémoire sur les GPU d’entraînement. C’est le mode par défaut.

from trl import RLOOConfig

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

Option 2 : Mode serveur

Dans ce mode, vLLM s’exécute dans un processus distinct (et utilise des GPU séparés) et communique avec l’entraîneur via HTTP. C’est idéal si vous disposez de GPU dédiés à l’inférence.

1. Démarrer le serveur vLLM :

   Copié

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

2. Activer le mode serveur dans votre script d’entraînement :

   Copié

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

Assurez-vous que le serveur utilise des GPU différents de ceux du formateur, sinon vous risquez de rencontrer des erreurs NCCL. Vous pouvez spécifier les GPU à utiliser avec la variable d’environnement CUDA_VISIBLE_DEVICES.

En fonction de la taille du modèle et des besoins globaux en mémoire GPU pour l’entraînement, vous devrez peut-être ajuster le paramètre vllm_gpu_memory_utilization dans RLOOConfig afin d’éviter une sous-utilisation ou des erreurs de mémoire insuffisante.

Nous mettons à disposition un HF Space pour vous aider à estimer l'utilisation recommandée de la mémoire GPU en fonction de la configuration de votre modèle et des paramètres de votre expérience. Utilisez-le simplement comme suit pour obtenir la recommandation de vllm_gpu_memory_utilization :

Si la valeur recommandée ne fonctionne pas dans votre environnement, nous suggérons d'ajouter une petite marge (par exemple, +0.05 ou +0.1) à la valeur recommandée pour garantir la stabilité.

Si vous rencontrez toujours des erreurs de mémoire insuffisante, définissez vllm_enable_sleep_mode sur True et les paramètres vllm ainsi que le cache seront déchargés pendant l'étape d'optimisation. Pour plus d'informations, consultez Réduire l'utilisation de la mémoire avec le mode veille de vLLM.

Par défaut, RLOO utilise MASTER_ADDR=localhost et MASTER_PORT=12345 pour vLLM, mais vous pouvez remplacer ces valeurs en définissant les variables d’environnement en conséquence.

Pour plus d'informations, consultez Accélérer l'entraînement avec vLLM.

RLOO à grande échelle : entraîner un modèle de 70B+ sur plusieurs nœuds

Lors de l'entraînement de grands modèles comme Qwen2.5-72B, plusieurs optimisations clés sont nécessaires pour rendre l'entraînement efficace et scalable sur plusieurs GPU et nœuds. Elles incluent :

• DeepSpeed ZeRO Stage 3 : ZeRO exploite le parallélisme de données pour distribuer les états du modèle (poids, gradients, états de l’optimiseur) sur plusieurs GPU et CPU, réduisant ainsi les besoins en mémoire et en calcul sur chaque appareil. Comme les grands modèles ne peuvent pas tenir sur un seul GPU, l’utilisation de ZeRO Stage 3 est requise pour entraîner ce type de modèles. Pour plus de détails, voir DeepSpeed Integration.
• Accelerate : Accelerate est une bibliothèque qui simplifie l’entraînement distribué sur plusieurs GPU et nœuds. Elle fournit une API simple pour lancer l’entraînement distribué et gère les complexités de l’entraînement distribué, telles que le parallélisme de données, l’accumulation de gradients et le chargement distribué des données. Pour plus de détails, voir Distributing Training.
• vLLM : Reportez-vous à la section précédente pour savoir comment utiliser vLLM afin d’accélérer la génération.

Voici un exemple de script SLURM pour entraîner un modèle de 70B avec RLOO sur plusieurs nœuds. Ce script entraîne un modèle sur 4 nœuds et utilise le 5e nœud pour la génération propulsée par 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_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,
        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()

Utiliser une fonction de récompense personnalisée

Le RLOOTrainer prend en charge l'utilisation de fonctions de récompense personnalisées à la place des modèles de récompense denses. Pour garantir la compatibilité, votre fonction de récompense doit satisfaire aux exigences suivantes :

Les fonctions de récompense peuvent être des fonctions Python synchrones ou des coroutines asynchrones async def. Lorsque vous fournissez plusieurs fonctions de récompense asynchrones, elles sont attendues de manière concurrente (exécutées en parallèle via asyncio.gather) afin que leur latence se superpose.

1. Arguments d'entrée :

   • La fonction doit accepter les éléments suivants en tant qu'arguments nommés :

     • prompts (contient les prompts),
     • completions (contient les complétions générées),
     • completion_ids (contient les complétions tokenisées),
     • trainer_state (TrainerState) : l'état actuel du formateur. Il peut être utilisé pour implémenter des fonctions de récompense dynamiques, comme l'apprentissage par curriculum, où la récompense est ajustée en fonction de la progression de l'entraînement.
     • log_extra : un appelable log_extra(column: str, values: list) pour ajouter des colonnes supplémentaires au tableau des complétions. Voir l’exemple 6.
     • log_metric : un appelable log_metric(name: str, value: float) pour journaliser des métriques scalaires sous forme de graphiques aux côtés de kl, entropy, etc. Voir l’exemple 6.
     • Tous les noms de colonnes (sauf prompt) que le jeu de données peut contenir. Par exemple, si le jeu de données contient une colonne nommée ground_truth, la fonction sera appelée avec ground_truth en tant qu'argument par mot-clé.

     La façon la plus simple de respecter cette exigence est d'utiliser **kwargs dans la signature de la fonction.

   • Selon le format du jeu de données, l'entrée variera :

     • Pour le format standard, prompts et completions seront des listes de chaînes de caractères.
     • Pour le format conversationnel, prompts et completions seront des listes de dictionnaires de messages.

2. Valeur de retour : La fonction doit retourner une liste de flottants. Chaque flottant représente la récompense correspondant à une seule complétion.

Exemple 1 : Récompenser les complétions plus longues

Voici un exemple de fonction de récompense pour un format standard qui récompense les complétions plus longues :

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]

Vous pouvez le tester comme suit :

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

Exemple 1.1 : Récompenser les complétions plus longues (en fonction du nombre de caractères)

Identique à l'exemple précédent, mais cette fois la fonction de récompense est basée sur le nombre de caractères plutôt que sur les jetons.

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]

Vous pouvez le tester comme suit :

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

Exemple 2 : Récompenser les complétions avec un format spécifique

Voici un exemple de fonction de récompense qui vérifie si l’achèvement respecte un format spécifique. Cet exemple s’inspire de la fonction format reward utilisée dans l’article DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. Elle est conçue pour un format conversationnel, où les invites et les achèvements sont constitués de messages structurés.

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]

Vous pouvez tester cette fonction comme suit :

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

Exemple 3 : Récompenser les complétions sur la base d'une référence

Voici un exemple de fonction de récompense qui vérifie si la complétion est correcte. Cet exemple s’inspire de la fonction accuracy reward utilisée dans l’article DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. Cet exemple est conçu pour le format standard, où le jeu de données contient une colonne nommée 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)]

Vous pouvez tester cette fonction comme suit :

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

Exemple 4 : Fonctions de récompense multi-tâches

Voici un exemple d'utilisation de plusieurs fonctions de récompense dans le RLOOTrainer. Dans cet exemple, nous définissons deux fonctions de récompense spécifiques à la tâche : math_reward_func et coding_reward_func. La fonction math_reward_func attribue une récompense aux problèmes de mathématiques en fonction de leur exactitude, tandis que la fonction coding_reward_func attribue une récompense aux problèmes de codage en fonction du fait que la solution fonctionne.

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

Dans cet exemple, math_reward_func et coding_reward_func sont conçus pour fonctionner avec un jeu de données mixte contenant à la fois des problèmes de mathématiques et de codage. La colonne task du jeu de données est utilisée pour déterminer quelle fonction de récompense appliquer à chaque problème. S'il n'existe pas de fonction de récompense pertinente pour un échantillon du jeu de données, la fonction de récompense renvoie None, et le RLOOTrainer poursuit avec les fonctions et tâches valides. Cela permet au RLOOTrainer de gérer plusieurs fonctions de récompense avec des domaines d'application différents.

Notez que le RLOOTrainer ignorera les récompenses None renvoyées par les fonctions de récompense et ne tiendra compte que des récompenses renvoyées par les fonctions pertinentes. Cela garantit que le modèle est entraîné sur les tâches pertinentes et ignore les tâches pour lesquelles il n’existe pas de fonction de récompense pertinente.

Exemple 5 : Fonctions de récompense asynchrones

Les fonctions de récompense personnalisées peuvent également être définies sous forme de coroutines async def. Cela est utile si votre récompense dépend d'une E/S lente (par exemple, l'appel d'un service distant). Lorsque vous passez plusieurs fonctions de récompense asynchrones, RLOOTrainer les exécute de manière concurrente afin que leur latence se superpose.

Voici un exemple minimal d'une fonction de récompense asynchrone qui simule une opération limitée par l'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]

Exemple 6 : Journalisation de colonnes et de métriques supplémentaires

Voici un exemple de fonction de récompense qui journalise des colonnes supplémentaires dans le tableau des complétions et des métriques scalaires sous forme de graphiques.

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

Passer la fonction de récompense à l'entraîneur

Pour utiliser votre fonction de récompense personnalisée, transmettez-la au RLOOTrainer comme suit :

from trl import RLOOTrainer

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

Vous pouvez passer plusieurs fonctions de récompense sous forme de liste ; cette liste peut inclure à la fois des fonctions synchrones et asynchrones :

from trl import RLOOTrainer

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

et la récompense sera calculée comme la somme des récompenses de chaque fonction, ou comme la somme pondérée si reward_weights est fourni dans la configuration.

Notez que RLOOTrainer prend en charge plusieurs fonctions de récompense de types différents. Consultez la documentation des paramètres pour plus de détails.

Entraînement des modèles de vision et de langage

RLOO prend en charge l’entraînement de modèles Vision-Language (VLM) sur des jeux de données multimodaux contenant à la fois du texte et des images.

Modèles pris en charge

Testé avec :

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

La compatibilité avec tous les VLM n'est pas garantie. Si vous pensez qu'un modèle devrait être pris en charge, n'hésitez pas à ouvrir un ticket sur GitHub — ou mieux encore, à soumettre une pull request avec les modifications nécessaires.

Démarrage rapide

Utilisez rloo_vlm.py pour l’ajustement fin d’un VLM. Exemple de commande pour l’entraînement sur lmms-lab/multimodal-open-r1-8k-verified :

accelerate launch \
  --config_file=examples/accelerate_configs/deepspeed_zero3.yaml \
  examples/rloo_visual_math/rloo_visual_math.py \
  --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

Conseils de configuration

• Utiliser LoRA sur les couches de projection vision-langage
• Activer la quantification 4 bits pour réduire l'utilisation de la mémoire
• Les VLM sont gourmands en mémoire — commencez par des tailles de lot plus petites
• La plupart des modèles sont compatibles avec vLLM (modes server et colocate)

Format du jeu de données

Chaque exemple d'entraînement doit inclure :

• prompt : texte formaté via le modèle de discussion du processeur
• image/images : Image PIL ou liste d'Images PIL

Le formateur gère automatiquement la conversion des images en tenseurs via le processeur d'images du modèle.

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

Références

1. Article RLOO
2. Retour aux bases : réexamen de l’optimisation de type REINFORCE pour l’apprentissage à partir des retours humains dans les LLM
3. Article - REINFORCE++ : une approche simple et efficace pour l’alignement des grands modèles de langage
4. Article de blogue - Remettre le RL dans le RLHF
5. Article de blog – Déchiffrer le RLHF et ses variantes : progrès et retours d’expérience en ingénierie pratique
6. YouTube - RLOO : une optimisation rentable pour l'apprentissage à partir des retours humains sur les LLM

Guide de migration depuis l'ancienne implémentation (0.21 et versions antérieures)

Avec la publication de la version 0.22.0, nous avons repensé le RLOOTrainer pour qu'il soit plus aligné avec les autres entraîneurs en ligne de la bibliothèque, comme GRPOTrainer. Cette nouvelle implémentation introduit plusieurs modifications aux paramètres de configuration et à la structure générale de l'entraîneur. Voici un résumé des principales modifications pour RLOOConfig :

Voici un résumé des principales modifications apportées à RLOOTrainer :