fix: easier environment setup; pin trl, transformers

README.md

@@ -45,23 +45,47 @@ We will use the DeepSeek-R1 [tech report](https://github.com/deepseek-ai/DeepSee
 
 **Note: Libraries rely on CUDA 12.1. Double check your system if you get segmentation faults.**
 
-To run the code in this project, first, create a Python virtual environment using e.g. `uv`.
-To install `uv`, follow the [UV Installation Guide](https://docs.astral.sh/uv/getting-started/installation/).
+First, create a virtual environment, either with Conda or `uv`.
+
+### Using `conda`
+
+Use your existing venv or create a new one with:
+
+```shell
+conda create -n openr1 python=3.11
+```
 
+### Using `uv`
+
+To install `uv`, follow the [UV Installation Guide](https://docs.astral.sh/uv/getting-started/installation/).
 
 ```shell
 uv venv openr1 --python 3.11 && source openr1/bin/activate && uv pip install --upgrade pip
 ```
 
-Next, install vLLM:
+### Dependencies
+
+From here, the easiest way to get started is to run:
 
 ```shell
-uv pip install vllm>=0.7.0
+./setup.sh
+```
 
-# For CUDA 12.1
-pip install vllm>=0.7.0 --extra-index-url https://download.pytorch.org/whl/cu121
+OR, do it manually:
+<details>
+<summary>Bash</summary>
+
+```shell
+# Install pinned vllm
+pip install "vllm>=0.7.1" --extra-index-url https://download.pytorch.org/whl/cu121
+
+# nvJitLink
 export LD_LIBRARY_PATH=$(python -c "import site; print(site.getsitepackages()[0] + '/nvidia/nvjitlink/lib')"):$LD_LIBRARY_PATH
+
+# Install dependencies
+pip install -e ".[dev]"
 ```
+</details>
 
 This will also install PyTorch `v2.5.1` and it is **very important** to use this version since the vLLM binaries are compiled for it. You can then install the remaining dependencies for your specific use case via `pip install -e .[LIST OF MODES]`. For most contributors, we recommend:
 

setup.py

@@ -45,26 +45,26 @@
     "bitsandbytes>=0.43.0",
     "ruff>=0.9.0",
     "datasets>=3.2.0",
-    "deepspeed==0.15.4",
     "distilabel[vllm,ray,openai]>=1.5.2",
     "einops>=0.8.0",
     "flake8>=6.0.0",
     "hf_transfer>=0.1.4",
     "huggingface-hub[cli]>=0.19.2,<1.0",
     "isort>=5.12.0",
-    "liger_kernel==0.5.2",
-    "lighteval @ git+https://github.com/huggingface/lighteval.git@0e462692436e1f0575bdb4c6ef63453ad9bde7d4#egg=lighteval[math]",
-    "math-verify>=0.3.3",  # Used for math verification in grpo
+    "lighteval @ git+https://github.com/huggingface/lighteval.git@cb35beae9f1bf8133f840de1ea5ad840e37c1e07#egg=lighteval[math]",
     "packaging>=23.0",
     "parameterized>=0.9.0",
     "pytest",
     "safetensors>=0.3.3",
     "sentencepiece>=0.1.99",
-    "torch>=2.5.1",
-    "transformers @ git+https://github.com/huggingface/transformers.git@main",
-    "trl @ git+https://github.com/huggingface/trl.git@main",
-    "vllm>=0.7.1",
+    "transformers>=4.48.2",
     "wandb>=0.19.1",
+    "deepspeed==0.15.4",
+    "liger_kernel==0.5.2",
+    "math-verify==0.5.1",  # Used for math verification in grpo
+    "torch==2.5.1",
+    "trl==0.14.0",
+    "vllm==0.7.1",
 ]
 
 # this is a lookup table with items like:

setup.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+
+REQUIRED_PREFIX="12.1"
+CUDA_FULL_VERSION=$(nvcc --version | grep "release" | awk '{print $6}' | cut -c2-)
+
+# Only stable for CUDA 12.1 right now.
+if [[ "$CUDA_FULL_VERSION" != ${REQUIRED_PREFIX}* ]]; then
+  echo "Error: CUDA version starting with $REQUIRED_PREFIX is required, but found CUDA version $CUDA_FULL_VERSION."
+  exit 1
+fi
+
+echo "---"
+echo "Using CUDA $CUDA_FULL_VERSION"
+
+echo "---"
+echo "Setting up vllm and nvJitLink."
+pip install "vllm>=0.7.1" --extra-index-url https://download.pytorch.org/whl/cu121 2>&1
+export LD_LIBRARY_PATH=$(python -c "import site; print(site.getsitepackages()[0] + '/nvidia/nvjitlink/lib')"):$LD_LIBRARY_PATH
+
+echo "---"
+echo "Installing dependencies."
+pip install -U -e ".[dev]"
+
+echo "---"
+echo "Finished setup. GRPO example (expects 8 accelerators):
+
+    HF_HUB_ENABLE_HF_TRANSFER=1 ACCELERATE_LOG_LEVEL=info accelerate launch \\
+        --config_file recipes/accelerate_configs/zero3.yaml \\
+        --num_processes=7 \\
+        src/open_r1/grpo.py \\
+        --config recipes/qwen/Qwen2.5-1.5B-Instruct/grpo/confg_full.yaml
+"