claude-opus-5

<persisted-output>
Output too large (37.2KB). Full output saved to: /logs/agent/sessions/projects/-workspace/3476b637-41e1-4437-8c9a-f53302e59a44/tool-results/bgrymjl8d.txt

Preview (first 2KB):
156 bash -c  rm -f /tmp/ri-exec-bbc4fda293b04c1ebd2462ad564faff4.pid setsid bash -c 'set -o pipefail; export PATH="$HOME/.local/bin:/usr/local/bin:$PATH"; claude --verbose --output-format=stream-json --permission-mode=bypassPermissions --disallowedTools WebSearch WebFetch --print -- '"'"'# Task: Improve On-Policy Self-Distillation (OPSD) on AIME 2024  ## Context You are given the code repository for "Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models" (arXiv 2601.18734v3) at `/workspace/repo`, and the paper at `/workspace/paper.pdf`.  The repo implements **OPSD**, the paper'"'"'"'"'"'"'"'"'s main approach: a single model acts as both student (sees only the problem) and teacher (additionally sees the ground-truth solution), and training minimizes a per-token forward-KL between the two distributions along the student'"'"'"'"'"'"'"'"'s own rollouts, with a fixed initial-policy teacher.  A baseline checkpoint trained with the paper'"'"'"'"'"'"'"'"'s released recipe (`/workspace/train_baseline.sh`, the 100-step Qwen3-1.7B configuration whose AIME24 number the paper publishes) is provided read-only at `/workspace/baseline/checkpoint/`, with its full training logs at `/workspace/baseline/logs/`. This is the paper'"'"'"'"'"'"'"'"'s native 4×H100 recipe (`num_processes 4`, `per_device 4`, `gradient_accumulation 2` → global batch 32, colocate vLLM util 0.6), evaluated data-parallel across the 4 GPUs (4 single-GPU vLLM engines, TP=1 each). Re-trained 5 times at the fixed scoring seed, this baseline scores **54.22 ± 1.48** (avg@12 accuracy) on AIME 2024.  ## Your goal Propose and implement a research idea — a **method** improvement to OPSD — that beats the paper'"'"'"'"'"'"'"'"'s own OPSD recipe **within a fixed training budget of 100 optimizer steps** (see "Hard constraints"). You are scored on the **absolute AIME24 avg@12 accuracy** of your re-trained method; **higher is better**, so the more your re-trained method'"'"'"'"'"'"'"'"'s avg@12 exceeds 
...
</persisted-output>

[stdout]
156 bash -c  rm -f /tmp/ri-exec-bbc4fda293b04c1ebd2462ad564faff4.pid setsid bash -c 'set -o pipefail; export PATH="$HOME/.local/bin:/usr/local/bin:$PATH"; claude --verbose --output-format=stream-json --permission-mode=bypassPermissions --disallowedTools WebSearch WebFetch --print -- '"'"'# Task: Improve On-Policy Self-Distillation (OPSD) on AIME 2024  ## Context You are given the code repository for "Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models" (arXiv 2601.18734v3) at `/workspace/repo`, and the paper at `/workspace/paper.pdf`.  The repo implements **OPSD**, the paper'"'"'"'"'"'"'"'"'s main approach: a single model acts as both student (sees only the problem) and teacher (additionally sees the ground-truth solution), and training minimizes a per-token forward-KL between the two distributions along the student'"'"'"'"'"'"'"'"'s own rollouts, with a fixed initial-policy teacher.  A baseline checkpoint trained with the paper'"'"'"'"'"'"'"'"'s released recipe (`/workspace/train_baseline.sh`, the 100-step Qwen3-1.7B configuration whose AIME24 number the paper publishes) is provided read-only at `/workspace/baseline/checkpoint/`, with its full training logs at `/workspace/baseline/logs/`. This is the paper'"'"'"'"'"'"'"'"'s native 4×H100 recipe (`num_processes 4`, `per_device 4`, `gradient_accumulation 2` → global batch 32, colocate vLLM util 0.6), evaluated data-parallel across the 4 GPUs (4 single-GPU vLLM engines, TP=1 each). Re-trained 5 times at the fixed scoring seed, this baseline scores **54.22 ± 1.48** (avg@12 accuracy) on AIME 2024.  ## Your goal Propose and implement a research idea — a **method** improvement to OPSD — that beats the paper'"'"'"'"'"'"'"'"'s own OPSD recipe **within a fixed training budget of 100 optimizer steps** (see "Hard constraints"). You are scored on the **absolute AIME24 avg@12 accuracy** of your re-trained method; **higher is better**, so the more your re-trained method'"'"'"'"'"'"'"'"'s avg@12 exceeds the baseline (**54.22**), the better you do. Simply training longer/bigger is not available to you — the budget is fixed — so improvements must come from the method itself (the loss, the teacher, data selection/weighting, the on-policy scheme, etc.).  ## Deliverables Your contribution is a **training method, not a checkpoint.** You produce it by: 1. **Editing the OPSD training code in `/workspace/repo`** — the loss, the    teacher construction, data selection/weighting, the collator, the on-policy    sampling, hyper-parameters in code, etc. This is where your research idea    lives. `opsd_train.py` must remain the training entrypoint. 2. *(optional)* **`/workspace/submission/recipe.env`** — whitelisted    hyper-parameter overrides, one `KEY=VALUE` per line (see the whitelist under    "Hard constraints"). An absent/empty recipe reproduces the baseline    hyper-parameters on top of whatever code changes you made.  There is **no checkpoint to submit.** After your session, the scorer re-trains your method from the frozen base under the fixed budget and evaluates it (see "Official evaluation"); any checkpoint you leave behind is ignored.  Develop and test exactly as the scorer will, with: `SEED=42 OUT_DIR=/workspace/runs/try1 bash /workspace/train_capped.sh` then `bash /workspace/dev_eval.sh /workspace/runs/try1/capped_seed42/checkpoint-100`. `train_capped.sh` is the dev mirror of the scorer'"'"'"'"'"'"'"'"'s training (same fixed budget, your repo code, your recipe.env); do not modify it — the scorer uses its own trusted copy, so edits only make your dev numbers diverge from the official ones.  ## Resources and budget - Hardware: 4×H100 for 6 hours total (your session ends when the budget   expires). This budget covers everything you run: diagnostic experiments,   training, and your own intermediate evaluations. The official scoring run   happens after your session and is not charged to you. For calibration: one   baseline training run takes ~19 min and one full dev evaluation ~27 min on   the 4 GPUs (LoRA submissions are merged into the base weights before   generation, so they evaluate at full-model speed). - Time remaining is tracked for you: read `/workspace/.timer/remaining_secs`   (and `/workspace/.timer/elapsed_secs`) at any point instead of doing   wall-clock arithmetic yourself. - Base model: Qwen3-1.7B, pre-downloaded at `/opt/models/Qwen3-1.7B`   (revision-pinned; do not modify these files). - Training data: `siyanzhao/Openthoughts_math_30k_opsd` (the paper'"'"'"'"'"'"'"'"'s dataset,   revision-pinned, available offline via   `load_dataset("siyanzhao/Openthoughts_math_30k_opsd")`). You may filter or   re-weight examples from this dataset. You may NOT add external data, and   there is no API budget for synthetic data generation. - Dev evaluation: `bash /workspace/dev_eval.sh <checkpoint_dir|base>` runs the   same evaluation as the official one (same code, same frozen settings, same   sampling seed). Use it for your own checks; it consumes your budget. Cheaper   partial evals: `VAL_N=4 NUM_PROBLEMS=15 bash /workspace/dev_eval.sh ...`   (the official run always uses all 30 problems and 12 samples). - **Evaluation noise (read this before you trust a number).** The benchmark is   30 problems; the metric is an average over 12 samples each. Partial evals are   cheap but noisy: a 15-problem / avg@4 partial has a standard error of roughly   **±4–5 points** and has repeatedly misled past attempts by 5+ points; only the   full 30-problem / avg@12 eval on the fixed sampling seed is meaningful.   Training itself is **not seed-deterministic**: re-running the *same* recipe at   the *same* seed still moves avg@12 by **~1.5 points (1σ)** (the on-policy   rollouts are sampled with vLLM, plus CUDA nondeterminism). The scorer absorbs   this by re-training your method **5 times at the fixed seed and averaging** the   avg@12 (see "How you are scored"), which shrinks the *scored* noise to **~0.66   (1σ)** — but any single dev train+eval *you* run carries the full ~1.5. So do   not trust a single dev run: a margin under ~1.5 on one run may be noise. Budget   for **repeated dev train+evals of your final candidate** and target a gain that   clearly survives the ~1.5 single-run spread. - Internet access: none. Everything needed (repo, paper, model, datasets,   Python environment) is already inside the container.  ## Official evaluation (how your method is scored) After your session ends, the scorer **re-trains your method from the frozen Qwen3-1.7B base** using its own trusted copy of `train_capped.sh`: your `/workspace/repo` code, your `recipe.env` method knobs, and a **fixed budget of 100 optimizer steps at global batch 32** (see "Hard constraints"). It does this **5 times at the same fixed seed (42)**, evaluates each resulting checkpoint on AIME 2024 (30 problems) with these frozen settings — temperature=1.0, top_p=0.95, top_k disabled, min_p=0, presence_penalty=0, max_new_tokens=38912, Qwen3 thinking mode enabled, 12 samples per problem (metric: avg@12 accuracy), fixed sampling seed, vLLM with the frozen Qwen3-1.7B tokenizer/chat template, data-parallel across the 4 GPUs (4 single-GPU engines, tensor_parallel_size=1 each; problems sharded round-robin and merged) — and **averages the 5 avg@12 scores**. The training and evaluation code is fixed and runs outside your session; only your method (repo code + recipe.env) matters. Because your method is re-trained 5 times at the fixed seed and the avg@12 is averaged (not your submitted checkpoint), the scored noise is ~0.66 (1σ); aim for a method whose gain clearly exceeds that rather than a lucky run.  ## Hard constraints - **Fixed training budget — you cannot change it.** The scorer always trains for   exactly **100 optimizer steps** at **global batch 32** (num_processes 4 ×   per_device_train_batch_size 4 × gradient_accumulation_steps 2) from the frozen   base. These budget knobs — `max_steps`, `num_train_epochs`,   `per_device_train_batch_size`, `gradient_accumulation_steps`, `num_processes`,   `model_name_or_path`, and the dataset — are set by the scorer, and any value   you put for them in `recipe.env` (or hard-code to exceed 100 steps: extra   epochs, resuming/continuing a checkpoint, multi-stage training that sums to   >100 steps) is ignored or disallowed. Your improvement must come from the   **method within this fixed budget**, not from training longer or bigger. - **Train fresh from the frozen base.** Your code must actually train the 100   steps from `/opt/models/Qwen3-1.7B`; do not load, cache, or copy in   externally-pretrained or previously-trained weights. - **`recipe.env` whitelist.** Only these method knobs are honored (others are   ignored): `learning_rate`, `lora_r`, `lora_alpha`, `lora_dropout`, `beta`,   `jsd_token_clip`, `top_k_loss`, `temperature`, `top_p`, `top_k`, `lmbda`,   `max_completion_length` (≤ 4096), `max_grad_norm`, `weight_decay`,   `warmup_ratio`, `lr_scheduler_type`, `ema_decay`, and the boolean flags   `fixed_teacher`, `use_ema_teacher`, `use_tinker_loss`, `reason_first`,   `teacher_thinking`, `student_thinking` (set to `true`/`false`). - Do not modify the benchmark data (`/opt/eval_assets/aime24.json`), the frozen   base model files (`/opt/models/Qwen3-1.7B`), or the training-dataset files (HF   cache under `/opt/hf_cache`). Integrity is verified before scoring; any   mismatch scores 0. You may filter/re-weight the training dataset **in code**. - No inference-time changes: your improvement must live in the trained model,   evaluated under the same frozen inference settings as the baseline (no extra   samples, no ensembling/voting, no decoding, prompt, or chat-template changes —   the official eval always uses the frozen base tokenizer and chat template). - Train the designated model only: Qwen3-1.7B. LoRA rank must be ≤ 256.  ## How you are scored Your method is scored on the **absolute AIME24 avg@12 accuracy** of its re-trains. The scorer trains your method **5 times at the same fixed seed (42)** under the fixed 100-step budget, evaluates each on AIME24 (avg@12), and reports the **mean of the 5 avg@12 scores**. **Higher is better.** The verifier reports that number directly — it does not subtract a baseline or otherwise transform it.  For reference, the benchmark records two fixed anchors alongside the score: a **baseline** of **54.22** avg@12 (the OPSD released recipe, measured the same way — 5 same-seed re-trains of the unmodified recipe, averaged) and a **theoretical best** of **100%** (a perfect AIME24). These anchors are used only by a separate downstream normalization step to place scores on a common scale; they do not change what the verifier measures.  The 5-run average has **~0.66 (1σ)** noise (the per-run avg@12 spread is ~1.48), so target a gain that clearly exceeds it rather than a lucky run. A method that fails to train, a training-longer/bigger workaround, or any tampering with the frozen inputs is marked invalid (no score) and excluded. '"'"' 2>&1 </dev/null | tee /logs/agent/claude-code.txt' & child="$!" echo "$child" > /tmp/ri-exec-bbc4fda293b04c1ebd2462ad564faff4.pid wait "$child" rc="$?" rm -f /tmp/ri-exec-bbc4fda293b04c1ebd2462ad564faff4.pid exit "$rc" 
158 bash -c set -o pipefail; export PATH="$HOME/.local/bin:/usr/local/bin:$PATH"; claude --verbose --output-format=stream-json --permission-mode=bypassPermissions --disallowedTools WebSearch WebFetch --print -- '# Task: Improve On-Policy Self-Distillation (OPSD) on AIME 2024  ## Context You are given the code repository for "Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models" (arXiv 2601.18734v3) at `/workspace/repo`, and the paper at `/workspace/paper.pdf`.  The repo implements **OPSD**, the paper'"'"'s main approach: a single model acts as both student (sees only the problem) and teacher (additionally sees the ground-truth solution), and training minimizes a per-token forward-KL between the two distributions along the student'"'"'s own rollouts, with a fixed initial-policy teacher.  A baseline checkpoint trained with the paper'"'"'s released recipe (`/workspace/train_baseline.sh`, the 100-step Qwen3-1.7B configuration whose AIME24 number the paper publishes) is provided read-only at `/workspace/baseline/checkpoint/`, with its full training logs at `/workspace/baseline/logs/`. This is the paper'"'"'s native 4×H100 recipe (`num_processes 4`, `per_device 4`, `gradient_accumulation 2` → global batch 32, colocate vLLM util 0.6), evaluated data-parallel across the 4 GPUs (4 single-GPU vLLM engines, TP=1 each). Re-trained 5 times at the fixed scoring seed, this baseline scores **54.22 ± 1.48** (avg@12 accuracy) on AIME 2024.  ## Your goal Propose and implement a research idea — a **method** improvement to OPSD — that beats the paper'"'"'s own OPSD recipe **within a fixed training budget of 100 optimizer steps** (see "Hard constraints"). You are scored on the **absolute AIME24 avg@12 accuracy** of your re-trained method; **higher is better**, so the more your re-trained method'"'"'s avg@12 exceeds the baseline (**54.22**), the better you do. Simply training longer/bigger is not available to you — the budget is fixed — so improvements must come from the method itself (the loss, the teacher, data selection/weighting, the on-policy scheme, etc.).  ## Deliverables Your contribution is a **training method, not a checkpoint.** You produce it by: 1. **Editing the OPSD training code in `/workspace/repo`** — the loss, the    teacher construction, data selection/weighting, the collator, the on-policy    sampling, hyper-parameters in code, etc. This is where your research idea    lives. `opsd_train.py` must remain the training entrypoint. 2. *(optional)* **`/workspace/submission/recipe.env`** — whitelisted    hyper-parameter overrides, one `KEY=VALUE` per line (see the whitelist under    "Hard constraints"). An absent/empty recipe reproduces the baseline    hyper-parameters on top of whatever code changes you made.  There is **no checkpoint to submit.** After your session, the scorer re-trains your method from the frozen base under the fixed budget and evaluates it (see "Official evaluation"); any checkpoint you leave behind is ignored.  Develop and test exactly as the scorer will, with: `SEED=42 OUT_DIR=/workspace/runs/try1 bash /workspace/train_capped.sh` then `bash /workspace/dev_eval.sh /workspace/runs/try1/capped_seed42/checkpoint-100`. `train_capped.sh` is the dev mirror of the scorer'"'"'s training (same fixed budget, your repo code, your recipe.env); do not modify it — the scorer uses its own trusted copy, so edits only make your dev numbers diverge from the official ones.  ## Resources and budget - Hardware: 4×H100 for 6 hours total (your session ends when the budget   expires). This budget covers everything you run: diagnostic experiments,   training, and your own intermediate evaluations. The official scoring run   happens after your session and is not charged to you. For calibration: one   baseline training run takes ~19 min and one full dev evaluation ~27 min on   the 4 GPUs (LoRA submissions are merged into the base weights before   generation, so they evaluate at full-model speed). - Time remaining is tracked for you: read `/workspace/.timer/remaining_secs`   (and `/workspace/.timer/elapsed_secs`) at any point instead of doing   wall-clock arithmetic yourself. - Base model: Qwen3-1.7B, pre-downloaded at `/opt/models/Qwen3-1.7B`   (revision-pinned; do not modify these files). - Training data: `siyanzhao/Openthoughts_math_30k_opsd` (the paper'"'"'s dataset,   revision-pinned, available offline via   `load_dataset("siyanzhao/Openthoughts_math_30k_opsd")`). You may filter or   re-weight examples from this dataset. You may NOT add external data, and   there is no API budget for synthetic data generation. - Dev evaluation: `bash /workspace/dev_eval.sh <checkpoint_dir|base>` runs the   same evaluation as the official one (same code, same frozen settings, same   sampling seed). Use it for your own checks; it consumes your budget. Cheaper   partial evals: `VAL_N=4 NUM_PROBLEMS=15 bash /workspace/dev_eval.sh ...`   (the official run always uses all 30 problems and 12 samples). - **Evaluation noise (read this before you trust a number).** The benchmark is   30 problems; the metric is an average over 12 samples each. Partial evals are   cheap but noisy: a 15-problem / avg@4 partial has a standard error of roughly   **±4–5 points** and has repeatedly misled past attempts by 5+ points; only the   full 30-problem / avg@12 eval on the fixed sampling seed is meaningful.   Training itself is **not seed-deterministic**: re-running the *same* recipe at   the *same* seed still moves avg@12 by **~1.5 points (1σ)** (the on-policy   rollouts are sampled with vLLM, plus CUDA nondeterminism). The scorer absorbs   this by re-training your method **5 times at the fixed seed and averaging** the   avg@12 (see "How you are scored"), which shrinks the *scored* noise to **~0.66   (1σ)** — but any single dev train+eval *you* run carries the full ~1.5. So do   not trust a single dev run: a margin under ~1.5 on one run may be noise. Budget   for **repeated dev train+evals of your final candidate** and target a gain that   clearly survives the ~1.5 single-run spread. - Internet access: none. Everything needed (repo, paper, model, datasets,   Python environment) is already inside the container.  ## Official evaluation (how your method is scored) After your session ends, the scorer **re-trains your method from the frozen Qwen3-1.7B base** using its own trusted copy of `train_capped.sh`: your `/workspace/repo` code, your `recipe.env` method knobs, and a **fixed budget of 100 optimizer steps at global batch 32** (see "Hard constraints"). It does this **5 times at the same fixed seed (42)**, evaluates each resulting checkpoint on AIME 2024 (30 problems) with these frozen settings — temperature=1.0, top_p=0.95, top_k disabled, min_p=0, presence_penalty=0, max_new_tokens=38912, Qwen3 thinking mode enabled, 12 samples per problem (metric: avg@12 accuracy), fixed sampling seed, vLLM with the frozen Qwen3-1.7B tokenizer/chat template, data-parallel across the 4 GPUs (4 single-GPU engines, tensor_parallel_size=1 each; problems sharded round-robin and merged) — and **averages the 5 avg@12 scores**. The training and evaluation code is fixed and runs outside your session; only your method (repo code + recipe.env) matters. Because your method is re-trained 5 times at the fixed seed and the avg@12 is averaged (not your submitted checkpoint), the scored noise is ~0.66 (1σ); aim for a method whose gain clearly exceeds that rather than a lucky run.  ## Hard constraints - **Fixed training budget — you cannot change it.** The scorer always trains for   exactly **100 optimizer steps** at **global batch 32** (num_processes 4 ×   per_device_train_batch_size 4 × gradient_accumulation_steps 2) from the frozen   base. These budget knobs — `max_steps`, `num_train_epochs`,   `per_device_train_batch_size`, `gradient_accumulation_steps`, `num_processes`,   `model_name_or_path`, and the dataset — are set by the scorer, and any value   you put for them in `recipe.env` (or hard-code to exceed 100 steps: extra   epochs, resuming/continuing a checkpoint, multi-stage training that sums to   >100 steps) is ignored or disallowed. Your improvement must come from the   **method within this fixed budget**, not from training longer or bigger. - **Train fresh from the frozen base.** Your code must actually train the 100   steps from `/opt/models/Qwen3-1.7B`; do not load, cache, or copy in   externally-pretrained or previously-trained weights. - **`recipe.env` whitelist.** Only these method knobs are honored (others are   ignored): `learning_rate`, `lora_r`, `lora_alpha`, `lora_dropout`, `beta`,   `jsd_token_clip`, `top_k_loss`, `temperature`, `top_p`, `top_k`, `lmbda`,   `max_completion_length` (≤ 4096), `max_grad_norm`, `weight_decay`,   `warmup_ratio`, `lr_scheduler_type`, `ema_decay`, and the boolean flags   `fixed_teacher`, `use_ema_teacher`, `use_tinker_loss`, `reason_first`,   `teacher_thinking`, `student_thinking` (set to `true`/`false`). - Do not modify the benchmark data (`/opt/eval_assets/aime24.json`), the frozen   base model files (`/opt/models/Qwen3-1.7B`), or the training-dataset files (HF   cache under `/opt/hf_cache`). Integrity is verified before scoring; any   mismatch scores 0. You may filter/re-weight the training dataset **in code**. - No inference-time changes: your improvement must live in the trained model,   evaluated under the same frozen inference settings as the baseline (no extra   samples, no ensembling/voting, no decoding, prompt, or chat-template changes —   the official eval always uses the frozen base tokenizer and chat template). - Train the designated model only: Qwen3-1.7B. LoRA rank must be ≤ 256.  ## How you are scored Your method is scored on the **absolute AIME24 avg@12 accuracy** of its re-trains. The scorer trains your method **5 times at the same fixed seed (42)** under the fixed 100-step budget, evaluates each on AIME24 (avg@12), and reports the **mean of the 5 avg@12 scores**. **Higher is better.** The verifier reports that number directly — it does not subtract a baseline or otherwise transform it.  For reference, the benchmark records two fixed anchors alongside the score: a **baseline** of **54.22** avg@12 (the OPSD released recipe, measured the same way — 5 same-seed re-trains of the unmodified recipe, averaged) and a **theoretical best** of **100%** (a perfect AIME24). These anchors are used only by a separate downstream normalization step to place scores on a common scale; they do not change what the verifier measures.  The 5-run average has **~0.66 (1σ)** noise (the per-run avg@12 spread is ~1.48), so target a gain that clearly exceeds it rather than a lucky run. A method that fails to train, a training-longer/bigger workaround, or any tampering with the frozen inputs is marked invalid (no score) and excluded. ' 2>&1 </dev/null | tee /logs/agent/claude-code.txt
159 claude --verbose --output-format=stream-json --permission-mode=bypassPermissions --disallowedTools WebSearch WebFetch --print -- # Task: Improve On-Policy Self-Distillation (OPSD) on AIME 2024  ## Context You are given the code repository for "Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models" (arXiv 2601.18734v3) at `/workspace/repo`, and the paper at `/workspace/paper.pdf`.  The repo implements **OPSD**, the paper's main approach: a single model acts as both student (sees only the problem) and teacher (additionally sees the ground-truth solution), and training minimizes a per-token forward-KL between the two distributions along the student's own rollouts, with a fixed initial-policy teacher.  A baseline checkpoint trained with the paper's released recipe (`/workspace/train_baseline.sh`, the 100-step Qwen3-1.7B configuration whose AIME24 number the paper publishes) is provided read-only at `/workspace/baseline/checkpoint/`, with its full training logs at `/workspace/baseline/logs/`. This is the paper's native 4×H100 recipe (`num_processes 4`, `per_device 4`, `gradient_accumulation 2` → global batch 32, colocate vLLM util 0.6), evaluated data-parallel across the 4 GPUs (4 single-GPU vLLM engines, TP=1 each). Re-trained 5 times at the fixed scoring seed, this baseline scores **54.22 ± 1.48** (avg@12 accuracy) on AIME 2024.  ## Your goal Propose and implement a research idea — a **method** improvement to OPSD — that beats the paper's own OPSD recipe **within a fixed training budget of 100 optimizer steps** (see "Hard constraints"). You are scored on the **absolute AIME24 avg@12 accuracy** of your re-trained method; **higher is better**, so the more your re-trained method's avg@12 exceeds the baseline (**54.22**), the better you do. Simply training longer/bigger is not available to you — the budget is fixed — so improvements must come from the method itself (the loss, the teacher, data selection/weighting, the on-policy scheme, etc.).  ## Deliverables Your contribution is a **training method, not a checkpoint.** You produce it by: 1. **Editing the OPSD training code in `/workspace/repo`** — the loss, the    teacher construction, data selection/weighting, the collator, the on-policy    sampling, hyper-parameters in code, etc. This is where your research idea    lives. `opsd_train.py` must remain the training entrypoint. 2. *(optional)* **`/workspace/submission/recipe.env`** — whitelisted    hyper-parameter overrides, one `KEY=VALUE` per line (see the whitelist under    "Hard constraints"). An absent/empty recipe reproduces the baseline    hyper-parameters on top of whatever code changes you made.  There is **no checkpoint to submit.** After your session, the scorer re-trains your method from the frozen base under the fixed budget and evaluates it (see "Official evaluation"); any checkpoint you leave behind is ignored.  Develop and test exactly as the scorer will, with: `SEED=42 OUT_DIR=/workspace/runs/try1 bash /workspace/train_capped.sh` then `bash /workspace/dev_eval.sh /workspace/runs/try1/capped_seed42/checkpoint-100`. `train_capped.sh` is the dev mirror of the scorer's training (same fixed budget, your repo code, your recipe.env); do not modify it — the scorer uses its own trusted copy, so edits only make your dev numbers diverge from the official ones.  ## Resources and budget - Hardware: 4×H100 for 6 hours total (your session ends when the budget   expires). This budget covers everything you run: diagnostic experiments,   training, and your own intermediate evaluations. The official scoring run   happens after your session and is not charged to you. For calibration: one   baseline training run takes ~19 min and one full dev evaluation ~27 min on   the 4 GPUs (LoRA submissions are merged into the base weights before   generation, so they evaluate at full-model speed). - Time remaining is tracked for you: read `/workspace/.timer/remaining_secs`   (and `/workspace/.timer/elapsed_secs`) at any point instead of doing   wall-clock arithmetic yourself. - Base model: Qwen3-1.7B, pre-downloaded at `/opt/models/Qwen3-1.7B`   (revision-pinned; do not modify these files). - Training data: `siyanzhao/Openthoughts_math_30k_opsd` (the paper's dataset,   revision-pinned, available offline via   `load_dataset("siyanzhao/Openthoughts_math_30k_opsd")`). You may filter or   re-weight examples from this dataset. You may NOT add external data, and   there is no API budget for synthetic data generation. - Dev evaluation: `bash /workspace/dev_eval.sh <checkpoint_dir|base>` runs the   same evaluation as the official one (same code, same frozen settings, same   sampling seed). Use it for your own checks; it consumes your budget. Cheaper   partial evals: `VAL_N=4 NUM_PROBLEMS=15 bash /workspace/dev_eval.sh ...`   (the official run always uses all 30 problems and 12 samples). - **Evaluation noise (read this before you trust a number).** The benchmark is   30 problems; the metric is an average over 12 samples each. Partial evals are   cheap but noisy: a 15-problem / avg@4 partial has a standard error of roughly   **±4–5 points** and has repeatedly misled past attempts by 5+ points; only the   full 30-problem / avg@12 eval on the fixed sampling seed is meaningful.   Training itself is **not seed-deterministic**: re-running the *same* recipe at   the *same* seed still moves avg@12 by **~1.5 points (1σ)** (the on-policy   rollouts are sampled with vLLM, plus CUDA nondeterminism). The scorer absorbs   this by re-training your method **5 times at the fixed seed and averaging** the   avg@12 (see "How you are scored"), which shrinks the *scored* noise to **~0.66   (1σ)** — but any single dev train+eval *you* run carries the full ~1.5. So do   not trust a single dev run: a margin under ~1.5 on one run may be noise. Budget   for **repeated dev train+evals of your final candidate** and target a gain that   clearly survives the ~1.5 single-run spread. - Internet access: none. Everything needed (repo, paper, model, datasets,   Python environment) is already inside the container.  ## Official evaluation (how your method is scored) After your session ends, the scorer **re-trains your method from the frozen Qwen3-1.7B base** using its own trusted copy of `train_capped.sh`: your `/workspace/repo` code, your `recipe.env` method knobs, and a **fixed budget of 100 optimizer steps at global batch 32** (see "Hard constraints"). It does this **5 times at the same fixed seed (42)**, evaluates each resulting checkpoint on AIME 2024 (30 problems) with these frozen settings — temperature=1.0, top_p=0.95, top_k disabled, min_p=0, presence_penalty=0, max_new_tokens=38912, Qwen3 thinking mode enabled, 12 samples per problem (metric: avg@12 accuracy), fixed sampling seed, vLLM with the frozen Qwen3-1.7B tokenizer/chat template, data-parallel across the 4 GPUs (4 single-GPU engines, tensor_parallel_size=1 each; problems sharded round-robin and merged) — and **averages the 5 avg@12 scores**. The training and evaluation code is fixed and runs outside your session; only your method (repo code + recipe.env) matters. Because your method is re-trained 5 times at the fixed seed and the avg@12 is averaged (not your submitted checkpoint), the scored noise is ~0.66 (1σ); aim for a method whose gain clearly exceeds that rather than a lucky run.  ## Hard constraints - **Fixed training budget — you cannot change it.** The scorer always trains for   exactly **100 optimizer steps** at **global batch 32** (num_processes 4 ×   per_device_train_batch_size 4 × gradient_accumulation_steps 2) from the frozen   base. These budget knobs — `max_steps`, `num_train_epochs`,   `per_device_train_batch_size`, `gradient_accumulation_steps`, `num_processes`,   `model_name_or_path`, and
[metadata] {"noOutputExpected": false, "persistedOutputPath": "/logs/agent/sessions/projects/-workspace/3476b637-41e1-4437-8c9a-f53302e59a44/tool-results/bgrymjl8d.txt", "persistedOutputSize": 38100}

0, 67919 MiB
1, 66345 MiB
2, 66305 MiB
3, 68141 MiB
9

[stdout]
0, 67919 MiB
1, 66345 MiB
2, 66305 MiB
3, 68141 MiB
9
[metadata] {"noOutputExpected": false}

pkill: refusing to run — this pattern matches the Claude CLI process (PID 159). Narrow the pattern, or target your own children with `pkill -P $$ ...`.
8
0, 66785 MiB
1, 66265 MiB
2, 66403 MiB
3, 66919 MiB

[stdout]
pkill: refusing to run — this pattern matches the Claude CLI process (PID 159). Narrow the pattern, or target your own children with `pkill -P $$ ...`.
8
0, 66785 MiB
1, 66265 MiB
2, 66403 MiB
3, 66919 MiB
[metadata] {"noOutputExpected": false}

(Bash completed with no output)

[metadata] {"noOutputExpected": false}

(Bash completed with no output)

[metadata] {"noOutputExpected": false}

The file /workspace/submission/recipe.env has been updated successfully. (file state is current in your context — no need to Read it back)

[metadata] {"type": "update", "filePath": "/workspace/submission/recipe.env", "content": "# OPSD method recipe — whitelisted hyper-parameter overrides (KEY=VALUE per line).\n#\n# The method改 lives in /workspace/repo; no hyper-parameter overrides are set yet.\n", "structuredPatch": [{"oldStart": 1, "oldLines": 8, "newStart": 1, "newLines": 3, "lines": [" # OPSD method recipe — whitelisted hyper-parameter overrides (KEY=VALUE per line).", " #", "-# The method itself lives in /workspace/repo (direction-preserving token-influence", "-# clipping + padding-hole-free position ids). The only knob set here is the LR", "-# schedule: the released OPSD recipe (train_baseline.sh / the reference run's logs)", "-# leaves lr_scheduler_type at the HF default, i.e. linear decay to 0 over the 100", "-# steps, which train_capped.sh would otherwise override to `constant`.", "-lr_scheduler_type=linear", "+# The method改 lives in /workspace/repo; no hyper-parameter overrides are set yet."]}], "originalFile": "# OPSD method recipe — whitelisted hyper-parameter overrides (KEY=VALUE per line).\n#\n# The method itself lives in /workspace/repo (direction-preserving token-influence\n# clipping + padding-hole-free position ids). The only knob set here is the LR\n# schedule: the released OPSD recipe (train_baseline.sh / the reference run's logs)\n# leaves lr_scheduler_type at the HF default, i.e. linear decay to 0 over the 100\n# steps, which train_capped.sh would otherwise override to `constant`.\nlr_scheduler_type=linear\n", "userModified": false}

The file /workspace/submission/recipe.env has been updated successfully. (file state is current in your context — no need to Read it back)

[metadata] {"type": "update", "filePath": "/workspace/submission/recipe.env", "content": "# OPSD method recipe — whitelisted hyper-parameter overrides (KEY=VALUE per line).\n#\n# The method itself lives in /workspace/repo. No hyper-parameter overrides are\n# set yet; train_capped.sh's own defaults are the paper's recipe.\n", "structuredPatch": [{"oldStart": 1, "oldLines": 3, "newStart": 1, "newLines": 4, "lines": [" # OPSD method recipe — whitelisted hyper-parameter overrides (KEY=VALUE per line).", " #", "-# The method改 lives in /workspace/repo; no hyper-parameter overrides are set yet.", "+# The method itself lives in /workspace/repo. No hyper-parameter overrides are", "+# set yet; train_capped.sh's own defaults are the paper's recipe."]}], "originalFile": "# OPSD method recipe — whitelisted hyper-parameter overrides (KEY=VALUE per line).\n#\n# The method改 lives in /workspace/repo; no hyper-parameter overrides are set yet.\n", "userModified": false}

=== [e2] train start 04:56:34  (remaining 17493s)
[modes] teacher_thinking=False student_thinking=False
[modes] teacher_thinking=False student_thinking=False

[stdout]
=== [e2] train start 04:56:34  (remaining 17493s)
[modes] teacher_thinking=False student_thinking=False
[modes] teacher_thinking=False student_thinking=False
[metadata] {"noOutputExpected": false}

File "<string>", line 6, in <module>
ModuleNotFoundError: No module named 'PyPDF2'

[stdout]
  File "<string>", line 6, in <module>
ModuleNotFoundError: No module named 'PyPDF2'
[metadata] {"noOutputExpected": false}

Traceback (most recent call last):
  File "<string>", line 2, in <module>
ModuleNotFoundError: No module named 'fitz'

[stdout]
Traceback (most recent call last):
  File "<string>", line 2, in <module>
ModuleNotFoundError: No module named 'fitz'
[metadata] {"noOutputExpected": false}

README.md:122:OPSD can also run in non-thinking setting where both the Qwen student and teacher are enabled_thinking=False during training (`--student_thinking False --teacher_thinking False`) and evaluated with non-thinking inference (`--no_thinking`), with faster evaluation time than thinking mode.
README.md:138:#### Qwen3-8B (`--jsd_token_clip 1e-7`)
README.md:182:#### Qwen3-4B (`--jsd_token_clip 1e-6`)
README.md:229:#### Qwen3-1.7B (`--jsd_token_clip 1e-6`)
README.md:283:| `--max_completion_length` | — | Student generation length for distillation. We use 1024 in our main experiments. |
README.md:285:| `--jsd_token_clip` | 0.05 | Clip the JSD loss for each token to a maximum value. This can improve stability. | 
scripts/run_opsd_1b.sh:16:    --max_completion_length 1024 \
scripts/run_opsd_1b.sh:34:    --lmbda 1 \
scripts/run_opsd_1b.sh:36:    --jsd_token_clip 0.05 \
scripts/run_grpo.sh:20:    --max_completion_length 16000 \
scripts/run_opsd_8b_nonthink.sh:16:    --max_completion_length 1024 \
scripts/run_opsd_8b_nonthink.sh:34:    --lmbda 1 \
scripts/run_opsd_8b_nonthink.sh:37:    --teacher_thinking False \
scripts/run_opsd_8b_nonthink.sh:38:    --jsd_token_clip 1e-7 \
scripts/run_opsd_4b_nonthink.sh:16:    --max_completion_length 1024 \
scripts/run_opsd_4b_nonthink.sh:34:    --lmbda 1 \
scripts/run_opsd_4b_nonthink.sh:37:    --teacher_thinking False \
scripts/run_opsd_4b_nonthink.sh:38:    --jsd_token_clip 1e-6 \
scripts/run_opsd_8b.sh:16:    --max_completion_length 1024 \
scripts/run_opsd_8b.sh:34:    --lmbda 1 \
scripts/run_opsd_8b.sh:36:    --jsd_token_clip 0.06 \
scripts/run_opsd_4b.sh:16:    --max_completion_length 1024 \
scripts/run_opsd_4b.sh:34:    --lmbda 1 \
scripts/run_opsd_4b.sh:36:    --jsd_token_clip 0.05 \

[stdout]
README.md:122:OPSD can also run in non-thinking setting where both the Qwen student and teacher are enabled_thinking=False during training (`--student_thinking False --teacher_thinking False`) and evaluated with non-thinking inference (`--no_thinking`), with faster evaluation time than thinking mode.
README.md:138:#### Qwen3-8B (`--jsd_token_clip 1e-7`)
README.md:182:#### Qwen3-4B (`--jsd_token_clip 1e-6`)
README.md:229:#### Qwen3-1.7B (`--jsd_token_clip 1e-6`)
README.md:283:| `--max_completion_length` | — | Student generation length for distillation. We use 1024 in our main experiments. |
README.md:285:| `--jsd_token_clip` | 0.05 | Clip the JSD loss for each token to a maximum value. This can improve stability. | 
scripts/run_opsd_1b.sh:16:    --max_completion_length 1024 \
scripts/run_opsd_1b.sh:34:    --lmbda 1 \
scripts/run_opsd_1b.sh:36:    --jsd_token_clip 0.05 \
scripts/run_grpo.sh:20:    --max_completion_length 16000 \
scripts/run_opsd_8b_nonthink.sh:16:    --max_completion_length 1024 \
scripts/run_opsd_8b_nonthink.sh:34:    --lmbda 1 \
scripts/run_opsd_8b_nonthink.sh:37:    --teacher_thinking False \
scripts/run_opsd_8b_nonthink.sh:38:    --jsd_token_clip 1e-7 \
scripts/run_opsd_4b_nonthink.sh:16:    --max_completion_length 1024 \
scripts/run_opsd_4b_nonthink.sh:34:    --lmbda 1 \
scripts/run_opsd_4b_nonthink.sh:37:    --teacher_thinking False \
scripts/run_opsd_4b_nonthink.sh:38:    --jsd_token_clip 1e-6 \
scripts/run_opsd_8b.sh:16:    --max_completion_length 1024 \
scripts/run_opsd_8b.sh:34:    --lmbda 1 \
scripts/run_opsd_8b.sh:36:    --jsd_token_clip 0.06 \
scripts/run_opsd_4b.sh:16:    --max_completion_length 1024 \
scripts/run_opsd_4b.sh:34:    --lmbda 1 \
scripts/run_opsd_4b.sh:36:    --jsd_token_clip 0.05 \
[metadata] {"noOutputExpected": false}

</td>
<td>

| Step | Avg@12 |
|---|---|
| Base | 23.1% |
| 25 | 24.7% |
| 50 | 27.8% |
| 75 | 26.9% |
| 100 | 29.2% |

</td>
</tr>
</table>
</div>

> **Evaluation settings:** temperature=1.0, thinking mode enabled, max new tokens=38912, top-p=none, top-k disabled, min-p=0, presence penalty=0, num samples=12


## Non-Thinking Mode

OPSD can also run in non-thinking setting where both the Qwen student and teacher are enabled_thinking=False during training (`--student_thinking False --teacher_thinking False`) and evaluated with non-thinking inference (`--no_thinking`), with faster evaluation time than thinking mode.

Training:
```bash
bash scripts/run_opsd_4b_nonthink.sh
bash scripts/run_opsd_8b_nonthink.sh
```

Evaluation:
```bash
cd eval
bash run_eval_nonthink.sh
```

### Evaluation Results with Non-Thinking Mode across Models

#### Qwen3-8B (`--jsd_token_clip 1e-7`)

<div align="center">
<table>
<tr>
<th align="center">AIME24</th>
<th align="center">AIME25</th>
<th align="center">HMMT25</th>
</tr>
<tr>
<td>

| Step | Avg@12 |
| 75 | 13.1% |
| 100 | 16.4% |
| 150 | 14.4% |

</td>
</tr>
</table>
</div>

#### Qwen3-1.7B (`--jsd_token_clip 1e-6`)

<div align="center">
<table>
<tr>
<th align="center">AIME24</th>
<th align="center">AIME25</th>
<th align="center">HMMT25</th>
</tr>
<tr>
<td>

| Step | Avg@12 |
|---|---|
| Base | 11.9% |
| 50 | 15.0% |
| 75 | 13.9% |
| 100 | 12.5% |

</td>
<td>

| Step | Avg@12 |
|---|---|
| Base | 9.2% |
| 50 | 6.2% |
| 75 | 8.3% |
| 100 | 8.1% |

</td>
<td>

| Step | Avg@12 |
|---|---|
| Base | 5.0% |
| 25 | 7.2% |
| 50 | 5.8% |
| 75 | 5.0% |

</td>
</tr>
</table>
</div>

> **Evaluation settings:** temperature=1.0, non-thinking mode, num samples=12.



## Key OPSD arguments

| Argument | Default | Description |
|---|---|---|
| `--fixed_teacher` | `False` | Fix the teacher to the initial policy (step 0). Requires --use_peft. Note ❗ If you disable PEFT, the teacher will keep updating at every training step, which may make training unstable. Our main results use the fixed teacher, which is currently implemented with LoRA adapter weights. |
| `--use_tinker_loss` | `False` | Use sampled-token policy-gradient objective instead of full-vocabulary JSD. More memory efficient. Currently no clipped implemented for this variant, could be unstable. |
| `--max_completion_length` | — | Student generation length for distillation. We use 1024 in our main experiments. |
| `--beta` | — | Interpolation weight for the JSD mixture distribution. Beta=0 means forward KL and 1 means reverse KL. |
| `--jsd_token_clip` | 0.05 | Clip the JSD loss for each token to a maximum value. This can improve stability. | 
| `--reason_first` | `False` | Prepend an explicit rationalization to the teacher context before distillation. |
| `--run_config` | `None` | Custom name suffix for the output directory and WandB run. |

### SFT Baseline

See [`scripts/run_sft.sh`](scripts/run_sft.sh).

### GRPO Baseline

See [`scripts/run_grpo.sh`](scripts/run_grpo.sh).

### Acknowledgements
Our implementation builds on [TRL GOLD Trainer](https://huggingface.co/docs/trl/gold_trainer). We sincerely thank [@simran135](https://github.com/simran135) and [@beanie00](https://github.com/beanie00) for identifying the prompt template bugs and the zero-2 issue, respectively!

## Citation

[stdout]
</td>
<td>

| Step | Avg@12 |
|---|---|
| Base | 23.1% |
| 25 | 24.7% |
| 50 | 27.8% |
| 75 | 26.9% |
| 100 | 29.2% |

</td>
</tr>
</table>
</div>

> **Evaluation settings:** temperature=1.0, thinking mode enabled, max new tokens=38912, top-p=none, top-k disabled, min-p=0, presence penalty=0, num samples=12


## Non-Thinking Mode

OPSD can also run in non-thinking setting where both the Qwen student and teacher are enabled_thinking=False during training (`--student_thinking False --teacher_thinking False`) and evaluated with non-thinking inference (`--no_thinking`), with faster evaluation time than thinking mode.

Training:
```bash
bash scripts/run_opsd_4b_nonthink.sh
bash scripts/run_opsd_8b_nonthink.sh
```

Evaluation:
```bash
cd eval
bash run_eval_nonthink.sh
```

### Evaluation Results with Non-Thinking Mode across Models

#### Qwen3-8B (`--jsd_token_clip 1e-7`)

<div align="center">
<table>
<tr>
<th align="center">AIME24</th>
<th align="center">AIME25</th>
<th align="center">HMMT25</th>
</tr>
<tr>
<td>

| Step | Avg@12 |
| 75 | 13.1% |
| 100 | 16.4% |
| 150 | 14.4% |

</td>
</tr>
</table>
</div>

#### Qwen3-1.7B (`--jsd_token_clip 1e-6`)

<div align="center">
<table>
<tr>
<th align="center">AIME24</th>
<th align="center">AIME25</th>
<th align="center">HMMT25</th>
</tr>
<tr>
<td>

| Step | Avg@12 |
|---|---|
| Base | 11.9% |
| 50 | 15.0% |
| 75 | 13.9% |
| 100 | 12.5% |

</td>
<td>

| Step | Avg@12 |
|---|---|
| Base | 9.2% |
| 50 | 6.2% |
| 75 | 8.3% |
| 100 | 8.1% |

</td>
<td>

| Step | Avg@12 |
|---|---|
| Base | 5.0% |
| 25 | 7.2% |
| 50 | 5.8% |
| 75 | 5.0% |

</td>
</tr>
</table>
</div>

> **Evaluation settings:** temperature=1.0, non-thinking mode, num samples=12.



## Key OPSD arguments

| Argument | Default | Description |
|---|---|---|
| `--fixed_teacher` | `False` | Fix the teacher to the initial policy (step 0). Requires --use_peft. Note ❗ If you disable PEFT, the teacher will keep updating at every training step, which may make training unstable. Our main results use the fixed teacher, which is currently implemented with LoRA adapter weights. |
| `--use_tinker_loss` | `False` | Use sampled-token policy-gradient objective instead of full-vocabulary JSD. More memory efficient. Currently no clipped implemented for this variant, could be unstable. |
| `--max_completion_length` | — | Student generation length for distillation. We use 1024 in our main experiments. |
| `--beta` | — | Interpolation weight for the JSD mixture distribution. Beta=0 means forward KL and 1 means reverse KL. |
| `--jsd_token_clip` | 0.05 | Clip the JSD loss for each token to a maximum value. This can improve stability. | 
| `--reason_first` | `False` | Prepend an explicit rationalization to the teacher context before distillation. |
| `--run_config` | `None` | Custom name suffix for the output directory and WandB run. |

### SFT Baseline

See [`scripts/run_sft.sh`](scripts/run_sft.sh).

### GRPO Baseline

See [`scripts/run_grpo.sh`](scripts/run_grpo.sh).

### Acknowledgements
Our implementation builds on [TRL GOLD Trainer](https://huggingface.co/docs/trl/gold_trainer). We sincerely thank [@simran135](https://github.com/simran135) and [@beanie00](https://github.com/beanie00) for identifying the prompt template bugs and the zero-2 issue, respectively!

## Citation
[metadata] {"noOutputExpected": false}

# Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models


<p align="center">
<a href="https://arxiv.org/pdf/2601.18734v3"><img src="https://img.shields.io/badge/arXiv-2601.18734-b31b1b.svg"></a>
<a href="https://siyan-zhao.github.io/blog/2026/opsd/"><img src="https://img.shields.io/badge/Blog-Post-blue.svg"></a>
</p>

---
## Overview

**On-Policy Self-Distillation (OPSD)** trains a single model to act as both student and teacher by conditioning on different contexts — the student sees only the problem, while the teacher additionally sees the ground-truth solution — and performs token-level distribution matching along the student's own on-policy trajectories.


## Updates

- **Mar 18, 2026**: Released updated code. 

  (1) Fixed chat template and zero2 bugs (see [template issue](https://github.com/huggingface/trl/issues/5241)), we re-ran experiments with updated results (detailed results & ablations updated on arxiv/blog). The fixes yield improved OPSD performance, most notably on Qwen3-1.7B.


-  **Mar 3, 2026**: Initial code release.

## Installation


```bash
conda env create -f environment.yml
conda activate opsd
```

```bash
pip install flash-attn==2.8.3 --no-build-isolation
```
If you encounter difficulties installing flash-attn, you can check the version matching your CUDA and PyTorch versions from the [flash-attention releases page](https://github.com/Dao-AILab/flash-attention/releases).

The code uses `trl`'s experimental GOLD trainer as a base.

## Repository Structure

```
├── opsd_trainer.py          # OPSDTrainer: core self-distillation trainer
├── data_collator.py         # Data collator for self-distillation
├── opsd_train.py            # OPSD training entry point
├── sft_train.py             # SFT baseline training entry point
├── grpo_train.py            # GRPO baseline training entry point
├── accelerate.yaml          # Accelerate config (multi-GPU)
├── scripts/
│   ├── run_opsd.sh          # Example launch script for OPSD
│   ├── run_sft.sh           # Example launch script for SFT
│   └── run_grpo.sh          # Example launch script for GRPO
└── eval/
    ├── evaluate_math.py     # Evaluation script (vLLM)
    └── run_eval.sh          # Example evaluation script
```

## Quick Start

Reproduce results on Qwen3-1.7B (🚀 training only takes **~15 minutes** on 4×H100 and peaks within 100 steps):

```bash
bash scripts/run_opsd_1b.sh
```
Evaluation: (evaluation takes ~ 30-50 minutes on 4xh100 for each checkpoint) 
```bash
cd eval
bash run_eval.sh
```

### Evaluation Results across Tasks on Qwen3-1.7B

<div align="center">
<table>
<tr>
<th align="center">AIME24</th>
<th align="center">AIME25</th>
<th align="center">HMMT25</th>
</tr>
<tr>
<td>

| Step | Avg@12 |
|---|---|
| Base | 51.5% |
| 25 | 51.4% |
| 50 | 52.8% |
| 75 | 54.4% |
| 100 | 57.2% |

</td>
<td>

| Step | Avg@12 |
|---|---|
| Base | 36.7% |
| 25 | 42.5% |
| 50 | 43.9% |
| 75 | 40.6% |
| 100 | 41.1% |

[stdout]
# Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models


<p align="center">
<a href="https://arxiv.org/pdf/2601.18734v3"><img src="https://img.shields.io/badge/arXiv-2601.18734-b31b1b.svg"></a>
<a href="https://siyan-zhao.github.io/blog/2026/opsd/"><img src="https://img.shields.io/badge/Blog-Post-blue.svg"></a>
</p>

---
## Overview

**On-Policy Self-Distillation (OPSD)** trains a single model to act as both student and teacher by conditioning on different contexts — the student sees only the problem, while the teacher additionally sees the ground-truth solution — and performs token-level distribution matching along the student's own on-policy trajectories.


## Updates

- **Mar 18, 2026**: Released updated code. 

  (1) Fixed chat template and zero2 bugs (see [template issue](https://github.com/huggingface/trl/issues/5241)), we re-ran experiments with updated results (detailed results & ablations updated on arxiv/blog). The fixes yield improved OPSD performance, most notably on Qwen3-1.7B.


-  **Mar 3, 2026**: Initial code release.

## Installation


```bash
conda env create -f environment.yml
conda activate opsd
```

```bash
pip install flash-attn==2.8.3 --no-build-isolation
```
If you encounter difficulties installing flash-attn, you can check the version matching your CUDA and PyTorch versions from the [flash-attention releases page](https://github.com/Dao-AILab/flash-attention/releases).

The code uses `trl`'s experimental GOLD trainer as a base.

## Repository Structure

```
├── opsd_trainer.py          # OPSDTrainer: core self-distillation trainer
├── data_collator.py         # Data collator for self-distillation
├── opsd_train.py            # OPSD training entry point
├── sft_train.py             # SFT baseline training entry point
├── grpo_train.py            # GRPO baseline training entry point
├── accelerate.yaml          # Accelerate config (multi-GPU)
├── scripts/
│   ├── run_opsd.sh          # Example launch script for OPSD
│   ├── run_sft.sh           # Example launch script for SFT
│   └── run_grpo.sh          # Example launch script for GRPO
└── eval/
    ├── evaluate_math.py     # Evaluation script (vLLM)
    └── run_eval.sh          # Example evaluation script
```

## Quick Start

Reproduce results on Qwen3-1.7B (🚀 training only takes **~15 minutes** on 4×H100 and peaks within 100 steps):

```bash
bash scripts/run_opsd_1b.sh
```
Evaluation: (evaluation takes ~ 30-50 minutes on 4xh100 for each checkpoint) 
```bash
cd eval
bash run_eval.sh
```

### Evaluation Results across Tasks on Qwen3-1.7B

<div align="center">
<table>
<tr>
<th align="center">AIME24</th>
<th align="center">AIME25</th>
<th align="center">HMMT25</th>
</tr>
<tr>
<td>

| Step | Avg@12 |
|---|---|
| Base | 51.5% |
| 25 | 51.4% |
| 50 | 52.8% |
| 75 | 54.4% |
| 100 | 57.2% |

</td>
<td>

| Step | Avg@12 |
|---|---|
| Base | 36.7% |
| 25 | 42.5% |
| 50 | 43.9% |
| 75 | 40.6% |
| 100 | 41.1% |
[metadata] {"noOutputExpected": false}

accelerate launch \
    --config_file accelerate.yaml \
    --num_processes 4 \
    --gradient_accumulation_steps 2 \
    --main_process_port 12949 \
    opsd_train.py \
    --model_name_or_path /data0/shared/Qwen3-1.7B \
    --learning_rate 5e-6 \
    --max_grad_norm 0.1 \
    --per_device_train_batch_size 4 \
    --gradient_checkpointing \
    --gradient_accumulation_steps 2 \
    --output_dir  /data0/siyanz/opsd/ \
    --run_config qwen31b_gen1024_fixteacher_temp11_forwardbeta0_clip005 \
    --num_train_epochs 30 \
    --max_completion_length 1024 \
    --save_steps 25 \
    --logging_steps 2 \
    --attn_implementation flash_attention_2 \
    --torch_dtype bfloat16 \
    --max_length 20000 \
    --beta 0 \
    --use_vllm \
    --vllm_mode colocate \
    --vllm_gpu_memory_utilization 0.6 \
    --vllm_tensor_parallel_size 1 \
    --use_peft \
    --lora_r 64 \
    --lora_alpha 128 \
    --lora_target_modules q_proj k_proj v_proj o_proj gate_proj up_proj down_proj \
    --temperature 1.1 \
    --top_p 0.95 \
    --top_k 20 \
    --lmbda 1 \
    --fixed_teacher \
    --jsd_token_clip 0.05 \
    --wandb_project OPSD
===
40:    --learning_rate 5e-6 \
41:    --max_grad_norm 0.1 \
47:    --num_train_epochs 30 \
48:    --max_steps 100 \

[stdout]
accelerate launch \
    --config_file accelerate.yaml \
    --num_processes 4 \
    --gradient_accumulation_steps 2 \
    --main_process_port 12949 \
    opsd_train.py \
    --model_name_or_path /data0/shared/Qwen3-1.7B \
    --learning_rate 5e-6 \
    --max_grad_norm 0.1 \
    --per_device_train_batch_size 4 \
    --gradient_checkpointing \
    --gradient_accumulation_steps 2 \
    --output_dir  /data0/siyanz/opsd/ \
    --run_config qwen31b_gen1024_fixteacher_temp11_forwardbeta0_clip005 \
    --num_train_epochs 30 \
    --max_completion_length 1024 \
    --save_steps 25 \
    --logging_steps 2 \
    --attn_implementation flash_attention_2 \
    --torch_dtype bfloat16 \
    --max_length 20000 \
    --beta 0 \
    --use_vllm \
    --vllm_mode colocate \
    --vllm_gpu_memory_utilization 0.6 \
    --vllm_tensor_parallel_size 1 \
    --use_peft \
    --lora_r 64 \
    --lora_alpha 128 \
    --lora_target_modules q_proj k_proj v_proj o_proj gate_proj up_proj down_proj \
    --temperature 1.1 \
    --top_p 0.95 \
    --top_k 20 \
    --lmbda 1 \
    --fixed_teacher \
    --jsd_token_clip 0.05 \
    --wandb_project OPSD
===
40:    --learning_rate 5e-6 \
41:    --max_grad_norm 0.1 \
47:    --num_train_epochs 30 \
48:    --max_steps 100 \
[metadata] {"noOutputExpected": false}

20:        student_thinking=False,
21:        teacher_thinking=True,
26:        self.student_thinking = student_thinking
27:        self.teacher_thinking = teacher_thinking
33:            "Do NOT use <think> tags. Do NOT derive your own solution. "
72:                student_messages, tokenize=False, add_generation_prompt=True, enable_thinking=self.student_thinking
108:                    teacher_messages, tokenize=False, add_generation_prompt=True, enable_thinking=self.teacher_thinking
===
74:    from vllm import LLM, SamplingParams
176:        student_thinking: bool = False,
191:                student_thinking=student_thinking,
842:            return self._generate_teacher_reasoning_vllm(teacher_reasoning_prompts)
862:                        reasoning_outputs = model.generate(
866:                            return_dict_in_generate=True,
902:            generated_outputs = model.generate(
906:                return_dict_in_generate=True,
936:    def _generate_on_policy_outputs_vllm(self, inputs, generation_config, pad_token_id=None):
979:                completion_ids = self.vllm_client.generate(
1006:            sampling_params = SamplingParams(
1030:            all_outputs = self.vllm_engine.generate(
1124:    def _generate_teacher_reasoning_vllm(
1154:                completion_ids = self.vllm_client.generate(
1172:            sampling_params = SamplingParams(
1188:            all_outputs = self.vllm_engine.generate(
1448:            result = self._generate_on_policy_outputs_vllm(

[stdout]
20:        student_thinking=False,
21:        teacher_thinking=True,
26:        self.student_thinking = student_thinking
27:        self.teacher_thinking = teacher_thinking
33:            "Do NOT use <think> tags. Do NOT derive your own solution. "
72:                student_messages, tokenize=False, add_generation_prompt=True, enable_thinking=self.student_thinking
108:                    teacher_messages, tokenize=False, add_generation_prompt=True, enable_thinking=self.teacher_thinking
===
74:    from vllm import LLM, SamplingParams
176:        student_thinking: bool = False,
191:                student_thinking=student_thinking,
842:            return self._generate_teacher_reasoning_vllm(teacher_reasoning_prompts)
862:                        reasoning_outputs = model.generate(
866:                            return_dict_in_generate=True,
902:            generated_outputs = model.generate(
906:                return_dict_in_generate=True,
936:    def _generate_on_policy_outputs_vllm(self, inputs, generation_config, pad_token_id=None):
979:                completion_ids = self.vllm_client.generate(
1006:            sampling_params = SamplingParams(
1030:            all_outputs = self.vllm_engine.generate(
1124:    def _generate_teacher_reasoning_vllm(
1154:                completion_ids = self.vllm_client.generate(
1172:            sampling_params = SamplingParams(
1188:            all_outputs = self.vllm_engine.generate(
1448:            result = self._generate_on_policy_outputs_vllm(
[metadata] {"noOutputExpected": false}

def _generate_on_policy_outputs_vllm(self, inputs, generation_config, pad_token_id=None):
        """Generate on-policy outputs from student prompts using vLLM."""
        import time

        device = self.accelerator.device

        prompts_text_for_vllm = self.processing_class.batch_decode(
            inputs["student_prompts"],
            skip_special_tokens=False,
        )
        # Remove padding token text if it appears, as vLLM expects clean prompts
        if self.processing_class.pad_token:
            prompts_text_for_vllm = [
                p.replace(self.processing_class.pad_token, "") for p in prompts_text_for_vllm
            ]

        # Also decode prompts WITH special tokens for logging
        prompts_text_with_special = self.processing_class.batch_decode(
            inputs["student_prompts"],
            skip_special_tokens=False,
        )

        # system_prompt = "Please reason step by step, and put your final answer within \\boxed{}."
        # target_system_prompt = "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."
        # prompts_text = [p.replace(target_system_prompt, system_prompt) for p in prompts_text]
        # Add system prompt to prompts

        max_completion_length = generation_config.max_new_tokens
        temperature = generation_config.temperature
        # vLLM uses top_k=-1 for no top_k, transformers uses 0 or None.
        top_k = generation_config.top_k if generation_config.top_k and generation_config.top_k > 0 else -1
        # top_p, repetition_penalty, min_p, presence_penalty are not directly in generation_config, get from trainer args
        top_p = self.args.top_p if hasattr(self.args, "top_p") else 1.0
        repetition_penalty = self.args.repetition_penalty if hasattr(self.args, "repetition_penalty") else 1.0
        min_p = self.args.min_p if hasattr(self.args, "min_p") else 0.0
        presence_penalty = self.args.presence_penalty if hasattr(self.args, "presence_penalty") else 0.0

        # Start timing for vLLM generation
        start_time = time.time()

        if self.vllm_mode == "server":
            all_prompts_text = gather_object(prompts_text_for_vllm)
            if self.accelerator.is_main_process:
                completion_ids = self.vllm_client.generate(
                    prompts=all_prompts_text,
                    n=1,  # In GKD, we generate 1 completion per prompt from student
                    repetition_penalty=repetition_penalty,
                    temperature=temperature,
                    top_p=top_p,
                    top_k=top_k,
                    min_p=min_p,
                    max_tokens=max_completion_length,
                    presence_penalty=presence_penalty,
                    guided_decoding_regex=self.vllm_guided_decoding_regex,
                )
            else:
                completion_ids = [None] * len(all_prompts_text)
            completion_ids = broadcast_object_list(completion_ids, from_process=0)
            process_slice = slice(
                self.accelerator.process_index * len(prompts_text_for_vllm),
                (self.accelerator.process_index + 1) * len(prompts_text_for_vllm),
            )
            completion_ids = completion_ids[process_slice]
        elif self.vllm_mode == "colocate":
            if self.vllm_guided_decoding_regex:
                guided_decoding = GuidedDecodingParams(
                    backend="outlines", regex=self.vllm_guided_decoding_regex
                )
            else:
                guided_decoding = None
            sampling_params = SamplingParams(
                n=1,
                repetition_penalty=repetition_penalty,
                temperature=temperature,
                top_p=top_p,
                top_k=top_k,
                min_p=min_p,
                max_tokens=max_completion_length,
                presence_penalty=presence_penalty,
                guided_decoding=guided_decoding,
            )

            if hasattr(self, "vllm_tp_group") and self.vllm_tensor_parallel_size > 1:
                # Gather prompts from all ranks in the TP group and flatten.
                # Each rank starts with its own prompts; after gathering, all ranks see the full group set.
                orig_size = len(prompts_text_for_vllm)
                gathered_prompts = [None for _ in range(self.vllm_tensor_parallel_size)]
                torch.distributed.all_gather_object(
                    gathered_prompts, prompts_text_for_vllm, group=self.vllm_tp_group
                )
                all_prompts_text = [p for sublist in gathered_prompts for p in sublist]
            else:
                all_prompts_text = prompts_text_for_vllm

            all_outputs = self.vllm_engine.generate(
                all_prompts_text, sampling_params=sampling_params, use_tqdm=False
            )
            completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs]

            if hasattr(self, "vllm_tp_group") and self.vllm_tensor_parallel_size > 1:
                # Slice completions for this rank within its TP group.
                # Each rank generates all outputs — we keep only our share.
                local_rank_in_group = torch.distributed.get_rank(group=self.vllm_tp_group)
                tp_slice = slice(local_rank_in_group * orig_size, (local_rank_in_group + 1) * orig_size)
                completion_ids = completion_ids[tp_slice]

            if self.vllm_enable_sleep_mode:
                self.vllm_engine.sleep(level=2)
        else:
            raise ValueError(f"Unknown vllm_mode: {self.vllm_mode}")

        # Calculate and print vLLM generation statistics
        elapsed_time = time.time() - start_time
        total_completion_tokens = sum(len(ids) for ids in completion_ids)
        num_prompts = len(completion_ids)
        avg_completion_length = total_completion_tokens / num_prompts if num_prompts > 0 else 0
        tokens_per_sec = total_completion_tokens / elapsed_time if elapsed_time > 0 else 0
        print(
            f"vLLM generation done - elapsed time: {elapsed_time:.2f}s, prompts: {num_prompts}, total tokens: {total_completion_tokens}, avg length: {avg_completion_length:.1f}, speed: {tokens_per_sec:.1f} tok/s"
        )

        # We need to combine prompt and completion for new_input_ids
        # Tokenize prompts again to get prompt_ids on the correct device and format
        # Use prompts_text_for_vllm (without special tokens) for tokenization since vLLM expects clean text
        # Ensure add_special_tokens=False as vLLM typically handles prompts as raw text

[stdout]
    def _generate_on_policy_outputs_vllm(self, inputs, generation_config, pad_token_id=None):
        """Generate on-policy outputs from student prompts using vLLM."""
        import time

        device = self.accelerator.device

        prompts_text_for_vllm = self.processing_class.batch_decode(
            inputs["student_prompts"],
            skip_special_tokens=False,
        )
        # Remove padding token text if it appears, as vLLM expects clean prompts
        if self.processing_class.pad_token:
            prompts_text_for_vllm = [
                p.replace(self.processing_class.pad_token, "") for p in prompts_text_for_vllm
            ]

        # Also decode prompts WITH special tokens for logging
        prompts_text_with_special = self.processing_class.batch_decode(
            inputs["student_prompts"],
            skip_special_tokens=False,
        )

        # system_prompt = "Please reason step by step, and put your final answer within \\boxed{}."
        # target_system_prompt = "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."
        # prompts_text = [p.replace(target_system_prompt, system_prompt) for p in prompts_text]
        # Add system prompt to prompts

        max_completion_length = generation_config.max_new_tokens
        temperature = generation_config.temperature
        # vLLM uses top_k=-1 for no top_k, transformers uses 0 or None.
        top_k = generation_config.top_k if generation_config.top_k and generation_config.top_k > 0 else -1
        # top_p, repetition_penalty, min_p, presence_penalty are not directly in generation_config, get from trainer args
        top_p = self.args.top_p if hasattr(self.args, "top_p") else 1.0
        repetition_penalty = self.args.repetition_penalty if hasattr(self.args, "repetition_penalty") else 1.0
        min_p = self.args.min_p if hasattr(self.args, "min_p") else 0.0
        presence_penalty = self.args.presence_penalty if hasattr(self.args, "presence_penalty") else 0.0

        # Start timing for vLLM generation
        start_time = time.time()

        if self.vllm_mode == "server":
            all_prompts_text = gather_object(prompts_text_for_vllm)
            if self.accelerator.is_main_process:
                completion_ids = self.vllm_client.generate(
                    prompts=all_prompts_text,
                    n=1,  # In GKD, we generate 1 completion per prompt from student
                    repetition_penalty=repetition_penalty,
                    temperature=temperature,
                    top_p=top_p,
                    top_k=top_k,
                    min_p=min_p,
                    max_tokens=max_completion_length,
                    presence_penalty=presence_penalty,
                    guided_decoding_regex=self.vllm_guided_decoding_regex,
                )
            else:
                completion_ids = [None] * len(all_prompts_text)
            completion_ids = broadcast_object_list(completion_ids, from_process=0)
            process_slice = slice(
                self.accelerator.process_index * len(prompts_text_for_vllm),
                (self.accelerator.process_index + 1) * len(prompts_text_for_vllm),
            )
            completion_ids = completion_ids[process_slice]
        elif self.vllm_mode == "colocate":
            if self.vllm_guided_decoding_regex:
                guided_decoding = GuidedDecodingParams(
                    backend="outlines", regex=self.vllm_guided_decoding_regex
                )
            else:
                guided_decoding = None
            sampling_params = SamplingParams(
                n=1,
                repetition_penalty=repetition_penalty,
                temperature=temperature,
                top_p=top_p,
                top_k=top_k,
                min_p=min_p,
                max_tokens=max_completion_length,
                presence_penalty=presence_penalty,
                guided_decoding=guided_decoding,
            )

            if hasattr(self, "vllm_tp_group") and self.vllm_tensor_parallel_size > 1:
                # Gather prompts from all ranks in the TP group and flatten.
                # Each rank starts with its own prompts; after gathering, all ranks see the full group set.
                orig_size = len(prompts_text_for_vllm)
                gathered_prompts = [None for _ in range(self.vllm_tensor_parallel_size)]
                torch.distributed.all_gather_object(
                    gathered_prompts, prompts_text_for_vllm, group=self.vllm_tp_group
                )
                all_prompts_text = [p for sublist in gathered_prompts for p in sublist]
            else:
                all_prompts_text = prompts_text_for_vllm

            all_outputs = self.vllm_engine.generate(
                all_prompts_text, sampling_params=sampling_params, use_tqdm=False
            )
            completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs]

            if hasattr(self, "vllm_tp_group") and self.vllm_tensor_parallel_size > 1:
                # Slice completions for this rank within its TP group.
                # Each rank generates all outputs — we keep only our share.
                local_rank_in_group = torch.distributed.get_rank(group=self.vllm_tp_group)
                tp_slice = slice(local_rank_in_group * orig_size, (local_rank_in_group + 1) * orig_size)
                completion_ids = completion_ids[tp_slice]

            if self.vllm_enable_sleep_mode:
                self.vllm_engine.sleep(level=2)
        else:
            raise ValueError(f"Unknown vllm_mode: {self.vllm_mode}")

        # Calculate and print vLLM generation statistics
        elapsed_time = time.time() - start_time
        total_completion_tokens = sum(len(ids) for ids in completion_ids)
        num_prompts = len(completion_ids)
        avg_completion_length = total_completion_tokens / num_prompts if num_prompts > 0 else 0
        tokens_per_sec = total_completion_tokens / elapsed_time if elapsed_time > 0 else 0
        print(
            f"vLLM generation done - elapsed time: {elapsed_time:.2f}s, prompts: {num_prompts}, total tokens: {total_completion_tokens}, avg length: {avg_completion_length:.1f}, speed: {tokens_per_sec:.1f} tok/s"
        )

        # We need to combine prompt and completion for new_input_ids
        # Tokenize prompts again to get prompt_ids on the correct device and format
        # Use prompts_text_for_vllm (without special tokens) for tokenization since vLLM expects clean text
        # Ensure add_special_tokens=False as vLLM typically handles prompts as raw text
[metadata] {"noOutputExpected": false}

148
vLLM generation done - elapsed time: 3.74s, prompts: 4, total tokens: 3413, avg length: 853.2, speed: 912.6 tok/s
vLLM generation done - elapsed time: 3.79s, prompts: 4, total tokens: 4096, avg length: 1024.0, speed: 1079.6 tok/s
'loss': -0.0021, 'grad_norm': 0.049206171184778214, 'learning_rate': 5e-06, 'on_policy_loss': -0.0021, 'epoch': 0.02
'loss': -0.002, 'grad_norm': 0.050117623060941696, 'learning_rate': 5e-06, 'on_policy_loss': -0.002, 'epoch': 0.02
17273

[stdout]
148
vLLM generation done - elapsed time: 3.74s, prompts: 4, total tokens: 3413, avg length: 853.2, speed: 912.6 tok/s
vLLM generation done - elapsed time: 3.79s, prompts: 4, total tokens: 4096, avg length: 1024.0, speed: 1079.6 tok/s
'loss': -0.0021, 'grad_norm': 0.049206171184778214, 'learning_rate': 5e-06, 'on_policy_loss': -0.0021, 'epoch': 0.02
'loss': -0.002, 'grad_norm': 0.050117623060941696, 'learning_rate': 5e-06, 'on_policy_loss': -0.002, 'epoch': 0.02
17273
[metadata] {"noOutputExpected": false}

chained
=== [e2] train start 04:56:34  (remaining 17493s)

[stdout]
chained
=== [e2] train start 04:56:34  (remaining 17493s)
[metadata] {"noOutputExpected": false}

=== [e2] train rc=0 05:12:42
/workspace/runs/e2/capped_seed42/checkpoint-100
=== [e2] eval start 05:12:42
'loss': -0.0131, 'grad_norm': 0.04136570915579796
16501

[stdout]
=== [e2] train rc=0 05:12:42
/workspace/runs/e2/capped_seed42/checkpoint-100
=== [e2] eval start 05:12:42
'loss': -0.0131, 'grad_norm': 0.04136570915579796
16501
[metadata] {"noOutputExpected": false}

=== peak mem in log:

[stdout]
=== peak mem in log:
[metadata] {"noOutputExpected": false}