A concise library of reinforcement learning for large language models.
This is the right library for you if you are tired with complicated abstractions.
We deliver a clear implementation within 1K lines.
You can simply launch the training with torchrun as you do in supervised fine-tuning.
Despite the simplicity, you should be able to scale up to moderate-sized, e.g., 32B, language models with
- Model partition via Fully Sharded Data Parallelism (and Tensor Parallelism upcoming!)
- Efficient sequence parallelism via ZigZag Ring Attention
- Inference engine and KV cache partition via Tensor Parallelism
We also support
- Balanced sequence packing for higher throughput
- Multi-turn rollout with SGLang async inference engine
RL2 is a production-ready library! Check our wandb report on LIMO, SkyworkRM, UltraFeedback, OpenReasonerZero, and SearchR1.
git clone https://github.com/ChenmienTan/RL2.git
cd RL2
pip install -e .
Hugging Face dataset and various file types, including JSON, JSONL, CSV, Parquet, and Arrow, are accepted. The data for PPO should be in the following format
[
{
"messages": [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "What is the capital of China?"}
],
"answer": "Beijing"
}
]
For SFT
[
{
"messages": [
{"role": "user", "content": "What is the capital of China?"},
{"role": "assistant", "content": "Beijing."}
]
}
]
For reward modeling and DPO
[
{
"messages": [
{"role": "user", "content": "What is the capital of China?"}
],
"chosen": "Beijing.",
"rejected": "Shanghai."
}
]
The reward function should be in the follwing format.
Specify the path to the Python script including the function via actor.rollout.env_path.
def reward_fn(messages, answer):
pred = parse_answer(messages[-1]["content"])
return float(is_equivalent(pred, answer))
If a reward model is used, it should be served outside of the training framework, e.g., using vLLM or SGLang, and be accessed in the reward function.
RL2 supports multi-turn rollout with function calling.
In this case, you should set rollout.max_turns > 1 and include function interact with the following format in the Python script including the reward function.
You should parse the called functions in past messages and return new messages including the results.
An empty list indicates no function is called.
def interact(messages):
queries = parse_query(messages[-1]["content])
results = [search(query) for query in queries]
return [
{"role": "tool", "content": result}
for result in results
]
Use torchrun to launch the training. For example, for single node
torchrun \
--nproc_per_node=<number of GPUs> \
-m RL2.trainer.ppo \
<args>
For multi nodes
torchrun \
--nnodes=<number of nodes> \
--node_rank=<rank of node> \
--nproc_per_node=<number of GPUs on a node> \
--master_addr=<address of master node> \
--master_port=<port of master node> \
-m RL2.trainer.ppo \
<args>
path: Hugging Face name or local path of dataset.max_length: The maximum length of a sequence.batch_size:batch_sizesamples will be used for an update.prompts_per_rollout:prompts_per_rolloutprompts will be used per rollout.responses_per_prompt:responses_per_prompttrajectories will be sampled for a prompt in rollout.
model_name: Hugging Face name or local path of model.optimizer_dir: The directory of optimizer state to be loaded.fsdp_size: The model parameters will be sharded across everyfsdp_sizeGPUs. Must be divisible by the total number of GPUs. Default to0, where a single copy of parameters is sharded across all GPUs.sp_size: The sequence will be sharded acrosssp_sizeGPUs. Must be divisible by the total number of GPUs.gradient_checkpointing: Whether to enable gradient checkpointing.max_length_per_device: The maximum length allowed for a single GPU at training. The length of any sequence cannot exceedsp_size * max_length_per_device.max_inference_length_per_device: The maximum length allowed for a single GPU at inference.update_per_rollout: The model will be updatedupdate_for_rollouttimes per rollout.clip: The clipping range of logp ratio or values.lr: The learning rate of optimizer.weight_decay: The coefficient of L2 regularization of optimizer.max_grad_norm: The norm of gradient will be clipped tomax_grad_normif it exceeds the value.warmup_ratio: The fraction of steps to warm up the optimizer.freeze_steps: The model will be freezed in the firstfreeze_stepssteps. Should only be enabled for actor whenadv.estimator=gaeto warmup critic.offload_model: Whether to offload model when not needed.offload_optimizer: Whether to offload optimizer when not needed. Notice that the optimization step will still run on GPUs, which differs from Adam offloading.save_dir: The directory of checkpoints to be saved.save_freq: A checkpoint will be saved everysave_freqsteps. Default toNone, where only a single checkpoint will be saved when the training is finished.save_optimizer: Whether to save the optimizer.
tp_size: The inference engine will be sharded acrosstp_sizeGPUs. Must be divisible by the total number of GPUs.gpu_memory_utilization: The fraction of memory reserved for inference engine.train_sampling_params: The sampling parameters for rollout in training. At leasttemperatureandmax_new_tokensshould be indicated.max_turns: The inference engine will generate at mostmax_turnstimes in a trajectory. Default to1, where the inference engine will only generate once and no function will be called.env_path: The path to the Python script containing functionreward_fnandinteract(ifmax_turns > 1).
coef: The coefficient of KL regularization.type: Ifreward, the KL estimator will be added into the reward of each action; ifloss, the KL estimator will be added into the loss of policy gradient. Should be set tolossfor GRPO.reward_estimator: The estimator used to compute KL reward. Whentype=loss, this will affect the value logged by wandb.loss_estimator: The estimator used to compute KL loss. Should be set tok3for GRPO.
estimator: Ifgae, generalized advantage estimator (and hence critic) will be used to estimate advantage; ifreinforce, normalized reward will be used to estimate advantage. We use token-level loss, so the RL algorithm will be Dr. GRPO ifnorm_var=False.gamma: The discount factor of future rewards. Will only be effective whenestimator=gae.lamda: The coefficient to tradeoff variance and bias of generalized advantage estimator. Will only be effective whenestimator=gae.norm_var: Whether to divide advantages by the standard error. Will only be effective whenestimator=reinforce. Should be set toTruefor GRPO.
project: The name of wandb project.experiment_name: The name of wandb experiment.n_epochs: The dataset will be iterated throughn_epochstimes.disable_wandb: Whether to disable wandb.
This project is built upon the basis of many remarkable projects, including but not limited to
- DeepSpeedChat for the proposal of hybrid engine
- RingFlashAttention for the support of ZigZag ring attention
- SGLang for the support of async inference engine
We also thank OpenRLHF and veRL for their pioneering work.
If you find this library useful, please cite in the following format
@misc{Tan2025RL2,
author={Chenmien Tan and Simon Yu and Lanbo Lin and Ze Zhang and Yuanwu Xu and Chenhao Jiang and Tianyuan Yang and Sicong Xie and Guannan Zhang},
title={RL2: Ray Less Reinforcement Learning},
note={GitHub repository},
howpublished={\url{https://github.com/ChenmienTan/RL2}},
year={2025}
}
We are Accio, the world's first B2B AI sourcing engine. Send us an email if you are interested in opportunities in agent and reinforcement learning.