This repository contains the training, inference, and scoring code for the paper Leveraging Instruction Tuning and Merging for Reasoning Model Adaptation.
The full experiments require Linux, NVIDIA GPUs, Docker with NVIDIA Container Toolkit, and Rust. macOS is suitable for development, tests, and inspecting results, but not for the CUDA/vLLM experiment runs.
Install system dependencies on Ubuntu:
sudo apt-get update
sudo apt-get install -y build-essential cargo docker.io git rustc
sudo usermod -aG docker "$USER"Verify GPU access from Docker:
docker run --rm --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smiInstall the locked Python environment and the generation image:
uv sync
docker pull vllm/vllm-openai:latestLog in to Hugging Face if a model is gated:
uv run hf auth loginText summarization scoring uses Gemini and requires a Google API key. Create a secrets file from the tracked template:
cp secret.sh.example secret.sh
# Edit secret.sh, then load it into the current shell:
source secret.shRun the CPU-only test suite:
uv run pytest -qThe checked-in experiment datasets are:
| File | Purpose |
|---|---|
datasets/training_set_filtered_code_max.jsonl |
Rust IFT training |
datasets/rust_dataset_true_passed_test750_max.jsonl |
Rust merge calibration |
datasets/mbpp_rs_final_dataset.jsonl |
Rust test |
datasets/TLDR_train_filtered_10k.jsonl |
summarization IFT training |
datasets/TLDR_validate_filtered_400.jsonl |
summarization merge calibration |
datasets/summary_test_set_cd_unique.jsonl |
SummEval test |
The paper runner deterministically creates the 2,048-prompt Tulu-3 subset used
for OPD from allenai/tulu-3-sft-mixture with seed 20260529.
The retained scripts under scripts/*_construction.py recreate the Rust test
set and the summarization train/test sets from their upstream datasets. They
are not run automatically because the checked-in JSONL files are the exact
filtered splits used by the paper.
Train a LoRA adapter:
uv run python -m src.training.sft \
--task rust \
--base-model open-thoughts/OpenThinker-7B \
--train-file datasets/training_set_filtered_code_max.jsonl \
--output-dir outputs/example/sft \
--epoch 5 \
--lr 5e-5 \
--batch-size 2 \
--gradient-accumulation-steps 8 \
--validation-size 0Add --kl-coefficient 5 --kl-mask-mode thinking for IFT+KL. Add
--full-finetune for the paper's full-finetuning ablation.
src.run_sft_auto_sweep provides the same trainer with multiple seeds,
learning rates, GPU allocation, checkpoint discovery, and resumable outputs.
For a LoRA IFT model, scaling the LoRA delta is exactly linear interpolation between the base and IFT weights. Search for the largest passing scale:
uv run python -m src.run_alpha_binary_search \
--base-model open-thoughts/OpenThinker-7B \
--peft-adapter outputs/example/sft \
--dataset-path datasets/rust_dataset_true_passed_test750_max.jsonl \
--output-dir outputs/example/search \
--prompt-task rust \
--steps 8 \
--threshold 90 \
--limit 750For full-model linear, SLERP, or TIES merging, use
src.run_merge_eval_pipeline. It searches candidates, saves the selected
merge, evaluates it, and resumes completed JSONL outputs.
uv run python scripts/prepare_opd_dataset.py \
--output datasets/on_policy_distill/tulu3_random_sft_seed20260529_2048.jsonl
uv run python -m src.training.on_policy_distill \
--base-model outputs/example/baked-sft \
--teacher-model open-thoughts/OpenThinker-7B \
--train-file datasets/on_policy_distill/tulu3_random_sft_seed20260529_2048.jsonl \
--task chat \
--messages-field messages \
--strip-final-assistant \
--output-dir outputs/example/opd \
--steps 25 \
--prompts-per-step 8 \
--samples-per-prompt 1 \
--lr 2e-4src.run_vllm_eval starts a vLLM Docker server, generates responses, stops the
server, and then scores the saved responses. The same command resumes partial
generation and scoring files.
Rust plus MATH-500:
uv run python -m src.run_vllm_eval \
--task both \
--base-model open-thoughts/OpenThinker-7B \
--lora-adapter outputs/example/sft \
--adapter-scale 0.25 \
--model-name example-merge \
--output-dir outputs/example/eval \
--gpu 0 \
--eval-dataset datasets/mbpp_rs_final_dataset.jsonl \
--seeds 0,1,2,3,4,5,6,7,8,9 \
--max-model-len 32768 \
--max-tokens 16384Summarization:
uv run python -m src.run_vllm_eval \
--task summarization \
--base-model open-thoughts/OpenThinker-7B \
--lora-adapter outputs/example/sft \
--adapter-scale 0.25 \
--model-name example-summary-merge \
--output-dir outputs/example/eval \
--gpu 0 \
--eval-dataset datasets/summary_test_set_cd_unique.jsonl \
--seeds 0,1,2,3,4 \
--max-model-len 32768 \
--max-tokens 16384Raw generations and scored rows are stored separately:
rust_generations.jsonl rust_scored.jsonl
math500_generations.jsonl math500_scored.jsonl
summarization_generations.jsonl summarization_scored.jsonl
Rust scoring compiles and runs each generated solution locally. MATH-500 uses symbolic answer extraction. Summarization sends each generated summary to the configured Gemini judge.
All reported experiments are described by exactly three configs:
| Config | Contents |
|---|---|
configs/paper/main.toml |
the eight model/task combinations, both seeds, merge searches, main evaluations, and merge-ratio grids |
configs/paper/baselines.toml |
IFT+KL and IFT+OPD |
configs/paper/ablations.toml |
OpenThinker hyperparameters, full fine-tuning, SLERP, and TIES |
Inspect the complete expanded command list without running it:
uv run python scripts/reproduce_paper.py --dry-runRun everything in dependency order:
uv run python scripts/reproduce_paper.pySelect a subset by experiment or stage:
uv run python scripts/reproduce_paper.py --only 'openthinker_*'
uv run python scripts/reproduce_paper.py --stage 'evaluate*'
uv run python scripts/reproduce_paper.py \
--config configs/paper/main.toml \
--only 'r1_qwen_text*'The runner prints every concrete command before execution. Training, generation, and scoring commands are append/resume-oriented, so rerunning the same selection continues from completed artifacts.
@inproceedings{
feng2026leveraging,
title={Leveraging Instruction Tuning and Merging for Reasoning Model Adaptation},
author={Yu-Du Feng and Niels M{\"u}ndler and Mark Vero and Martin Vechev},
booktitle={Decision-Making from Offline Datasets to Online Adaptation: Black-Box Optimization to Reinforcement Learning},
year={2026},
url={https://openreview.net/forum?id=Ja67cQ7t0G}
}