repository·main·Indexed 12 days ago
https://github.com/huggingface/open-r1A project for training and distilling reasoning models, including OpenR1 Distill 7B and OlympicCoder (7B and 32B). It features tools for dataset pass rate filtering, GRPO and SFT training configurations, and reward functions for enforcing reasoning formats (<think> and <answer> tags). The project also includes Piston workers for running IOI and CodeForces problems in isolated Slurm or Docker environments and SGLang integration for serving DeepSeek-R1.
The pass rate filtering feature allows you to filter datasets by generating and computing pass rates on variable tasks. The process involves chunking the dataset for computation.
Key scripts for this process:
scripts/pass_rate_filtering/compute_pass_rate.py: The core logic for computing pass rates.scripts/pass_rate_filtering/launch_filtering.sh: A shell script to launch the filtering process (currently hardcoded for DAPO).To initiate training for the OpenR1 Distill 7B model using Slurm, use the sbatch command with the slurm/train.slurm script. This requires setting the --model to OpenR1-Distill-7B, the --task to sft, and using the distill configuration with the zero3 accelerator.
sbatch --nodes=1 slurm/train.slurm --model OpenR1-Distill-7B --task sft --config distill --accelerator zero3To train the 7B version of the OlympicCoder model, run the Slurm batch job on a single node using the sft task, the v00.00 configuration, and the zero3 accelerator.
sbatch --nodes=1 slurm/train.slurm --model OlympicCoder-7B --task sft --config v00.00 --accelerator zero3To run a single worker locally using Docker, use the docker run command. You must map a local directory to /piston/packages to persist package installations and specify the desired port. The following command includes configuration for timeouts, file size limits, and networking restrictions.
docker run -d \
--name piston_worker \
-v /path/to/local/packages:/piston/packages \
-e PORT=2000 \
-e PISTON_COMPILE_TIMEOUT=60000 \
-e PISTON_RUN_TIMEOUT=60000 \
-e PISTON_OUTPUT_MAX_SIZE=1000000000 \
-e PISTON_MAX_FILE_SIZE=1000000000 \
-e PISTON_DISABLE_NETWORKING=true \
-e PISTON_REPO_URL=https://github.com/guipenedo/piston/releases/download/pkgs/index \
-p 2000:2000 \
--entrypoint /bin/bash \
ghcr.io/engineer-man/piston@sha256:63b5654156a89c5a2ad281aface21416615d62ec056d88efe8fcd307ce73575a \
-c "sed -i '/app.use(body_parser.urlencoded/c\ app.use(body_parser.urlencoded({ extended: true, limit: \"512mb\" }));' src/index.js && \
sed -i '/app.use(body_parser.json/c\ app.use(body_parser.json({ limit: \"512mb\" }));' src/index.js && \
node src"To serve DeepSeek-R1 on SLURM nodes using SGLang, create a dedicated Conda environment and install the required dependencies. Note that the CUDA version in the URLs (cu124) should be adjusted to match your specific environment if necessary.
Steps:
torch 2.5.1 with the appropriate CUDA index.sgl-kernel.sglang[all] with the specific FlashInfer wheels for your CUDA/Torch version.conda create -n sglang124 python=3.11
conda activate sglang124
pip install torch==2.5.1 --index-url https://download.pytorch.org/whl/cu124
pip install sgl-kernel --force-reinstall --no-deps
pip install "sglang[all]>=0.4.2.post4" --find-links https://flashinfer.ai/whl/cu124/torch2.5/flashinfer/Use sbatch to submit the slurm/serve_r1.slurm script to your cluster. You must specify the model checkpoint path and the Conda environment name used during setup.
Arguments:
-m: The path to the DeepSeek-R1 checkpoint (e.g., /fsx/deepseek-r1-checkpoint).-e: The name of the Conda environment to use (e.g., sglang124).sbatch slurm/serve_r1.slurm -m "/fsx/deepseek-r1-checkpoint" -e "sglang124"To deploy a fleet of Piston workers on a Slurm cluster, use the provided shell scripts. You may need to adapt the paths within launch_piston_workers.sh and launch_single_piston.sh before execution. Each worker is launched as a separate Slurm job named piston-worker-<port>, where <port> is the listening port.
slurm/piston/launch_piston_workers.sh (number of workers to launch)Piston workers require specific packages to run IOI or CodeForces problems. After launching a worker, you must send a POST request to the /api/v2/packages endpoint to trigger the installation. Once installed, packages are stored in a shared mounted directory, so subsequent workers will already have them available.
# For IOI
curl -X POST http://<worker-ip>:<port>/api/v2/packages -H "Content-Type: application/json" -d '{"language": "cms_ioi", "version": "1.0.0"}'
# For CodeForces
curl -X POST http://<worker-ip>:<port>/api/v2/packages -H "Content-Type: application/json" -d '{"language": "codeforces", "version": "1.0.0"}'To train the 32B version of the OlympicCoder model, you must use 16 nodes. For this larger model, use the fsdp accelerator instead of zero3.
Note: For the 32B model, it is necessary to use FSDP1 and paged AdamW 8-bit to accommodate the largest possible context size.
sbatch --nodes=16 slurm/train.slurm --model OlympicCoder-32B --task sft --config v00.00 --accelerator fsdpBecause the filtering script chunks the dataset by default, you must manually merge the chunks back together. To do this, iterate through the generated configuration names (e.g., gen-{start}-{end}) and the filtered configuration names (e.g., filt-0.1-0.6-{start}-{end}) using load_dataset with the appropriate revision and config_name, then use concatenate_datasets to combine them. Finally, use push_to_hub to save the merged datasets back to the Hugging Face Hub.
Note: The example below uses the gen revision for the original datasets and the pass_rate revision for the filtered datasets.
from datasets import load_dataset, concatenate_datasets
name = "open-r1/DAPO-Math-17k-Processed-R1-Distill-Qwen-Math-7B-Merges-v00.02-v01.02-0.3-0.7-filter"
gen_datasets = []
filt_datasets = []
for start in range(0,17400,200):
end = start + 200
if start == 17200:
end = 17398
gen_config_name = f"gen-{start}-{end}"
gen_dataset = load_dataset(name, gen_config_name, revision="gen", split="train")
gen_datasets.append(gen_dataset)
filt_config_name = f"filt-0.1-0.6-{start}-{end}"
filt_dataset = load_dataset(name, filt_config_name, revision="pass_rate", split="train")
filt_datasets.append(filt_dataset)
gen_dataset = concatenate_datasets(gen_datasets)
gen_dataset.push_to_hub(name, config_name="gen", split="train")
print(gen_dataset)
filt_dataset = concatenate_datasets(filt_datasets)
filt_dataset.push_to_hub(name, config_name="default", split="train")
print(filt_dataset)The project uses a provider pattern to abstract code execution across different sandboxed environments.
CodeExecutionProvider (Interface): An abstract base class defining the execute_scripts(scripts: List[str], languages: List[str]) -> List[float] contract. All providers must implement this.E2BProvider: Implements execution via E2B sandboxes.MorphProvider: Implements execution via MorphCloud.execute_scripts is called, the provider manages the lifecycle of sandboxes (creation, execution, and destruction/killing) and handles concurrency using asyncio.Semaphore to respect the num_parallel limit.List[float]. They attempt to parse the numerical result from the sandbox output (stdout or text) to serve as a reward signal.The client automatically determines the task_type for the execution based on the provided file list:
manager.cpp is found within the graders/ directory, the task_type is set to Communication. This enables multi-process testing with specific parameters like num_processes and user_io.manager.cpp is present, the task_type is set to Batch.Problem ID Extraction:
The problem_id is derived from the filename of any .cpp file located in the graders/ directory (excluding grader.cpp, manager.cpp, or stub.cpp).