prime-rl
is developed and tested on NVIDIA A100, H100, H200, and B200; likely compatible with other Ampere, Hopper and Blackwell-class GPUs. If your installation fails, please create an issue.
Quick Installation (Recommended)
curl -sSL https://raw.githubusercontent.com/PrimeIntellect-ai/prime-rl/main/scripts/install.sh | bash
Manual Installation
- Clone the repository
git clone [email protected]:PrimeIntellect-ai/prime-rl.git
cd prime-rl
- Install uv
curl -LsSf https://astral.sh/uv/install.sh | sh
source $HOME/.local/bin/env
- Synchronize the environment
uv sync && uv sync --all-extras
Validate your environment setup
- Check that environment uses Python 3.12
uv run python -V
- Check that
flash-attn
is installed
uv run python -c "import flash_attn"
- Check that you can run SFT trainer in debug model (*this requires 1 GPU)
uv run sft @ configs/debug/sft.toml
- Check that you can run the RL trainer debug mode (this requires 1 GPU)
uv run trainer @ configs/debug/train.toml
- Check that you can run the orchestrator against an inference server (this requires 1 GPU)
uv run inference @ configs/debug/infer.toml
uv run orchestrator @ configs/debug/orch.toml
- Check that you can run a simple SFT warmup (this requires 1 GPU)
uv run sft @ configs/reverse_text/sft.toml
- Check that you can run a toy RL run (this requires 2 GPUs)
uv run rl \
--trainer @ configs/reverse_text/train.toml \
--orchestrator @ configs/reverse_text/orch.toml \
--inference @ configs/reverse_text/infer.toml
- If you want to log your runs to W&B (
wandb
), log in
uv run wandb login
# Or set `export WANDB_API_KEY=...`
- If you require gated/ private models or datasets from HuggingFace, log in
uv run huggingface-cli login
# Or set `export HF_TOKEN=...`
- You may want to increase the number of open files to prevent errors like
Too many open files
.
ulimit -n 32000
- We provide a convenient tmux layout script to start a run and view logs. To start the session simply run
bash scripts/tmux.sh
The default session name is prime-rl
and outputs directory is outputs
. They are configurable, as described in bash scripts/tmux.sh -h
.
We provide a convenience endpoint rl
for single-node RL experiments. It configures and starts the trainer, orchestrator and, optionally, an inference server. It allows setting shared configs conveniently (e.g. the model name or async level should be the same across all modules) and dispatches and monitors subprocesses. To stream the logs from each module we use file logging. By default, only the trainer logs will be displayed on the main process (use the tmux layout script to conveniently view all logs).
To check all available configuration options, run uv run rl --help
.
Reverse Text
Train a tiny model (PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT
) that has undergone SFT warmup (see below) to learn to reverse a small chunk of text. Training is extremely quick because we allow a maximum context of 128 tokens. We use this run in CI.
uv run rl \
--trainer @ configs/reverse_text/train.toml \
--orchestrator @ configs/reverse_text/orch.toml \
--inference @ configs/reverse_text/infer.toml
With two small GPUs (e.g. RTX 3090/4090), this experiment should finish in less than 5 minutes.
Hendrycks Math
Train a small model (willcb/DeepSeek-R1-Distill-Qwen-1.5B
) on high-school level math questions. It is recommended to have at least 2xA100-80GB GPUs or more for this experiment.
On eight GPUs, run the following command to run the experiment.
uv run rl \
--trainer @ configs/hendrycks_math/1b/train.toml \
--orchestrator @ configs/hendrycks_math/1b/orch.toml \
--inference @ configs/hendrycks_math/1b/infer.toml \
--trainer-gpus 2 --inference-gpus 6
NB: This setup requires 8 GPUs - 2 are used for the FSDP trainer, 6 are used for inference.
INTELLECT-2 Math
Train a small model (willcb/DeepSeek-R1-Distill-Qwen-1.5B
) on complex math questions from the INTELLECT-2 dataset.
uv run rl \
--trainer @ configs/intellect_math/1b/train.toml \
--orchestrator @ configs/intellect_math/1b/orch.toml \
--inference @ configs/intellect_math/1b/infer.toml \
--trainer-gpus 2 --inference-gpus 6
NB: This setup requires 8 GPUs - 2 are used for the FSDP trainer, 6 are used for inference.
TBD
For small models/ quick ablations, it can be more efficient to parallelize experiments within a node (e.g. split your GPUs to run two experiments in parallel). Because the trainer communicates with the orchestrator via a shared file system, and the orchestrator communicates with the inference engine via an OAI-compatible API, the connection points have to be uniquely set. For example, if you have access to 4 GPUs you can run two 2 GPU training runs in parallel as follows:
Start the first experiment in a tmux session exp-1
with outputs directory outputs
. Specify it both in the tmux script, as well as in the start command (will use the first 2 GPUs)
bash scripts/tmux.sh -s exp-1 -o outputs1
# Start the first experiment
uv run rl \
--trainer @ configs/reverse_text/train.toml \
--orchestrator @ configs/reverse_text/orch.toml \
--inference @ configs/reverse_text/infer.toml \
--outputs-dir outputs1
For the second experiment, start a second tmux session named exp-2
with outputs directory outputs2
. In addition, specify a new server port for the inference engine and orchestrator (will use the first 2 GPUs)
bash scripts/tmux.sh -s exp-2 -o outputs2
# Start the second experiment
CUDA_VISIBLE_DEVICES=2,3 uv run rl \
--trainer @ configs/reverse_text/train.toml \
--orchestrator @ configs/reverse_text/orch.toml \
--inference @ configs/reverse_text/infer.toml \
--inference.server.port 8001 \
--orchestrator.client.port 8001 \
--outputs-dir outputs2
We have built-in support for SFT using the sft
entrypoint. SFT often proves useful prior to RL training to get the model in-distribution and have higher initial reward that it can hillclimb from.
To check all available configuration options, run uv run sft --help
.
Reverse Text
Finetune PrimeIntellect/Qwen3-0.6B
(Qwen/Qwen3-0.6B
but with Qwen-2.5 chat template) to reverse a tiny chunk of text, used as a warmup model train in reverse-text
environment. We use this run in CI.
uv run sft @ configs/reverse_text/sft.toml
If you want to train on multiple GPUs or nodes, use torchrun
. For example, to do the same run as above on 4 colocated GPUs.
uv run torchrun --nproc-per-node 4 src/prime_rl/trainer/sft/train.py @ configs/reverse_text/sft.toml
We provide a convenience endpoint for running a full evaluation suite of common benchmarks such as AIME, MATH-500 or LiveCodeBench against your model using the eval
entrypoint.
uv run inference --model.name Qwen/Qwen3-0.6B --max-model-len 2048
uv run eval --model.name Qwen/Qwen3-0.6B --environment-ids math500,aime2024,aime2025
To check all available configuration options, run uv run eval --help
.
TBD
For now, development is only possible on CUDA-enabled devices.
- Install pre-commit hooks
uv run pre-commit install
Sources
We support the following sources for configuration, in this order of precedence:
-
Command-line arguments: You can pass (nested) arguments as
--key.subkey value
to the script. For example, to set the model name you can run--model.name
-
Config files: You can pass
.toml
config files (defined in theconfigs
directory) using the@
prefix. For example, to use thedebug.toml
config file, you can runuv run inference @ configs/debug/infer.toml
. (If you leave a space between the@
and the config file, you will get shell path auto-completions.) -
Environment variables: You can set environment variables to override the config values. All environment variables must be prefixed with
PRIME_
and use the__
delimiter to nest the keys. For example, to set the model name you can runexport PRIME_MODEL__NAME=Qwen/Qwen3-0.6B
. -
Defaults: For almost all config arguments, we have a default value which will be used if no other source is provided.
In general we recommend setting configurations via config files to define reproducible experiments and use command-line arguments to override the config values to run variants of the same experiment. Environment variables are usually only used in production settings to communicate with the Prime Protocol worker. In most cases, you should not need to use environment variables.
The precedence order will be important if multiple sources try to configure the same argument. For example, in the following command, all sources will define a model name
# qwen8b.toml
[model]
name = "Qwen/Qwen3-8B"
# qwen14b.toml
[model]
name = "Qwen/Qwen-14B"
PRIME_MODEL__NAME=Qwen/Qwen3-4B uv run inference @qwen8b.toml @qwen14b.toml --model.name Qwen/Qwen3-32B
In this example, the CLI argument --model.name Qwen/Qwen3-32B
will take precendence and the script will use Qwen/Qwen3-32B
as the model name. If the CLI argument wasn't set, then the second config file would take precedence and the script would use Qwen/Qwen-14B
as the model name. If the second config file wasn't set, then the first config file would take precedence and the script would use Qwen/Qwen3-8B
as the model name. Finally, if the first config file wasn't set, then the environment variable would take precedence and the script would use Qwen/Qwen-4B
as the model name. If the environment variable wasn't set, then the default value would be used and the script would use Qwen/Qwen3-0.6B
as the model name.
For development purposes it is useful start the inference server once and keep it alive across experiments to avoid suffering the vLLM startup time repeatedly. The recommended workflow is as follows:
- Start the preconfigured tmux session using the tmux script
bash scripts/tmux.sh
- Start the inference server in the
Inference
pane.
uv run inference @ configs/reverse_text/infer.toml
- Start the trainer and orchestrator in the
Trainer
pane.
uv run rl \
--trainer @ configs/reverse_text/train.toml \
--orchestrator @ configs/reverse_text/orch.toml
To kill the tmux session when you're done:
tmux kill-session -t prime-rl
prime-rl
supports Environment modules built with verifiers
(repo) for training tasks. All of our current research environments live in a separate Prime Environments repository.
To add a new training or evaluation environment, please follow the instructions in the Prime Environments repository.
To then use it as part of prime-rl
, install the newly pushed environment via the Environment Hub.
To install your Environment module temporarily within prime-rl
, do:
uv run prime env install custom-environment
To persist your Environment module installation in the package-wide pyproject.toml
, do:
uv add --optional vf "custom-environment @ https://hub.primeintellect.ai/your-username/custom-environment/@latest/custom-environment-0.1.3-py2.py3-none-any.whl"
For quick API-based testing post-installation, do:
uv run vf-eval custom-environment # -h for config options; defaults to gpt-4.1-mini, 5 prompts, 3 rollouts each
For training, create trainer
/inference
/orchestrator
config files following the aforementioned examples, then set id = custom-environment
in the [environment]
section of your orchestrator
config (along with any desired Environment-level args in [environment.args]
).
For any serious run we recommend logging to W&B. Since it is disabled by default, you have to set up W&B. First, make sure that you are logged in.
uv run wandb login
# Or set `export WANDB_API_KEY=...`
Both the trainer and orchestrator can log to W&B as separate runs using the --monitor.wandb
subconfig. You can set the project (--monitor.wandb.project
, defaults to prime-rl
), run name (--monitor.wandb.name
, defaults to None
which will make W&B generate a name randomly), run ID (--monitor.wandb.id
, defaults to None
), the log directory (--monitor.wandb.dir
, defaults to logs
) and whether or not to run in offline mode (--monitor.wandb.offline
, defaults to False
).
First, start your inference server
uv run inference @ configs/reverse_text/infer.toml
Then, start the trainer and orchestrator with logging enabled.
CUDA_VISIBLE_DEVICES=1 uv run trainer @ configs/reverse_text/train.toml --monitor.wandb.project example-project --monitor.wandb.name trainer
uv run orchestrator @ configs/reverse_text/orch.toml --monitor.wandb.project example-project --monitor.wandb.name orchestrator
Usually it will be more convenient to use the rl
entrypoint. To setup W&B concisely, you can specify shared configs using the --wandb
subconfig, e.g. the project (--wandb.project
), run name (--wandb.name
), directory (--wandb.dir
) and offline mode (--wandb.offline
). It will automatically share these configs to the trainer and orchestrator. For the run name, it will automatically suffix the specified name with -trainer
and -orchestrator
to clearly distinguish those runs.
uv run rl \
--trainer @ configs/reverse_text/train.toml \
--orchestrator @ configs/reverse_text/orch.toml \
--inference @ configs/reverse_text/infer.toml \
--wandb.project example-project \
--wandb.name example-run
We support logging samples (e.g. prompt, completion, reward, advantage for selected rollouts) and distributions (e.g. reward, advantage, entropy distributions) as W&B tables using the monitor.wandb.log-extras
subconfig. On the orchestrator you can log activate logging samples (--monitor.wandb.log-extras.samples
) and distributions (--monitor.wandb.log-extras.samples
). On the trainer you can only log distributions (--monitor.wandb.log-extras.distributions
). On both, you can specify the logging step interval using --monitor.wandb.log-extras.interval
. To log all extras on trainer and orchestrator every 10 steps,
uv run rl \
--trainer @ configs/reverse_text/train.toml \
--orchestrator @ configs/reverse_text/orch.toml \
--inference @ configs/reverse_text/infer.toml \
--wandb.project example-project \
--wandb.name example-run \
--trainer.monitor.wandb.log-extras.distributions \
--trainer.monitor.wandb.log-extras.interval 10 \
--orchestrator.monitor.wandb.log-extras.samples \
--orchestrator.monitor.wandb.log-extras.distributions \
--orchestrator.monitor.wandb.log-extras.interval 10
Our codebase supports checkpointing. Because of the trainer/ orchestrator design, as well as the natural asynchrony checkpointing is non-standard.
- Trainer (
src/prime_rl/trainer/ckpt.py
): Checkpoints FSDP model shard, optimizer state and progress (training step, total samples, total tokens) - Orchestrator (
src/prime_rl/orchestrator/ckpt.py
): Checkpoints orchestrator progress
NB: Each run with asynchrony level async_level
and some checkpoint step x
, requires weight checkpoints in the step range [x-async_level, x]
. Currently we do not duplicate weight checkpoints into the checkpoints
directory but simply keep them around in weights
, by keeping the trainer from cleaning up weight checkpoints that are required for resuming training. This way, the orchestrator only needs to checkpoint its progress (read: step) to load the correct weights into the inference engine upon resuming.
The default checkpoint directory is checkpoints
and each checkpoint step will live in a subdirectory enumerated by the step, i.e. checkpoints/step_{step}
. The trainer checkpoint is called trainer.pt
for single GPU workloads, else trainer_{local_rank}.pt
. The orchestrator checkpoint is called orchestrator.pt
. Thus, this is a typical directory structure:
checkpoints
├── step_10
│ ├── orchestrator.pt
│ └── trainer.pt
├── step_25
│ ├── orchestrator.pt
│ └── trainer.pt
└── step_30
├── orchestrator.pt
└── trainer.pt
Checkpointing is configured by the CheckpointConfig
, with the config key --ckpt
. One can specify the interval (--ckpt.interval
, defaults to 50
), whether to save checkpoints asynchronoously (--ckpt.save-async
, defaults to False
), and how many recent step checkpoints to keep on disk (--ckpt.keep
, defaults to None
which means no cleanup).
By default, runs do no write checkpoints to save disk space. To checkpoint every 10 steps on our debug RL run, run the following command
CUDA_VISIBLE_DEVICES=1 uv run trainer @ configs/reverse_text/train.toml --ckpt.interval 10
To resume a run use the --ckpt.resume-step
flag. To resume from the checkpoint step 10 from the previous command, run the following command
CUDA_VISIBLE_DEVICES=1 uv run trainer @ configs/reverse_text/train.toml --ckpt.resume-step 10
Because we save progress information, resuming from a checkpoint is fully W&B compatible. By default, resuming from a checkpoint, will simply create a new run. To resume the same W&B run, you'd have to pass the same W&B run ID for both the trainer and the orchestrator, e.g.
CUDA_VISIBLE_DEVICES=1 uv run trainer @ configs/reverse_text/train.toml \
--monitor.wandb.project <project> \
--ckpt.resume-step 10 \
--monitor.wandb.id <trainer-run-id> \
You also need to restart the orchestrator from a checkpoint, the api is the same as the trainer, e.g.
uv run orchestrator @ configs/reverse_text/orch.toml \
--monitor.wandb.project <project> \
--ckpt.resume-step 10 \
--monitor.wandb.id <orchestrator-run-id>
If you started your run using rl.py
, you can resume the same run by passing the same W&B run ID for both the trainer and the orchestrator, e.g.
uv run rl \
--trainer @ configs/reverse_text/train.toml \
--orchestrator @ configs/reverse_text/orch.toml \
--ckpt.resume-step 10 \
--trainer.monitor.wandb.id <trainer-run-id> \
--orchestrator.monitor.wandb.id <orchestrator-run-id>
You don't need to restart the inference server if started manually, the orchestrator will automatically send the right checkpoint to the inference server when resuming.
We provide a convenient way to benchmark the performance of the inference engine and trainer using the --bench
flag. It will run each module in isolation for a few steps and log performance statistics to the console and, optionally, W&B.
Inference
To benchmark inference, first spin up the inference server with an experiment configuration
uv run inference @ configs/reverse_text/infer.toml
Then, start the orchestrator with the matching configuration file in benchmark mode
uv run orchestrator @ configs/reverse_text/orch.toml --bench
Trainer
To benchmark the RL trainer, simply run the trainer against a fake data loader with batch certain specifications.
uv run trainer @ configs/reverse_text/train.toml --bench --data.fake.micro_batch_size 8 --data.fake.batch_size 128 --data.fake.seq_len 128
RL
You can benchmark both the RL trainer and inference at the same time with the rl.py
entrypoint. Note, that the benchmarking is still decoupled.
uv run rl \
--trainer @ configs/reverse_text/train.toml \
--orchestrator @ configs/reverse_text/orch.toml \
--inference @ configs/reverse_text/infer.toml \
--bench
SFT
Benchmark the SFT trainer against fixed
or variable
length fake data by specifyin --data.fake.type
uv run sft --bench --data.fake.type fixed --data.micro-batch-size 8 --data.batch-size 8 --data.seq-len 128
Run the full test suite
uv run pytest -v
To run unit tests, run
uv run pytest tests/unit -v
To run integration tests, run
uv run pytest tests/integration -v
To run CPU-only tests, use the inverse of the gpu
marker:
uv run pytest -v -m "not gpu"
To run fast tests, use the inverse of the slow
marker:
uv run pytest -v -m "not slow"
At each step n
, all artifacts (e.g., checkpoint, rollout, gradient) are tagged with the same step number.
- Step 0:
- Uses checkpoint 0 on rollout 0 to compute gradient 0.
- Then computes checkpoint 1 as:
ckpt 1 = ckpt 0 - grad 0
In general, the model used for generating rollouts at step n
is from ckpt[n - async_level]
.
- When async_level = 0, the rollout and gradient are based on the same model version. This is equivalent to synchronous on-policy training.
This project is licensed under the Apache 2.0 license, as found in the License file.
If you find our work useful, feel free to cite it using
@misc{primeintellect2025prime-rl,
author = {Prime Intellect},
title = {PRIME-RL},
url = {https://github.com/PrimeIntellect-ai/prime-rl},
year = {2025}
}