RLOO Trainer

개요

TRL은 Arash Ahmadian, Chris Cremer, Matthias Gallé, Marzieh Fadaee, Julia Kreutzer, Ahmet Üstün 및 Sara Hooker의 논문 Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs에 설명된 대로 언어 모델 학습을 위한 RLOO Trainer를 지원합니다.

해당 논문의 초록은 다음과 같습니다:

인간 피드백 기반 강화 학습(Reinforcement Learning from Human Feedback, RLHF) 형태의 AI 정렬은 점점 더 고성능 대규모 언어 모델의 핵심 요소로 간주되고 있습니다. 최근 문헌에서는 근접 정책 최적화(Proximal Policy Optimization, PPO)를 RLHF의 RL 부분에 대한 표준적인 방법으로 제시하고 있습니다. 그러나 PPO는 높은 계산 비용과 민감한 하이퍼파라미터 튜닝을 수반합니다. 우리는 PPO 개발로 이어진 대부분의 동기 부여 원칙이 RLHF에서는 실질적인 우려 사항이 아니라고 주장하며, 성능을 유지하거나 오히려 향상시키면서 계산 비용이 덜 드는 방법을 옹호합니다. 우리는 RL 맥락에서 인간 선호도 기반 정렬의 수식화를 재검토합니다. 단순성을 지침으로 삼아, 우리는 PPO의 많은 구성 요소가 RLHF 맥락에서 불필요하며, PPO와 DPO, RAFT와 같은 신생 “RL-free” 방법론 모두보다 훨씬 단순한 REINFORCE 스타일의 최적화 변형이 더 우수한 성능을 낸다는 것을 보여줍니다. 우리의 연구는 LLM 정렬 특성에 대한 신중한 적응을 통해 낮은 비용으로 온라인 RL 최적화의 이점을 얻을 수 있음을 시사합니다.

이 포스트 트레이닝 방법은 Costa Huang이 기여했으며, 이후 Shirin Yamani가 리팩토링했습니다.

빠른 시작

이 예제는 RLOO 방법을 사용하여 모델을 학습하는 방법을 보여줍니다. DeepMath-103K 데이터셋의 프롬프트를 사용하여 Qwen 0.5B Instruct 모델을 학습합니다. 데이터셋의 데이터는 여기서 확인할 수 있습니다:

아래는 모델을 학습하는 스크립트입니다. accuracy_reward에는 math-verify가 필요하며, pip install trl[math_verify]로 설치할 수 있습니다.

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

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

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

다음 명령을 사용하여 스크립트를 실행합니다:

torchrun --nproc_per_node 8 train_rloo.py

RLOO 방법 자세히 살펴보기

RLOO는 온라인 학습 알고리즘으로, 학습 중에 학습된 모델 자체에서 생성된 데이터를 사용하여 반복적으로 성능을 향상시킵니다. RLOO 목적 함수의 핵심 직관은 모델이 참조 정책(reference policy)에 가깝도록 유지하면서 생성된 completions의 이점을 최대화하는 것입니다. RLOO의 작동 방식을 이해하려면 다음 네 가지 주요 단계로 나눌 수 있습니다: completions 생성, 이점(advantage) 계산, KL 발산 추정, 손실(loss) 계산.

완성 결과 생성

각 학습 단계에서 프롬프트 배치(batches)를 샘플링하고 각 프롬프트에 대해 $G$ 개의 completion을 생성합니다(기호는 $o_i$).

보상 계산

RLOO에서 보상(reward)은 두 가지 구성 요소로 이루어져 있습니다. 보상 모델(또는 보상 함수)이 제공하는 보상과, 정책(policy)이 고정된 참조 정책(reference policy)에서 너무 벗어나지 않도록 억제하는 KL 페널티입니다.

1. 각 $G$ 생성된 시퀀스 $o_i = (o_{i,1}, \dots, o_{i,T})$ 에 대해 쿼리 $q$, 에 조건부화된 각 시퀀스에 대해 보상 모델 $R(o_i, q)$.
2. 동시에, 현재 정책 $\pi_\theta$ 과 고정된 참조 정책 $\pi_{\text{ref}}$ 사이의 시퀀스 전체에 걸친 KL 발산을 추정합니다. 시퀀스 $o_i$ 에 대한 KL 추정값은 다음과 같습니다: $$\mathbb{D}_{\mathrm{KL}}\!\left[\pi_\theta\|\pi_{\mathrm{ref}}\right] = \sum_{t=1}^T \log \frac{\pi_\theta(o_{i,t} \mid q, o_{i,<t})}{\pi_{\mathrm{ref}}(o_{i,t} \mid q, o_{i,<t})}.$$

시퀀스에 할당되는 최종 보상 $o_i$ 은 다음과 같습니다: $$r_i = R(o_i, q) - \beta \, \mathbb{D}_{\mathrm{KL}}\!\left[\pi_\theta \|\pi_{\mathrm{ref}}\right],$$

여기서 $\beta > 0$ KL 패널티의 강도를 제어합니다.

순수 온라인 설정(num_iterations = 1, 기본값)에서는 데이터가 현재 정책(policy)에 의해 생성됩니다. 이 경우 KL 패널티는 현재 정책을 직접 사용하여 계산됩니다.

더 일반적인 설정(예: 배치당 여러 개의 gradient step)에서는 데이터가 이전 스냅샷 $\pi_{\text{old}}$. 이 정책 기준으로 정의된 KL을 사용하여, 벌칙이 샘플링 분포와 일치하도록 합니다:

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

동일하게, 샘플링된 시퀀스 $o$, 에 대해 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})}.$

이점(advantage) 계산

각 completion에 대한 보상이 계산되면, 현재 샘플을 제외한 동일한 배치 내 다른 모든 샘플의 평균 보상을 baseline으로 계산합니다. 이 baseline은 정책 기울기 추정치의 분산을 줄이는 데 사용됩니다. 각 completion의 advantage는 해당 completion의 자체 보상에서 이 leave-one-out baseline을 뺀 차이로 구해집니다.

정식적으로, G개의 completion 배치에 대해 completion의 baseline은 다음과 같습니다:$$b_i = \frac{1}{G-1} \sum_{j \neq i} r_j$$

그리고 각 completion의 advantage는 해당 reward와 baseline의 차이로 계산됩니다: $$A_i = r_i - b_i$$

손실 계산

REINFORCE 손실은 다음과 같이 간단히 정의됩니다: $$\mathcal{L}_{\text{RLOO}}(\theta) = - \frac{1}{G} \sum_{i=1}^G \hat{A}_i \, \log \pi_\theta(o_i \mid q)$$

실제로 동일한 배치에서 여러 단계의 기울기 업데이트를 수행하면, 현재 파라미터에 대해 동작이 사실상 오프-폴리시(off-policy)가 됩니다. 이를 보정하기 위해 중요도 샘플링(importance sampling) 비율을 도입합니다. 샘플링과 기울기 단계 사이에 정책이 변경될 때 과도하게 큰 업데이트를 방지하기 위해 이 비율을 클리핑(clip)합니다: $$\mathcal{L}_{\text{RLOO}}(\theta) = - \frac{1}{G} \sum_{i=1}^G \min \left( \frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)} \hat{A}_i, \, \text{clip}\left(\frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)}, 1-\epsilon, 1+\epsilon\right) \hat{A}_i \right)$$

완전 온라인, 단일 단계 설정(기본값)에서 $\frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)} = 1$ 이는 표준 REINFORCE로 축소됩니다.

GRPO와 달리 RLOO는 KL 항을 통해 기울기를 역전파하지 않습니다. 여기서 KL은 순수하게 보상 성형(reward shaping)에 사용되며, no_grad 아래에서 계산되어 스칼라 보상 $r_i$, 에 포함되며, 이는 분리된(detached) 이점(advantage)이 됩니다. $\hat{A}_i$.

즉, 정책으로 전달되는 경로의 유일한 경로는 중요도 비율(importance-ratio) 항입니다 $\frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)}$, 그리고 KL 항은 정책으로 역전파되지 않습니다.

반면, GRPO 는 현재 정책 $\pi_\theta$, 을 통해 실시간으로 계산되는 명시적이고 미분 가능한 KL 항을 목적 함수에 추가하여, 그 기울기가 비율 항과 KL 페널티 양쪽을 통해 흐르도록 합니다.

기록된 지표

학습 및 평가 중에 다음 메트릭을 기록합니다:

• num_tokens: 프롬프트와 완성 결과를 모두 포함하여 지금까지 처리된 토큰의 총 개수.
• step_time: 학습 단계당 평균 소요 시간(초 단위, 생성 포함).
• completions/mean_length: 생성된 completions의 평균 길이.
• completions/min_length: 생성된 completions의 최소 길이입니다.
• completions/max_length: 생성된 completions의 최대 길이.
• completions/mean_terminated_length: EOS로 종료된 생성된 completions의 평균 길이입니다.
• completions/min_terminated_length: EOS로 종료되는 생성된 completions의 최소 길이입니다.
• completions/max_terminated_length: EOS로 종료되는 생성된 completions의 최대 길이입니다.
• completions/clipped_ratio: 잘림(클립)된 completions의 비율.
• rewards/{reward_func_name}/mean: 특정 보상 함수의 평균 보상입니다.
• rewards/{reward_func_name}/std: 특정 보상 함수의 보상의 표준편차입니다.
• reward: 함수 간 보상을 합산한 후(가중치는 reward_weights로 적용) 계산된 전체 평균 보상입니다.
• reward_std: 함수 전체에 걸쳐 합산된 보상(가중치 reward_weights 적용)의 표준편차로, 전체 배치에 대해 계산됩니다.
• frac_reward_zero_std: 생성 배치 내에서 보상 표준편차가 0인 샘플의 비율입니다. 이는 해당 프롬프트에 대한 다양성이 적다는 것을 의미하며(모든 답변이 정답이거나 오답임)
• entropy: 생성된 완성본 전체에 걸친 토큰 예측의 평균 엔트로피입니다. (mask_truncated_completions=True인 경우, 마스킹된 시퀀스의 토큰은 제외됩니다.)
• aux_loss: Mixture-of-Experts 모델의 로드 밸런싱 보조 손실입니다. router_aux_loss_coef로 스케일링되어 손실에 추가되기 전의 값입니다. 모델이 MoE 모델이고 router_aux_loss_coef가 0이 아닌 경우에만 기록됩니다.
• kl: 생성된 completions에 대해 모델과 reference model 사이의 평균 KL 발산을 계산한 값입니다. beta가 0이 아닌 경우에만 기록됩니다.
• clip_ratio/region_mean: RLOO 목적 함수가 신뢰 영역(trust region) 내에 머물도록 클리핑되는 시퀀스 확률의 비율: $\text{clip}\left( r_{i}(\theta), 1 - \epsilon_\mathrm{low}, 1 + \epsilon_\mathrm{high} \right)\,, \quad r_{i}(\theta) = \frac{\pi_\theta(o_{i} \mid q)}{\pi_{\theta_{\text{old}}}(o_{i} \mid q)}$. 값이 높을수록 더 많은 샘플이 클리핑되어 정책 $\pi_\theta$ 이 변경될 수 있는 범위가 제한됩니다.
• clip_ratio/low_mean: trust region의 하한에서 클리핑된 시퀀스 확률의 평균 비율: $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/low_min: trust region 하한에서 completion별 가장 작은 clip 지표: $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/high_mean: trust region의 상한에서 클리핑된 시퀀스 확률의 평균 비율: $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.
• clip_ratio/high_max: trust region의 상한에서 completion별 최대 clip 지표: $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.

커스터마이징

vLLM 기반 생성으로 학습 속도 향상

온라인 방식으로 학습할 때 생성 과정이 주요 병목 현상이 되는 경우가 많습니다. 생성을 가속화하려면 LLM용 고처리량, 저지연 추론 엔진인 vLLM을 사용할 수 있습니다. 이를 활성화하려면 먼저 패키지를 설치해야 합니다.

pip install trl[vllm]

학습 중 vLLM을 사용하는 두 가지 방식을 지원합니다: 서버 모드와 콜로케이트 모드입니다.

옵션 1: Colocate 모드

이 모드에서는 vLLM이 트레이너 프로세스 내에서 실행되며 학습 모델과 GPU 메모리를 공유합니다. 이를 통해 별도의 서버를 시작할 필요가 없으며 GPU 활용도를 높일 수 있지만, 학습 GPU에서 메모리 경쟁이 발생할 수 있습니다. 이 모드가 기본 모드입니다.

from trl import RLOOConfig

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

옵션 2: 서버 모드

이 모드에서는 vLLM이 별도의 프로세스(그리고 별도의 GPU)에서 실행되며 HTTP를 통해 trainer와 통신합니다. 추론에 전용 GPU가 있는 경우 이상적입니다.

1. vLLM 서버 시작:

   복사됨

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

2. 학습 스크립트에서 서버 모드를 활성화합니다:

   복사됨

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

서버가 trainer와 다른 GPU를 사용하도록 확인하세요. 그렇지 않으면 NCCL 오류가 발생할 수 있습니다. 사용할 GPU는 CUDA_VISIBLE_DEVICES 환경 변수로 지정할 수 있습니다.

모델 크기와 학습에 필요한 전체 GPU 메모리 요구 사항에 따라, 리소스 과소 사용 또는 메모리 부족 오류를 방지하기 위해 RLOOConfig의 vllm_gpu_memory_utilization 매개변수를 조정해야 할 수 있습니다.

모델 구성 및 실험 설정에 기반하여 권장 GPU 메모리 사용률을 추정하는 데 도움이 되는 HF Space를 제공합니다. 아래처럼 사용하면 vllm_gpu_memory_utilization 권장 값을 얻을 수 있습니다:

권장 값이 사용자의 환경에서 작동하지 않는 경우, 안정성을 보장하기 위해 권장 값에 작은 버퍼(예: +0.05 또는 +0.1)를 추가하는 것을 권장합니다.

여전히 out-of-memory 오류가 발생하면 vllm_enable_sleep_mode을 True로 설정하세요. 그러면 최적화 단계에서 vllm 매개변수와 캐시가 offload됩니다. 자세한 내용은 vLLM Sleep Mode를 사용한 메모리 사용량 감소를 참조하세요.

기본적으로 RLOO는 vLLM에 대해 MASTER_ADDR=localhost와 MASTER_PORT=12345를 사용하지만, 환경 변수를 적절히 설정하여 이 값들을 덮어쓸 수 있습니다.

자세한 내용은 vLLM을 활용한 학습 속도 향상을 참조하세요.

대규모 RLOO: 여러 노드에서 70B+ 모델 학습

Qwen2.5-72B와 같은 대규모 모델을 학습할 때는 여러 GPU와 노드에 걸쳐 학습을 효율적이고 확장 가능하게 만들기 위해 몇 가지 핵심 최적화가 필요합니다. 여기에는 다음이 포함됩니다:

• DeepSpeed ZeRO Stage 3: ZeRO는 데이터 병렬성을 활용하여 모델 상태(가중치, 기울기, 옵티마이저 상태)를 여러 GPU와 CPU에 분산 배치하여 각 장치의 메모리 및 계산 요구 사항을 줄입니다. 대규모 모델은 단일 GPU에 적재할 수 없으므로, 이러한 모델을 학습하려면 ZeRO Stage 3 사용이 필수적입니다. 자세한 내용은 DeepSpeed Integration을 참조하세요.
• torchrun: PyTorch의 런처는 모든 노드에서 GPU당 하나의 프로세스를 시작하고 공통 rendezvous를 통해 연결합니다. 자세한 내용은 학습 분산을 참조하세요.
• vLLM: 생성 속도를 높이기 위해 vLLM을 사용하는 방법은 이전 섹션을 참고하세요.

아래는 여러 노드에서 RLOO를 사용하여 70B 모델을 학습하기 위한 SLURM 스크립트 예시입니다. 이 스크립트는 4개의 노드에서 모델을 학습하고, 5번째 노드에서는 vLLM 기반 생성을 수행합니다.

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

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

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

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

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

wait

import argparse

from datasets import load_dataset
from trl import RLOOTrainer, RLOOConfig

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

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

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

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

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

if __name__=="__main__":
    main()

사용자 정의 보상 함수 사용

RLOOTrainer는 밀도 보상 모델 대신 사용자 정의 보상 함수를 사용할 수 있습니다. 호환성을 보장하기 위해 보상 함수는 다음 요구 사항을 충족해야 합니다:

Reward 함수는 동기 Python callable이거나 비동기 async def 코루틴일 수 있습니다. 여러 비동기 reward 함수를 제공하면 asyncio.gather를 통해 병렬로 실행되어 대기되므로, 지연 시간이 겹치게 됩니다.

1. 입력 인자:

   • 이 함수는 다음을 키워드 인수로 받아야 합니다:

     • prompts(프롬프트를 포함합니다),
     • completions (생성된 completions을 포함),
     • completion_ids (토큰화된 completions을 포함),
     • trainer_state (TrainerState): 트레이너의 현재 상태입니다. 학습 진행도에 따라 보상이 조정되는 커리큘럼 학습과 같은 동적 보상 함수를 구현하는 데 사용할 수 있습니다.
     • log_extra: completions 테이블에 추가 열을 추가하기 위한 callable log_extra(column: str, values: list). 예제 6을 참조하세요.
     • log_metric: kl, entropy 등과 함께 스칼라 메트릭을 플롯으로 기록하기 위한 콜러블 log_metric(name: str, value: float). 예제 6을 참조하세요.
     • 데이터셋에 있을 수 있는 모든 열 이름(prompt 제외). 예를 들어, 데이터셋에 ground_truth라는 이름의 열이 포함된 경우, 해당 함수는 ground_truth를 키워드 인수로 받아 호출됩니다.

     이 요구 사항을 충족하는 가장 쉬운 방법은 함수 시그니처에 **kwargs를 사용하는 것입니다.

   • 데이터셋 형식에 따라 입력이 달라집니다:

     • 표준 형식의 경우, prompts와 completions는 문자열 목록입니다.
     • 대화 형식의 경우, prompts와 completions는 메시지 딕셔너리의 목록입니다.

2. 반환 값: 함수는 float의 리스트를 반환해야 합니다. 각 float는 단일 completion에 해당하는 보상을 나타냅니다.

예제 1: 더 긴 완성본을 보상

아래는 더 긴 완성본을 보상하는 표준 형식의 보상 함수 예시입니다:

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

다음과 같이 테스트할 수 있습니다:

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

예제 1.1: 더 긴 완성본을 보상합니다 (문자 수 기준)

이전 예제와 동일하지만, 이번에는 보상 함수가 토큰이 아닌 문자 수를 기반으로 합니다.

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

다음과 같이 테스트할 수 있습니다:

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

예제 2: 특정 형식으로 보상 완료

아래는 완성물이 특정 형식을 갖는지 확인하는 보상 함수의 예시입니다. 이 예시는 DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning 논문에서 사용된 format reward 함수에서 영감을 받았습니다. 프롬프트와 완성물이 구조화된 메시지로 구성된 대화형 형식을 위해 설계되었습니다.

import re

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

다음과 같이 이 함수를 테스트할 수 있습니다:

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

예제 3: 참조를 기반으로 보상 완료

아래는 완성물이 올바른지 확인하는 보상 함수의 예시입니다. 이 예시는 DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning 논문에서 사용된 accuracy reward 함수에서 영감을 받았습니다. 이 예시는 standard format를 위해 설계되었으며, 데이터셋에는 ground_truth라는 이름의 열이 포함되어 있습니다.

import re

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

다음과 같이 이 함수를 테스트할 수 있습니다:

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

예제 4: 다중 작업 보상 함수

아래는 RLOOTrainer에서 여러 보상 함수를 사용하는 예시입니다. 이 예시에서는 두 가지 작업별 보상 함수인 math_reward_func와 coding_reward_func를 정의합니다. math_reward_func는 수학 문제의 정답 여부 기준으로 보상을 부여하고, coding_reward_func는 코딩 문제의 솔루션이 정상적으로 작동하는지 여부 기준으로 보상을 부여합니다.

from datasets import Dataset
from trl import RLOOTrainer

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

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

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

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

trainer.train()

이 예시에서 math_reward_func와 coding_reward_func는 수학 문제와 코딩 문제를 모두 포함하는 혼합 데이터셋과 함께 작동하도록 설계되었습니다. 데이터셋의 task 컬럼은 각 문제에 어떤 보상 함수를 적용할지 결정하는 데 사용됩니다. 데이터셋의 샘플에 해당하는 보상 함수가 없으면 보상 함수는 None을 반환하고, RLOOTrainer는 유효한 함수와 작업을 계속 처리합니다. 이를 통해 RLOOTrainer는 서로 다른 적용 범위를 가진 여러 보상 함수를 처리할 수 있습니다.

RLOOTrainer는 보상 함수가 반환하는 None 보상을 무시하고 관련 함수가 반환한 보상만 고려한다는 점에 유의하세요. 이를 통해 모델은 관련 작업으로 학습되고, 관련 보상 함수가 없는 작업은 무시됩니다.

예제 5: 비동상 보상 함수

커스텀 보상 함수는 async def 코루틴으로 정의할 수도 있습니다. 보상이 느린 I/O(예: 원격 서비스 호출)에 의존하는 경우 유용합니다. 여러 개의 async 보상 함수를 전달하면 RLOOTrainer는 이를 병렬로 실행하여 지연 시간이 겹치게 합니다.

다음은 I/O 바운드 작업을 시뮬레이션하는 비동기 보상 함수의 최소 예시입니다:

import asyncio

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

예제 6: 추가 열 및 지표 로깅

completions 테이블에 추가 열을 기록하고 스칼라 메트릭을 플롯으로 기록하는 보상 함수의 예시는 다음과 같습니다.

import re

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

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

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

    return rewards

트레이너에 보상 함수 전달

사용자 정의 보상 함수를 사용하려면 다음과 같이 RLOOTrainer에 전달하세요:

from trl import RLOOTrainer

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

여러 개의 보상 함수를 리스트로 전달할 수 있으며, 이 리스트에는 동기 함수와 비동기 함수가 모두 포함될 수 있습니다:

from trl import RLOOTrainer

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

보상은 각 함수의 보상 합으로 계산되며, 설정에 reward_weights가 제공된 경우 가중 합으로 계산됩니다.

RLOOTrainer는 서로 다른 유형의 여러 보상 함수를 지원합니다. 자세한 내용은 매개변수 문서를 참조하세요.

비전 언어 모델 학습

RLOO는 텍스트와 이미지를 모두 포함하는 멀티모달 데이터셋에서 Vision-Language Models(VLM)을 학습하는 것을 지원합니다.

지원되는 모델

테스트 환경:

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

모든 VLM과의 호환성이 보장되지는 않습니다. 특정 모델이 지원되어야 한다고 생각하시면 GitHub에서 이슈를 열어 주세요. 더 좋다면, 필요한 변경 사항이 포함된 pull request를 제출해 주세요.

빠른 시작

rloo_vlm.py를 사용하여 VLM을 파인튜닝합니다. lmms-lab/multimodal-open-r1-8k-verified로 학습하는 명령어 예시:

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

구성 팁

• 시각-언어 투영 레이어에 LoRA 사용
• 메모리 사용량을 줄이기 위해 4-bit 양자화를 활성화
• VLM은 메모리 집약적입니다 — 작은 배치 크기부터 시작하세요
• 대부분의 모델은 vLLM(server 및 colocate 모드)과 호환됩니다

데이터셋 형식

각 학습 샘플에는 다음을 포함해야 합니다:

• prompt: 프로세서의 채팅 템플릿을 통해 포맷된 텍스트
• image/images: PIL Image 또는 PIL Image 목록

트레이너는 모델의 이미지 프로세서를 통해 이미지에서 텐서로의 변환을 자동으로 처리합니다.

RLOOTrainer

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

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

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

RLOOConfig

참고

1. RLOO 논문
2. 기본으로 돌아가기: LLM에서 인간 피드백을 통한 학습을 위한 REINFORCE 스타일 최적화 재검토
3. 논문 - REINFORCE++: 대규모 언어 모델 정렬을 위한 단순하고 효율적인 접근법
4. 블로그 게시물 - RLHF에 RL을 다시 적용하기
5. 블로그 포스트 - RLHF와 그 변형 파헤치기: 진행 상황과 실무 엔지니어링 인사이트
6. Youtube - RLOO: LLM에서 인간 피드백을 학습하기 위한 비용 효율적 최적화

구버전 구현(0.21 이하)에서 마이그레이션 가이드

0.22.0 버전 출시와 함께 RLOOTrainer가 리팩토링되어, GRPOTrainer와 같은 라이브러리 내 다른 온라인 트레이너와 더 일관되도록 개선되었습니다. 이 새로운 구현은 트레이너의 구성 매개변수와 전체 구조에 여러 가지 변경 사항을 도입합니다. 아래는 RLOOConfig의 주요 변경 사항 요약입니다:

다음은 RLOOTrainer의 주요 변경 사항 요약입니다: