basetenlabs/relex-replication

basetenlabs/relex-replication

Description: Minimal code to replicate our RELEX (rank-1 RL extrapolation) experiments on MATH, IFEval, Knights & Knaves, and function calling

Language: Python

Stars: 0

Forks: 0

Open issues: 0

Created: 2026-10-02T20:53:59Z

Pushed: 2026-10-05T00:35:36Z

Default branch: main

Fork: no

Archived: no

README:

RELEX replication

Replicates and stress-tests RELEX from You Only Need Minimal RLVR Training: Extrapolating LLMs via Rank-1 Trajectories (authors' code) on MATH with Qwen2.5-Math-1.5B, and on MATH, Knights & Knaves (kk), IFEval (ifeval) and function calling (fc: xLAM training, BFCL evaluation) with Qwen3-4B-Base and Qwen3-8B-Base.

Setup and training

Run commands from this repository's root. The setup environment is Python 3.12, Torch 2.13 with CUDA 13.0.3, and vllm/vllm-openai:v0.28.0 (image digest sha256:61fc8a896b0a4fbbbdc063bc4b0dbc25ce98e02b5050c24aeb7830ac02039b14). Use that image on Linux GPU hosts, then install the requirements. The repository is run directly from source; pip install . is not the supported setup.

python -m pip install -r requirements.txt
export PYTHONHASHSEED=1 VLLM_WORKER_MULTIPROC_METHOD=spawn TOKENIZERS_PARALLELISM=false
export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1

For CPU tests and extrapolation (including on macOS), install requirements-cpu.txt instead and run python -m pytest -q. macOS requires version 14 or newer. Training and generation require CUDA GPUs; CPU tests do not validate those paths. The original results used Torch 2.10, CUDA 12.9 and vLLM 0.19.0; the upgraded GPU environment has not yet been validated against those runs. Data and tokenizers download automatically at their pinned revisions. xLAM is gated: accept its terms on Hugging Face and set HF_TOKEN before function-calling training.

| Recipe | Config | Environment | Training hardware | Snapshot weights alone | |---|---|---|---|---| | Qwen2.5-Math-1.5B | configs/qwen25_math_1.5b.json | math | 2 nodes × 8 H200 | ~230 GiB | | Qwen3-4B-Base | configs/qwen3_4b.json | math, kk, ifeval, fc | 2 nodes × 8 H200 | ~600 GiB | | Qwen3-8B-Base | configs/qwen3_8b.json | kk, ifeval, fc | 1 node × 8 B200 | ~920 GiB |

Allow additional disk for the base model, two full optimizer checkpoints, download caches and derived models. For the 8B runs, plan at least 1.5 TiB per experiment. Each environment is a separate training run. The two-node runs require a shared output filesystem at the same path on both nodes, plus network connectivity to node 0. Keep the reference GPU count to preserve data sharding.

For 1.5B, set these variables on both nodes, using machine rank 0 on the first and 1 on the second. For 4B, change C, E, and RUN to the chosen row above and a distinct output path.

C=configs/qwen25_math_1.5b.json
E=math
RUN=/shared/runs/1.5b-math
# Set MAIN_PROCESS_IP to node 0's reachable address and MACHINE_RANK to 0 or 1.
: "${MAIN_PROCESS_IP:?set the address of node 0}"
: "${MACHINE_RANK:?set 0 on node 0 or 1 on node 1}"
accelerate launch --multi_gpu --num_machines 2 --num_processes 16 \
--machine_rank "$MACHINE_RANK" --main_process_ip "$MAIN_PROCESS_IP" --main_process_port 29500 \
--mixed_precision bf16 --dynamo_backend no \
-m relex_replication.train --config "$C" --env "$E" --output "$RUN"

For 8B, run once on its eight-B200 node:

C=configs/qwen3_8b.json
E=kk # kk, ifeval, or fc
RUN=/mnt/checkpoints/runs/8b-$E
accelerate launch --multi_gpu --num_machines 1 --num_processes 8 --machine_rank 0 \
--mixed_precision bf16 --dynamo_backend no \
-m relex_replication.train --config "$C" --env "$E" --output "$RUN"

Training writes base_model/, trajectory/global_step_N/, resume/checkpoint-N/ and the resolved recipe in run.json. To resume, repeat the same launch command with --resume "$RUN/resume/checkpoint-N", using a completed checkpoint and the original --output. Optimizer, scheduler and Trainer RNG state resume; vLLM sampling is not guaranteed bitwise identical after a restart. Fresh training cannot overwrite an existing trajectory.

Qwen2.5-Math-1.5B experiments

A fresh trajectory reruns the method; reproducing the recorded numbers requires the original saved weights. The original 1.5B producer is Baseten project qklp483, job w5j4g73, checkpoint prefix relex_qwen25_math15b_method/ (access required). Retain its base_model/ and trajectory/global_step_N/ layout and point RUN there. Original Qwen3 producers are 4B: K&K w7r5463, IFEval wx2zn73, FC 32o07vw; 8B: K&K wp1rv9w, IFEval q498j1w, FC q8zem8q. Those archives may require exporting their checkpoints into this layout before using --run. The reference repository's source manifests describe their archive paths.

After training or restoring the 1.5B trajectory:

RUN=/shared/runs/1.5b-math
python -m relex_replication.extrapolate --run "$RUN" --prefix 2 --output models/1.5b-relex2
python -m relex_replication.extrapolate --run "$RUN" --alpha 1000 --output models/1.5b-alpha1000

# MATH evaluation uses one GPU. Choose the GPU with CUDA_VISIBLE_DEVICES if needed.
python -m relex_replication.evaluate --model "$RUN/base_model" --output evals/1.5b-base
python -m relex_replication.evaluate --model "$RUN/trajectory/global_step_500" --output evals/1.5b-step500
python -m relex_replication.evaluate --model models/1.5b-alpha1000 --output evals/1.5b-alpha1000
python -m relex_replication.evaluate --model models/1.5b-relex2 --output evals/1.5b-relex2
python -m relex_replication.evaluate --model relex-rlvr/RELEX-Qwen2.5-Math-1.5B \
--revision 458878c19dd5abbdbf89de91f7267ad249fca00a --output evals/1.5b-released-relex
python -m relex_replication.analyze --base evals/1.5b-base --target evals/1.5b-relex2

For the original full sweeps, repeat the build/evaluate pair at prefixes 2 5 10 25 50 75 and alphas 0 1 10 30 100 300 500 1000 3000, with a distinct model and evaluation directory per setting. Extrapolation runs on CPU and needs the base plus every checkpoint through the selected prefix. A build requires a new output directory and currently restarts from scratch if interrupted. Evaluation resumes completed batches only when the model hashes and resolved settings match. It writes evaluation.json, raw generations.jsonl, scores.jsonl and summary.json.

| Model | MATH test...