A concise library of post-training for large language models.
This is the right library for you if you want to learn reinforcement learning for large language models or have a quick test for your own algorithm. We deliver a clear implementation without complicated abstractions.
Despite the simplicity, you should be able to scale up with 3D (DP/CP/TP) parallelism in FSDP backend and 5D parallelism (DP/CP/PP/TP/EP) in Megatron backend. We also support
RL2 is a production-ready library! It achieves comparable performance with other popular LLM RL libraries.
Also check our wandb report on OpenThoughts, SkyworkRM, UltraFeedback, TinyZero, LetterCounting, and SearchR1.
From PyPI:
pip install rl-square
Using Pre-built Docker Images:
# FSDP backend only
docker pull rlsquare/rl2:latest
# Megatron backend (includes FSDP)
docker pull rlsquare/rl2-megatron:latest
Hugging Face dataset and various file types, i.e., JSON, JSONL, CSV, Parquet, and Arrow, are accepted. All trainers support formats of both raw text and messages. The former is more flexible but may be model-specific.
[
{
"prompt": "The capital of China is",
"response": "Beijing."
}
]
[
{
"messages": [
{"role": "user", "content": "What is the capital of China?"},
{"role": "assistant", "content": "Beijing."}
]
}
]
Multi-turn is only supported by the latter format.
[
{
"prompt": "The capital of China is",
"chosen": "Beijing.",
"rejected": "Shanghai."
}
]
[
{
"chosen": [
{"role": "user", "content": "What is the capital of China?"},
{"role": "assistant", "content": "Beijing."}
],
"rejected": [
{"role": "user", "content": "What is the capital of China?"},
{"role": "assistant", "content": "Shanghai."}
]
}
]
[
{
"prompt": "The capital of China is",
"extra_info": {
"answer": "Beijing"
}
}
]
[
{
"messages": [
{"role": "user", "content": "What is the capital of China?"}
],
"extra_info": {
"answer": "Beijing"
}
}
]
In PPO, the language model interacts with the environment through a user-defined function step in the following format.
async def step(
state: str, action: str, extra_info: Dict
) -> Dict:
action_type = parse_action_type(action)
env_response = {
"next_state": None,
"reward": 0.0,
"score": 0.0,
"done": False,
"extra_info": extra_info
}
if action_type == "search":
query = parse_query(action)
passage = await search_result(query)
env_response["next_state"] = state + action + passage
elif action_type == "answer":
pred = parse_pred(action)
reward = float(is_equivalent(pred, extra_info["answer"]))
env["reward"] = reward
env["score"] = score
env_response["done"] = True
return env_response
state and action are the input and output of language model in the last turn and next_state is the input of language model in the next turn.
When state + action is a prefix of next_state, the two turns will be processed in a single sequence.reward is used to compute advantages (and subsequently update the model) while score is used to log the model performance.
Diverge values may be used when needed.done indicates whether to proceed to the next turn.extra_info contains everything not aforementioned, e.g., answer.The function should be included in a Python script where the path is specified by actor.rollout.env_path.
Use torchrun to launch the trainer. 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>
By default, i.e., ddp_size=1, tp_size=1, your model will be partitioned via ZeRO stage 3.
ddp_size specifies the number of model parameter copies.
Larger ddp_size leads to higher memory consumption and lower communication cost.
For large models, you may specify tp_size > 1 to enable tensor parallelism.
The product of ddp_size and tp_size should be a factor of the total number of GPUs.
For SFT, RM, and DPO, max_length is used to truncate sequences.
In RM and DPO, the chosen and rejected sequences will be packed together, so the actual sequence length can be up to twice of max_length.
For PPO, max_new_tokens is used to terminate generations.
The length of any sequence cannot exceed cp_size * tp_size * max_length_per_device.
The default algorithm is Dr. GRPO, where the loss is averaged at the token level and the advantage is not divided by the standard deviation.
kl.type=reward, kl.reward_estimator=k1, and adv.estimator=gaeactor.avg_level=sequence, kl.type=loss, kl.loss_estimator=k3, and adv.norm_var=trueThis project is built upon the basis of many remarkable projects, including but not limited to
We also thank OpenRLHF, veRL, and slime 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 Lei Wei 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 AI sourcing agent. We are always looking for talents in Hangzhou. Send us an email if you are interested in internship/full-time positions in post-training/agent.
Python
99.1%
A concise library of post-training for large language models.
This is the right library for you if you want to learn reinforcement learning for large language models or have a quick test for your own algorithm. We deliver a clear implementation without complicated abstractions.
Despite the simplicity, you should be able to scale up with 3D (DP/CP/TP) parallelism in FSDP backend and 5D parallelism (DP/CP/PP/TP/EP) in Megatron backend. We also support
RL2 is a production-ready library! It achieves comparable performance with other popular LLM RL libraries.
Also check our wandb report on OpenThoughts, SkyworkRM, UltraFeedback, TinyZero, LetterCounting, and SearchR1.
From PyPI:
pip install rl-square
Using Pre-built Docker Images:
# FSDP backend only
docker pull rlsquare/rl2:latest
# Megatron backend (includes FSDP)
docker pull rlsquare/rl2-megatron:latest
Hugging Face dataset and various file types, i.e., JSON, JSONL, CSV, Parquet, and Arrow, are accepted. All trainers support formats of both raw text and messages. The former is more flexible but may be model-specific.
[
{
"prompt": "The capital of China is",
"response": "Beijing."
}
]
[
{
"messages": [
{"role": "user", "content": "What is the capital of China?"},
{"role": "assistant", "content": "Beijing."}
]
}
]
Multi-turn is only supported by the latter format.
[
{
"prompt": "The capital of China is",
"chosen": "Beijing.",
"rejected": "Shanghai."
}
]
[
{
"chosen": [
{"role": "user", "content": "What is the capital of China?"},
{"role": "assistant", "content": "Beijing."}
],
"rejected": [
{"role": "user", "content": "What is the capital of China?"},
{"role": "assistant", "content": "Shanghai."}
]
}
]
[
{
"prompt": "The capital of China is",
"extra_info": {
"answer": "Beijing"
}
}
]
[
{
"messages": [
{"role": "user", "content": "What is the capital of China?"}
],
"extra_info": {
"answer": "Beijing"
}
}
]
In PPO, the language model interacts with the environment through a user-defined function step in the following format.
async def step(
state: str, action: str, extra_info: Dict
) -> Dict:
action_type = parse_action_type(action)
env_response = {
"next_state": None,
"reward": 0.0,
"score": 0.0,
"done": False,
"extra_info": extra_info
}
if action_type == "search":
query = parse_query(action)
passage = await search_result(query)
env_response["next_state"] = state + action + passage
elif action_type == "answer":
pred = parse_pred(action)
reward = float(is_equivalent(pred, extra_info["answer"]))
env["reward"] = reward
env["score"] = score
env_response["done"] = True
return env_response
state and action are the input and output of language model in the last turn and next_state is the input of language model in the next turn.
When state + action is a prefix of next_state, the two turns will be processed in a single sequence.reward is used to compute advantages (and subsequently update the model) while score is used to log the model performance.
Diverge values may be used when needed.done indicates whether to proceed to the next turn.extra_info contains everything not aforementioned, e.g., answer.The function should be included in a Python script where the path is specified by actor.rollout.env_path.
Use torchrun to launch the trainer. 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>
By default, i.e., ddp_size=1, tp_size=1, your model will be partitioned via ZeRO stage 3.
ddp_size specifies the number of model parameter copies.
Larger ddp_size leads to higher memory consumption and lower communication cost.
For large models, you may specify tp_size > 1 to enable tensor parallelism.
The product of ddp_size and tp_size should be a factor of the total number of GPUs.
For SFT, RM, and DPO, max_length is used to truncate sequences.
In RM and DPO, the chosen and rejected sequences will be packed together, so the actual sequence length can be up to twice of max_length.
For PPO, max_new_tokens is used to terminate generations.
The length of any sequence cannot exceed cp_size * tp_size * max_length_per_device.
The default algorithm is Dr. GRPO, where the loss is averaged at the token level and the advantage is not divided by the standard deviation.
kl.type=reward, kl.reward_estimator=k1, and adv.estimator=gaeactor.avg_level=sequence, kl.type=loss, kl.loss_estimator=k3, and adv.norm_var=trueThis project is built upon the basis of many remarkable projects, including but not limited to
We also thank OpenRLHF, veRL, and slime 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 Lei Wei 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 AI sourcing agent. We are always looking for talents in Hangzhou. Send us an email if you are interested in internship/full-time positions in post-training/agent.
Python
99.1%