Pelatih RLOO

Ikhtisar

TRL mendukung RLOO Trainer untuk melatih model bahasa, seperti yang dijelaskan dalam makalah Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs oleh Arash Ahmadian, Chris Cremer, Matthias Gallé, Marzieh Fadaee, Julia Kreutzer, Ahmet Üstün dan Sara Hooker.

Abstrak dari makalah tersebut adalah sebagai berikut:

AI alignment dalam bentuk Reinforcement Learning from Human Feedback (RLHF) semakin dianggap sebagai komponen penting untuk model bahasa besar berkinerja tinggi. Proximal Policy Optimization (PPO) telah diposisikan oleh literatur terbaru sebagai metode kanonik untuk bagian RL dari RLHF. Namun, metode ini melibatkan biaya komputasi yang tinggi dan penyetelan hiperparameter yang sensitif. Kami berpendapat bahwa sebagian besar prinsip motivasi yang mendorong pengembangan PPO kurang menjadi kekhawatiran praktis dalam RLHF, dan kami mengadvokasi metode yang kurang mahal secara komputasi namun tetap mempertahankan bahkan meningkatkan kinerja. Kami meninjau kembali formulasi alignment dari preferensi manusia dalam konteks RL. Dengan menjaga kesederhanaan sebagai prinsip panduan, kami menunjukkan bahwa banyak komponen PPO tidak diperlukan dalam konteks RLHF dan bahwa varian optimasi bergaya REINFORCE yang jauh lebih sederhana melampaui kinerja PPO maupun metode “RL-free” yang baru diusulkan seperti DPO dan RAFT. Pekerjaan kami menunjukkan bahwa adaptasi yang cermat terhadap karakteristik alignment LLM memungkinkan pemanfaatan optimasi RL online dengan biaya rendah.

Metode pasca-pelatihan ini disumbangkan oleh Costa Huang dan kemudian di-refactor oleh Shirin Yamani.

Mulai cepat

Contoh ini menunjukkan cara melatih model menggunakan metode RLOO. Kami melatih Qwen 0.5B Instruct model dengan prompt dari DeepMath-103K dataset. Anda dapat melihat data dalam dataset di sini:

Berikut adalah skrip untuk melatih model. accuracy_reward memerlukan math-verify, yang dapat Anda instal dengan 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()

Jalankan skrip menggunakan perintah berikut:

torchrun --nproc_per_node 8 train_rloo.py

Menggali lebih dalam ke metode RLOO

RLOO adalah algoritma pembelajaran online, yang berarti ia meningkatkan secara iteratif dengan menggunakan data yang dihasilkan oleh model terlatih itu sendiri selama pelatihan. Intuisi di balik tujuan RLOO adalah memaksimalkan keuntungan dari penyelesaian yang dihasilkan, sambil memastikan bahwa model tetap dekat dengan kebijakan referensi. Untuk memahami cara kerja RLOO, prosesnya dapat dipecah menjadi empat langkah utama: menghasilkan penyelesaian, menghitung keuntungan, mengestimasi divergensi KL, dan menghitung loss.

Membuat kelengkapan

Pada setiap langkah pelatihan, kami mengambil sampel batch prompt dan menghasilkan satu set $G$ komplemen untuk setiap prompt (dilambangkan sebagai $o_i$).

Menghitung reward

Dalam RLOO, reward terdiri dari dua komponen: reward yang diberikan oleh model reward (atau fungsi reward) dan penalti KL yang mencegah kebijakan menyimpang terlalu jauh dari kebijakan referensi yang tetap

1. Untuk setiap $G$ sekuens yang dihasilkan $o_i = (o_{i,1}, \dots, o_{i,T})$ yang dikondisikan pada query $q$, kami menghitung reward skalar menggunakan model reward $R(o_i, q)$.
2. Secara bersamaan, kami memperkirakan divergensi KL antara kebijakan saat ini $\pi_\theta$ dan kebijakan referensi tetap $\pi_{\text{ref}}$ pada sekuens. Estimasi KL untuk sekuens $o_i$ adalah: $$\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})}.$$

Hadiah akhir yang diberikan ke urutan $o_i$ adalah sebagai berikut: $$r_i = R(o_i, q) - \beta \, \mathbb{D}_{\mathrm{KL}}\!\left[\pi_\theta \|\pi_{\mathrm{ref}}\right],$$

di mana $\beta > 0$ mengontrol kekuatan penalti KL.

Dalam pengaturan murni online (num_iterations = 1, default), data dihasilkan oleh kebijakan saat ini. Dalam kasus ini, penalti KL dihitung secara langsung menggunakan kebijakan saat ini.

Dalam pengaturan yang lebih umum (misalnya, beberapa langkah gradien per batch), data justru dihasilkan oleh snapshot sebelumnya $\pi_{\text{old}}$. Untuk menjaga agar penalti konsisten dengan distribusi sampling, KL didefinisikan terhadap kebijakan ini:

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

Secara ekuivalen, untuk urutan yang disampel $o$, estimasi Monte Carlo adalah:

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

Menghitung keuntungan

Setelah reward untuk setiap completion dihitung, kita menghitung baseline sebagai rata-rata reward dari semua sampel lain dalam batch yang sama, dengan mengecualikan sampel saat ini. Baseline ini digunakan untuk mengurangi varians dari estimasi policy gradient. Advantage untuk setiap completion kemudian diperoleh sebagai selisih antara reward-nya sendiri dan baseline leave-one-out ini.

Secara formal, untuk batch berisi G penyelesaian, baseline untuk penyelesaian adalah:$$b_i = \frac{1}{G-1} \sum_{j \neq i} r_j$$

dan kemudian keuntungan untuk setiap penyelesaian dihitung sebagai selisih antara reward-nya dan baseline: $$A_i = r_i - b_i$$

Menghitung loss

Loss REINFORCE didefinisikan secara sederhana sebagai: $$\mathcal{L}_{\text{RLOO}}(\theta) = - \frac{1}{G} \sum_{i=1}^G \hat{A}_i \, \log \pi_\theta(o_i \mid q)$$

Dalam praktiknya, melakukan beberapa langkah gradien pada batch yang sama membuat aksi secara efektif menjadi off-policy relatif terhadap parameter saat ini. Untuk mengoreksi hal ini, kami memperkenalkan rasio importance sampling. Untuk mencegah pembaruan yang terlalu besar ketika kebijakan berubah antara pengambilan sampel dan langkah gradien, kami memotong (clip) rasio ini: $$\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)$$

Dalam pengaturan online penuh, satu langkah (default), $\frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)} = 1$ dan ini berkurang menjadi REINFORCE standar.

Berbeda dengan GRPO, RLOO tidak melakukan backpropagation gradien melalui istilah KL. Di sini, KL murni digunakan untuk reward shaping: ia dihitung di bawah no_grad dan digabungkan ke dalam reward skalar $r_i$, yang menjadi advantage yang terdetas $\hat{A}_i$.

Dengan kata lain, satu-satunya jalur yang diambil gradien ke dalam kebijakan adalah suku rasio kepentingan $\frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)}$, dan suku KL tidak melakukan backpropagation ke dalam kebijakan.

Sebaliknya, GRPO menambahkan suku KL eksplisit yang dapat dibedakan ke dalam objektifnya, yang dihitung secara langsung dengan kebijakan saat ini $\pi_\theta$, sehingga gradiennya mengalir melalui suku rasio dan penalti KL.

Metrik yang dicatat

Selama pelatihan dan evaluasi, kami mencatat metrik berikut:

• num_tokens: Total jumlah token yang diproses sejauh ini, termasuk prompt dan penyelesaian.
• step_time: Rata-rata waktu (dalam detik) yang dibutuhkan per langkah pelatihan (termasuk generasi).
• completions/mean_length: Rata-rata panjang dari kelengkapan yang dihasilkan.
• completions/min_length: Panjang minimum dari kelengkapan yang dihasilkan.
• completions/max_length: Panjang maksimum dari kelengkapan yang dihasilkan.
• completions/mean_terminated_length: Rata-rata panjang dari kelengkapan yang dihasilkan yang diakhiri dengan EOS.
• completions/min_terminated_length: Panjang minimum dari completion yang dihasilkan yang diakhiri dengan EOS.
• completions/max_terminated_length: Panjang maksimum dari completion yang dihasilkan yang diakhiri dengan EOS.
• completions/clipped_ratio: Rasio penyelesaian yang dipotong (clipped).
• rewards/{reward_func_name}/mean: Rata-rata reward dari fungsi reward tertentu.
• rewards/{reward_func_name}/std: Deviasi standar dari reward yang dihasilkan oleh fungsi reward tertentu.
• reward: Rata-rata keseluruhan reward setelah menjumlahkan reward di seluruh fungsi (diberi bobot oleh reward_weights).
• reward_std: Standar deviasi dari total reward di seluruh fungsi (diberat oleh reward_weights), dihitung pada batch penuh.
• frac_reward_zero_std: Proporsi sampel dalam batch generasi yang memiliki standar deviasi reward nol, yang mengindikasikan bahwa terdapat sedikit keberagaman untuk prompt tersebut (semua jawaban benar atau salah).
• entropy: Entropi rata-rata dari prediksi token di seluruh kelengkapan yang dihasilkan. (Jika mask_truncated_completions=True, token dari urutan yang dimasker dikecualikan.)
• aux_loss: kerugian bantu penyeimbangan beban dari model Mixture-of-Experts, sebelum dikalikan dengan router_aux_loss_coef dan ditambahkan ke kerugian. Dicatat hanya ketika model adalah model MoE dan router_aux_loss_coef tidak nol.
• kl: Rata-rata divergensi KL antara model dan model referensi, dihitung di atas kelengkapan yang dihasilkan. Dicatat hanya jika beta tidak nol.
• clip_ratio/region_mean: Rasio probabilitas sekuens di mana tujuan RLOO dipangkas agar tetap berada di dalam 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)}$. Nilai yang lebih tinggi berarti lebih banyak sampel yang dipangkas, yang membatasi seberapa besar kebijakan $\pi_\theta$ dapat berubah.
• clip_ratio/low_mean: Rata-rata rasio probabilitas urutan yang dipangkas pada batas bawah trust region: $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/low_min: Indikator clip per-penyelesaian terkecil pada batas bawah trust region: $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/high_mean: Rata-rata rasio probabilitas urutan yang dipangkas pada batas atas trust region: $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.
• clip_ratio/high_max: Indikator clip per-completion terbesar pada batas atas trust region: $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.

Kustomisasi

Percepat pelatihan dengan generasi yang didukung vLLM

Generasi sering menjadi bottleneck utama saat melatih dengan metode online. Untuk mempercepat generasi, Anda dapat menggunakan vLLM, mesin inferensi ber-throughput tinggi dan ber-latensi rendah untuk LLM. Untuk mengaktifkannya, pertama instal paketnya dengan

pip install trl[vllm]

Kami mendukung dua cara penggunaan vLLM selama pelatihan: mode server dan mode colocate.

Opsi 1: Mode Colocate

Pada mode ini, vLLM berjalan di dalam proses trainer dan berbagi memori GPU dengan model pelatihan. Hal ini menghindari peluncuran server terpisah dan dapat meningkatkan pemanfaatan GPU, tetapi dapat menyebabkan perebutan memori pada GPU pelatihan. Ini adalah mode default.

from trl import RLOOConfig

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

Opsi 2: Mode server

Pada mode ini, vLLM berjalan dalam proses terpisah (dan menggunakan GPU terpisah) serta berkomunikasi dengan trainer melalui HTTP. Ini ideal jika Anda memiliki GPU khusus untuk inferensi.

1. Mulai server vLLM:

   Disalin

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

2. Aktifkan mode server di skrip pelatihan Anda:

   Disalin

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

Pastikan server menggunakan GPU yang berbeda dari trainer, jika tidak Anda mungkin akan menemui error NCCL. Anda dapat menentukan GPU yang akan digunakan dengan variabel lingkungan CUDA_VISIBLE_DEVICES.

Tergantung pada ukuran model dan kebutuhan memori GPU keseluruhan untuk pelatihan, Anda mungkin perlu menyesuaikan parameter vllm_gpu_memory_utilization di RLOOConfig untuk menghindari underutilization atau error out-of-memory.

Kami menyediakan HF Space untuk membantu memperkirakan pemanfaatan memori GPU yang direkomendasikan berdasarkan konfigurasi model dan pengaturan eksperimen Anda. Cukup gunakan sebagai berikut untuk mendapatkan rekomendasi vllm_gpu_memory_utilization:

Jika nilai yang direkomendasikan tidak berfungsi di lingkungan Anda, kami menyarankan untuk menambahkan buffer kecil (misalnya, +0.05 atau +0.1) ke nilai yang direkomendasikan untuk memastikan stabilitas.

Jika Anda masih mengalami error out-of-memory, atur vllm_enable_sleep_mode ke True sehingga parameter vllm dan cache akan di-offload selama langkah optimasi. Untuk informasi lebih lanjut, lihat Mengurangi Penggunaan Memori dengan vLLM Sleep Mode.

Secara default, RLOO menggunakan MASTER_ADDR=localhost dan MASTER_PORT=12345 untuk vLLM, tetapi Anda dapat menimpa nilai-nilai ini dengan mengatur variabel lingkungan secara sesuai.

Untuk informasi lebih lanjut, lihat Mempercepat pelatihan dengan vLLM.

RLOO berskala besar: latih Model 70B+ di beberapa node

Saat melatih model besar seperti Qwen2.5-72B, Anda memerlukan beberapa optimasi kunci untuk membuat pelatihan efisien dan dapat diskalakan di beberapa GPU dan node. Optimasi ini meliputi:

• DeepSpeed ZeRO Stage 3: ZeRO memanfaatkan data parallelism untuk mendistribusikan state model (weights, gradients, optimizer states) ke beberapa GPU dan CPU, sehingga mengurangi kebutuhan memori dan komputasi pada setiap perangkat. Karena model besar tidak muat di satu GPU, penggunaan ZeRO Stage 3 diperlukan untuk melatih model semacam itu. Untuk detail lebih lanjut, lihat DeepSpeed Integration.
• torchrun: peluncur PyTorch memulai satu proses per GPU di setiap node dan menghubungkannya melalui rendezvous bersama. Untuk detail lebih lanjut, lihat Mendistribusikan Pelatihan.
• vLLM: Lihat bagian sebelumnya tentang cara menggunakan vLLM untuk mempercepat generasi.

Di bawah ini adalah contoh skrip SLURM untuk melatih model 70B dengan RLOO di beberapa node. Skrip ini melatih model di 4 node dan menggunakan node ke-5 untuk generasi yang didukung 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()

Menggunakan fungsi reward kustom

RLOOTrainer mendukung penggunaan fungsi reward kustom sebagai pengganti model reward padat. Untuk memastikan kompatibilitas, fungsi reward Anda harus memenuhi persyaratan berikut:

Fungsi reward dapat berupa callable Python sinkron atau coroutine async def asinkron. Ketika Anda menyediakan beberapa fungsi reward asinkron, fungsi-fungsi tersebut ditunggu secara bersamaan (dijalankan paralel melalui asyncio.gather) sehingga latensinya tumpang tindih.

1. Argumen input:

   • Fungsi harus menerima argumen berikut sebagai argumen kata kunci:

     • prompts (mengandung prompt),
     • completions (mengandung kelengkapan yang dihasilkan),
     • completion_ids (mengandung pelengkapan yang telah ditokenisasi),
     • trainer_state (TrainerState): Status terkini dari trainer. Hal ini dapat digunakan untuk mengimplementasikan fungsi reward dinamis, seperti pembelajaran kurikulum, di mana reward disesuaikan berdasarkan kemajuan pelatihan.
     • log_extra: callable log_extra(column: str, values: list) untuk menambahkan kolom tambahan ke tabel completions. Lihat Contoh 6.
     • log_metric: callable log_metric(name: str, value: float) untuk mencatat metrik skalar sebagai plot di samping kl, entropy, dll. Lihat Contoh 6.
     • Semua nama kolom (kecuali prompt) yang mungkin dimiliki dataset. Misalnya, jika dataset berisi kolom bernama ground_truth, fungsi tersebut akan dipanggil dengan ground_truth sebagai argumen kata kunci.

     Cara termudah untuk memenuhi persyaratan ini adalah dengan menggunakan **kwargs dalam tanda tangan fungsi.

   • Tergantung pada format dataset, input akan bervariasi:

     • Untuk format standar, prompts dan completions akan berupa daftar string.
     • Untuk format percakapan, prompts dan completions akan berupa daftar kamus pesan.

2. Nilai yang dikembalikan: Fungsi harus mengembalikan daftar float. Setiap float mewakili reward yang sesuai dengan satu completion.

Contoh 1: Memberi imbalan pada penyelesaian yang lebih panjang

Berikut adalah contoh fungsi reward untuk format standar yang memberikan reward lebih tinggi untuk penyelesaian yang lebih panjang:

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]

Anda dapat mengujinya sebagai berikut:

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

Contoh 1.1: Memberi imbalan untuk penyelesaian yang lebih panjang (berdasarkan jumlah karakter)

Sama seperti contoh sebelumnya, tetapi kali ini fungsi reward didasarkan pada jumlah karakter, bukan token.

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]

Anda dapat mengujinya sebagai berikut:

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

Contoh 2: Memberikan imbalan untuk penyelesaian dengan format tertentu

Berikut adalah contoh fungsi reward yang memeriksa apakah penyelesaian memiliki format tertentu. Contoh ini terinspirasi dari fungsi format reward yang digunakan dalam makalah DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. Fungsi ini dirancang untuk format percakapan, di mana prompt dan penyelesaian terdiri dari pesan terstruktur.

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]

Anda dapat menguji fungsi ini sebagai berikut:

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

Contoh 3: Memberikan imbalan pada penyelesaian berdasarkan referensi

Berikut adalah contoh fungsi reward yang memeriksa apakah penyelesaian benar. Contoh ini terinspirasi dari fungsi reward akurasi yang digunakan dalam makalah DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. Contoh ini dirancang untuk format standar, di mana dataset berisi kolom bernama 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)]

Anda dapat menguji fungsi ini sebagai berikut:

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

Contoh 4: Fungsi reward multi-tugas

Berikut adalah contoh penggunaan beberapa fungsi reward dalam RLOOTrainer. Dalam contoh ini, kami mendefinisikan dua fungsi reward yang spesifik untuk tugas: math_reward_func dan coding_reward_func. Fungsi math_reward_func memberikan reward untuk soal matematika berdasarkan kebenarannya, sedangkan coding_reward_func memberikan reward untuk soal pemrograman berdasarkan apakah solusinya berfungsi.

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

Pada contoh ini, math_reward_func dan coding_reward_func dirancang untuk bekerja dengan dataset campuran yang berisi soal matematika dan pemrograman. Kolom task dalam dataset digunakan untuk menentukan fungsi reward mana yang akan diterapkan pada setiap soal. Jika tidak ada fungsi reward yang relevan untuk sampel dalam dataset, fungsi reward akan mengembalikan None, dan RLOOTrainer akan melanjutkan dengan fungsi dan tugas yang valid. Hal ini memungkinkan RLOOTrainer untuk menangani beberapa fungsi reward dengan cakupan penerapan yang berbeda.

Perhatikan bahwa RLOOTrainer akan mengabaikan reward None yang dikembalikan oleh fungsi reward dan hanya mempertimbangkan reward yang dikembalikan oleh fungsi yang relevan. Hal ini memastikan bahwa model dilatih pada tugas yang relevan dan mengabaikan tugas yang tidak memiliki fungsi reward yang relevan.

Contoh 5: Fungsi reward asinkron

Fungsi reward kustom juga dapat didefinisikan sebagai korutin async def. Hal ini berguna jika reward Anda bergantung pada I/O yang lambat (misalnya, memanggil layanan jarak jauh). Ketika Anda meneruskan beberapa fungsi reward asinkron, RLOOTrainer akan mengeksekusinya secara bersamaan sehingga latensinya saling tumpang tindih.

Berikut adalah contoh minimal dari fungsi reward asinkron yang mensimulasikan operasi yang terikat 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]

Contoh 6: Pencatatan kolom dan metrik tambahan

Di bawah ini adalah contoh fungsi reward yang mencatat kolom tambahan ke tabel completions dan metrik skalar sebagai plot.

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

Menyalurkan fungsi reward ke trainer

Untuk menggunakan fungsi reward kustom Anda, berikan ke RLOOTrainer sebagai berikut:

from trl import RLOOTrainer

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

Anda dapat meneruskan beberapa fungsi reward sebagai daftar; daftar ini dapat mencakup fungsi sinkron dan asinkron:

from trl import RLOOTrainer

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

dan hadiah akan dihitung sebagai jumlah dari hadiah dari setiap fungsi, atau jumlah berbobot jika reward_weights disediakan dalam konfigurasi.

Perhatikan bahwa RLOOTrainer mendukung beberapa fungsi reward dengan tipe yang berbeda. Lihat dokumentasi parameter untuk detail lebih lanjut.

Pelatihan Model Bahasa Visual

RLOO mendukung pelatihan Vision-Language Models (VLMs) pada dataset multimodal yang berisi teks dan gambar.

Model yang Didukung

Diuji dengan:

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

Kompatibilitas dengan semua VLM tidak dijamin. Jika Anda yakin suatu model harus didukung, silakan buka issue di GitHub — atau lebih baik lagi, kirimkan pull request dengan perubahan yang diperlukan.

Mulai Cepat

Gunakan rloo_vlm.py untuk fine-tune VLM. Contoh perintah untuk pelatihan pada 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

Tips Konfigurasi

• Gunakan LoRA pada lapisan proyeksi bahasa-visi
• Aktifkan kuantisasi 4-bit untuk mengurangi penggunaan memori
• VLMs sangat intensif terhadap memori — mulailah dengan ukuran batch yang lebih kecil
• Sebagian besar model kompatibel dengan vLLM (mode server dan colocate)

Format Dataset

Setiap sampel pelatihan harus mencakup:

• prompt: Teks yang diformat melalui template chat prosesor
• image/images: Gambar PIL atau daftar Gambar PIL

Trainer secara otomatis menangani konversi gambar-ke-tensor melalui image processor model.

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

Referensi

1. Makalah RLOO
2. Paper Kembali ke Dasar: Mengunjungi Ulang Optimasi Gaya REINFORCE untuk Pembelajaran dari Umpan Balik Manusia pada LLM
3. Makalah - REINFORCE++: Pendekatan Sederhana dan Efisien untuk Menyelaraskan Model Bahasa Besar
4. Artikel Blog - Mengembalikan RL ke RLHF
5. Artikel Blog - Mengurai RLHF dan Variannya: Kemajuan dan Wawasan Teknik Praktis
6. Youtube - RLOO: Optimasi yang Efisien Biaya untuk Pembelajaran dari Umpan Balik Manusia pada LLM

Panduan Migrasi dari implementasi lama (0.21 dan sebelumnya)

Dengan rilis versi 0.22.0, kami telah merombak RLOOTrainer agar lebih selaras dengan trainer online lainnya di pustaka ini, seperti GRPOTrainer. Implementasi baru ini memperkenalkan beberapa perubahan pada parameter konfigurasi dan struktur keseluruhan trainer. Berikut adalah ringkasan perubahan utama untuk RLOOConfig:

Berikut adalah ringkasan perubahan utama untuk RLOOTrainer: