Trainer RLOO

Tổng quan

TRL hỗ trợ RLOO Trainer để huấn luyện các mô hình ngôn ngữ, như được mô tả trong bài báo Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs của Arash Ahmadian, Chris Cremer, Matthias Gallé, Marzieh Fadaee, Julia Kreutzer, Ahmet Üstün và Sara Hooker.

Tóm tắt từ bài báo như sau:

Việc căn chỉnh AI dưới hình thức Học tăng cường từ phản hồi của con người (Reinforcement Learning from Human Feedback - RLHF) ngày càng được coi là một thành phần thiết yếu cho các mô hình ngôn ngữ lớn có hiệu suất cao. Tối ưu hóa chính sách gần (Proximal Policy Optimization - PPO) đã được các tài liệu gần đây xác định là phương pháp chuẩn cho phần RL của RLHF. Tuy nhiên, phương pháp này liên quan đến cả chi phí tính toán cao và việc tinh chỉnh siêu tham số nhạy cảm. Chúng tôi cho rằng hầu hết các nguyên tắc động lực dẫn đến sự phát triển của PPO ít là mối lo ngại thực tế trong RLHF và ủng hộ một phương pháp ít tốn kém tính toán hơn, đồng thời duy trì và thậm chí tăng hiệu suất. Chúng tôi xem xét lại cách thức biểu diễn việc căn chỉnh từ sở thích của con người trong bối cảnh RL. Giữ sự đơn giản làm nguyên tắc chỉ đạo, chúng tôi chỉ ra rằng nhiều thành phần của PPO là không cần thiết trong bối cảnh RLHF và rằng các biến thể tối ưu hóa theo phong cách REINFORCE đơn giản hơn nhiều vượt trội hơn cả PPO và các phương pháp “không RL” mới được đề xuất như DPO và RAFT. Công việc của chúng tôi cho thấy việc thích ứng cẩn thận với các đặc điểm căn chỉnh của LLM cho phép tận dụng tối ưu hóa RL trực tuyến với chi phí thấp.

Phương pháp hậu huấn luyện này được đóng góp bởi Costa Huang và sau đó được tái cấu trúc bởi Shirin Yamani.

Bắt đầu nhanh

Ví dụ này minh họa cách huấn luyện một mô hình bằng phương pháp RLOO. Chúng ta huấn luyện mô hình Qwen 0.5B Instruct với các prompt từ dataset DeepMath-103K. Bạn có thể xem dữ liệu trong dataset tại đây:

Dưới đây là script để huấn luyện mô hình. accuracy_reward cần math-verify, bạn có thể cài đặt bằng 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()

Chạy script bằng lệnh sau:

torchrun --nproc_per_node 8 train_rloo.py

Tìm hiểu sâu hơn về phương pháp RLOO

RLOO là một thuật toán học trực tuyến, nghĩa là nó được cải thiện theo từng lần lặp bằng cách sử dụng dữ liệu do chính mô hình đã được huấn luyện tạo ra trong quá trình huấn luyện. Trực giác đằng sau mục tiêu của RLOO là tối đa hóa lợi thế của các phần hoàn thành được tạo ra, đồng thời đảm bảo rằng mô hình vẫn ở gần chính sách tham chiếu. Để hiểu cách RLOO hoạt động, nó có thể được chia thành bốn bước chính: tạo các phần hoàn thành, tính toán lợi thế, ước lượng phân kỳ KL và tính toán tổn thất.

Tạo các câu trả lời hoàn chỉnh

Ở mỗi bước huấn luyện, chúng ta lấy mẫu một lô prompt và tạo ra một tập hợp $G$ các completion cho mỗi prompt (ký hiệu là $o_i$).

Tính phần thưởng

Trong RLOO, phần thưởng bao gồm hai thành phần: phần thưởng do mô hình phần thưởng (hoặc hàm phần thưởng) cung cấp và một hình phạt KL nhằm ngăn chặn chính sách lệch quá xa so với một chính sách tham chiếu cố định

1. Với mỗi $G$ chuỗi được tạo ra $o_i = (o_{i,1}, \dots, o_{i,T})$ dựa trên truy vấn $q$, chúng ta tính một phần thưởng vô hướng bằng mô hình phần thưởng $R(o_i, q)$.
2. Đồng thời, chúng ta ước lượng độ phân kỳ KL giữa chính sách hiện tại $\pi_\theta$ và chính sách tham chiếu cố định $\pi_{\text{ref}}$ trên toàn bộ chuỗi. Ước lượng KL cho chuỗi $o_i$ là: $$\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})}.$$

Phần thưởng cuối cùng được gán cho chuỗi $o_i$ là: $$r_i = R(o_i, q) - \beta \, \mathbb{D}_{\mathrm{KL}}\!\left[\pi_\theta \|\pi_{\mathrm{ref}}\right],$$

trong đó $\beta > 0$ kiểm soát cường độ của hình phạt KL.

Trong cài đặt hoàn toàn trực tuyến (num_iterations = 1, mặc định), dữ liệu được tạo ra bởi chính sách hiện tại. Trong trường hợp này, hình phạt KL được tính trực tiếp bằng chính sách hiện tại.

Trong bối cảnh tổng quát hơn (ví dụ: nhiều bước gradient trên mỗi batch), dữ liệu được tạo ra bởi một snapshot trước đó $\pi_{\text{old}}$. Để giữ cho hình phạt nhất quán với phân phối lấy mẫu, KL được định nghĩa theo chính sách này:

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

Tương đương, với một chuỗi được lấy mẫu $o$, ước lượng Monte Carlo là:

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

Tính toán lợi thế

Sau khi phần thưởng cho mỗi kết quả hoàn thành đã được tính toán, chúng ta tính một đường cơ sở (baseline) là phần thưởng trung bình của tất cả các mẫu khác trong cùng một lô (batch), loại trừ mẫu hiện tại. Đường cơ sở này được sử dụng để giảm phương sai của ước lượng gradient chính sách. Lợi thế (advantage) cho mỗi kết quả hoàn thành sau đó được xác định bằng hiệu giữa phần thưởng của chính nó và đường cơ sở leave-one-out này.

Theo nghĩa chính thức, đối với một batch gồm G completion, baseline cho completion là:$$b_i = \frac{1}{G-1} \sum_{j \neq i} r_j$$

sau đó lợi thế cho mỗi completion được tính là hiệu giữa phần thưởng của nó và baseline: $$A_i = r_i - b_i$$

Tính toán tổn thất

Mất REINFORCE được định nghĩa đơn giản như sau: $$\mathcal{L}_{\text{RLOO}}(\theta) = - \frac{1}{G} \sum_{i=1}^G \hat{A}_i \, \log \pi_\theta(o_i \mid q)$$

Trong thực tế, thực hiện nhiều bước gradient trên cùng một batch khiến các hành động trở nên off-policy so với các tham số hiện tại. Để khắc phục điều này, chúng tôi đưa vào tỷ lệ importance sampling. Để ngăn chặn các cập nhật quá lớn khi chính sách thay đổi giữa bước lấy mẫu và bước gradient, chúng tôi cắt (clip) tỷ lệ này: $$\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)$$

Trong cài đặt trực tuyến hoàn toàn, một bước (mặc định), $\frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)} = 1$ và điều này rút gọn về REINFORCE tiêu chuẩn.

Khác với GRPO, RLOO không backpropagate gradient qua số hạng KL. Ở đây, KL chỉ được dùng để định hình phần thưởng: nó được tính dưới no_grad và được gộp vào phần thưởng vô hướng $r_i$, trở thành advantage đã được detach $\hat{A}_i$.

Nói cách khác, con đường duy nhất mà gradient đi vào policy là thông qua số hạng tỉ trọng quan trọng $\frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)}$, và số hạng KL không được backpropagate vào policy.

Ngược lại, GRPO thêm một số hạng KL rõ ràng, có thể tính đạo hàm vào hàm mục tiêu, được tính trực tiếp với chính sách hiện tại $\pi_\theta$, do đó gradient của nó chảy qua cả số hạng tỷ lệ lẫn hình phạt KL.

Các chỉ số được ghi log

Trong quá trình huấn luyện và đánh giá, chúng ta ghi lại các chỉ số sau:

• num_tokens: Tổng số token đã được xử lý cho đến nay, bao gồm cả prompt và completion.
• step_time: Thời gian trung bình (tính bằng giây) cho mỗi bước huấn luyện (bao gồm cả quá trình sinh).
• completions/mean_length: Độ dài trung bình của các câu trả lời được tạo ra.
• completions/min_length: Độ dài tối thiểu của các phần hoàn thành được tạo ra.
• completions/max_length: Độ dài tối đa của các phần hoàn thành được tạo ra.
• completions/mean_terminated_length: Độ dài trung bình của các phần hoàn thành được tạo ra kết thúc bằng EOS.
• completions/min_terminated_length: Độ dài tối thiểu của các phần hoàn thành được tạo ra kết thúc bằng EOS.
• completions/max_terminated_length: Độ dài tối đa của các phần hoàn thành được tạo ra kết thúc bằng EOS.
• completions/clipped_ratio: Tỷ lệ các câu trả lời bị cắt (clipped).
• rewards/{reward_func_name}/mean: Phần thưởng trung bình từ một hàm phần thưởng cụ thể.
• rewards/{reward_func_name}/std: Độ lệch chuẩn của phần thưởng từ một hàm phần thưởng cụ thể.
• reward: Phần thưởng trung bình tổng thể sau khi cộng phần thưởng giữa các hàm (được trọng số hóa bởi reward_weights).
• reward_std: Độ lệch chuẩn của phần thưởng tổng hợp trên các hàm (được trọng số hóa bởi reward_weights), được tính trên toàn bộ lô.
• frac_reward_zero_std: Tỷ lệ các mẫu trong lô sinh có độ lệch chuẩn phần thưởng bằng 0, cho thấy mức độ đa dạng thấp cho prompt đó (tất cả các câu trả lời đều đúng hoặc đều sai).
• entropy: Entropy trung bình của các dự đoán token trên các chuỗi hoàn thành được tạo ra. (Nếu mask_truncated_completions=True, các token của các chuỗi bị mask sẽ được loại trừ.)
• aux_loss: Tổn phụ trợ cân bằng tải của mô hình Mixture-of-Experts, trước khi được nhân với router_aux_loss_coef và cộng vào tổn. Chỉ được ghi log khi mô hình là mô hình MoE và router_aux_loss_coef khác không.
• kl: Trung bình của độ phân kỳ KL giữa mô hình và mô hình tham chiếu, được tính trên các phần hoàn thành đã tạo. Chỉ được ghi log nếu beta khác không.
• clip_ratio/region_mean: Tỷ lệ xác suất của các chuỗi mà mục tiêu RLOO bị cắt (clip) để giữ trong vùng tin cậy: $\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)}$. Giá trị cao hơn có nghĩa là nhiều mẫu bị cắt hơn, điều này hạn chế mức độ thay đổi của chính sách $\pi_\theta$ có thể thay đổi.
• clip_ratio/low_mean: Tỷ lệ trung bình của xác suất chuỗi bị cắt ở cận dưới của vùng tin cậy: $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/low_min: Chỉ báo clip nhỏ nhất cho mỗi completion ở cận dưới của vùng tin cậy: $r_{i,t}(\theta) < 1 - \epsilon_\mathrm{low}$.
• clip_ratio/high_mean: Tỷ lệ trung bình của xác suất chuỗi bị cắt ở cận trên của vùng tin cậy: $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.
• clip_ratio/high_max: Chỉ báo clip lớn nhất cho mỗi completion ở cận trên của vùng tin cậy: $r_{i,t}(\theta) > 1 + \epsilon_\mathrm{high}$.

Tùy chỉnh

Tăng tốc huấn luyện với sinh nội dung do vLLM hỗ trợ

Việc tạo sinh (generation) thường là nút thắt cổ chai chính khi huấn luyện bằng các phương pháp trực tuyến. Để tăng tốc quá trình tạo sinh, bạn có thể sử dụng vLLM, một bộ máy suy luận (inference engine) có thông lượng cao và độ trễ thấp dành cho LLM. Để bật tính năng này, trước tiên hãy cài đặt gói bằng

pip install trl[vllm]

Chúng tôi hỗ trợ hai cách sử dụng vLLM trong quá trình huấn luyện: chế độ máy chủ và chế độ colocate.

Tùy chọn 1: Chế độ Colocate

Trong chế độ này, vLLM chạy bên trong tiến trình trainer và chia sẻ bộ nhớ GPU với mô hình huấn luyện. Điều này giúp tránh việc khởi động một server riêng và có thể cải thiện mức sử dụng GPU, nhưng có thể dẫn đến tranh chấp bộ nhớ trên các GPU huấn luyện. Đây là chế độ mặc định.

from trl import RLOOConfig

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

Tùy chọn 2: Chế độ máy chủ

Trong chế độ này, vLLM chạy trong một tiến trình riêng (và sử dụng các GPU riêng) và giao tiếp với trainer qua HTTP. Đây là lựa chọn lý tưởng nếu bạn có các GPU dành riêng cho suy luận.

1. Khởi động máy chủ vLLM:

   Đã sao chép

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

2. Bật chế độ máy chủ trong script huấn luyện của bạn:

   Đã sao chép

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

Hãy đảm bảo rằng server đang sử dụng các GPU khác với trainer, nếu không bạn có thể gặp lỗi NCCL. Bạn có thể chỉ định các GPU cần sử dụng bằng biến môi trường CUDA_VISIBLE_DEVICES.

Tùy thuộc vào kích thước mô hình và yêu cầu bộ nhớ GPU tổng thể cho việc huấn luyện, bạn có thể cần điều chỉnh tham số vllm_gpu_memory_utilization trong RLOOConfig để tránh tình trạng khai thác không đủ hoặc lỗi hết bộ nhớ.

Chúng tôi cung cấp một HF Space để giúp ước lượng mức sử dụng bộ nhớ GPU được khuyến nghị dựa trên cấu hình mô hình và cài đặt thí nghiệm của bạn. Chỉ cần sử dụng nó như sau để nhận khuyến nghị vllm_gpu_memory_utilization:

Nếu giá trị khuyến nghị không hoạt động trong môi trường của bạn, chúng tôi đề xuất thêm một bộ đệm nhỏ (ví dụ: +0.05 hoặc +0.1) vào giá trị khuyến nghị để đảm bảo tính ổn định.

Nếu bạn vẫn gặp lỗi hết bộ nhớ, hãy đặt vllm_enable_sleep_mode thành True và các tham số cũng như bộ nhớ đệm của vllm sẽ được chuyển ra ngoài trong bước tối ưu hóa. Để biết thêm thông tin, xem Giảm mức sử dụng bộ nhớ với vLLM Sleep Mode.

Mặc định, RLOO sử dụng MASTER_ADDR=localhost và MASTER_PORT=12345 cho vLLM, nhưng bạn có thể ghi đè các giá trị này bằng cách đặt các biến môi trường tương ứng.

Để biết thêm thông tin, hãy xem Tăng tốc huấn luyện với vLLM.

RLOO ở quy mô lớn: huấn luyện Mô hình 70B+ trên nhiều nút

Khi huấn luyện các mô hình lớn như Qwen2.5-72B, bạn cần một số tối ưu hóa chính để đảm bảo quá trình huấn luyện hiệu quả và có khả năng mở rộng trên nhiều GPU và node. Những tối ưu hóa này bao gồm:

• DeepSpeed ZeRO Stage 3: ZeRO tận dụng song song hóa dữ liệu để phân phối trạng thái mô hình (trọng số, gradient, trạng thái optimizer) trên nhiều GPU và CPU, giảm thiểu yêu cầu về bộ nhớ và tính toán trên mỗi thiết bị. Vì các mô hình lớn không thể chứa vừa trên một GPU duy nhất, việc sử dụng ZeRO Stage 3 là bắt buộc để huấn luyện các mô hình như vậy. Để biết thêm chi tiết, xem DeepSpeed Integration.
• torchrun: trình khởi chạy của PyTorch sẽ khởi tạo một tiến trình cho mỗi GPU trên mọi nút và kết nối chúng thông qua một điểm hội tụ chung. Để biết thêm chi tiết, hãy xem Phân phối huấn luyện.
• vLLM: Xem phần trước để biết cách sử dụng vLLM để tăng tốc độ sinh.

Dưới đây là một ví dụ về script SLURM để huấn luyện một mô hình 70B với RLOO trên nhiều nút. Script này huấn luyện mô hình trên 4 nút và sử dụng nút thứ 5 cho việc tạo sinh được hỗ trợ bởi 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()

Sử dụng hàm phần thưởng tùy chỉnh

RLOOTrainer hỗ trợ sử dụng các hàm phần thưởng tùy chỉnh thay vì các mô hình phần thưởng dày đặc. Để đảm bảo tính tương thích, hàm phần thưởng của bạn phải đáp ứng các yêu cầu sau:

Các hàm reward có thể là các callable Python đồng bộ hoặc các coroutine async def bất đồng bộ. Khi bạn cung cấp nhiều hàm reward bất đồng bộ, chúng sẽ được chờ đồng thời (chạy song song thông qua asyncio.gather) để độ trễ của chúng chồng lên nhau.

1. Các đối số đầu vào:

   • Hàm phải nhận các đối số sau dưới dạng đối số từ khóa:

     • prompts (chứa các prompts),
     • completions (chứa các phần hoàn thành đã tạo),
     • completion_ids (chứa các completion đã được token hóa),
     • trainer_state (TrainerState): Trạng thái hiện tại của trainer. Có thể sử dụng để triển khai các hàm phần thưởng động, chẳng hạn như curriculum learning, trong đó phần thưởng được điều chỉnh dựa trên tiến độ huấn luyện.
     • log_extra: một callable log_extra(column: str, values: list) để thêm các cột bổ sung vào bảng completions. Xem Ví dụ 6.
     • log_metric: một hàm có thể gọi log_metric(name: str, value: float) để ghi log các chỉ số vô hướng dưới dạng biểu đồ bên cạnh kl, entropy, v.v. Xem Ví dụ 6.
     • Tất cả tên cột (ngoại trừ prompt) mà bộ dữ liệu có thể có. Ví dụ, nếu bộ dữ liệu chứa một cột có tên ground_truth, hàm sẽ được gọi với ground_truth như một đối số từ khóa.

     Cách dễ nhất để đáp ứng yêu cầu này là sử dụng **kwargs trong chữ ký hàm.

   • Tùy thuộc vào định dạng bộ dữ liệu, đầu vào sẽ khác nhau:

     • Đối với định dạng chuẩn, prompts và completions sẽ là các danh sách chuỗi.
     • Đối với định dạng hội thoại, prompts và completions sẽ là các danh sách từ điển tin nhắn.

2. Giá trị trả về: Hàm phải trả về một danh sách các số float. Mỗi số float đại diện cho phần thưởng tương ứng với một kết quả hoàn thành duy nhất.

Ví dụ 1: Phần thưởng cho các câu trả lời dài hơn

Dưới đây là một ví dụ về hàm phần thưởng cho định dạng chuẩn, phần thưởng cho các câu trả lời dài hơn:

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]

Bạn có thể kiểm tra nó như sau:

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

Ví dụ 1.1: Thưởng cho các phần hoàn thành dài hơn (dựa trên số ký tự)

Tương tự như ví dụ trước, nhưng lần này hàm phần thưởng được dựa trên số lượng ký tự thay vì 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]

Bạn có thể kiểm tra nó như sau:

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

Ví dụ 2: Thưởng cho các phần hoàn thành có định dạng cụ thể

Dưới đây là một ví dụ về hàm phần thưởng kiểm tra xem phần hoàn thành có đúng định dạng cụ thể hay không. Ví dụ này được lấy cảm hứng từ hàm phần thưởng định dạng được sử dụng trong bài báo DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. Nó được thiết kế cho định dạng hội thoại, trong đó các prompt và phần hoàn thành bao gồm các tin nhắn có cấu trúc.

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]

Bạn có thể kiểm tra hàm này như sau:

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

Ví dụ 3: Hoàn thành phần thưởng dựa trên một tham chiếu

Dưới đây là một ví dụ về hàm phần thưởng kiểm tra xem câu trả liệu có chính xác hay không. Ví dụ này được lấy cảm hứng từ hàm phần thưởng độ chính xác được sử dụng trong bài báo DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. Ví dụ này được thiết kế cho định dạng chuẩn, trong đó tập dữ liệu chứa một cột có tên là 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)]

Bạn có thể kiểm tra hàm này như sau:

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

Ví dụ 4: Các hàm phần thưởng đa nhiệm

Dưới đây là một ví dụ về việc sử dụng nhiều hàm phần thưởng trong RLOOTrainer. Trong ví dụ này, chúng ta định nghĩa hai hàm phần thưởng dành riêng cho từng tác vụ: math_reward_func và coding_reward_func. Hàm math_reward_func cấp phần thưởng cho các bài toán dựa trên tính chính xác của chúng, trong khi coding_reward_func cấp phần thưởng cho các bài toán lập trình dựa trên việc giải pháp có chạy được hay không.

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

Trong ví dụ này, math_reward_func và coding_reward_func được thiết kế để hoạt động với một tập dữ liệu hỗn hợp chứa cả các bài toán toán học và lập trình. Cột task trong tập dữ liệu được dùng để xác định hàm thưởng nào sẽ được áp dụng cho từng bài toán. Nếu không có hàm thưởng phù hợp cho một mẫu trong tập dữ liệu, hàm thưởng sẽ trả về None, và RLOOTrainer sẽ tiếp tục với các hàm và tác vụ hợp lệ. Điều này cho phép RLOOTrainer xử lý nhiều hàm thưởng với phạm vi áp dụng khác nhau.

Lưu ý rằng RLOOTrainer sẽ bỏ qua các phần thưởng None được trả về bởi các hàm phần thưởng và chỉ xem xét các phần thưởng được trả về bởi các hàm liên quan. Điều này đảm bảo rằng mô hình được huấn luyện trên các tác vụ liên quan và bỏ qua các tác vụ không có hàm phần thưởng liên quan.

Ví dụ 5: Các hàm phần thưởng bất đồng bộ

Các hàm phần thưởng tùy chỉnh cũng có thể được định nghĩa dưới dạng coroutine async def. Điều này hữu ích nếu phần thưởng của bạn phụ thuộc vào I/O chậm (ví dụ: gọi một dịch vụ từ xa). Khi bạn truyền nhiều hàm phần thưởng async, RLOOTrainer sẽ thực thi chúng đồng thời để độ trễ của chúng chồng lên nhau.

Dưới đây là một ví dụ tối giản về một hàm thưởng bất đồng bộ mô phỏng một thao tác bị giới hạn bởi 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]

Ví dụ 6: Ghi log các cột và chỉ số bổ sung

Dưới đây là một ví dụ về hàm phần thưởng ghi log các cột bổ sung vào bảng completions và các chỉ số vô hướng dưới dạng biểu đồ.

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

Truyền hàm phần thưởng cho trình huấn luyện

Để sử dụng hàm phần thưởng tùy chỉnh của bạn, hãy truyền nó vào RLOOTrainer như sau:

from trl import RLOOTrainer

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

Bạn có thể truyền nhiều hàm phần thưởng dưới dạng một danh sách; danh sách này có thể bao gồm cả các hàm đồng bộ và bất đồng bộ:

from trl import RLOOTrainer

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

và phần thưởng sẽ được tính là tổng các phần thưởng từ mỗi hàm, hoặc tổng có trọng số nếu reward_weights được cung cấp trong cấu hình.

Lưu ý rằng RLOOTrainer hỗ trợ nhiều hàm phần thưởng có các kiểu khác nhau. Xem tài liệu tham số để biết thêm chi tiết.

Huấn luyện Mô hình Ngôn ngữ Thị giác

RLOO hỗ trợ huấn luyện các Mô hình Ngôn ngữ - Thị giác (VLM) trên các bộ dữ liệu đa phương thức chứa cả văn bản và hình ảnh.

Các Mô hình Được hỗ trợ

Đã kiểm tra với:

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

Không đảm bảo tương thích với tất cả các VLM. Nếu bạn cho rằng một model nên được hỗ trợ, hãy mở một issue trên GitHub — hoặc tốt hơn nữa, gửi một pull request kèm theo các thay đổi cần thiết.

Bắt đầu nhanh

Sử dụng rloo_vlm.py để tinh chỉnh một VLM. Lệnh ví dụ để huấn luyện trên 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

Mẹo cấu hình

• Sử dụng LoRA trên các lớp chiếu ngôn ngữ-thị giác
• Bật lượng tử hóa 4-bit để giảm mức sử dụng bộ nhớ
• VLMs đòi hỏi nhiều bộ nhớ — hãy bắt đầu với kích thước batch nhỏ hơn
• Hầu hết các model đều tương thích với vLLM (chế độ server và colocate)

Định dạng Dữ liệu

Mỗi mẫu huấn luyện nên bao gồm:

• prompt: Văn bản được định dạng qua chat template của processor
• image/images: Ảnh PIL hoặc danh sách các ảnh PIL

Trainer tự động xử lý việc chuyển đổi ảnh sang tensor thông qua bộ xử lý ảnh của 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

Tham khảo

1. Bài báo RLOO
2. Bài viết Quay về cơ bản: Xem xét lại tối ưu hóa kiểu REINFORCE để học từ phản hồi của con người trong LLM
3. Bài báo - REINFORCE++: Một cách tiếp cận đơn giản và hiệu quả để căn chỉnh các mô hình ngôn ngữ lớn
4. Bài viết trên blog - Đưa RL trở lại RLHF
5. Bài viết trên blog - Giải mã RLHF và các biến thể của nó: Tiến bộ và những hiểu biết thực tiễn về kỹ thuật
6. Youtube - RLOO: Tối ưu hóa tiết kiệm chi phí cho việc học từ phản hồi của con người trong LLM

Hướng dẫn chuyển đổi từ bản triển khai cũ (0.21 trở xuống)

Với việc phát hành phiên bản 0.22.0, chúng tôi đã làm mới RLOOTrainer để phù hợp hơn với các trainer online khác trong thư viện, chẳng hạn như GRPOTrainer. Bản triển khai mới này giới thiệu một số thay đổi đối với các tham số cấu hình và cấu trúc tổng thể của trainer. Dưới đây là tóm tắt các thay đổi chính đối với RLOOConfig:

Dưới đây là tóm tắt các thay đổi chính đối với RLOOTrainer: