+
+
+
+Figure 1: We train Qwen2.5-Math-1.5B/7B on the DAPO + MATH dataset and evaluate them on five math reasoning benchmarks: MATH500, AMC, Gaokao, Minerva, and Olympiad Bench. Compared to the baseline method (Dynamic Sampling), our approach (GRESO) reduces rollout overhead by up to 2x while achieving comparable training performance, improving the efficiency of rollout scaling. +
+ + + +## Getting Started +Our implementation is based on [volcengine/verl](https://github.com/volcengine/verl) +. + +### 1. Environment Setup + + + +```bash +conda create -n greso python==3.11 +conda activate greso + +# Install verl +# exit the project folder to install verl +cd .. +git clone https://github.com/volcengine/verl.git +cd verl +pip3 install -e . + +# Install the latest stable version of vLLM +pip3 install vllm==0.8.2 + +# Install flash-attn +pip3 install flash-attn --no-build-isolation + +## misc +pip install wandb IPython matplotlib ipdb latex2sympy2-extended math-verify torchdata pylatexenc +``` + +### 2. Download & Preprocess Data + +You can download the dataset using the following command: + +```bash +# cd the project folder +conda activate greso +export PYTHONPATH="$PYTHONPATH:$(pwd)" + +bash train-scripts/generate_dataset.sh +``` + +### 3. Training + +Train Qwen Math 1.5b with GRESO on 4xH100: +```bash +bash train-scripts/math_qwen_1_5b_dm_greso.sh +``` + +Train Qwen Math 7b with GRESO on 8xH100: +```bash +bash train-scripts/math_qwen_7b_dm_greso.sh +``` + +See more scripts in `train-scripts` folder. diff --git a/docker/Dockerfile.megatron b/docker/Dockerfile.megatron new file mode 100644 index 0000000..dd35521 --- /dev/null +++ b/docker/Dockerfile.megatron @@ -0,0 +1,9 @@ +FROM verlai/verl:vemlp-th2.4.0-cu124-vllm0.6.3-ray2.10-te1.7-v0.0.3 + +RUN pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable + +RUN cd /opt/nvidia && git clone --single-branch --branch core_r0.11.0 https://github.com/NVIDIA/Megatron-LM.git Megatron-LM + +# only config pip index with https://pypi.tuna.tsinghua.edu.cn/simple if needed +# unset for now +RUN cd /opt/nvidia/Megatron-LM && pip3 install --no-deps -e . \ No newline at end of file diff --git a/docker/Dockerfile.ngc.vllm b/docker/Dockerfile.ngc.vllm new file mode 100644 index 0000000..3a43a25 --- /dev/null +++ b/docker/Dockerfile.ngc.vllm @@ -0,0 +1,47 @@ +# docker buildx build --platform linux/x86_64 -t "verlai/verl:ngc-th2.4.0-cu124-vllm0.6.3-ray2.4-te1.7-v0.0.6" -f docker/Dockerfile.ngc.vllm . --builder cloud-verlai-verl-builder --progress=plain --push +FROM nvcr.io/nvidia/pytorch:24.05-py3 + +# uninstall nv-pytorch fork +RUN pip3 uninstall pytorch-quantization \ + pytorch-triton \ + torch \ + torch-tensorrt \ + torchvision \ + xgboost transformer_engine flash_attn \ + apex megatron-core -y + +RUN pip3 install torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cu124 + +# =============== Megatron dependencies (optional) ================= +# install apex, set MAX_JOBS to avoid OOMs +RUN MAX_JOBS=4 pip3 install -v --disable-pip-version-check --no-cache-dir --no-build-isolation \ + --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" \ + git+https://github.com/NVIDIA/apex +# =============== End of Megatron dependencies (optional) ================= + +RUN pip3 install --no-cache-dir \ + accelerate \ + codetiming \ + datasets \ + dill \ + hydra-core \ + numpy \ + 'pandas' \ + 'peft' \ + 'pyarrow>=15.0.0' \ + 'pybind11' \ + 'pylatexenc' \ + 'ray>=2.10' \ + 'tensordict<0.6' \ + 'transformers' \ + 'vllm==0.6.3.post1' \ + 'wandb' + +# full dependencies +RUN pip3 install pytest yapf py-spy pyext liger-kernel + +# =============== Megatron dependencies (optional) ================= +# install Transformer Engine, which requires FA 2.5.8. Do it in a separate step for docker cache +RUN MAX_JOBS=4 NINJA_FLAGS="-j4" pip3 install flash-attn==2.5.8 --no-cache-dir --no-build-isolation +RUN MAX_JOBS=1 NINJA_FLAGS="-j1" TE_BUILD_WITH_NINJA=0 pip3 install git+https://github.com/eric-haibin-lin/TransformerEngine.git@v1.7.0 +# =============== End of Megatron dependencies (optional) ================= diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm new file mode 100644 index 0000000..6a0d20e --- /dev/null +++ b/docker/Dockerfile.rocm @@ -0,0 +1,45 @@ +# Build the docker in the repo dir: +# docker build -f docker/Dockerfile.rocm -t verl-rocm:03.04.2015 . +# docker images # you can find your built docker + + +FROM rocm/vllm:rocm6.2_mi300_ubuntu20.04_py3.9_vllm_0.6.4 + +# Set working directory +# WORKDIR $PWD/app + +# Set environment variables +ENV PYTORCH_ROCM_ARCH="gfx90a;gfx942" + +# Install vllm +RUN pip uninstall -y vllm && \ + rm -rf vllm && \ + git clone -b v0.6.3 https://github.com/vllm-project/vllm.git && \ + cd vllm && \ + MAX_JOBS=$(nproc) python3 setup.py install && \ + cd .. && \ + rm -rf vllm + +# Copy the entire project directory +COPY . . + +# Install dependencies +RUN pip install "tensordict<0.6" --no-deps && \ + pip install accelerate \ + codetiming \ + datasets \ + dill \ + hydra-core \ + liger-kernel \ + numpy \ + pandas \ + peft \ + "pyarrow>=15.0.0" \ + pylatexenc \ + "ray[data,train,tune,serve]" \ + torchdata \ + transformers \ + wandb \ + orjson \ + pybind11 && \ + pip install -e . --no-deps \ No newline at end of file diff --git a/docker/Dockerfile.vemlp.vllm.te b/docker/Dockerfile.vemlp.vllm.te new file mode 100644 index 0000000..a03d7c1 --- /dev/null +++ b/docker/Dockerfile.vemlp.vllm.te @@ -0,0 +1,41 @@ +# docker buildx build --platform linux/x86_64 -t "verlai/verl:$TAG" -f docker/$FILE . + +# the one in docker.io is an alias for the one veturbo +# FROM vemlp-cn-beijing.cr.volces.com/veturbo/pytorch:2.4-cu124 +FROM docker.io/haibinlin/verl:v0.0.5-th2.4.0-cu124-base + +# only config pip index with https://pypi.tuna.tsinghua.edu.cn/simple if needed +# unset for now +RUN pip3 config unset global.index-url + +# transformers 4.47.0 contains the following bug: +# AttributeError: 'Gemma2Attention' object has no attribute '_flash_attn_uses_top_left_mask' +RUN pip3 install --no-cache-dir \ + torch==2.4.0 \ + accelerate \ + codetiming \ + dill \ + hydra-core \ + numpy \ + pybind11 \ + tensordict \ + "transformers <= 4.46.0" + +RUN pip3 install --no-cache-dir flash-attn==2.7.0.post2 --no-build-isolation + +# vllm depends on ray, and veRL does not support ray > 2.37 +RUN pip3 install --no-cache-dir vllm==0.6.3 ray==2.10 + +# install apex +RUN MAX_JOBS=4 pip3 install -v --disable-pip-version-check --no-cache-dir --no-build-isolation \ + --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" \ + git+https://github.com/NVIDIA/apex + +# install Transformer Engine +# - flash-attn pinned to 2.5.3 by TransformerEngine, switch to eric-haibin-lin/TransformerEngine.git@v1.7.0 to relax version req +# - install with: MAX_JOBS=1 NINJA_FLAGS="-j1" TE_BUILD_WITH_NINJA=0 to avoid OOM +# - cudnn is required by TransformerEngine +# RUN CUDNN_PATH=/opt/conda/lib/python3.11/site-packages/nvidia/cudnn \ +# pip3 install git+https://github.com/eric-haibin-lin/TransformerEngine.git@v1.7.0 +RUN MAX_JOBS=1 NINJA_FLAGS="-j1" pip3 install flash-attn==2.5.3 --no-cache-dir --no-build-isolation +RUN MAX_JOBS=1 NINJA_FLAGS="-j1" pip3 install git+https://github.com/NVIDIA/TransformerEngine.git@v1.7 diff --git a/docs/index.html b/docs/index.html new file mode 100644 index 0000000..6253477 --- /dev/null +++ b/docs/index.html @@ -0,0 +1,370 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + ++ TL;DR: + GRESO is a lightweight pre-rollout filter that skips uninformative prompts using reward dynamics, saving RL training time without hurting accuracy. +
+
+ Introduction
+ + Reinforcement learning, such as PPO and GRPO, has powered recent breakthroughs in LLM reasoning. + Scaling rollout to sample more prompts enables models to selectively use higher-quality data for training, which can stabilize RL training and improve model performance, but at the cost of significant computational overhead. + In this paper, we first show that a substantial portion of this overhead can be avoided by skipping uninformative prompts before rollout. + Our analysis of reward dynamics reveals a strong temporal consistency in prompt value: prompts that are uninformative in one epoch of training are likely to remain uninformative in near future epochs. + Based on these insights, we propose GRESO (GRPO with Efficient Selective Rollout), an online, lightweight pre-rollout filtering algorithm that predicts and skips uninformative prompts using reward training dynamics. + By evaluating GRESO on a broad range of math reasoning benchmarks and models, like Qwen2.5-Math-1.5B/7B and DeepSeek-R1-Distill-Qwen-1.5B, we show that GRESO achieves up to + 2.4× wall-clock time speedup in rollout and up to + 2.0× speedup in total training time without accuracy degradation. +
+ + ++ Scaling computational resources to sample responses for more prompts at this rollout stage can enhance reinforcement learning, which allows models to selectively utilize higher-quality data and thus train models with better converged performance. However, scaling up rollouts introduces significant computational overhead, as rollout remains a major bottleneck in RL training. For instance, as shown in the right figures, filtering out uninformative examples1 and resampling to fill the batch with effective data,also known as Dynamic Sampling (DS), can improve model performance, but it comes at the cost of significantly increased rollout overhead. +
+ ++ Motivated by this challenge, we aim to investigate the following research question in this work: +
+ ++ How can we perform more selective rollouts—focusing on sampling more valuable prompts—to make this scaling more efficient? ++
+ 1 In GRPO, many examples yield identical rewards across all responses, resulting in zero advantage and thus contributing no learning signal during training.
+ +
+ + In this paper, we aim to design an efficient selective rollout strategy for LLM RL to make rollout scaling more efficient. + We begin by analyzing the training dynamics of prompts across epochs and observe a strong temporal consistency across different training epochs. + In particular, prompts that yield zero advantage in one epoch are more likely to do so in future epochs as well. + This temporal correlation suggests that historical reward dynamics can be leveraged to predict and preemptively skip zero-variance examples before rollout. + Building on these observations, we propose GRESO (GRPO with Efficient Selective Rollout), an online efficient pre-rollout filtering algorithm that reduces rollout cost by selectively skipping prompts predicted to be zero-variance. + Instead of performing filtering after rollout, GRESO estimates a skipping probability for each prompt based on its reward dynamics during training prior to the rollout stage, significantly reducing prompt selection overhead and making the rollout scaling more efficient. +
+ ++ As shown in the above figure, compared to the baseline method (Dynamic Sampling), our approach (GRESO) reduces rollout overhead by up to 2x while achieving comparable training performance, improving the efficiency of rollout scaling. + (We train Qwen2.5-Math-1.5B/7B on the DAPO + MATH dataset and evaluate them on five math reasoning benchmarks: MATH500, AMC, Gaokao, Minerva, and Olympiad Bench.) +
+ +
+ GRESO: GRPO with Efficient Selective Rollout
+ 
+
+ Training data typically exhibits strong temporal correlations across epochs.
+ We hypothesize that zero-variance prompts in GRPO training similarly have such strong correlations in their training dynamics, enabling opportunities for more efficient identification of these prompts prior to the rollout stage.
+ To test this hypothesis, we conduct a study on the temporal correlation of zero-variance prompts in GRPO training.
+ Specifically, we train Qwen2.5-Math-7B with GRPO and measure two probabilities to study the temporal correlation of zero-variance prompts:
+ 1) P(Previous|Current): The probability that a prompt identified as zero-variance in the current epoch was also zero-variance in any previous epoch.
+ 2) P(Current|Previous): The probability that a prompt identified as zero-variance in any previous epoch remains zero-variance in the current epoch.
+
+ The results shown in the above Figure (a) indicate that zero-variance prompts exhibit strong temporal correlations throughout training. + We have two key observations: +
+P(Previous|Current) curve shows that the majority of zero-variance prompts in a given epoch (e.g., over 90%) were also identified as zero-variance in earlier epochs.P(Current|Previous) curve shows that approximately 20% of prompts previously labeled as zero-variance become effective prompts that contribute to training again. This suggests that, rather than statically pruning zero-variance prompts, it is beneficial to retain some degree of exploration to help retain potentially valuable prompts.+ Based on the above observations, as shown in the above Figure (b), + we propose a probabilistic filtering strategy based on reward dynamics observed during training. Rather than deterministically discarding prompts that previously yielded identical (i.e., zero-variance) rewards, we assign each prompt a probability of being filtered—this probability increases with the number of recent consecutive rollouts where the prompt showed zero-variance. Concretely, for each prompt, we track how many times in a row it has produced zero-variance responses in the most recent epochs. The more consecutive times this occurs, the more likely the prompt will be skipped in the current rollout. However, we always retain a minimum exploration probability, ensuring that even frequently zero-variance prompts have a small chance of being re-sampled. This approach skips uninformative prompts to exploit training efficiency, while occasionally revisiting them to maintain exploration. +
+
Comparable Performance with Fewer Rollouts
+ + To verify the effectiveness of GRESO, we present a comprehensive evaluation of GRESO and Dynamic Sampling (DS), which filters out zero-variance examples and resamples to fill the batch with effective data, across six math reasoning benchmarks, using three different model settings in the above table. + The models are trained on either the DAPO + MATH dataset (DM) or the Open R1 subset (OR1). + We report both the performance and the number of rollouts from the checkpoint that achieves the best average performance across six benchmarks. + Across all training settings, GRESO achieves comparable accuracy as DS, while significantly reducing the number of rollout samples—achieving up to 3.35× fewer rollouts. + For example, on Qwen2.5-Math-7B trained on the DM dataset, GRESO achieves a comparable average accuracy to DS (57.5% vs. 57.8%), while reducing the number of rollouts from 13.1M to 6.3M. + These results demonstrate that GRESO maintains performance while substantially lowering the cost on rollouts. + Similar improvements are observed across other evaluation settings. +
++ To better understand the efficiency of our proposed methods, we report the detailed end-to-end training time breakdown for different stages: rollout, actor model update, and other overheads (e.g., reference model and advantage calculation). + Qwen2.5-Math-1.5B is trained on 4×H100 GPUs, while the other two models are trained on 8×H100 GPUs. + The right table compares the training time breakdown between GRESO and Dynamic Sampling for models trained on the DAPO + MATH dataset. + For all three models, GRESO significantly reduces rollout time—achieving up to 2.4× speedup in rollout and 2.0× speedup in total training time compared to DS. + For instance, on Qwen2.5-Math-7B, GRESO reduces rollout time from 155.9 hours to 65.5 hours, cutting overall training time from 178.0 to 88.3 hours. +
+
+ + Selection Dynamics of different prompts in GRESO. Each row is a prompt, and each column is an epoch. + + We present a case study illustrating how GRESO selects or skips prompts over training epochs. We observe that very easy prompts tend to remain easy throughout training; although frequently skipped, GRESO still occasionally selects them to ensure a minimal level of exploration. For prompts of moderate difficulty, as the model becomes stronger over time, these prompts gradually become easier and are increasingly skipped. In contrast, some hard prompts become solvable~(i.e., effective prompts) in later epochs or even easy prompts. However, certain hard prompts remain unsolved throughout training. +
+
+ @article{zheng2025greso,
+ title={Act Only When It Pays: Efficient Reinforcement Learning for LLM Reasoning via Selective Rollouts},
+ author = {Zheng, Haizhong and Zhou, Yang and Bartoldson, Brian R. and Kailkhura, Bhavya and Lai, Fan and Zhao, Jiawei and Chen, Beidi},
+ journal={arXiv preprint arXiv:2404.11912},
+ year={2025}
+}| Python version: | \n", + "3.9.2 | \n", + "
| Ray version: | \n", + "2.10.0 | \n", + "