Skip to content

Commit b806e10

Browse files
authored
Bump vLLM and TRL (#665)
* Bump vLLM and TRL * Fix Makefile
1 parent a6b4f66 commit b806e10

3 files changed

Lines changed: 15 additions & 38 deletions

File tree

Makefile

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,11 @@ check_dirs := src tests
88

99
# dev dependencies
1010
install:
11-
uv venv openr1 --python 3.11 && . openr1/bin/activate && uv pip install --upgrade pip
12-
uv pip install vllm==0.8.4
13-
uv pip install setuptools
14-
uv pip install flash-attn --no-build-isolation
11+
uv venv openr1 --python 3.11
12+
. openr1/bin/activate && uv pip install --upgrade pip && \
13+
uv pip install vllm==0.8.5.post1 && \
14+
uv pip install setuptools && \
15+
uv pip install flash-attn --no-build-isolation && \
1516
GIT_LFS_SKIP_SMUDGE=1 uv pip install -e ".[dev]"
1617

1718
style:

README.md

Lines changed: 9 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ uv venv openr1 --python 3.11 && source openr1/bin/activate && uv pip install --u
6969
Next, install vLLM and FlashAttention:
7070

7171
```shell
72-
uv pip install vllm==0.8.4
72+
uv pip install vllm==0.8.5.post1
7373
uv pip install setuptools && uv pip install flash-attn --no-build-isolation
7474
```
7575

@@ -140,8 +140,9 @@ By default, these scripts will push each model to your Hugging Face Hub username
140140
# Change the base model to a smaller variant
141141
accelerate launch --config_file recipes/accelerate_configs/zero3.yaml src/open_r1/sft.py \
142142
--config recipes/OpenR1-Distill-7B/sft/config_distill.yaml \
143-
--model_name_or_path Qwen/Qwen3-0.6-Base \
144-
--hub_model_id OpenR1-Distill-0.6B
143+
--model_name_or_path Qwen/Qwen3-0.6B-Base \
144+
--hub_model_id OpenR1-Distill-0.6B \
145+
--output_dir data/OpenR1-Distill-0.6B
145146
```
146147

147148
If you also wish to override the Weights and Biases default settings, you can do so as follows:
@@ -214,43 +215,18 @@ You can adjust the YAML config to train on a different base model or dataset.
214215

215216
### GRPO
216217

217-
We use TRL's [vLLM backend](https://huggingface.co/docs/trl/speeding_up_training?vllm+examples=GRPO#vllm-for-fast-generation-in-online-methods) to scale training to large models across multiple nodes. For single-node training of smol models across 8 GPUs, first spin up the vLLM server to run on e.g. 1 GPU as follows:
218+
We use TRL's [vLLM backend](https://huggingface.co/docs/trl/speeding_up_training?vllm+examples=GRPO#vllm-for-fast-generation-in-online-methods) to scale training to large models across multiple nodes. For single-node training of smol models across 8 GPUs, use `vllm_mode="colocate"` to run vLLM in the same process as the training script:
218219

219220
```shell
220-
CUDA_VISIBLE_DEVICES=0 trl vllm-serve --model deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B
221-
```
222-
223-
Once the server is up, run training on the remaining GPUs as follows:
224-
225-
```shell
226-
CUDA_VISIBLE_DEVICES=1,2,3,4,5,6,7 ACCELERATE_LOG_LEVEL=info \
227-
accelerate launch --config_file recipes/accelerate_configs/zero2.yaml --num_processes 7 \
228-
src/open_r1/grpo.py --config recipes/DeepSeek-R1-Distill-Qwen-1.5B/grpo/config_demo.yaml
221+
ACCELERATE_LOG_LEVEL=info \
222+
accelerate launch --config_file recipes/accelerate_configs/zero3.yaml \
223+
src/open_r1/grpo.py --config recipes/DeepSeek-R1-Distill-Qwen-1.5B/grpo/config_demo.yaml \
224+
--vllm_mode colocate
229225
```
230226

231227
> [!WARNING]
232228
> The chat template used in the distilled DeepSeek models omits the contents of the reasoning block within the `<think>` and `</think>` tags. It also prefills the assistant response with `<think>` which interferes with the format reward function. To handle that, it is important to override the chat template as done in e.g. [recipes/DeepSeek-R1-Distill-Qwen-1.5B/grpo/config_demo.yaml](./recipes/DeepSeek-R1-Distill-Qwen-1.5B/grpo/config_demo.yaml).
233229
234-
To increase the throughput with data parallel on e.g. 2 GPUs, run:
235-
236-
```shell
237-
CUDA_VISIBLE_DEVICES=0,1 trl vllm-serve --model deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --data_parallel_size 2
238-
```
239-
240-
Then run training on the remaining GPUs as follows:
241-
242-
```shell
243-
CUDA_VISIBLE_DEVICES=2,3,4,5,6,7 ACCELERATE_LOG_LEVEL=info \
244-
accelerate launch --config_file recipes/accelerate_configs/zero2.yaml --num_processes 6 \
245-
src/open_r1/grpo.py --config recipes/DeepSeek-R1-Distill-Qwen-1.5B/grpo/config_demo.yaml
246-
```
247-
248-
For larger models, use tensor parallelism:
249-
250-
```shell
251-
CUDA_VISIBLE_DEVICES=0,1 trl vllm-serve --model deepseek-ai/DeepSeek-R1-Distill-Qwen-14B --tensor_parallel_size 2
252-
```
253-
254230
For multi-node training on N+1 nodes, with 1 node running the vLLM server and N nodes running training, we provide an example Slurm script. For example, to run the above example on 1+1 nodes with data parallelism, run:
255231

256232
```shell

setup.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@
6969
"sentencepiece>=0.1.99",
7070
"torch==2.6.0",
7171
"transformers==4.52.3",
72-
"trl[vllm] @ git+https://github.com/huggingface/trl.git@9ac614fb081e17805f7f62ab3f5f7036bdefe7b0", # Support for activation offload: https://github.com/huggingface/trl/pull/2954
72+
"trl[vllm]==0.18.0",
7373
"wandb>=0.19.1",
7474
"async-lru>=2.0.5",
7575
"aiofiles>=24.1.0",

0 commit comments

Comments
 (0)