diff --git a/es-fine-tuning-paper/.gitignore b/es-fine-tuning-paper/.gitignore
new file mode 100644
index 0000000..ec5ff4a
--- /dev/null
+++ b/es-fine-tuning-paper/.gitignore
@@ -0,0 +1,27 @@
+# Created by venv; see https://docs.python.org/3/library/venv.html
+bin
+include
+lib
+
+# Evaluation results and datasets (too large for git)
+src/evals/
+src/data/
+src/countdown/
+
+# Python cache
+__pycache__/
+*.pyc
+*.pyo
+*.pyd
+.Python
+
+# Jupyter
+.ipynb_checkpoints/
+
+# macOS
+.DS_Store
+
+# Environment
+.env
+.venv
+venv/
\ No newline at end of file
diff --git a/es-fine-tuning-paper/LICENSE.txt b/es-fine-tuning-paper/LICENSE.txt
new file mode 100644
index 0000000..9c42e4e
--- /dev/null
+++ b/es-fine-tuning-paper/LICENSE.txt
@@ -0,0 +1,225 @@
+License
+The following license governs the use of es-fine-tuning-paper Software in academic
+and educational environments. Commercial use requires a commercial license from
+Cognizant Technology Solutions Corp, www.cognizant.com.
+
+ACADEMIC PUBLIC LICENSE
+Written by Andras Varga (license text is in public domain)
+Adapted by Karl Mutch (license text is in public domain)
+Adapted by Garrett Bingham (license text is in public domain)
+Adapted by Xin Qiu (license text is in public domain)
+
+Preamble
+This license contains the terms and conditions of using es-fine-tuning-paper Software
+in noncommercial settings: at academic institutions for teaching and research
+use, and at non-profit research organizations. You will find that this license
+provides noncommercial users of es-fine-tuning-paper Software with rights that are
+similar to the well-known GNU General Public License, yet it retains the
+possibility for es-fine-tuning-paper Software authors to financially support the
+development by selling commercial licenses. In fact, if you intend to use
+es-fine-tuning-paper Software in a “for-profit” environment, where research is
+conducted to develop or enhance a product, is used in a commercial service
+offering, or when a commercial company uses es-fine-tuning-paper Software to
+participate in a research project (for example government-funded or EU-funded
+research projects), then you need to obtain a commercial license for
+es-fine-tuning-paper Software. In that case, please contact the Author to inquire
+about commercial licenses.
+
+What are the rights given to noncommercial users? Similarly to GPL, you have the
+right to use the software, to distribute copies, to receive source code, to
+ehange the software and distribute your modifications or the modified software.
+Also similarly to the GPL, if you distribute verbatim or modified copies of
+this software, they must be distributed under this license.
+
+By modeling the GPL, this license guarantees that you’re safe when using
+es-fine-tuning-paper Software in your work, for teaching or research. This license
+guarantees that es-fine-tuning-paper Software will remain available free of charge
+for nonprofit use. You can modify es-fine-tuning-paper Software to your purposes,
+and you can also share your modifications. Even in the unlikely case of the
+authors abandoning es-fine-tuning-paper Software entirely, this license permits
+anyone to continue developing it from the last release, and to create further
+releases under this license.
+
+We believe that the combination of noncommercial open-source and commercial
+licensing will be beneficial for the whole user community, because income from
+commercial licenses will enable faster development and a higher level of
+software quality, while further enjoying the informal, open communication and
+collaboration channels of open source development.
+
+The precise terms and conditions for using, copying, distribution and
+modification follow.
+
+TERMS AND CONDITIONS FOR USE, COPYING, DISTRIBUTION AND MODIFICATION
+
+Definitions
+“Program” means a copy of es-fine-tuning-paper Software, which is said to be
+distributed under this Academic Public License.
+
+“Work based on the Program” means either the Program or any derivative work
+under copyright law: that is to say, a work containing the Program or a portion
+of it, either verbatim or with modifications and/or translated into another
+language. (Hereinafter, translation is included without limitation in the term
+“modification”.)
+
+“Using the Program” means any act of creating executables that contain or
+directly use libraries that are part of the Program, running any of the tools
+that are part of the Program, or creating works based on the Program.
+
+Each licensee is addressed as “you”.
+
+§1. Permission is hereby granted to use the Program free of charge for any
+noncommercial purpose, including teaching and research at universities, colleges
+and other educational institutions, research at non-profit research
+institutions, and personal non-profit purposes. For using the Program for
+commercial purposes, including but not restricted to consulting activities,
+design of commercial hardware or software networking products, and a commercial
+entity participating in research projects, you have to contact the Author for an
+appropriate license. Permission is also granted to use the Program for a
+reasonably limited period of time for the purpose of evaluating its usefulness
+for a particular purpose.
+
+§2. You may copy and distribute verbatim copies of the Program’s source code as
+you receive it, in any medium, provided that you conspicuously and appropriately
+publish on each copy an appropriate copyright notice and disclaimer of warranty;
+keep intact all the notices that refer to this License and to the absence of any
+warranty; and give any other recipients of the Program a copy of this License
+along with the Program.
+
+§3. You may modify your copy or copies of the Program or any portion of it, thus
+forming a work based on the Program, and copy and distribute such modifications
+or work under the terms of Section 2 above, provided that you also meet all of
+these conditions:
+
+a) You must cause the modified files to carry prominent notices stating that
+you changed the files and the date of any change.
+
+b) You must cause any work that you distribute or publish, that in whole or in
+part contains or is derived from the Program or any part thereof, to be
+licensed as a whole at no charge to all third parties under the terms of this
+License.
+
+These requirements apply to the modified work as a whole. If identifiable
+sections of that work are not derived from the Program, and can be reasonably
+considered independent and separate works in themselves, then this License, and
+its terms, do not apply to those sections when you distribute them as separate
+works. But when you distribute the same sections as part of a whole which is a
+work based on the Program, the distribution of the whole must be on the terms of
+this License, whose regulations for other licensees extend to the entire whole,
+and thus to each and every part regardless of who wrote it. (If the same,
+independent sections are distributed as part of a package that is otherwise
+reliant on, or is based on the Program, then the distribution of the whole
+package, including but not restricted to the independent section, must be on the
+unmodified terms of this License, regadless of who the author of the included
+sections was.)
+
+Thus, it is not the intent of this section to claim rights or contest your
+rights to work written entirely by you; rather, the intent is to exercise the
+right to control the distribution of derivative or collective works based or
+reliant on the Program.
+
+In addition, mere aggregation of another work not based on the Program with the
+Program (or with a work based on the Program) on a volume of storage or
+distribution medium does not bring the other work under the scope of this
+License.
+
+§4. You may copy and distribute the Program (or a work based on it, under
+Section 3) in object code or executable form under the terms of Sections 2 and 3
+above provided that you also do one of the following:
+
+a) Accompany it with the complete corresponding machine-readable source code,
+which must be distributed under the terms of Sections 2 and 3 above on a medium
+customarily used for software interchange; or,
+
+b) Accompany it with a written offer, valid for at least three years, to give
+any third party, for a charge no more than your cost of physically performing
+source distribution, a complete machine-readable copy of the corresponding
+source code, to be distributed under the terms of Sections 2 and 3 above on a
+medium customarily used for software interchange; or,
+
+c) Accompany it with the information you received as to the offer to distribute
+corresponding source code. (This alternative is allowed only for noncommercial
+distribution and only if you received the program in object code or executable
+form with such an offer, in accord with Subsection b) above.)
+
+The source code for a work means the preferred form of the work for making
+modifications to it. For an executable work, complete source code means all the
+source code for all modules it contains, plus any associated interface
+definition files, plus the scripts used to control compilation and installation
+of the executable. However, as a special exception, the source code distributed
+need not include anything that is normally distributed (in either source or
+binary form) with the major components (compiler, kernel, and so on) of the
+operating system on which the executable runs, unless that component itself
+accompanies the executable.
+
+If distribution of executable or object code is made by offering access to copy
+from a designated place, then offering equivalent access to copy the source code
+from the same place counts as distribution of the source code, even though third
+parties are not compelled to copy the source along with the object code.
+
+§5. You may not copy, modify, sublicense, or distribute the Program except as
+expressly provided under this License. Any attempt otherwise to copy, modify,
+sublicense or distribute the Program is void, and will automatically terminate
+your rights under this License. However, parties who have received copies, or
+rights, from you under this License will not have their licenses terminated so
+long as such parties remain in full compliance.
+
+§6. You are not required to accept this License, since you have not signed it.
+Nothing else grants you permission to modify or distribute the Program or its
+derivative works; law prohibits these actions if you do not accept this License.
+Therefore, by modifying or distributing the Program (or any work based on the
+Program), you indicate your acceptance of this License and all its terms and
+conditions for copying, distributing or modifying the Program or works based on
+it, to do so.
+
+§7. Each time you redistribute the Program (or any work based on the Program),
+the recipient automatically receives a license from the original licensor to
+copy, distribute or modify the Program subject to these terms and conditions.
+You may not impose any further restrictions on the recipients’ exercise of the
+rights granted herein. You are not responsible for enforcing compliance by third
+parties to this License.
+
+§8. If, as a consequence of a court judgment or allegation of patent
+infringement or for any other reason (not limited to patent issues), conditions
+are imposed on you (whether by court order, agreement or otherwise) that
+contradict the conditions of this License, they do not excuse you from the
+conditions of this License. If you cannot distribute so as to satisfy
+simultaneously your obligations under this License and any other pertinent
+obligations, then as a consequence you may not distribute the Program at all.
+For example, if a patent license would not permit royalty-free redistribution of
+the Program by all those who receive copies directly or indirectly through you,
+then the only way you could satisfy both it and this License would be to refrain
+entirely from distribution of the Program.
+
+If any portion of this section is held invalid or unenforceable under any
+particular circumstance, the balance of the section is intended to apply and the
+section as a whole is intended to apply in other circumstances.
+
+§9. If the distribution and/or use of the Program are restricted in certain
+countries either by patents or by copyrighted interfaces, the original copyright
+holder who places the Program under this License may add an explicit
+geographical distribution limitation excluding those countries, so that
+distribution is permitted only in or among countries not thus excluded. In such
+case, this License incorporates the limitation as if written in the body of this
+License.
+
+NO WARRANTY
+
+§10. BECAUSE THE PROGRAM IS LICENSED FREE OF CHARGE, THERE IS NO WARRANTY FOR
+THE PROGRAM, TO THE EXTENT PERMITTED BY APPLICABLE LAW. EXCEPT WHEN OTHERWISE
+STATED IN WRITING THE COPYRIGHT HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM
+“AS IS” WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING,
+BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
+PARTICULAR PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE
+PROGRAM IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
+ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
+
+§11. IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED ON IN WRITING WILL
+ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MAY MODIFY AND/OR REDISTRIBUTE THE
+PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY GENERAL,
+SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE USE OR INABILITY
+TO USE THE PROGRAM INCLUDING BUT NOT LIMITED TO LOSS OF DATA OR DATA BEING
+RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD PARTIES OR A FAILURE OF
+THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), EVEN IF SUCH HOLDER OR OTHER
+PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.
+
+END OF TERMS AND CONDITIONS
diff --git a/es-fine-tuning-paper/README.md b/es-fine-tuning-paper/README.md
new file mode 100644
index 0000000..572c1ea
--- /dev/null
+++ b/es-fine-tuning-paper/README.md
@@ -0,0 +1,228 @@
+# Evolution Strategies vs GRPO for LLM Fine-Tuning
+
+This repository contains the implementation and experimental code for comparing **Evolution Strategies (ES)** and **Group Relative Policy Optimization (GRPO)** on mathematical reasoning tasks.
+
+> **Paper:** "Evolution Strategies at Scale: LLM Fine-Tuning Beyond Reinforcement Learning"
+> **alphaXiv:** https://alphaXiv.org/abs/2509.24372
+
+## Highlights
+
+- **ES achieves 89% accuracy on GSM8K with only 10% of training data** (vs 85.5% for GRPO) - a 3.5-point advantage
+- **GRPO dominates with more data**: 90.9% vs 86.5% on GSM8K at 40% data
+- - **GRPO** works well with base models (Llama compared to Qwen) than ES
+- Direct comparison across data regimes (10%, 40%, 70%, 100%) and model types (base vs instruct)
+
+📖 **See our [blog](https://pr-2747.storybook.alphaxiv.org/iframe.html?id=blog--post-es-for-fine-tuning-llms&viewMode=story) for detailed documentation.**
+
+
+## Key Results
+
+| Method | GSM8K (10%) | GSM8K (40%) | Countdown (10%) | Countdown (40%) |
+|--------|-------------|-------------|-----------------|-----------------|
+| **ES (Qwen-3B-Instruct)** | **89.0%** | 86.5% | **36.0%** | 35.0% |
+| **GRPO (Qwen-3B-Instruct)** | 85.5% | **90.9%** | 34.1% | **39.6%** |
+
+💡 **Key Takeaway:** ES excels with limited data (10%), while GRPO dominates with more data (40%+). On GSM8K, ES achieves a 3.5-point advantage at 10% data. On Countdown, results are mixed, with ES slightly ahead at 10% but GRPO pulling ahead at 40%.
+
+## Repository Structure
+
+```
+├── src/
+│ ├── scripts/
+│ │ ├── es/ # Evolution Strategies training
+│ │ ├── grpo/ # GRPO training scripts
+│ │ ├── evaluation/ # Model evaluation
+│ │ └── data_prep/ # Data preparation
+│ ├── rewards/ # Task-specific reward functions
+│ ├── utils/ # Utilities
+│ ├── data/ # Training/test datasets
+│ ├── evals/ # Evaluation results
+│ └── countdown/ # Countdown task implementation
+├── BLOG.md # Detailed technical blog post
+└── assets/ # Visualizations and figures
+```
+
+
+## Setup
+
+### Prerequisites
+- **Python:** 3.11 (for ES) or 3.10+ (for GRPO)
+- **Conda:** Anaconda or Miniconda (for ES environment management)
+- **Docker:** Required for GRPO training (verl-docker-run.sh)
+- **GPU:** CUDA-capable GPU(s), 80GB+ memory recommended for 3B models
+
+### Quick Setup
+
+**The `speedrun.sh` script handles all environment setup automatically:**
+
+```bash
+# For ES training - automatically creates conda env with all dependencies
+./speedrun.sh --method es --task gsm8k --train-split 0.1
+
+# For GRPO training - automatically sets up Docker container
+./speedrun.sh --method grpo --task gsm8k --train-split 0.1
+```
+
+
+## Quick Start
+
+### Unified Training Script (Recommended)
+
+Use `speedrun.sh` for automated training with proper environment setup:
+
+#### Evolution Strategies (ES)
+
+```bash
+# GSM8K with ES (10% data, 8 perturbations)
+./speedrun.sh --method es --task gsm8k --train-split 0.1 --num-samples 700
+
+# Countdown with ES (40% data)
+./speedrun.sh --method es --task countdown --train-split 0.4 --num-samples 800
+
+# ES with custom population size and iterations
+./speedrun.sh --method es --task gsm8k --population-size 30 --num-iterations 200
+```
+
+**ES Environment Setup:** The script automatically creates and configures a conda environment (`es-debug`) with:
+- Python 3.11
+- vLLM 0.11.0 (CUDA 12.9)
+- Transformers 4.57 (critical for compatibility)
+- Required dependencies (tensorboard, pandas, uv)
+
+#### Group Relative Policy Optimization (GRPO)
+
+```bash
+# GSM8K with GRPO (10% data)
+./speedrun.sh --method grpo --task gsm8k --train-split 0.1
+
+# Countdown with GRPO (40% data)
+./speedrun.sh --method grpo --task countdown --train-split 0.4
+
+# Run both ES and GRPO for comparison(./speedrun.sh --method both --task gsm8k --train-split 0.1
+```
+
+**GRPO Environment:** Uses Docker container via `verl-docker-run.sh` (automatically set up by speedrun.sh)
+
+## 📚 Training Details
+
+### Evolution Strategies (ES)
+- **Training:** Full-parameter fine-tuning (no LoRA except for 100% dataset)
+- **Population size:** N=8 or N=30 perturbations
+- **Learning rate (α):** 0.0005
+- **Noise std (σ):** 0.001
+- **Inference:** vLLM for fast parallel evaluation
+- **Speed:** [~10X speedup](https://github.com/VsonicV/es-fine-tuning-paper) using authors' accelerated implementation
+
+### GRPO (Group Relative Policy Optimization)
+- **Algorithm:** On-policy RL with group-based advantage estimation
+- **Training:** Full-parameter fine-tuning (LoRA only for 100% dataset experiments)
+ - LoRA rank: 64, LoRA alpha: 32
+- **Rollouts:** N=8 rollouts per prompt
+- **KL divergence penalty:** coef=0.001 (low_var_kl) to prevent drift from reference policy
+- **Gradient checkpointing:** Enabled for memory efficiency
+- **Distributed training:** FSDP (Fully Sharded Data Parallel)
+ - 8 GPUs per node
+ - Reference model parameter offload enabled
+ - Actor parameter offload disabled for faster training
+- **Framework:** VERL (Volcano Engine Reinforcement Learning)
+
+**Task-Specific Hyperparameters:**
+
+| Parameter | GSM8K | Countdown |
+|-----------|-------|:-----------|
+| Learning Rate | 3×10⁻⁶ | 1×10⁻⁶ |
+| Batch Size | 32 | 128 |
+| Max Prompt Length | 512 | 256 |
+| Max Response Length | 1024 | 1024 |
+| Rollouts (N) | 8 | 8 |
+| KL Coef | 0.001 | 0.001 |
+| LoRA Rank/Alpha | 64/32 | 64/32 |
+| Save Frequency | 23 steps | 100 steps |
+
+## Evaluation
+
+### Unified Evaluation Script (Recommended)
+
+Evaluate both ES and GRPO models with a single command:
+
+```bash
+# Evaluate ES model on GSM8K
+./evaluation.sh --method es --task gsm8k --train-split 0.1
+
+# Evaluate GRPO model on Countdown
+./evaluation.sh --method grpo --task countdown --train-split 0.4
+
+# Evaluate both methods and generate comparison
+./evaluation.sh --method both --task gsm8k --train-split 0.1
+```
+
+**Options:**
+- `--method`: `es`, `grpo`, or `both`
+- `--task`: `gsm8k` or `countdown`
+- `--train-split`: Training data fraction used (e.g., 0.1, 0.4)
+- `--batch-size`: Batch size for evaluation (default: 32)
+- `--num-gpus`: Number of GPUs to use (default: 4)
+- `--checkpoint-dir`: Custom checkpoint directory (optional)
+
+The script automatically:
+- Detects checkpoint locations
+- Uses vLLM for fast ES evaluation
+- Merges FSDP checkpoints for GRPO
+- Saves results to `./src/evals/`
+
+
+## Data Preparation
+
+**Data preparation is handled automatically by `speedrun.sh`.** It detects the task and train split, then prepares the data accordingly.
+
+To skip automatic data prep (if data already exists):
+```bash
+./speedrun.sh --method es --task gsm8k --train-split 0.1 --skip-data-prep
+```
+
+
+Manual data preparation (optional)
+
+### Prepare GSM8K dataset
+```bash
+bash ./src/scripts/data_prep/prepare_gsm8k_data.sh \
+ --local_dir ./src/data/gsm8k-0.1 \
+ --train_split 0.1 \
+ --test_samples 200
+```
+
+### Prepare Countdown dataset
+```bash
+bash ./src/scripts/data_prep/prepare_countdown_data.sh \
+ --local_dir ./src/data/countdown-0.4 \
+ --train_split 0.4 \
+ --test_samples 200
+```
+
+
+## Results & Analysis
+
+Detailed results and analysis are available in:
+- **[BLOG](https://pr-2747.storybook.alphaxiv.org/iframe.html?id=blog--post-es-for-fine-tuning-llms&viewMode=story)** - Technical blog post with full experimental details
+- **[eval results + rollouts/](https://huggingface.co/datasets/alphaXiv/es-grpo-results/tree/main)** - Raw evaluation results (JSON files)
+- **[assets/](assets/)** - Visualization charts
+
+## Hardware Requirements
+
+**Our Setup:**
+- **Platform:** Lambda Labs Lambda Stack 22.04
+- **GPUs:** 8× A100 (80GB)
+
+
+## Acknowledgments
+
+- Thanks to the authors of the original paper for their detailed implementation and insights. [Their work](https://alphaXiv.org/abs/2509.24372) provided a strong foundation for our experiments and analysis.
+
+## Contact & Discussions
+
+- **Issues:** Report bugs or request features via [GitHub Issues](https://github.com/alphaXiv/paper-implementations/issues)
+- **Discussions:** Join the ES fine-tuning forum in [Discussions](https://github.com/alphaXiv/paper-implementations/discussions)
+
+## 📄 License
+
+See [LICENSE.txt](LICENSE.txt) for details.
diff --git a/es-fine-tuning-paper/assets/algorithm1-es.png b/es-fine-tuning-paper/assets/algorithm1-es.png
new file mode 100644
index 0000000..0e2b0da
Binary files /dev/null and b/es-fine-tuning-paper/assets/algorithm1-es.png differ
diff --git a/es-fine-tuning-paper/assets/es_reward_cal.png b/es-fine-tuning-paper/assets/es_reward_cal.png
new file mode 100644
index 0000000..ad466bb
Binary files /dev/null and b/es-fine-tuning-paper/assets/es_reward_cal.png differ
diff --git a/es-fine-tuning-paper/assets/grpo-working.png b/es-fine-tuning-paper/assets/grpo-working.png
new file mode 100644
index 0000000..41e7c6d
Binary files /dev/null and b/es-fine-tuning-paper/assets/grpo-working.png differ
diff --git a/es-fine-tuning-paper/assets/rl-vs_es.png b/es-fine-tuning-paper/assets/rl-vs_es.png
new file mode 100644
index 0000000..84e009e
Binary files /dev/null and b/es-fine-tuning-paper/assets/rl-vs_es.png differ
diff --git a/es-fine-tuning-paper/assets/table1_data_efficiency.png b/es-fine-tuning-paper/assets/table1_data_efficiency.png
new file mode 100644
index 0000000..5b92432
Binary files /dev/null and b/es-fine-tuning-paper/assets/table1_data_efficiency.png differ
diff --git a/es-fine-tuning-paper/assets/table2_base_models.png b/es-fine-tuning-paper/assets/table2_base_models.png
new file mode 100644
index 0000000..5865544
Binary files /dev/null and b/es-fine-tuning-paper/assets/table2_base_models.png differ
diff --git a/es-fine-tuning-paper/assets/table3_population_scaling.png b/es-fine-tuning-paper/assets/table3_population_scaling.png
new file mode 100644
index 0000000..a1b5e8e
Binary files /dev/null and b/es-fine-tuning-paper/assets/table3_population_scaling.png differ
diff --git a/es-fine-tuning-paper/assets/table4_summary.png b/es-fine-tuning-paper/assets/table4_summary.png
new file mode 100644
index 0000000..b7b869d
Binary files /dev/null and b/es-fine-tuning-paper/assets/table4_summary.png differ
diff --git a/es-fine-tuning-paper/evaluation.sh b/es-fine-tuning-paper/evaluation.sh
new file mode 100755
index 0000000..c5b52c5
--- /dev/null
+++ b/es-fine-tuning-paper/evaluation.sh
@@ -0,0 +1,337 @@
+#!/bin/bash
+# Evaluation runner for ES and GRPO trained models
+# Runs evaluations on test sets and aggregates results
+
+set -e
+
+# Colors for output
+RED='\033[0;31m'
+GREEN='\033[0;32m'
+YELLOW='\033[1;33m'
+BLUE='\033[0;34m'
+NC='\033[0m' # No Color
+
+# Default values
+METHOD="" # "es", "grpo", or "both"
+TASK="" # "gsm8k" or "countdown"
+TRAIN_SPLIT="" # Training split used (e.g., 0.1, 0.4)
+MODEL_NAME="Qwen/Qwen2.5-3B-Instruct"
+MODEL_TYPE="instruct"
+BATCH_SIZE=32
+NUM_GPUS=4
+CHECKPOINT_DIR=""
+EVAL_SPLIT="test"
+HF_TOKEN_VALUE=""
+
+# Print usage
+usage() {
+ cat << EOF
+${GREEN}Evaluation Runner for ES and GRPO Models${NC}
+Usage: $0 --method --task --train-split [OPTIONS]
+
+${YELLOW}Required:${NC}
+ --method METHOD Evaluation method: 'es', 'grpo', or 'both'
+ --task TASK Task to evaluate: 'gsm8k' or 'countdown'
+ --train-split FRACTION Train split used during training (e.g., 0.1, 0.4)
+
+${YELLOW}Optional:${NC}
+ --model MODEL_NAME Model name (default: Qwen/Qwen2.5-3B-Instruct)
+ --model-type TYPE Model type: 'instruct' or 'base' (default: instruct)
+ --checkpoint-dir DIR Override checkpoint directory auto-detection
+ --batch-size N Batch size for evaluation (default: 32)
+ --num-gpus N Number of GPUs to use (default: 4)
+ --eval-split SPLIT Split to evaluate on (default: test)
+ --hf-token TOKEN HuggingFace token (or set HF_TOKEN env var)
+
+${YELLOW}Examples:${NC}
+ # Evaluate ES model on GSM8K test set
+ $0 --method es --task gsm8k --train-split 0.1
+
+ # Evaluate GRPO model on Countdown
+ $0 --method grpo --task countdown --train-split 0.4
+
+ # Evaluate both methods
+ $0 --method both --task gsm8k --train-split 0.1
+
+ # Custom checkpoint directory
+ $0 --method es --task gsm8k --train-split 0.1 --checkpoint-dir ./my_checkpoints
+
+EOF
+ exit 1
+}
+
+# Parse arguments
+while [[ $# -gt 0 ]]; do
+ case $1 in
+ --method)
+ METHOD="$2"
+ shift 2
+ ;;
+ --task)
+ TASK="$2"
+ shift 2
+ ;;
+ --train-split)
+ TRAIN_SPLIT="$2"
+ shift 2
+ ;;
+ --model)
+ MODEL_NAME="$2"
+ shift 2
+ ;;
+ --model-type)
+ MODEL_TYPE="$2"
+ shift 2
+ ;;
+ --checkpoint-dir)
+ CHECKPOINT_DIR="$2"
+ shift 2
+ ;;
+ --batch-size)
+ BATCH_SIZE="$2"
+ shift 2
+ ;;
+ --num-gpus)
+ NUM_GPUS="$2"
+ shift 2
+ ;;
+ --eval-split)
+ EVAL_SPLIT="$2"
+ shift 2
+ ;;
+ --hf-token)
+ HF_TOKEN_VALUE="$2"
+ shift 2
+ ;;
+ --help)
+ usage
+ ;;
+ *)
+ echo -e "${RED}Error: Unknown option $1${NC}"
+ usage
+ ;;
+ esac
+done
+
+# Validate required arguments
+if [ -z "$METHOD" ]; then
+ echo -e "${RED}Error: --method is required${NC}"
+ usage
+fi
+
+if [ "$METHOD" != "es" ] && [ "$METHOD" != "grpo" ] && [ "$METHOD" != "both" ]; then
+ echo -e "${RED}Error: --method must be 'es', 'grpo', or 'both'${NC}"
+ usage
+fi
+
+if [ -z "$TASK" ]; then
+ echo -e "${RED}Error: --task is required${NC}"
+ usage
+fi
+
+if [ "$TASK" != "gsm8k" ] && [ "$TASK" != "countdown" ]; then
+ echo -e "${RED}Error: --task must be 'gsm8k' or 'countdown'${NC}"
+ usage
+fi
+
+if [ -z "$TRAIN_SPLIT" ]; then
+ echo -e "${RED}Error: --train-split is required${NC}"
+ usage
+fi
+
+# Set environment variables
+if [ -n "$HF_TOKEN_VALUE" ]; then
+ export HF_TOKEN="$HF_TOKEN_VALUE"
+fi
+
+if [ -z "$HF_TOKEN" ]; then
+ echo -e "${YELLOW}Warning: HF_TOKEN not set. Some models may not be accessible.${NC}"
+fi
+
+# Get script directory
+SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
+cd "$SCRIPT_DIR"
+
+# Set data directory
+if [ "$TASK" == "gsm8k" ]; then
+ DATA_DIR="./src/data/gsm8k-$TRAIN_SPLIT"
+else
+ DATA_DIR="./src/data/countdown-$TRAIN_SPLIT"
+fi
+
+echo -e "${BLUE}╔════════════════════════════════════════════════════════════╗${NC}"
+echo -e "${BLUE}║ Evaluation Runner ║${NC}"
+echo -e "${BLUE}╚════════════════════════════════════════════════════════════╝${NC}"
+echo ""
+echo -e "${GREEN}Configuration:${NC}"
+echo -e " Method: ${YELLOW}$METHOD${NC}"
+echo -e " Task: ${YELLOW}$TASK${NC}"
+echo -e " Train Split: ${YELLOW}$TRAIN_SPLIT${NC}"
+echo -e " Eval Split: ${YELLOW}$EVAL_SPLIT${NC}"
+echo -e " Model: ${YELLOW}$MODEL_NAME${NC}"
+echo -e " Model Type: ${YELLOW}$MODEL_TYPE${NC}"
+echo -e " Batch Size: ${YELLOW}$BATCH_SIZE${NC}"
+echo -e " Num GPUs: ${YELLOW}$NUM_GPUS${NC}"
+echo -e " Data Dir: ${YELLOW}$DATA_DIR${NC}"
+echo ""
+
+# Function to evaluate ES model
+evaluate_es() {
+ echo -e "${BLUE}────────────────────────────────────────────${NC}"
+ echo -e "${BLUE}Evaluating ES Model${NC}"
+ echo -e "${BLUE}────────────────────────────────────────────${NC}"
+
+ # Auto-detect checkpoint directory if not provided
+ if [ -z "$CHECKPOINT_DIR" ]; then
+ if [ "$TASK" == "gsm8k" ]; then
+ ES_CHECKPOINT="./checkpoints/es_gsm8k_${MODEL_TYPE}_split${TRAIN_SPLIT}"
+ else
+ ES_CHECKPOINT="./checkpoints/es_countdown_split${TRAIN_SPLIT}"
+ fi
+ else
+ ES_CHECKPOINT="$CHECKPOINT_DIR"
+ fi
+
+ echo -e "Checkpoint: ${YELLOW}$ES_CHECKPOINT${NC}"
+
+ # Check if checkpoint exists
+ if [ ! -d "$ES_CHECKPOINT" ]; then
+ echo -e "${RED}Error: ES checkpoint directory not found: $ES_CHECKPOINT${NC}"
+ return 1
+ fi
+
+ # Find the latest checkpoint
+ LATEST_CHECKPOINT=$(ls -d "$ES_CHECKPOINT"/iteration_* 2>/dev/null | sort -V | tail -n 1)
+
+ if [ -z "$LATEST_CHECKPOINT" ]; then
+ echo -e "${YELLOW}Warning: No iteration checkpoints found, using base checkpoint${NC}"
+ LATEST_CHECKPOINT="$ES_CHECKPOINT"
+ else
+ echo -e "Latest checkpoint: ${YELLOW}$(basename $LATEST_CHECKPOINT)${NC}"
+ fi
+
+ # Determine output file
+ MODEL_BASENAME=$(echo "$MODEL_NAME" | tr '/' '_' | tr '[:upper:]' '[:lower:]')
+ OUTPUT_FILE="./src/evals/es-evals/${MODEL_BASENAME}_${MODEL_TYPE}_eval_results_${TASK}_es_${TRAIN_SPLIT}.json"
+ mkdir -p "$(dirname "$OUTPUT_FILE")"
+
+ echo -e "Output: ${YELLOW}$OUTPUT_FILE${NC}"
+ echo ""
+
+ # Run evaluation using vLLM
+ if [ "$TASK" == "gsm8k" ]; then
+ python ./src/scripts/evaluation/eval_gsm8k_vllm.py \
+ --model_path "$LATEST_CHECKPOINT" \
+ --data_file "$DATA_DIR/${EVAL_SPLIT}.parquet" \
+ --output_file "$OUTPUT_FILE" \
+ --batch_size "$BATCH_SIZE" \
+ --num_gpus "$NUM_GPUS"
+ else
+ python ./src/scripts/evaluation/eval_countdown_vllm.py \
+ --model_path "$LATEST_CHECKPOINT" \
+ --data_file "$DATA_DIR/${EVAL_SPLIT}.parquet" \
+ --output_file "$OUTPUT_FILE" \
+ --batch_size "$BATCH_SIZE" \
+ --num_gpus "$NUM_GPUS"
+ fi
+
+ echo -e "${GREEN}✓ ES Evaluation complete${NC}"
+ echo -e " Results saved to: ${YELLOW}$OUTPUT_FILE${NC}"
+ echo ""
+}
+
+# Function to evaluate GRPO model
+evaluate_grpo() {
+ echo -e "${BLUE}────────────────────────────────────────────${NC}"
+ echo -e "${BLUE}Evaluating GRPO Model${NC}"
+ echo -e "${BLUE}────────────────────────────────────────────${NC}"
+
+ # Auto-detect checkpoint directory if not provided
+ if [ -z "$CHECKPOINT_DIR" ]; then
+ if [ "$TASK" == "gsm8k" ]; then
+ GRPO_CHECKPOINT="./checkpoints/verl_grpo_gsm8k_${MODEL_TYPE}"
+ else
+ GRPO_CHECKPOINT="./checkpoints/verl_grpo_countdown"
+ fi
+ else
+ GRPO_CHECKPOINT="$CHECKPOINT_DIR"
+ fi
+
+ echo -e "Checkpoint: ${YELLOW}$GRPO_CHECKPOINT${NC}"
+
+ # Check if checkpoint exists
+ if [ ! -d "$GRPO_CHECKPOINT" ]; then
+ echo -e "${RED}Error: GRPO checkpoint directory not found: $GRPO_CHECKPOINT${NC}"
+ return 1
+ fi
+
+ # Determine output file
+ MODEL_BASENAME=$(echo "$MODEL_NAME" | tr '/' '_' | tr '[:upper:]' '[:lower:]')
+ OUTPUT_FILE="./src/evals/${MODEL_BASENAME}_${MODEL_TYPE}_eval_results_${TASK}_${TRAIN_SPLIT}.json"
+ mkdir -p "$(dirname "$OUTPUT_FILE")"
+
+ echo -e "Output: ${YELLOW}$OUTPUT_FILE${NC}"
+ echo ""
+
+ # Run evaluation script (this will merge checkpoints and evaluate)
+ if [ "$TASK" == "gsm8k" ]; then
+ bash ./src/scripts/evaluation/evaluate_gsm8k.sh \
+ --checkpoint_dir "$GRPO_CHECKPOINT" \
+ --data_file "$DATA_DIR/${EVAL_SPLIT}.parquet" \
+ --output_file "$OUTPUT_FILE" \
+ --batch_size "$BATCH_SIZE"
+ else
+ bash ./src/scripts/evaluation/evaluate_countdown.sh \
+ --checkpoint_dir "$GRPO_CHECKPOINT" \
+ --data_file "$DATA_DIR/${EVAL_SPLIT}.parquet" \
+ --output_file "$OUTPUT_FILE" \
+ --batch_size "$BATCH_SIZE"
+ fi
+
+ echo -e "${GREEN}✓ GRPO Evaluation complete${NC}"
+ echo -e " Results saved to: ${YELLOW}$OUTPUT_FILE${NC}"
+ echo ""
+}
+
+# Run evaluations
+if [ "$METHOD" == "es" ]; then
+ evaluate_es
+elif [ "$METHOD" == "grpo" ]; then
+ evaluate_grpo
+elif [ "$METHOD" == "both" ]; then
+ evaluate_es
+ evaluate_grpo
+
+ # Generate comparison summary
+ echo -e "${BLUE}────────────────────────────────────────────${NC}"
+ echo -e "${BLUE}Generating Comparison Summary${NC}"
+ echo -e "${BLUE}────────────────────────────────────────────${NC}"
+
+ # This would call a Python script to compare results
+ # python ./src/scripts/evaluation/compare_results.py \
+ # --es_results "$OUTPUT_FILE_ES" \
+ # --grpo_results "$OUTPUT_FILE_GRPO" \
+ # --output "./src/evals/comparison_${TASK}_${TRAIN_SPLIT}.json"
+
+ echo -e "${YELLOW}Note: Run generate_charts.py to visualize comparison${NC}"
+ echo ""
+fi
+
+echo ""
+echo -e "${GREEN}═══════════════════════════════════════${NC}"
+echo -e "${GREEN}✓ Evaluation Complete!${NC}"
+echo -e "${GREEN}═══════════════════════════════════════${NC}"
+echo -e "Method: ${YELLOW}$METHOD${NC}"
+echo -e "Task: ${YELLOW}$TASK${NC}"
+echo -e "Train Split: ${YELLOW}$TRAIN_SPLIT${NC}"
+echo ""
+echo -e "${YELLOW}Results location:${NC}"
+echo -e " ${YELLOW}./src/evals/${NC}"
+echo ""
+echo -e "${YELLOW}Next steps:${NC}"
+echo -e "1. Review results in evaluation output files"
+echo -e "2. Generate visualizations:"
+echo -e " ${YELLOW}python src/scripts/generate_charts.py${NC}"
+echo -e "3. Upload results to HuggingFace:"
+echo -e " ${YELLOW}python upload_inference_results.py${NC}"
+echo -e "${GREEN}═══════════════════════════════════════${NC}"
diff --git a/es-fine-tuning-paper/pyvenv.cfg b/es-fine-tuning-paper/pyvenv.cfg
new file mode 100644
index 0000000..83c2108
--- /dev/null
+++ b/es-fine-tuning-paper/pyvenv.cfg
@@ -0,0 +1,5 @@
+home = /opt/miniconda3/bin
+include-system-site-packages = false
+version = 3.13.2
+executable = /opt/miniconda3/bin/python3.13
+command = /opt/miniconda3/bin/python -m venv /Users/yuvrajsingh9886/Desktop/paper-implementations/es-fine-tuning-paper
diff --git a/es-fine-tuning-paper/requirement.txt b/es-fine-tuning-paper/requirement.txt
new file mode 100644
index 0000000..0723247
--- /dev/null
+++ b/es-fine-tuning-paper/requirement.txt
@@ -0,0 +1,6 @@
+transformers>=4.30.0
+accelerate>=0.20.0
+numpy>=1.21.0
+psutil>=5.8.0
+pandas>=2.3.0
+datasets
\ No newline at end of file
diff --git a/es-fine-tuning-paper/speedrun.sh b/es-fine-tuning-paper/speedrun.sh
new file mode 100755
index 0000000..a122db9
--- /dev/null
+++ b/es-fine-tuning-paper/speedrun.sh
@@ -0,0 +1,482 @@
+#!/bin/bash
+# Speedrun script for ES and GRPO training on GSM8K and Countdown tasks
+# Uses Docker environment setup via verl-docker-run.sh
+# Handles data preparation and training execution with proper environment setup
+
+set -e
+
+# Colors for output
+RED='\033[0;31m'
+GREEN='\033[0;32m'
+YELLOW='\033[1;33m'
+BLUE='\033[0;34m'
+NC='\033[0m' # No Color
+
+# Default values
+METHOD="" # "es" or "grpo"
+TASK="" # "gsm8k" or "countdown"
+MODEL_NAME="Qwen/Qwen2.5-3B-Instruct"
+MODEL_TYPE="instruct"
+TRAIN_SPLIT=0.1
+NUM_SAMPLES=700 # For ES training
+TEST_SAMPLES=200
+POPULATION_SIZE=8 # For ES
+NUM_ITERATIONS=100 # For ES
+CUDA_DEVICES="0,1,2,3"
+NUM_ENGINES=4 # For ES vLLM engines
+SKIP_DATA_PREP=false
+SKIP_DOCKER_SETUP=false
+HF_TOKEN_VALUE=""
+WANDB_API_KEY_VALUE=""
+
+# Print usage
+usage() {
+ cat << EOF
+${GREEN}ES vs GRPO Speedrun Script${NC}
+Usage: $0 --method --task [OPTIONS]
+
+${YELLOW}Required:${NC}
+ --method METHOD Training method: 'es' or 'grpo' or 'both'
+ --task TASK Task to run: 'gsm8k' or 'countdown'
+
+${YELLOW}Optional - General:${NC}
+ --model MODEL_NAME Model name (default: Qwen/Qwen2.5-3B-Instruct)
+ --model-type TYPE Model type: 'instruct' or 'base' (default: instruct)
+ --train-split FRACTION Fraction of data for training (default: 0.1)
+ --test-samples N Number of samples reserved for test (default: 200)
+ --skip-data-prep Skip data preparation step
+ --skip-docker Skip Docker environment setup
+ --hf-token TOKEN HuggingFace token (or set HF_TOKEN env var)
+ --wandb-key KEY Weights & Biases API key (or set WANDB_API_KEY env var)
+
+${YELLOW}Optional - ES Specific:${NC}
+ --num-samples N Number of training samples for ES (default: 700)
+ --population-size N ES population size (default: 8)
+ --num-iterations N ES iterations (default: 100)
+ --cuda-devices DEVICES CUDA devices for ES (default: 0,1,2,3)
+ --num-engines N Number of vLLM engines for ES (default: 4)
+
+${YELLOW}Examples:${NC}
+ # Run ES training on GSM8K with 10% data
+ $0 --method es --task gsm8k --train-split 0.1 --num-samples 700
+
+ # Run GRPO training on GSM8K with 40% data
+ $0 --method grpo --task gsm8k --train-split 0.4
+
+ # Run both ES and GRPO on Countdown
+ $0 --method both --task countdown --train-split 0.4
+
+ # Run with base model
+ $0 --method grpo --task gsm8k --model-type base --model Qwen/Qwen2.5-3B
+
+ # Skip Docker setup (if already running)
+ $0 --method es --task gsm8k --skip-docker
+
+EOF
+ exit 1
+}
+
+# Parse arguments
+while [[ $# -gt 0 ]]; do
+ case $1 in
+ --method)
+ METHOD="$2"
+ shift 2
+ ;;
+ --task)
+ TASK="$2"
+ shift 2
+ ;;
+ --model)
+ MODEL_NAME="$2"
+ shift 2
+ ;;
+ --model-type)
+ MODEL_TYPE="$2"
+ shift 2
+ ;;
+ --train-split)
+ TRAIN_SPLIT="$2"
+ shift 2
+ ;;
+ --num-samples)
+ NUM_SAMPLES="$2"
+ shift 2
+ ;;
+ --test-samples)
+ TEST_SAMPLES="$2"
+ shift 2
+ ;;
+ --population-size)
+ POPULATION_SIZE="$2"
+ shift 2
+ ;;
+ --num-iterations)
+ NUM_ITERATIONS="$2"
+ shift 2
+ ;;
+ --cuda-devices)
+ CUDA_DEVICES="$2"
+ shift 2
+ ;;
+ --num-engines)
+ NUM_ENGINES="$2"
+ shift 2
+ ;;
+ --skip-data-prep)
+ SKIP_DATA_PREP=true
+ shift
+ ;;
+ --skip-docker)
+ SKIP_DOCKER_SETUP=true
+ shift
+ ;;
+ --hf-token)
+ HF_TOKEN_VALUE="$2"
+ shift 2
+ ;;
+ --wandb-key)
+ WANDB_API_KEY_VALUE="$2"
+ shift 2
+ ;;
+ --help)
+ usage
+ ;;
+ *)
+ echo -e "${RED}Error: Unknown option $1${NC}"
+ usage
+ ;;
+ esac
+done
+
+# Validate required arguments
+if [ -z "$METHOD" ]; then
+ echo -e "${RED}Error: --method is required${NC}"
+ usage
+fi
+
+if [ "$METHOD" != "es" ] && [ "$METHOD" != "grpo" ] && [ "$METHOD" != "both" ]; then
+ echo -e "${RED}Error: --method must be 'es', 'grpo', or 'both'${NC}"
+ usage
+fi
+
+if [ -z "$TASK" ]; then
+ echo -e "${RED}Error: --task is required${NC}"
+ usage
+fi
+
+if [ "$TASK" != "gsm8k" ] && [ "$TASK" != "countdown" ]; then
+ echo -e "${RED}Error: --task must be 'gsm8k' or 'countdown'${NC}"
+ usage
+fi
+
+# Validate model type
+if [ "$MODEL_TYPE" != "instruct" ] && [ "$MODEL_TYPE" != "base" ]; then
+ echo -e "${RED}Error: --model-type must be 'instruct' or 'base'${NC}"
+ usage
+fi
+
+# Set environment variables
+if [ -n "$HF_TOKEN_VALUE" ]; then
+ export HF_TOKEN="$HF_TOKEN_VALUE"
+fi
+
+if [ -n "$WANDB_API_KEY_VALUE" ]; then
+ export WANDB_API_KEY="$WANDB_API_KEY_VALUE"
+fi
+
+# Check for required environment variables
+if [ -z "$HF_TOKEN" ]; then
+ echo -e "${RED}Error: HF_TOKEN not set. Use --hf-token or set HF_TOKEN environment variable${NC}"
+ exit 1
+fi
+
+if [ "$METHOD" == "grpo" ] || [ "$METHOD" == "both" ]; then
+ if [ -z "$WANDB_API_KEY" ]; then
+ echo -e "${YELLOW}Warning: WANDB_API_KEY not set. GRPO training may fail without it.${NC}"
+ echo -e "${YELLOW}Use --wandb-key or set WANDB_API_KEY environment variable${NC}"
+ fi
+fi
+
+# Get script directory
+SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
+cd "$SCRIPT_DIR"
+
+# Function to setup ES environment
+setup_es_environment() {
+ echo -e "${GREEN}[ES Setup] Checking ES training environment...${NC}"
+
+ # Check if conda is available
+ if ! command -v conda &> /dev/null; then
+ echo -e "${RED}Error: conda not found. Please install Anaconda/Miniconda first.${NC}"
+ exit 1
+ fi
+
+ # Check if environment exists
+ if ! conda env list | grep -q "^es-debug "; then
+ echo -e "${YELLOW}ES environment not found. Creating 'es-debug' environment...${NC}"
+
+ # Create conda environment with Python 3.11
+ conda create -n es-debug python=3.11 -y
+ if [ $? -ne 0 ]; then
+ echo -e "${RED}Error: Failed to create conda environment${NC}"
+ exit 1
+ fi
+
+ echo -e "${GREEN}✓ Created es-debug environment${NC}"
+
+ # Activate and install dependencies
+ echo -e "${YELLOW}Installing dependencies...${NC}"
+ eval "$(conda shell.bash hook)"
+ conda activate es-debug
+
+ # Install base packages
+ pip install uv tensorboard pandas
+
+ # Install vLLM with CUDA 12.9 support
+ echo -e "${YELLOW}Installing vLLM 0.11.0 (this may take a few minutes)...${NC}"
+ uv pip install vllm==0.11.0 --torch-backend=cu129
+
+ # CRITICAL: Install transformers 4.57 (HF 5.x has breaking changes)
+ echo -e "${YELLOW}Installing transformers==4.57 (critical for compatibility)...${NC}"
+ pip install transformers==4.57
+
+ echo -e "${GREEN}✓ ES environment setup complete${NC}"
+ else
+ echo -e "${GREEN}✓ ES environment 'es-debug' already exists${NC}"
+
+ # Activate environment
+ eval "$(conda shell.bash hook)"
+ conda activate es-debug
+
+ # Verify critical packages
+ if ! python -c "import vllm" 2>/dev/null; then
+ echo -e "${YELLOW}vLLM not found, installing...${NC}"
+ uv pip install vllm==0.11.0 --torch-backend=cu129
+ fi
+
+ if ! python -c "import transformers; assert transformers.__version__.startswith('4.57')" 2>/dev/null; then
+ echo -e "${YELLOW}Transformers version mismatch, installing 4.57...${NC}"
+ pip install transformers==4.57
+ fi
+ fi
+
+ echo -e "${GREEN}✓ ES environment ready${NC}"
+ echo ""
+}
+
+echo -e "${BLUE}╔════════════════════════════════════════════════════════════╗${NC}"
+echo -e "${BLUE}║ ES vs GRPO Training Speedrun ║${NC}"
+echo -e "${BLUE}╚════════════════════════════════════════════════════════════╝${NC}"
+echo ""
+echo -e "${GREEN}Configuration:${NC}"
+echo -e " Method: ${YELLOW}$METHOD${NC}"
+echo -e " Task: ${YELLOW}$TASK${NC}"
+echo -e " Model: ${YELLOW}$MODEL_NAME${NC}"
+echo -e " Model Type: ${YELLOW}$MODEL_TYPE${NC}"
+echo -e " Train Split: ${YELLOW}$TRAIN_SPLIT${NC}"
+echo -e " Test Samples: ${YELLOW}$TEST_SAMPLES${NC}"
+if [ "$METHOD" == "es" ] || [ "$METHOD" == "both" ]; then
+ echo -e " ES Samples: ${YELLOW}$NUM_SAMPLES${NC}"
+ echo -e " Population: ${YELLOW}$POPULATION_SIZE${NC}"
+ echo -e " Iterations: ${YELLOW}$NUM_ITERATIONS${NC}"
+ echo -e " CUDA Devices: ${YELLOW}$CUDA_DEVICES${NC}"
+fi
+echo ""
+
+# Step 0: Docker environment setup (for GRPO)
+if [ "$METHOD" == "grpo" ] || [ "$METHOD" == "both" ]; then
+ if [ "$SKIP_DOCKER_SETUP" = false ]; then
+ echo -e "${GREEN}[Step 0] Setting up Docker environment for GRPO...${NC}"
+ if [ ! -f "./verl-docker-run.sh" ]; then
+ echo -e "${RED}Error: verl-docker-run.sh not found${NC}"
+ exit 1
+ fi
+ bash ./verl-docker-run.sh
+ echo -e "${GREEN}✓ Docker environment ready${NC}"
+ echo ""
+ else
+ echo -e "${GREEN}[Step 0] Skipping Docker setup (using existing container)${NC}"
+ # Verify container is running
+ if ! sudo docker ps | grep -q "verl-es-fine-tuning-paper"; then
+ echo -e "${RED}Error: Docker container 'verl-es-fine-tuning-paper' not running${NC}"
+ echo -e "${YELLOW}Run without --skip-docker to set it up${NC}"
+ exit 1
+ fi
+ echo -e "${GREEN}✓ Container verified${NC}"
+ echo ""
+ fi
+else
+ echo -e "${GREEN}[Step 0] Skipping Docker setup (ES-only mode)${NC}"
+ echo ""
+fi
+
+# Step 1: Data Preparation
+if [ "$SKIP_DATA_PREP" == "false" ]; then
+ echo -e "${GREEN}[Step 1] Data Preparation${NC}"
+
+ if [ "$TASK" == "gsm8k" ]; then
+ DATA_DIR="./src/data/gsm8k-$TRAIN_SPLIT"
+ echo -e "Preparing GSM8K data..."
+ echo -e "Output directory: ${YELLOW}$DATA_DIR${NC}"
+
+ bash ./src/scripts/data_prep/prepare_gsm8k_data.sh \
+ --local_dir "$DATA_DIR" \
+ --train_split "$TRAIN_SPLIT" \
+ --test_samples "$TEST_SAMPLES"
+
+ elif [ "$TASK" == "countdown" ]; then
+ DATA_DIR="./src/data/countdown-$TRAIN_SPLIT"
+ echo -e "Preparing Countdown data..."
+ echo -e "Output directory: ${YELLOW}$DATA_DIR${NC}"
+
+ bash ./src/scripts/data_prep/prepare_countdown_data.sh \
+ --local_dir "$DATA_DIR" \
+ --train_split "$TRAIN_SPLIT" \
+ --test_samples "$TEST_SAMPLES"
+ fi
+
+ echo -e "${GREEN}✓ Data preparation complete${NC}"
+ echo ""
+else
+ echo -e "${YELLOW}[Step 1] Skipping data preparation${NC}"
+ if [ "$TASK" == "gsm8k" ]; then
+ DATA_DIR="./src/data/gsm8k-$TRAIN_SPLIT"
+ elif [ "$TASK" == "countdown" ]; then
+ DATA_DIR="./src/data/countdown-$TRAIN_SPLIT"
+ fi
+ echo -e "Expected data directory: ${YELLOW}$DATA_DIR${NC}"
+ echo ""
+fi
+
+# Step 2: Training
+echo -e "${GREEN}[Step 2] Training${NC}"
+
+# ES Training
+if [ "$METHOD" == "es" ] || [ "$METHOD" == "both" ]; then
+ # Setup ES environment first
+ setup_es_environment
+
+ echo -e "${BLUE}Running ES Training...${NC}"
+
+ if [ "$TASK" == "gsm8k" ]; then
+ ES_SCRIPT="./src/scripts/es/es_fine_tuning_gsm8k_accl.py"
+ echo -e "ES Script: ${YELLOW}$ES_SCRIPT${NC}"
+
+ CUDA_VISIBLE_DEVICES=$CUDA_DEVICES python "$ES_SCRIPT" \
+ --model_name "$MODEL_NAME" \
+ --data_dir "$DATA_DIR" \
+ --num_samples "$NUM_SAMPLES" \
+ --population_size "$POPULATION_SIZE" \
+ --num_iterations "$NUM_ITERATIONS" \
+ --num_engines "$NUM_ENGINES" \
+ --output_dir "./checkpoints/es_gsm8k_${MODEL_TYPE}_split${TRAIN_SPLIT}"
+
+ elif [ "$TASK" == "countdown" ]; then
+ ES_SCRIPT="./src/scripts/es/es_fine_tuning_countdown_accl.py"
+ echo -e "ES Script: ${YELLOW}$ES_SCRIPT${NC}"
+
+ CUDA_VISIBLE_DEVICES=$CUDA_DEVICES python "$ES_SCRIPT" \
+ --model_name "$MODEL_NAME" \
+ --data_dir "$DATA_DIR" \
+ --num_samples "$NUM_SAMPLES" \
+ --population_size "$POPULATION_SIZE" \
+ --num_iterations "$NUM_ITERATIONS" \
+ --num_engines "$NUM_ENGINES" \
+ --output_dir "./checkpoints/es_countdown_split${TRAIN_SPLIT}"
+ fi
+
+ echo -e "${GREEN}✓ ES Training complete${NC}"
+ echo ""
+fi
+
+# GRPO Training
+if [ "$METHOD" == "grpo" ] || [ "$METHOD" == "both" ]; then
+ echo -e "${BLUE}Running GRPO Training (in Docker)...${NC}"
+
+ if [ "$TASK" == "gsm8k" ]; then
+ if [ "$MODEL_TYPE" == "base" ]; then
+ GRPO_SCRIPT="./src/scripts/grpo/grpo-gsm8k-base.sh"
+ echo -e "GRPO Script: ${YELLOW}$GRPO_SCRIPT${NC}"
+
+ # Update data paths in script
+ sudo docker exec verl-es-fine-tuning-paper bash -c \
+ "cd /workspace/es-fine-tuning-paper && \
+ sed -i \"s|data.train_files=.*train.parquet|data.train_files=$DATA_DIR/train.parquet|g\" $GRPO_SCRIPT && \
+ sed -i \"s|data.val_files=.*validation.parquet|data.val_files=$DATA_DIR/validation.parquet|g\" $GRPO_SCRIPT"
+
+ else # instruct
+ GRPO_SCRIPT="./src/scripts/grpo/grpo-gsm8k.sh"
+ echo -e "GRPO Script: ${YELLOW}$GRPO_SCRIPT${NC}"
+
+ # Update data paths in script
+ sudo docker exec verl-es-fine-tuning-paper bash -c \
+ "cd /workspace/es-fine-tuning-paper && \
+ sed -i \"s|data.train_files=.*train.parquet|data.train_files=$DATA_DIR/train.parquet|g\" $GRPO_SCRIPT && \
+ sed -i \"s|data.val_files=.*validation.parquet|data.val_files=$DATA_DIR/validation.parquet|g\" $GRPO_SCRIPT"
+ fi
+
+ # Run GRPO training with environment variables
+ sudo docker exec \
+ -e HF_TOKEN="$HF_TOKEN" \
+ -e WANDB_API_KEY="$WANDB_API_KEY" \
+ verl-es-fine-tuning-paper bash -c \
+ "cd /workspace/es-fine-tuning-paper && bash $GRPO_SCRIPT"
+
+ elif [ "$TASK" == "countdown" ]; then
+ GRPO_SCRIPT="./src/scripts/grpo/grpo-countdown-custom.sh"
+ echo -e "GRPO Script: ${YELLOW}$GRPO_SCRIPT${NC}"
+
+ # Update data paths in script
+ sudo docker exec verl-es-fine-tuning-paper bash -c \
+ "cd /workspace/es-fine-tuning-paper && \
+ sed -i \"s|data.train_files=.*train.parquet|data.train_files=$DATA_DIR/train.parquet|g\" $GRPO_SCRIPT && \
+ sed -i \"s|data.val_files=.*validation.parquet|data.val_files=$DATA_DIR/validation.parquet|g\" $GRPO_SCRIPT"
+
+ # Run GRPO training with environment variables
+ sudo docker exec \
+ -e HF_TOKEN="$HF_TOKEN" \
+ -e WANDB_API_KEY="$WANDB_API_KEY" \
+ verl-es-fine-tuning-paper bash -c \
+ "cd /workspace/es-fine-tuning-paper && bash $GRPO_SCRIPT"
+ fi
+
+ echo -e "${GREEN}✓ GRPO Training complete${NC}"
+ echo ""
+fi
+
+echo ""
+echo -e "${GREEN}═══════════════════════════════════════${NC}"
+echo -e "${GREEN}✓ All Training Complete!${NC}"
+echo -e "${GREEN}═══════════════════════════════════════${NC}"
+echo -e "Method: ${YELLOW}$METHOD${NC}"
+echo -e "Task: ${YELLOW}$TASK${NC}"
+echo -e "Model Type: ${YELLOW}$MODEL_TYPE${NC}"
+echo -e "Data Dir: ${YELLOW}$DATA_DIR${NC}"
+echo ""
+
+# Show checkpoint locations
+if [ "$METHOD" == "es" ] || [ "$METHOD" == "both" ]; then
+ if [ "$TASK" == "gsm8k" ]; then
+ ES_CHECKPOINT="./checkpoints/es_gsm8k_${MODEL_TYPE}_split${TRAIN_SPLIT}"
+ else
+ ES_CHECKPOINT="./checkpoints/es_countdown_split${TRAIN_SPLIT}"
+ fi
+ echo -e "ES Checkpoints: ${YELLOW}$ES_CHECKPOINT${NC}"
+fi
+
+if [ "$METHOD" == "grpo" ] || [ "$METHOD" == "both" ]; then
+ if [ "$TASK" == "gsm8k" ]; then
+ GRPO_CHECKPOINT="./checkpoints/verl_grpo_gsm8k_${MODEL_TYPE}"
+ else
+ GRPO_CHECKPOINT="./checkpoints/verl_grpo_countdown"
+ fi
+ echo -e "GRPO Checkpoints: ${YELLOW}$GRPO_CHECKPOINT${NC}"
+fi
+
+echo ""
+echo -e "${YELLOW}Next steps:${NC}"
+echo -e "1. Run evaluations:"
+echo -e " ${YELLOW}bash evaluation.sh --method $METHOD --task $TASK --train-split $TRAIN_SPLIT${NC}"
+echo -e "2. View results in ${YELLOW}./src/evals/${NC}"
+echo -e "${GREEN}═══════════════════════════════════════${NC}"
diff --git a/es-fine-tuning-paper/src/countdown/__pycache__/countdown_task.cpython-313.pyc b/es-fine-tuning-paper/src/countdown/__pycache__/countdown_task.cpython-313.pyc
new file mode 100644
index 0000000..1799755
Binary files /dev/null and b/es-fine-tuning-paper/src/countdown/__pycache__/countdown_task.cpython-313.pyc differ
diff --git a/es-fine-tuning-paper/src/countdown/countdown_task.py b/es-fine-tuning-paper/src/countdown/countdown_task.py
new file mode 100644
index 0000000..a213ff6
--- /dev/null
+++ b/es-fine-tuning-paper/src/countdown/countdown_task.py
@@ -0,0 +1,96 @@
+import re
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import pandas as pd
+from torch.utils.data import Dataset
+
+
+def format_reward_function(response: str, end_token: Optional[str] = None) -> float:
+ """
+ Checks if the response follows the format ......
+ """
+ # Strip end token if present
+ if end_token and response.endswith(end_token):
+ response = response[: -len(end_token)]
+
+ think_regex = r".*?<\/think>"
+ answer_regex = r".*?<\/answer>"
+ full_format_regex = r"^.*?<\/think>\n.*?<\/answer>$"
+
+ think_match = re.search(think_regex, response, re.DOTALL)
+ answer_match = re.search(answer_regex, response, re.DOTALL)
+ full_format_match = re.match(full_format_regex, response, re.DOTALL)
+
+ if full_format_match:
+ return 1.0
+
+ reward = 0.0
+
+ if think_match:
+ reward += 0.1
+
+ if answer_match:
+ reward += 0.5
+
+ return reward
+
+def answer_reward_function(
+ response: str, numbers: List[int] = None, target: int = None
+) -> float:
+ # modified
+ """
+ Checks if the last ... uses all numbers exactly once and evaluates to the target.
+ Returns 1.0 if the last one is correct, else 0.0.
+ """
+ answer_regex = r"(.*?)<\/answer>"
+ all_matches = re.findall(answer_regex, response, re.DOTALL)
+
+ if not all_matches:
+ return 0.0
+
+ # Only check the last answer
+ answer_content = all_matches[-1]
+
+ allowed_chars = r"^[0-9+\-*/() ]+$"
+
+ if not answer_content:
+ return 0.0
+ if not re.match(allowed_chars, answer_content):
+ return 0.0
+
+ # Check numbers used
+ used_numbers = [int(n) for n in re.findall(r"\d+", answer_content)]
+ if sorted(used_numbers) != sorted(numbers):
+ return 0.0
+
+ # Try evaluating
+ try:
+ result = eval(answer_content, {"__builtins__": None}, {})
+ if abs(float(result) - float(target)) < 1e-5:
+ return 1.0
+ except:
+ return 0.0
+
+ return 0.0
+
+
+def reward_function(
+ response: str,
+ numbers: List[int] = None,
+ target: int = None,
+ end_token: str = None,
+) -> Dict[str, Any]:
+ """Reward function for Countdown Tasks.
+
+ Total reward = 0.1 * format_reward + answer_reward
+ """
+ format_reward = format_reward_function("" + response, end_token)
+ answer_reward = answer_reward_function(response, numbers, target)
+ return {
+ "reward": format_reward * 0.1 + answer_reward,
+ "reward_info": {
+ "format_reward": format_reward,
+ "answer_reward": answer_reward,
+ },
+ }
\ No newline at end of file
diff --git a/es-fine-tuning-paper/src/rewards/__pycache__/gsm8k_reward.cpython-313.pyc b/es-fine-tuning-paper/src/rewards/__pycache__/gsm8k_reward.cpython-313.pyc
new file mode 100644
index 0000000..3edb693
Binary files /dev/null and b/es-fine-tuning-paper/src/rewards/__pycache__/gsm8k_reward.cpython-313.pyc differ
diff --git a/es-fine-tuning-paper/src/rewards/countdown_reward.py b/es-fine-tuning-paper/src/rewards/countdown_reward.py
new file mode 100644
index 0000000..5c0abbf
--- /dev/null
+++ b/es-fine-tuning-paper/src/rewards/countdown_reward.py
@@ -0,0 +1,95 @@
+"""
+Alternative: Custom reward function for Countdown Task
+This file can be used with VERL's custom_reward_function config
+
+Usage in config:
+custom_reward_function:
+ path: /path/to/countdown_reward.py
+ name: countdown_reward_function
+ reward_kwargs: {}
+"""
+
+import re
+from typing import Any, Dict, List, Optional
+from verl import DataProto
+
+
+def format_reward_function(response: str, end_token: Optional[str] = None) -> float:
+ """Checks if the response follows the format ......"""
+ if end_token and response.endswith(end_token):
+ response = response[: -len(end_token)]
+
+ think_regex = r".*?<\/think>"
+ answer_regex = r".*?<\/answer>"
+ full_format_regex = r"^.*?<\/think>\n.*?<\/answer>$"
+
+ think_match = re.search(think_regex, response, re.DOTALL)
+ answer_match = re.search(answer_regex, response, re.DOTALL)
+ full_format_match = re.match(full_format_regex, response, re.DOTALL)
+
+ if full_format_match:
+ return 1.0
+
+ reward = 0.0
+ if think_match:
+ reward += 0.1
+ if answer_match:
+ reward += 0.5
+
+ return reward
+
+
+def answer_reward_function(response: str, numbers: List[int], target: int) -> float:
+ """Checks if the answer uses all numbers exactly once and evaluates to the target"""
+ answer_regex = r"(.*?)<\/answer>"
+ all_matches = re.findall(answer_regex, response, re.DOTALL)
+
+ if not all_matches:
+ return 0.0
+
+ answer_content = all_matches[-1].strip()
+
+ if not answer_content:
+ return 0.0
+
+ allowed_chars = r"^[0-9+\-*/() ]+$"
+ if not re.match(allowed_chars, answer_content):
+ return 0.0
+
+ used_numbers = [int(n) for n in re.findall(r"\d+", answer_content)]
+ if sorted(used_numbers) != sorted(numbers):
+ return 0.0
+
+ try:
+ result = eval(answer_content, {"__builtins__": None}, {})
+ if abs(float(result) - float(target)) < 1e-5:
+ return 1.0
+ except:
+ pass
+
+ return 0.0
+
+
+def countdown_reward_function(data_source: str, solution_str: str, ground_truth: Dict[str, Any],
+ extra_info: Optional[Dict] = None) -> Dict[str, Any]:
+ """
+ Custom reward function for countdown task compatible with VERL's custom_reward_function mechanism.
+
+ This function signature matches what VERL expects from compute_score functions.
+ """
+ numbers = ground_truth.get("numbers", [])
+ target = ground_truth.get("target", 0)
+
+ if isinstance(target, str):
+ target = float(target) # Handle decimal targets from division
+
+ format_reward = format_reward_function("" + solution_str)
+ answer_reward = answer_reward_function(solution_str, numbers, target)
+
+ total_reward = format_reward * 0.1 + answer_reward
+
+ return {
+ "score": total_reward,
+ "format_reward": format_reward,
+ "answer_reward": answer_reward,
+ }
diff --git a/es-fine-tuning-paper/src/rewards/gsm8k_reward.py b/es-fine-tuning-paper/src/rewards/gsm8k_reward.py
new file mode 100644
index 0000000..056e7e7
--- /dev/null
+++ b/es-fine-tuning-paper/src/rewards/gsm8k_reward.py
@@ -0,0 +1,143 @@
+"""
+Reward function for GSM8K Task
+This file provides a reward function for evaluating model responses on GSM8K math problems.
+
+The reward function checks:
+1. Format: Whether the response contains the expected "#### " format
+2. Correctness: Whether the extracted answer matches the ground truth
+
+Usage:
+ from src.rewards.gsm8k_reward import reward_function
+ reward_dict = reward_function(response, ground_truth="42")
+"""
+
+import re
+from typing import Any, Dict, Optional
+
+
+def extract_solution(solution_str: str, method: str = "flexible") -> Optional[str]:
+ """
+ Extract the numerical answer from GSM8K response format.
+
+ Args:
+ solution_str: The model's response string
+ method: "strict" requires "#### " format, "flexible" extracts any number
+
+ Returns:
+ The extracted answer as a string, or None if no answer found
+ """
+ assert method in ["strict", "flexible"]
+
+ if method == "strict":
+ # Require the "#### " format (tests formatting)
+ solution = re.search(r"#### (\-?[0-9\.\,]+)", solution_str)
+ if solution is None:
+ return None
+ final_answer = solution.group(1).replace(",", "").replace("$", "")
+ return final_answer
+
+ elif method == "flexible":
+ # Find all numbers in the response and take the last one
+ answer = re.findall(r"(\-?[0-9\.\,]+)", solution_str)
+ if len(answer) == 0:
+ return None
+
+ invalid_str = ["", "."]
+ final_answer = None
+ # Find the last valid number
+ for final_answer in reversed(answer):
+ if final_answer not in invalid_str:
+ final_answer = final_answer.replace(",", "").replace("$", "")
+ break
+
+ return final_answer
+
+
+def format_reward_function(response: str, end_token: Optional[str] = None) -> float:
+ """
+ Checks if the response follows the GSM8K format with "#### "
+
+ Args:
+ response: The model's response string
+ end_token: Optional end token to strip from response
+
+ Returns:
+ 1.0 if proper format is found, 0.0 otherwise
+ """
+ # Strip end token if present
+ if end_token and response.endswith(end_token):
+ response = response[: -len(end_token)]
+
+ # Check for "#### " format
+ format_regex = r"#### (\-?[0-9\.\,]+)"
+ format_match = re.search(format_regex, response)
+
+ if format_match:
+ return 1.0
+
+ return 0.0
+
+
+def answer_reward_function(response: str, ground_truth: str) -> float:
+ """
+ Checks if the extracted answer matches the ground truth.
+
+ Args:
+ response: The model's response string
+ ground_truth: The correct answer as a string
+
+ Returns:
+ 1.0 if answer is correct, 0.0 otherwise
+ """
+ # Extract answer using flexible method (more forgiving)
+ answer = extract_solution(response, method="flexible")
+
+ if answer is None:
+ return 0.0
+
+ # Compare numerical values
+ try:
+ answer_float = float(answer)
+ gt_float = float(ground_truth.replace(",", "").replace("$", ""))
+
+ # Use small epsilon for floating point comparison
+ if abs(answer_float - gt_float) < 1e-5:
+ return 1.0
+ except (ValueError, AttributeError):
+ # Fallback to string comparison if conversion fails
+ if answer == ground_truth.replace(",", "").replace("$", ""):
+ return 1.0
+
+ return 0.0
+
+
+def reward_function(
+ response: str,
+ ground_truth: str,
+ end_token: Optional[str] = None,
+) -> Dict[str, Any]:
+ """
+ Reward function for GSM8K Tasks.
+
+ Total reward = 0.1 * format_reward + answer_reward
+
+ Args:
+ response: The model's response string
+ ground_truth: The correct answer
+ end_token: Optional end token to strip from response
+
+ Returns:
+ Dictionary containing:
+ - reward: Total reward (float between 0 and 1.1)
+ - reward_info: Breakdown of format_reward and answer_reward
+ """
+ format_reward = format_reward_function(response, end_token)
+ answer_reward = answer_reward_function(response, ground_truth)
+
+ return {
+ "reward": format_reward * 0.1 + answer_reward,
+ "reward_info": {
+ "format_reward": format_reward,
+ "answer_reward": answer_reward,
+ },
+ }
diff --git a/es-fine-tuning-paper/src/scripts/data_prep/grpo_data_countdown.py b/es-fine-tuning-paper/src/scripts/data_prep/grpo_data_countdown.py
new file mode 100644
index 0000000..f8bad47
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/data_prep/grpo_data_countdown.py
@@ -0,0 +1,235 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Preprocess the Countdown dataset to parquet format for VERL GRPO training
+Reserves last N samples for final evaluation
+"""
+
+import argparse
+import json
+import os
+import re
+from typing import Any, Dict, List, Optional
+
+import pandas as pd
+
+from verl.utils.hdfs_io import copy, makedirs
+
+
+SYSTEM_MESSAGE = (
+ "You are a helpful assistant. You first think about the reasoning process "
+ "in your mind and then provide the user with the answer."
+)
+USER_TEMPLATE = (
+ "Using the numbers {numbers}, create an equation that equals {target}. "
+ "You can use basic arithmetic operations (+, -, *, /) and each number can only be used once. "
+ "Show your work in tags. "
+ "And return the final answer in tags, for example (1 + 2) / 3 ."
+)
+RESPONSE_PROMPT = "Let me solve this step by step.\n"
+
+
+def format_reward_function(response: str, end_token: Optional[str] = None) -> float:
+ """
+ Checks if the response follows the format ......
+ """
+ # Strip end token if present
+ if end_token and response.endswith(end_token):
+ response = response[: -len(end_token)]
+
+ think_regex = r".*?<\/think>"
+ answer_regex = r".*?<\/answer>"
+ full_format_regex = r"^.*?<\/think>\n.*?<\/answer>$"
+
+ think_match = re.search(think_regex, response, re.DOTALL)
+ answer_match = re.search(answer_regex, response, re.DOTALL)
+ full_format_match = re.match(full_format_regex, response, re.DOTALL)
+
+ if full_format_match:
+ return 1.0
+
+ reward = 0.0
+
+ if think_match:
+ reward += 0.1
+
+ if answer_match:
+ reward += 0.5
+
+ return reward
+
+
+def answer_reward_function(
+ response: str, numbers: List[int] = None, target: int = None
+) -> float:
+ """
+ Checks if the answer uses all numbers exactly once and evaluates to the target
+ """
+ answer_regex = r"(.*?)<\/answer>"
+ answer_match = re.search(answer_regex, response, re.DOTALL)
+ if not answer_match:
+ return 0.0
+
+ answer_content = answer_match.group(1)
+ if not answer_content:
+ return 0.0
+
+ allowed_chars = r"^[0-9+\-*/() ]+$"
+ if not re.match(allowed_chars, answer_content):
+ return 0.0
+
+ # Check if the answer uses all numbers exactly once
+ used_numbers = [int(n) for n in re.findall(r"\d+", answer_content)]
+ if sorted(used_numbers) != sorted(numbers):
+ return 0.0
+
+ # Check if the answer evaluates to the target
+ try:
+ result = eval(answer_content, {"__builtins__": None}, {})
+ if abs(float(result) - float(target)) < 1e-5:
+ return 1.0
+ except:
+ pass
+
+ return 0.0
+
+
+def reward_function(
+ response: str,
+ numbers: List[int] = None,
+ target: int = None,
+ end_token: str = None,
+) -> Dict[str, Any]:
+ """Reward function for Countdown Tasks.
+
+ Total reward = 0.1 * format_reward + answer_reward
+ """
+ format_reward = format_reward_function("" + response, end_token)
+ answer_reward = answer_reward_function(response, numbers, target)
+ return {
+ "reward": format_reward * 0.1 + answer_reward,
+ "reward_info": {
+ "format_reward": format_reward,
+ "answer_reward": answer_reward,
+ },
+ }
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--local_dir", default="./src/data/countdown-0.4")
+ parser.add_argument("--json_file", default="./src/data/countdown-full/countdown.json")
+ parser.add_argument("--hdfs_dir", default=None)
+ parser.add_argument("--train_split", type=float, default=0.4,
+ help="Fraction of available data to use for training")
+ parser.add_argument("--test_samples", type=int, default=200,
+ help="Number of samples to reserve for final evaluation")
+
+ args = parser.parse_args()
+
+ data_source = "countdown_task"
+
+ # Load the countdown JSON data
+ with open(args.json_file, 'r') as f:
+ countdown_data = json.load(f)
+
+ print(f"Loaded {len(countdown_data)} countdown tasks")
+
+ # Process each example to VERL format
+ processed_data = []
+ for idx, example in enumerate(countdown_data):
+ numbers = example["numbers"]
+ target = float(example["target"]) # Use float to handle decimal targets
+ solution = example.get("solution", "")
+
+ # Create the user message
+ user_message = USER_TEMPLATE.format(numbers=numbers, target=target)
+
+ data = {
+ "data_source": data_source,
+ "prompt": [
+ {
+ "role": "system",
+ "content": SYSTEM_MESSAGE,
+ },
+ {
+ "role": "user",
+ "content": user_message,
+ },
+ {
+ "role": "assistant",
+ "content": RESPONSE_PROMPT,
+ }
+ ],
+ "ability": "math",
+ "reward_model": {
+ "style": "rule",
+ "ground_truth": {
+ "numbers": numbers,
+ "target": target,
+ }
+ },
+ "extra_info": {
+ "index": idx,
+ "numbers": numbers,
+ "target": target,
+ "solution": solution,
+ },
+ }
+ processed_data.append(data)
+
+ # Convert to DataFrame for easier manipulation
+ df = pd.DataFrame(processed_data)
+
+ total_samples = len(df)
+ reserved_test_size = min(args.test_samples, total_samples)
+
+ # Reserve last N samples for final evaluation
+ available_data = df.iloc[:total_samples - reserved_test_size]
+ reserved_test = df.iloc[total_samples - reserved_test_size:]
+
+ # Split available data into train and validation
+ available_size = len(available_data)
+ train_size = max(1, int(available_size * args.train_split))
+
+ train_df = available_data.iloc[:train_size]
+ val_df = available_data.iloc[train_size:]
+
+ local_dir = args.local_dir
+ hdfs_dir = args.hdfs_dir
+
+ # Save to parquet files
+ os.makedirs(local_dir, exist_ok=True)
+ train_df.to_parquet(os.path.join(local_dir, "train.parquet"))
+ val_df.to_parquet(os.path.join(local_dir, "validation.parquet"))
+ reserved_test.to_parquet(os.path.join(local_dir, "test.parquet"))
+
+ print(f"\n{'='*60}")
+ print(f"Countdown Data Split Summary ({args.train_split*100:.0f}% of available data)")
+ print(f"{'='*60}")
+ print(f"Total countdown tasks: {total_samples} samples")
+ print(f"\nData splits:")
+ print(f" Available for training/validation: {available_size} samples")
+ print(f" Train: {len(train_df)} samples ({args.train_split*100:.0f}% of available)")
+ print(f" Validation: {len(val_df)} samples ({(1-args.train_split)*100:.0f}% of available)")
+ print(f" Test (reserved): {len(reserved_test)} samples (LAST {reserved_test_size} from dataset)")
+ print(f"\nFiles saved to: {local_dir}")
+ print(f" - train.parquet: {len(train_df)} samples")
+ print(f" - validation.parquet: {len(val_df)} samples")
+ print(f" - test.parquet: {len(reserved_test)} samples (RESERVED FOR FINAL EVAL ONLY)")
+ print(f"{'='*60}\n")
+
+ if hdfs_dir is not None:
+ makedirs(hdfs_dir)
+ copy(src=local_dir, dst=hdfs_dir)
diff --git a/es-fine-tuning-paper/src/scripts/data_prep/grpo_data_gsm8k.py b/es-fine-tuning-paper/src/scripts/data_prep/grpo_data_gsm8k.py
new file mode 100644
index 0000000..4843052
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/data_prep/grpo_data_gsm8k.py
@@ -0,0 +1,140 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Preprocess the GSM8K dataset to parquet format for VERL GRPO training
+Reserves last 200 samples from TEST dataset for final evaluation
+"""
+
+import argparse
+import os
+import re
+
+import datasets
+
+from verl.utils.hdfs_io import copy, makedirs
+
+
+def extract_solution(solution_str):
+ solution = re.search("#### (\\-?[0-9\\.\\,]+)", solution_str)
+ assert solution is not None
+ final_solution = solution.group(0)
+ final_solution = final_solution.split("#### ")[1].replace(",", "")
+ return final_solution
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--local_dir", default="./src/data/gsm8k-0.1")
+ parser.add_argument("--hdfs_dir", default=None)
+ parser.add_argument("--train_split", type=float, default=0.1,
+ help="Fraction of train dataset to use for training")
+ parser.add_argument("--test_samples", type=int, default=200,
+ help="Number of samples to reserve from TEST dataset for final evaluation")
+
+ args = parser.parse_args()
+
+ data_source = "openai/gsm8k"
+
+ dataset = datasets.load_dataset(data_source, "main")
+
+ train_dataset = dataset["train"]
+ test_dataset = dataset["test"]
+
+ instruction_following = 'Let\'s think step by step and output the final answer after "####".'
+
+ # add a row to each data item that represents a unique id
+ def make_map_fn(split):
+ def process_fn(example, idx):
+ question_raw = example.pop("question")
+ question = question_raw + " " + instruction_following
+ answer_raw = example.pop("answer")
+ solution = extract_solution(answer_raw)
+
+ data = {
+ "data_source": data_source,
+ "prompt": [
+ {
+ "role": "user",
+ "content": question,
+ }
+ ],
+ "ability": "math",
+ "reward_model": {"style": "rule", "ground_truth": solution},
+ "extra_info": {
+ "split": split,
+ "index": idx,
+ "answer": answer_raw,
+ "question": question_raw,
+ },
+ }
+ return data
+ return process_fn
+
+ # Process both datasets
+ train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True)
+ test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True)
+
+ # Reserve last N samples from TEST dataset for final evaluation
+ test_total = len(test_dataset)
+ reserved_test_size = min(args.test_samples, test_total)
+
+ # Split test dataset: everything except last N samples can be used for validation
+ test_for_val = test_dataset.select(range(test_total - reserved_test_size))
+
+ # Reserved test set (last N samples from test dataset)
+ reserved_test = test_dataset.select(range(test_total - reserved_test_size, test_total))
+
+ # Create training data from train dataset
+ train_total = len(train_dataset)
+ train_size = max(1, int(train_total * args.train_split))
+ train_data = train_dataset.select(range(train_size))
+
+ # Combine remaining train data with available test data for validation
+ train_for_val = train_dataset.select(range(train_size, train_total))
+ available_data = datasets.concatenate_datasets([train_for_val, test_for_val])
+
+ # Use combined data as validation
+ val_data = available_data
+
+ local_dir = args.local_dir
+ hdfs_dir = args.hdfs_dir
+
+ # Create directory if it doesn't exist
+ os.makedirs(local_dir, exist_ok=True)
+
+ # Save to parquet files
+ train_data.to_parquet(os.path.join(local_dir, "train.parquet"))
+ val_data.to_parquet(os.path.join(local_dir, "validation.parquet"))
+ reserved_test.to_parquet(os.path.join(local_dir, "test.parquet"))
+
+ print(f"\n{'='*60}")
+ print(f"GSM8K Data Split Summary ({args.train_split*100:.0f}% of train dataset)")
+ print(f"{'='*60}")
+ print(f"Original GSM8K train dataset: {train_total} samples")
+ print(f"Original GSM8K test dataset: {test_total} samples")
+ print(f"\nData splits:")
+ print(f" Train: {len(train_data)} samples ({args.train_split*100:.0f}% of train dataset)")
+ print(f" Validation: {len(val_data)} samples (remaining train + test except last {reserved_test_size})")
+ print(f" - From train: {len(train_for_val)} samples")
+ print(f" - From test: {len(test_for_val)} samples")
+ print(f" Test (reserved): {len(reserved_test)} samples (LAST {reserved_test_size} from test dataset)")
+ print(f"\nFiles saved to: {local_dir}")
+ print(f" - train.parquet: {len(train_data)} samples")
+ print(f" - validation.parquet: {len(val_data)} samples")
+ print(f" - test.parquet: {len(reserved_test)} samples (RESERVED FOR FINAL EVAL ONLY)")
+ print(f"{'='*60}\n")
+
+ if hdfs_dir is not None:
+ makedirs(hdfs_dir)
+ copy(src=local_dir, dst=hdfs_dir)
diff --git a/es-fine-tuning-paper/src/scripts/data_prep/prepare_base_tokenizers.sh b/es-fine-tuning-paper/src/scripts/data_prep/prepare_base_tokenizers.sh
new file mode 100755
index 0000000..2b2baa8
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/data_prep/prepare_base_tokenizers.sh
@@ -0,0 +1,49 @@
+#!/bin/bash
+set -e
+
+# Script to prepare custom tokenizers for base models (Llama and Qwen)
+# These tokenizers add the simple chat template: "Question: {input} Answer: Let's think step by step."
+
+echo "=========================================="
+echo "Preparing Base Model Tokenizers"
+echo "=========================================="
+
+# Create tokenizers directory
+mkdir -p ./tokenizers
+
+# Check if HF_TOKEN is set
+if [ -z "$HF_TOKEN" ]; then
+ echo "Warning: HF_TOKEN environment variable is not set"
+ echo "You may need to login to HuggingFace to access Llama models"
+ echo "Run: export HF_TOKEN=your_huggingface_token"
+fi
+
+# Prepare Llama 3.2 3B base model tokenizer
+echo ""
+echo "Creating tokenizer for Llama-3.2-3B (base)..."
+python3 base_model_tokenizer.py \
+ --model_path meta-llama/Llama-3.2-3B \
+ --save_path ./tokenizers/llama-3.2-3b-base-chat \
+ --test
+
+echo ""
+echo "✓ Llama tokenizer created at: ./tokenizers/llama-3.2-3b-base-chat"
+
+# Prepare Qwen2.5 3B base model tokenizer
+echo ""
+echo "Creating tokenizer for Qwen2.5-3B (base)..."
+python3 base_model_tokenizer.py \
+ --model_path Qwen/Qwen2.5-3B \
+ --save_path ./tokenizers/qwen2.5-3b-base-chat \
+ --test
+
+echo ""
+echo "✓ Qwen tokenizer created at: ./tokenizers/qwen2.5-3b-base-chat"
+
+echo ""
+echo "=========================================="
+echo "Tokenizer Preparation Complete!"
+echo "=========================================="
+echo "Use these tokenizer paths in your training config:"
+echo " Llama: ./tokenizers/llama-3.2-3b-base-chat"
+echo " Qwen: ./tokenizers/qwen2.5-3b-base-chat"
diff --git a/es-fine-tuning-paper/src/scripts/data_prep/prepare_countdown_data.sh b/es-fine-tuning-paper/src/scripts/data_prep/prepare_countdown_data.sh
new file mode 100755
index 0000000..35af53f
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/data_prep/prepare_countdown_data.sh
@@ -0,0 +1,62 @@
+#!/bin/bash
+# Script to prepare Countdown data for GRPO training with reserved test set
+
+set -e
+
+echo "=========================================="
+echo "Preparing Countdown data for GRPO training"
+echo "=========================================="
+
+# Default parameters
+LOCAL_DIR="./src/data/countdown-0.4"
+JSON_FILE="./src/data/countdown-full/countdown.json"
+TRAIN_SPLIT=0.4
+TEST_SAMPLES=200
+
+# Install dependencies
+pip install pandas
+
+# Parse command line arguments
+while [[ $# -gt 0 ]]; do
+ case $1 in
+ --local_dir)
+ LOCAL_DIR="$2"
+ shift 2
+ ;;
+ --json_file)
+ JSON_FILE="$2"
+ shift 2
+ ;;
+ --train_split)
+ TRAIN_SPLIT="$2"
+ shift 2
+ ;;
+ --test_samples)
+ TEST_SAMPLES="$2"
+ shift 2
+ ;;
+ *)
+ echo "Unknown option: $1"
+ echo "Usage: $0 [--local_dir DIR] [--json_file FILE] [--train_split FRACTION] [--test_samples N]"
+ exit 1
+ ;;
+ esac
+done
+
+echo "Configuration:"
+echo " Input JSON file: $JSON_FILE"
+echo " Output directory: $LOCAL_DIR"
+echo " Train split: $TRAIN_SPLIT (${TRAIN_SPLIT%.*}0% of available data)"
+echo " Test samples (reserved): $TEST_SAMPLES"
+echo ""
+
+# Run the data preparation script
+python3 grpo_data_countdown.py \
+ --local_dir "$LOCAL_DIR" \
+ --json_file "$JSON_FILE" \
+ --train_split "$TRAIN_SPLIT" \
+ --test_samples "$TEST_SAMPLES"
+
+echo ""
+echo "Data preparation complete!"
+echo "You can now run GRPO training with: ./grpo-countdown.sh"
diff --git a/es-fine-tuning-paper/src/scripts/data_prep/prepare_gsm8k_data.sh b/es-fine-tuning-paper/src/scripts/data_prep/prepare_gsm8k_data.sh
new file mode 100755
index 0000000..db5050b
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/data_prep/prepare_gsm8k_data.sh
@@ -0,0 +1,55 @@
+#!/bin/bash
+# Script to prepare GSM8K data for GRPO training with reserved test set
+
+set -e
+
+echo "=========================================="
+echo "Preparing GSM8K data for GRPO training"
+echo "=========================================="
+
+# Default parameters
+LOCAL_DIR="./src/data/gsm8k-0.4"
+TRAIN_SPLIT=0.4
+TEST_SAMPLES=200
+
+#install datasets
+pip install datasets
+
+# Parse command line arguments
+while [[ $# -gt 0 ]]; do
+ case $1 in
+ --local_dir)
+ LOCAL_DIR="$2"
+ shift 2
+ ;;
+ --train_split)
+ TRAIN_SPLIT="$2"
+ shift 2
+ ;;
+ --test_samples)
+ TEST_SAMPLES="$2"
+ shift 2
+ ;;
+ *)
+ echo "Unknown option: $1"
+ echo "Usage: $0 [--local_dir DIR] [--train_split FRACTION] [--test_samples N]"
+ exit 1
+ ;;
+ esac
+done
+
+echo "Configuration:"
+echo " Output directory: $LOCAL_DIR"
+echo " Train split: $TRAIN_SPLIT (${TRAIN_SPLIT%.*}0% of available data)"
+echo " Test samples (reserved): $TEST_SAMPLES"
+echo ""
+
+# Run the data preparation script
+python3 grpo_data_gsm8k.py \
+ --local_dir "$LOCAL_DIR" \
+ --train_split "$TRAIN_SPLIT" \
+ --test_samples "$TEST_SAMPLES"
+
+echo ""
+echo "Data preparation complete!"
+echo "You can now run GRPO training with: ./grpo-gsm8k.sh"
diff --git a/es-fine-tuning-paper/src/scripts/es/es-gsm8k.sh b/es-fine-tuning-paper/src/scripts/es/es-gsm8k.sh
new file mode 100755
index 0000000..1b6841a
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/es/es-gsm8k.sh
@@ -0,0 +1,99 @@
+#!/bin/bash
+set -x
+
+# ES Fine-tuning Script for GSM8K
+# This script runs Evolution Strategies fine-tuning on the GSM8K dataset
+
+echo "Starting ES Fine-tuning for GSM8K..."
+
+# Check if HF_TOKEN is set for gated models (e.g., Llama)
+if [ -z "$HF_TOKEN" ]; then
+ echo "Warning: HF_TOKEN environment variable is not set"
+ echo "If using gated models (e.g., Llama), you may encounter authentication errors"
+ echo "Set with: export HF_TOKEN=your_huggingface_token"
+else
+ echo "HF_TOKEN found, logging in to HuggingFace..."
+ huggingface-cli login --token "$HF_TOKEN" 2>&1 | grep -v "Token is valid"
+ echo "✓ Logged in to HuggingFace"
+fi
+echo ""
+
+# Default hyperparameters
+SIGMA=${SIGMA:-0.001}
+ALPHA=${ALPHA:-0.0005}
+POPULATION_SIZE=${POPULATION_SIZE:-8}
+NUM_ENGINES=${NUM_ENGINES:-8}
+NUM_ITERATIONS=${NUM_ITERATIONS:-100}
+NUM_TRAIN_SAMPLES=${NUM_TRAIN_SAMPLES:-200}
+CUDA_DEVICES=${CUDA_DEVICES:-"0,1,2,3,4,5,6,7"}
+MODEL_NAME=${MODEL_NAME:-"meta-llama/Llama-3.2-3B"}
+EXPERIMENT_DIR=${EXPERIMENT_DIR:-"es-ft-gsm8k-experiment-exp3"}
+DATA_PATH="src/data/gsm8k-0.1/train.parquet"
+TOKENIZER_PATH=${TOKENIZER_PATH:-"./src/tokenizers/llama-3.2-3b-base-chat"}
+
+
+echo "Configuration:"
+echo " Model: $MODEL_NAME"
+if [ -n "$TOKENIZER_PATH" ]; then
+ echo " Tokenizer: $TOKENIZER_PATH (custom)"
+else
+ echo " Tokenizer: Using model's default tokenizer"
+fi
+echo " Sigma: $SIGMA"
+echo " Alpha: $ALPHA"
+echo " Population Size: $POPULATION_SIZE"
+echo " Number of Engines: $NUM_ENGINES"
+echo " Number of Iterations: $NUM_ITERATIONS"
+echo " Training Samples: $NUM_TRAIN_SAMPLES"
+echo " CUDA Devices: $CUDA_DEVICES"
+echo " Experiment Directory: $EXPERIMENT_DIR"
+echo ""
+
+# Check if data exists
+if [ ! -f $DATA_PATH ]; then
+ echo "Error: GSM8K training data not found at $DATA_PATH"
+ echo "Please prepare the data first using grpo_data_gsm8k.py"
+ exit 1
+fi
+
+# Run ES fine-tuning
+CMD="python3 es_fine_tuning_gsm8k_accl.py \
+ --model_name \"$MODEL_NAME\" \
+ --sigma $SIGMA \
+ --alpha $ALPHA \
+ --population_size $POPULATION_SIZE \
+ --num_engines $NUM_ENGINES \
+ --num_iterations $NUM_ITERATIONS \
+ --num_train_samples $NUM_TRAIN_SAMPLES \
+ --cuda_devices \"$CUDA_DEVICES\" \
+ --experiment_dir \"$EXPERIMENT_DIR\" \
+ --data_path \"$DATA_PATH\" \
+ --verbose"
+
+# Add tokenizer path only if specified
+if [ -n "$TOKENIZER_PATH" ]; then
+ CMD="$CMD --tokenizer_path \"$TOKENIZER_PATH\""
+fi
+
+eval $CMD
+
+# Check exit status
+if [ $? -eq 0 ]; then
+ echo ""
+ echo "ES Fine-tuning completed successfully!"
+ echo "Model saved to: $EXPERIMENT_DIR/gsm8k_nccl_*/model_saves/final_model_iteration_$NUM_ITERATIONS"
+else
+ echo ""
+ echo "ES Fine-tuning failed with exit code $?"
+ exit 1
+fi
+
+# Optional: Evaluate the final model on test set
+# Uncomment the following lines to run evaluation after training
+# echo ""
+# echo "Evaluating final model on test set..."
+# python3 src/evaluate_model.py \
+# --model_path "$EXPERIMENT_DIR/gsm8k_nccl_*/model_saves/final_model_iteration_$NUM_ITERATIONS" \
+# --test_file src/data/gsm8k-0.1/test.parquet \
+# --task_type gsm8k \
+# --output_file gsm8k_es_eval_results.json
diff --git a/es-fine-tuning-paper/src/scripts/es/es_fine-tuning_countdown_accl.py b/es-fine-tuning-paper/src/scripts/es/es_fine-tuning_countdown_accl.py
new file mode 100644
index 0000000..2ea8a05
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/es/es_fine-tuning_countdown_accl.py
@@ -0,0 +1,348 @@
+import argparse
+from datetime import datetime
+import gc
+import json
+import os
+import random
+import shutil
+import signal
+import sys
+import time
+
+import numpy as np
+import ray
+from ray.util.placement_group import placement_group, remove_placement_group
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+import torch
+import wandb
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from vllm import LLM, SamplingParams
+from vllm.utils import get_ip, get_open_port
+
+from src.countdown.countdown_task import reward_function
+
+# Default Hyperparameters
+SIGMA = 0.001
+ALPHA = 0.0005
+POPULATION_SIZE = 30
+NUM_ENGINES = 4
+NUM_ITERATIONS = 1000
+EXPERIMENT_DIR = "es-ft-experiment"
+
+def parse_args():
+ parser = argparse.ArgumentParser(
+ description="ES Fine-tuning for Countdown Task with multi-engine NCCL sync"
+ )
+ parser.add_argument("--model_name", type=str, default="Qwen/Qwen2.5-3B-Instruct")
+ parser.add_argument("--sigma", type=float, default=SIGMA)
+ parser.add_argument("--alpha", type=float, default=ALPHA)
+ parser.add_argument("--population_size", type=int, default=POPULATION_SIZE)
+ parser.add_argument("--num_engines", type=int, default=NUM_ENGINES)
+ parser.add_argument("--num_iterations", type=int, default=NUM_ITERATIONS)
+ parser.add_argument("--experiment_dir", type=str, default=EXPERIMENT_DIR)
+ parser.add_argument(
+ "--num_train_samples",
+ type=int,
+ default=200,
+ help="Number of examples to use from the training dataset (default: 200)",
+ )
+ parser.add_argument("--cuda_devices", type=str, default="0,1,2,3")
+ parser.add_argument('--verbose', action='store_true', help='Print verbose logs')
+ parser.add_argument(
+ "--global_seed",
+ type=int,
+ help="Global random seed",
+ )
+ args = parser.parse_args()
+ # Optional: scope host visibility; vLLM actors will ignore it and pick device from PG
+ os.environ["CUDA_VISIBLE_DEVICES"] = args.cuda_devices
+
+ # set global random seed
+ if args.global_seed is not None:
+ random.seed(args.global_seed)
+ np.random.seed(args.global_seed)
+ torch.manual_seed(args.global_seed)
+ torch.cuda.manual_seed_all(args.global_seed)
+
+ return args
+
+class ESNcclLLM(LLM):
+ def __init__(self, *args, **kwargs):
+ # Let Ray/PG determine the actual visible device in the actor
+ os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+ os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
+ super().__init__(*args, **kwargs)
+
+def launch_engines(num_engines, model_name):
+ # Strict 1-GPU isolation via PGs
+ pgs = [placement_group([{"GPU": 1, "CPU": 0}], lifetime="detached") for _ in range(num_engines)]
+ ray.get([pg.ready() for pg in pgs])
+
+ strategies = [
+ PlacementGroupSchedulingStrategy(
+ placement_group=pg,
+ placement_group_capture_child_tasks=True,
+ placement_group_bundle_index=0,
+ )
+ for pg in pgs
+ ]
+
+ engines = [
+ ray.remote(num_cpus=0, num_gpus=0, scheduling_strategy=strategy)(ESNcclLLM).remote(
+ model=model_name,
+ tensor_parallel_size=1,
+ distributed_executor_backend="ray",
+ worker_extension_cls="utils.worker_extn.WorkerExtension",
+ dtype="float16",
+ enable_prefix_caching=False,
+ enforce_eager=False,
+ )
+ for strategy in strategies
+ ]
+ return engines, pgs
+
+def evaluate_countdown_handle(llm, task_datas):
+ prompts = [d["context"] for d in task_datas]
+ sampling_params = SamplingParams(
+ temperature=0.0,
+ seed=42,
+ max_tokens=1024,
+ )
+ handle = llm.generate.remote(prompts, sampling_params, use_tqdm=False)
+ return handle, time.time()
+
+def _postprocess_outputs(outputs, task_datas):
+ rewards = []
+ avg_rewards = []
+ for output, data in zip(outputs, task_datas):
+ response = output.outputs[0].text
+ r = reward_function(response, data["numbers"], data["target"])
+ rewards.append(r)
+ avg_rewards.append(r["reward"])
+ return {
+ "rewards": rewards,
+ "avg_reward": float(np.mean(avg_rewards)) if avg_rewards else 0.0,
+ }
+
+def main(args):
+ # Ensure local Ray
+ os.environ.pop("RAY_ADDRESS", None)
+ os.environ.pop("RAY_HEAD_IP", None)
+ os.environ.pop("RAY_GCS_SERVER_ADDRESS", None)
+ ray.init(address="local", include_dashboard=False, ignore_reinit_error=True)
+
+ # Logging
+ run_name = f"countdown_nccl_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+ logging_dir = f"{args.experiment_dir}/{run_name}"
+ wandb.init(
+ project="es-fine-tuning-countdown",
+ name=run_name,
+ config={
+ "model_name": args.model_name,
+ "sigma": args.sigma,
+ "alpha": args.alpha,
+ "population_size": args.population_size,
+ "num_engines": args.num_engines,
+ "num_iterations": args.num_iterations,
+ "num_train_samples": args.num_train_samples,
+ "global_seed": args.global_seed,
+ }
+ )
+
+ # Prepare an HF checkpoint for vLLM to load
+ model_saves_dir = f"{logging_dir}/model_saves"
+ os.makedirs(model_saves_dir, exist_ok=True)
+
+ base_model = AutoModelForCausalLM.from_pretrained(
+ args.model_name, torch_dtype=torch.float16
+ ).to("cpu")
+ tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+
+ base_model_path = f"{model_saves_dir}/base_model"
+ if os.path.exists(base_model_path):
+ shutil.rmtree(base_model_path)
+ os.makedirs(base_model_path, exist_ok=True)
+ tokenizer.save_pretrained(base_model_path)
+ base_model.save_pretrained(base_model_path)
+ del base_model
+ gc.collect()
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+
+ # Load data
+ data_path = "src/data/countdown-full/countdown.json"
+ with open(data_path, "r") as f:
+ task_datas = json.load(f)
+ task_datas = task_datas[: args.num_train_samples]
+
+ # Launch engines
+ engines, pgs = launch_engines(args.num_engines, base_model_path)
+
+ # Init inter-engine communicator once
+ master_address = get_ip()
+ master_port = get_open_port()
+ ray.get([
+ engines[i].collective_rpc.remote(
+ "init_inter_engine_group", args=(master_address, master_port, i, args.num_engines)
+ )
+ for i in range(args.num_engines)
+ ])
+
+ def cleanup():
+ for llm in engines:
+ try:
+ ray.kill(llm)
+ except Exception:
+ pass
+ for pg in pgs:
+ try:
+ remove_placement_group(pg)
+ except Exception:
+ pass
+ ray.shutdown()
+
+ def sig_handler(sig, frame):
+ cleanup()
+ sys.exit(0)
+
+ signal.signal(signal.SIGINT, sig_handler)
+ signal.signal(signal.SIGTERM, sig_handler)
+
+ # Engines start with identical weights (loaded from the same HF checkpoint)
+ # For each iteration:
+ # - Explore: per-seed add noise -> eval -> subtract noise (GPU-only)
+ # - Compute ES update on engine 0 only
+ # - Broadcast weights from engine 0 to all engines (NCCL)
+ for i in range(args.num_iterations):
+ print(f"\n\n=== Generation {i} ===")
+ total_iter_start = time.time()
+
+ # Random seeds for population
+ seeds = [random.randint(0, 1_000_000) for _ in range(args.population_size)]
+ seeds_perf = {}
+
+ # Round-robin scheduling
+ seed_iter = iter(seeds)
+ inflight = {}
+ results_this_gen = []
+
+ # Kick off an eval on each engine
+ for eng_idx, llm in enumerate(engines):
+ try:
+ seed = next(seed_iter)
+ except StopIteration:
+ break
+ # Add exploration noise
+ ray.get(llm.collective_rpc.remote("perturb_self_weights", args=(seed, args.sigma, False)))
+ handle, start_ts = evaluate_countdown_handle(llm, task_datas)
+ inflight[handle] = {
+ "engine": llm,
+ "engine_idx": eng_idx,
+ "seed": seed,
+ "start_ts": start_ts,
+ }
+
+ while inflight:
+ done, _ = ray.wait(list(inflight.keys()), num_returns=1)
+ h = done[0]
+ meta = inflight.pop(h)
+
+ outputs = ray.get(h)
+ metrics = _postprocess_outputs(outputs, task_datas)
+ elapsed = time.time() - meta["start_ts"]
+
+ seeds_perf[meta["seed"]] = metrics
+ results_this_gen.append(
+ {"seed": meta["seed"], "avg_reward": metrics["avg_reward"], "time": elapsed}
+ )
+
+ llm = meta["engine"]
+ # Remove exploration noise
+ ray.get(llm.collective_rpc.remote("restore_self_weights", args=(meta["seed"], args.sigma)))
+
+ # Schedule next seed on this engine
+ try:
+ next_seed = next(seed_iter)
+ except StopIteration:
+ continue
+
+ ray.get(llm.collective_rpc.remote("perturb_self_weights", args=(next_seed, args.sigma, False)))
+ handle, start_ts = evaluate_countdown_handle(llm, task_datas)
+ inflight[handle] = {
+ "engine": llm,
+ "engine_idx": meta["engine_idx"],
+ "seed": next_seed,
+ "start_ts": start_ts,
+ }
+ if args.verbose:
+ print(f"Scheduled seed {next_seed} on engine {meta['engine_idx']}")
+
+ # Normalize rewards
+ all_avg_rewards = [v["avg_reward"] for v in seeds_perf.values()]
+ mean_reward = float(np.mean(all_avg_rewards)) if all_avg_rewards else 0.0
+ std_reward = float(np.std(all_avg_rewards)) if all_avg_rewards else 0.0
+ min_reward = float(np.min(all_avg_rewards)) if all_avg_rewards else 0.0
+ max_reward = float(np.max(all_avg_rewards)) if all_avg_rewards else 0.0
+
+ print(f"Mean reward: {mean_reward}, std: {std_reward}, min: {min_reward}, max: {max_reward}")
+ for k in seeds_perf:
+ seeds_perf[k]["norm_reward"] = (seeds_perf[k]["avg_reward"] - mean_reward) / (std_reward + 1e-8)
+ if args.verbose:
+ print(f"Seed {k} normalized reward: {seeds_perf[k]['norm_reward']}")
+
+ wandb.log({
+ "reward/mean": mean_reward,
+ "reward/std": std_reward,
+ "reward/min": min_reward,
+ "reward/max": max_reward,
+ "iteration": i
+ })
+
+ # Compute ES update ONLY on engine 0 (baseline is already current weights)
+ per_seed_coeffs = [
+ (seed, (args.alpha / args.population_size) * float(seeds_perf[seed]["norm_reward"]))
+ for seed in seeds
+ ]
+
+ perturb_start = time.time()
+ handles = []
+ for seed, coeff in per_seed_coeffs:
+ # Use sigma_or_scale=1.0 so the applied scale is `coeff`
+ handles.append(engines[0].collective_rpc.remote("perturb_self_weights", args=(seed, coeff, False)))
+ ray.get(handles)
+ if args.verbose:
+ print(f"Applied perturbations in {time.time() - perturb_start}s")
+ wandb.log({"time/perturbation_application": time.time() - perturb_start, "iteration": i})
+
+ # Broadcast updated weights from engine 0 to all engines (avoid CPU copies)
+ broadcast_start = time.time()
+ ray.get([e.collective_rpc.remote("broadcast_all_weights", args=(0,)) for e in engines])
+ if args.verbose:
+ print(f"Broadcasted updated weights in {time.time() - broadcast_start}s")
+ wandb.log({"time/broadcast": time.time() - broadcast_start, "iteration": i})
+
+ # Logging per-result and timing
+ if args.verbose:
+ for res_idx, res in enumerate(results_this_gen):
+ print(f"IDX:{res_idx} Seed {res['seed']} avg_reward: {res['avg_reward']}, time: {res['time']}s")
+ total_iter_end = time.time()
+ wandb.log({"time/iteration": total_iter_end - total_iter_start, "iteration": i})
+ print(f"wall clock time for iteration {i}: {total_iter_end - total_iter_start}s")
+ print(f"=== Generation {i} finished ===\n")
+
+ # Save final model weights (all engines are in sync; save from engine 0)
+ final_model_path = f"{model_saves_dir}/final_model_iteration_{args.num_iterations}"
+ os.makedirs(final_model_path, exist_ok=True)
+ ray.get(
+ engines[0].collective_rpc.remote(
+ "save_self_weights_to_disk", args=(f"{final_model_path}/pytorch_model.pth",)
+ )
+ )
+ print(f"Final model weights saved to {final_model_path}.")
+
+ wandb.finish()
+ cleanup()
+
+if __name__ == "__main__":
+ args = parse_args()
+ main(args)
diff --git a/es-fine-tuning-paper/src/scripts/es/es_fine_tuning_gsm8k_accl.py b/es-fine-tuning-paper/src/scripts/es/es_fine_tuning_gsm8k_accl.py
new file mode 100644
index 0000000..0287ca7
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/es/es_fine_tuning_gsm8k_accl.py
@@ -0,0 +1,371 @@
+import argparse
+from datetime import datetime
+import gc
+import json
+import os
+import random
+import shutil
+import signal
+import sys
+import time
+
+import numpy as np
+import ray
+from ray.util.placement_group import placement_group, remove_placement_group
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+import torch
+from torch.utils.tensorboard import SummaryWriter
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from vllm import LLM, SamplingParams
+from vllm.utils import get_ip, get_open_port
+
+from src.rewards.gsm8k_reward import reward_function
+
+# Default Hyperparameters
+SIGMA = 0.001
+ALPHA = 0.0005
+POPULATION_SIZE = 30
+NUM_ENGINES = 4
+NUM_ITERATIONS = 1000
+EXPERIMENT_DIR = "es-ft-experiment"
+DATA_DIR="src/data"
+
+def parse_args():
+ parser = argparse.ArgumentParser(
+ description="ES Fine-tuning for GSM8K Task with multi-engine NCCL sync"
+ )
+ parser.add_argument("--model_name", type=str, default="Qwen/Qwen2.5-3B-Instruct")
+ parser.add_argument("--sigma", type=float, default=SIGMA)
+ parser.add_argument("--alpha", type=float, default=ALPHA)
+ parser.add_argument("--population_size", type=int, default=POPULATION_SIZE)
+ parser.add_argument("--num_engines", type=int, default=NUM_ENGINES)
+ parser.add_argument("--num_iterations", type=int, default=NUM_ITERATIONS)
+ parser.add_argument("--experiment_dir", type=str, default=EXPERIMENT_DIR)
+ parser.add_argument("--data_path", type=str, default=DATA_DIR)
+ parser.add_argument(
+ "--num_train_samples",
+ type=int,
+ default=200,
+ help="Number of training examples to use from the GSM8K dataset (default: 200)",
+ )
+ parser.add_argument("--cuda_devices", type=str, default="0,1,2,3")
+ parser.add_argument('--verbose', action='store_true', help='Print verbose logs')
+ parser.add_argument(
+ "--tokenizer_path",
+ type=str,
+ default=None,
+ help="Path to custom tokenizer (e.g., src/tokenizers/qwen2.5-3b-base-chat)",
+ )
+ parser.add_argument(
+ "--global_seed",
+ type=int,
+ help="Global random seed",
+ )
+ args = parser.parse_args()
+ # Optional: scope host visibility; vLLM actors will ignore it and pick device from PG
+ os.environ["CUDA_VISIBLE_DEVICES"] = args.cuda_devices
+
+ # set global random seed
+ if args.global_seed is not None:
+ random.seed(args.global_seed)
+ np.random.seed(args.global_seed)
+ torch.manual_seed(args.global_seed)
+ torch.cuda.manual_seed_all(args.global_seed)
+
+ return args
+
+class ESNcclLLM(LLM):
+ def __init__(self, *args, **kwargs):
+ # Let Ray/PG determine the actual visible device in the actor
+ os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+ os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
+ super().__init__(*args, **kwargs)
+
+def launch_engines(num_engines, model_name, tokenizer_path=None):
+ # Strict 1-GPU isolation via PGs
+ pgs = [placement_group([{"GPU": 1, "CPU": 0}], lifetime="detached") for _ in range(num_engines)]
+ ray.get([pg.ready() for pg in pgs])
+
+ strategies = [
+ PlacementGroupSchedulingStrategy(
+ placement_group=pg,
+ placement_group_capture_child_tasks=True,
+ placement_group_bundle_index=0,
+ )
+ for pg in pgs
+ ]
+
+ # Prepare vLLM kwargs
+ vllm_kwargs = {
+ "model": model_name,
+ "tensor_parallel_size": 1,
+ "distributed_executor_backend": "ray",
+ "worker_extension_cls": "utils.worker_extn.WorkerExtension",
+ "dtype": "float16",
+ "enable_prefix_caching": False,
+ "enforce_eager": False,
+ }
+
+ # Add tokenizer path if provided
+ if tokenizer_path:
+ vllm_kwargs["tokenizer"] = tokenizer_path
+
+ engines = [
+ ray.remote(num_cpus=0, num_gpus=0, scheduling_strategy=strategy)(ESNcclLLM).remote(
+ **vllm_kwargs
+ )
+ for strategy in strategies
+ ]
+ return engines, pgs
+
+def evaluate_gsm8k_handle(llm, task_datas):
+ prompts = [d["context"] for d in task_datas]
+ sampling_params = SamplingParams(
+ temperature=0.0,
+ seed=42,
+ max_tokens=1024,
+ )
+ handle = llm.generate.remote(prompts, sampling_params, use_tqdm=False)
+ return handle, time.time()
+
+def _postprocess_outputs(outputs, task_datas):
+ rewards = []
+ avg_rewards = []
+ for output, data in zip(outputs, task_datas):
+ response = output.outputs[0].text
+ r = reward_function(response, data["ground_truth"])
+ rewards.append(r)
+ avg_rewards.append(r["reward"])
+ return {
+ "rewards": rewards,
+ "avg_reward": float(np.mean(avg_rewards)) if avg_rewards else 0.0,
+ }
+
+def main(args):
+ # Ensure local Ray
+ os.environ.pop("RAY_ADDRESS", None)
+ os.environ.pop("RAY_HEAD_IP", None)
+ os.environ.pop("RAY_GCS_SERVER_ADDRESS", None)
+ ray.init(address="local", include_dashboard=False, ignore_reinit_error=True)
+
+ # Logging
+ logging_dir = f"{args.experiment_dir}/gsm8k_nccl_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+ writer = SummaryWriter(log_dir=logging_dir)
+
+ # Prepare an HF checkpoint for vLLM to load
+ model_saves_dir = f"{logging_dir}/model_saves"
+ os.makedirs(model_saves_dir, exist_ok=True)
+
+ # Get HF token from environment if available
+ hf_token = os.environ.get("HF_TOKEN", None)
+
+ base_model = AutoModelForCausalLM.from_pretrained(
+ args.model_name, torch_dtype=torch.float16, token=hf_token
+ ).to("cpu")
+
+ # Load tokenizer from custom path if provided, otherwise from model
+ if args.tokenizer_path:
+ # Convert to absolute path to avoid HuggingFace validation issues
+ tokenizer = AutoTokenizer.from_pretrained(
+ args.tokenizer_path,
+ )
+ else:
+ print(f"Loading tokenizer from model: {args.model_name}")
+ tokenizer = AutoTokenizer.from_pretrained(args.model_name, token=hf_token)
+
+ base_model_path = f"{model_saves_dir}/base_model"
+ if os.path.exists(base_model_path):
+ shutil.rmtree(base_model_path)
+ os.makedirs(base_model_path, exist_ok=True)
+ tokenizer.save_pretrained(base_model_path)
+ base_model.save_pretrained(base_model_path)
+ del base_model
+ gc.collect()
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+
+ # Load GSM8K data
+ import datasets
+ data_path = args.data_path
+ ds = datasets.load_dataset('parquet', data_files={'train': data_path})
+
+ # Convert to format expected by evaluation function
+ task_datas = []
+ for item in ds['train']:
+ # Extract prompt text from the prompt field
+ prompt_content = item['prompt'][0]['content']
+ ground_truth = item['reward_model']['ground_truth']
+ task_datas.append({
+ "context": prompt_content,
+ "ground_truth": ground_truth,
+ })
+
+ task_datas = task_datas[: args.num_train_samples]
+
+ # Launch engines with custom tokenizer path
+ engines, pgs = launch_engines(args.num_engines, base_model_path, tokenizer_path=base_model_path)
+
+ # Init inter-engine communicator once
+ master_address = get_ip()
+ master_port = get_open_port()
+ ray.get([
+ engines[i].collective_rpc.remote(
+ "init_inter_engine_group", args=(master_address, master_port, i, args.num_engines)
+ )
+ for i in range(args.num_engines)
+ ])
+
+ def cleanup():
+ for llm in engines:
+ try:
+ ray.kill(llm)
+ except Exception:
+ pass
+ for pg in pgs:
+ try:
+ remove_placement_group(pg)
+ except Exception:
+ pass
+ ray.shutdown()
+
+ def sig_handler(sig, frame):
+ cleanup()
+ sys.exit(0)
+
+ signal.signal(signal.SIGINT, sig_handler)
+ signal.signal(signal.SIGTERM, sig_handler)
+
+ # Engines start with identical weights (loaded from the same HF checkpoint)
+ # For each iteration:
+ # - Explore: per-seed add noise -> eval -> subtract noise (GPU-only)
+ # - Compute ES update on engine 0 only
+ # - Broadcast weights from engine 0 to all engines (NCCL)
+ for i in range(args.num_iterations):
+ print(f"\n\n=== Generation {i} ===")
+ total_iter_start = time.time()
+
+ # Random seeds for population
+ seeds = [random.randint(0, 1_000_000) for _ in range(args.population_size)]
+ seeds_perf = {}
+
+ # Round-robin scheduling
+ seed_iter = iter(seeds)
+ inflight = {}
+ results_this_gen = []
+
+ # Kick off an eval on each engine
+ for eng_idx, llm in enumerate(engines):
+ try:
+ seed = next(seed_iter)
+ except StopIteration:
+ break
+ # Add exploration noise
+ ray.get(llm.collective_rpc.remote("perturb_self_weights", args=(seed, args.sigma, False)))
+ handle, start_ts = evaluate_gsm8k_handle(llm, task_datas)
+ inflight[handle] = {
+ "engine": llm,
+ "engine_idx": eng_idx,
+ "seed": seed,
+ "start_ts": start_ts,
+ }
+
+ while inflight:
+ done, _ = ray.wait(list(inflight.keys()), num_returns=1)
+ h = done[0]
+ meta = inflight.pop(h)
+
+ outputs = ray.get(h)
+ metrics = _postprocess_outputs(outputs, task_datas)
+ elapsed = time.time() - meta["start_ts"]
+
+ seeds_perf[meta["seed"]] = metrics
+ results_this_gen.append(
+ {"seed": meta["seed"], "avg_reward": metrics["avg_reward"], "time": elapsed}
+ )
+
+ llm = meta["engine"]
+ # Remove exploration noise
+ ray.get(llm.collective_rpc.remote("restore_self_weights", args=(meta["seed"], args.sigma)))
+
+ # Schedule next seed on this engine
+ try:
+ next_seed = next(seed_iter)
+ except StopIteration:
+ continue
+
+ ray.get(llm.collective_rpc.remote("perturb_self_weights", args=(next_seed, args.sigma, False)))
+ handle, start_ts = evaluate_gsm8k_handle(llm, task_datas)
+ inflight[handle] = {
+ "engine": llm,
+ "engine_idx": meta["engine_idx"],
+ "seed": next_seed,
+ "start_ts": start_ts,
+ }
+ if args.verbose:
+ print(f"Scheduled seed {next_seed} on engine {meta['engine_idx']}")
+
+ # Normalize rewards
+ all_avg_rewards = [v["avg_reward"] for v in seeds_perf.values()]
+ mean_reward = float(np.mean(all_avg_rewards)) if all_avg_rewards else 0.0
+ std_reward = float(np.std(all_avg_rewards)) if all_avg_rewards else 0.0
+ min_reward = float(np.min(all_avg_rewards)) if all_avg_rewards else 0.0
+ max_reward = float(np.max(all_avg_rewards)) if all_avg_rewards else 0.0
+
+ print(f"Mean reward: {mean_reward}, std: {std_reward}, min: {min_reward}, max: {max_reward}")
+ for k in seeds_perf:
+ seeds_perf[k]["norm_reward"] = (seeds_perf[k]["avg_reward"] - mean_reward) / (std_reward + 1e-8)
+ if args.verbose:
+ print(f"Seed {k} normalized reward: {seeds_perf[k]['norm_reward']}")
+
+ writer.add_scalar("reward/mean", mean_reward, i)
+ writer.add_scalar("reward/std", std_reward, i)
+ writer.add_scalar("reward/min", min_reward, i)
+ writer.add_scalar("reward/max", max_reward, i)
+
+ # Compute ES update ONLY on engine 0 (baseline is already current weights)
+ per_seed_coeffs = [
+ (seed, (args.alpha / args.population_size) * float(seeds_perf[seed]["norm_reward"]))
+ for seed in seeds
+ ]
+
+ perturb_start = time.time()
+ handles = []
+ for seed, coeff in per_seed_coeffs:
+ # Use sigma_or_scale=1.0 so the applied scale is `coeff`
+ handles.append(engines[0].collective_rpc.remote("perturb_self_weights", args=(seed, coeff, False)))
+ ray.get(handles)
+ if args.verbose:
+ print(f"Applied perturbations in {time.time() - perturb_start}s")
+ writer.add_scalar("time/perturbation_application", time.time() - perturb_start, i)
+
+ # Broadcast updated weights from engine 0 to all engines (avoid CPU copies)
+ broadcast_start = time.time()
+ ray.get([e.collective_rpc.remote("broadcast_all_weights", args=(0,)) for e in engines])
+ if args.verbose:
+ print(f"Broadcasted updated weights in {time.time() - broadcast_start}s")
+ writer.add_scalar("time/broadcast", time.time() - broadcast_start, i)
+
+ # Logging per-result and timing
+ if args.verbose:
+ for res_idx, res in enumerate(results_this_gen):
+ print(f"IDX:{res_idx} Seed {res['seed']} avg_reward: {res['avg_reward']}, time: {res['time']}s")
+ total_iter_end = time.time()
+ writer.add_scalar("time/iteration", total_iter_end - total_iter_start, i)
+ print(f"wall clock time for iteration {i}: {total_iter_end - total_iter_start}s")
+ print(f"=== Generation {i} finished ===\n")
+
+ # Save final model weights (all engines are in sync; save from engine 0)
+ final_model_path = f"{model_saves_dir}/final_model_iteration_{args.num_iterations}"
+ os.makedirs(final_model_path, exist_ok=True)
+ ray.get(
+ engines[0].collective_rpc.remote(
+ "save_self_weights_to_disk", args=(f"{final_model_path}/pytorch_model.pth",)
+ )
+ )
+ print(f"Final model weights saved to {final_model_path}.")
+
+ cleanup()
+
+if __name__ == "__main__":
+ args = parse_args()
+ main(args)
diff --git a/es-fine-tuning-paper/src/scripts/evaluation/eval_countdown_vllm.py b/es-fine-tuning-paper/src/scripts/evaluation/eval_countdown_vllm.py
new file mode 100644
index 0000000..361ded8
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/evaluation/eval_countdown_vllm.py
@@ -0,0 +1,376 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
+import json
+import time
+import argparse
+import tempfile
+from typing import List, Tuple, Dict, Any
+
+import numpy as np
+from vllm import LLM, SamplingParams
+
+import ray
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+from ray.util.placement_group import placement_group
+
+def parse_args():
+ parser = argparse.ArgumentParser(description='vLLM evaluation for ES models (Qwen/Llama/etc)')
+ parser.add_argument('--model_id', type=str, default="Qwen/Qwen2.5-3B-Instruct",
+ help='HF model name for vLLM')
+ parser.add_argument('--trained_model_path', type=str, required=True,
+ help='Path to the trained model directory')
+ parser.add_argument('--tokenizer_path', type=str, default=None,
+ help='Path to tokenizer (e.g. base_model or custom tokenizer saved during training). '
+ 'Use this when only weights were saved so eval matches training tokenization.')
+
+ # Data args (match the regular evaluator)
+ parser.add_argument('--train_data_path', type=str,
+ default='src/data/countdown-full/countdown.json',
+ help='Path to training data JSON file')
+ parser.add_argument('--eval_data_path', type=str,
+ default='src/data/countdown-full/countdown.json',
+ help='Path to evaluation data JSON file')
+ parser.add_argument('--eval_samples', type=int, default=100,
+ help='Number of evaluation samples to evaluate')
+ parser.add_argument('--eval_offset', type=int, default=-100,
+ help='Offset for evaluation data (negative means from end)')
+
+ # Generation args
+ parser.add_argument('--max_new_tokens', type=int, default=1024,
+ help='Maximum number of new tokens to generate')
+ parser.add_argument('--do_sample', action='store_true',
+ help='Whether to use sampling instead of greedy decoding')
+ parser.add_argument('--temperature', type=float, default=0.8,
+ help='Temperature for sampling')
+ parser.add_argument('--top_p', type=float, default=0.9,
+ help='Top-p for nucleus sampling')
+
+ # Batch/engine args
+ parser.add_argument('--batch_size', type=int, default=None,
+ help='Batch size for prompts per vLLM.generate call (default: min(32, dataset_size))')
+ parser.add_argument('--tensor_parallel_size', type=int, default=1,
+ help='Tensor parallelism for vLLM')
+ parser.add_argument('--dtype', type=str, default='float16',
+ choices=['float16', 'bfloat16', 'float32'],
+ help='Model dtype for vLLM')
+
+ # Output/verbosity
+ parser.add_argument('--output_dir', type=str, default=None,
+ help='Directory to save inference results (default: based on model name)')
+ parser.add_argument('--save_responses', action='store_true',
+ help='Save individual responses to file')
+ parser.add_argument('--verbose', action='store_true',
+ help='Print verbose output')
+ parser.add_argument('--show_examples', type=int, default=5,
+ help='Number of examples to show in detail')
+ return parser.parse_args()
+
+
+class ESNcclLLM(LLM):
+ def __init__(self, *args, **kwargs):
+ # Let Ray/PG determine the actual visible device in the actor
+ os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+ os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
+ super().__init__(*args, **kwargs)
+
+def launch_engines(num_engines, model_name, tokenizer_path=None):
+ # Strict 1-GPU isolation via PGs
+ pgs = [placement_group([{"GPU": 1, "CPU": 0}], lifetime="detached") for _ in range(num_engines)]
+ ray.get([pg.ready() for pg in pgs])
+
+ strategies = [
+ PlacementGroupSchedulingStrategy(
+ placement_group=pg,
+ placement_group_capture_child_tasks=True,
+ placement_group_bundle_index=0,
+ )
+ for pg in pgs
+ ]
+
+ engine_kwargs = dict(
+ model=model_name,
+ tensor_parallel_size=1,
+ distributed_executor_backend="ray",
+ worker_extension_cls="utils.worker_extn.WorkerExtension",
+ dtype="float16",
+ enable_prefix_caching=False,
+ enforce_eager=False,
+ )
+ if tokenizer_path is not None:
+ engine_kwargs["tokenizer"] = tokenizer_path
+
+ engines = [
+ ray.remote(num_cpus=0, num_gpus=0, scheduling_strategy=strategy)(ESNcclLLM).remote(
+ **engine_kwargs
+ )
+ for strategy in strategies
+ ]
+ return engines, pgs
+
+def load_data(data_path: str, num_samples: int = None, offset: int = 0) -> List[Tuple[str, str]]:
+ if not os.path.exists(data_path):
+ raise FileNotFoundError(f"Data file not found: {data_path}")
+ with open(data_path, 'r') as f:
+ data_json = json.load(f)
+
+ dataset = [(item['context'], item['target']) for item in data_json]
+
+ if offset < 0:
+ start_idx, end_idx = len(dataset) + offset, len(dataset)
+ else:
+ start_idx, end_idx = offset, len(dataset)
+ dataset = dataset[start_idx:end_idx]
+
+ if num_samples is not None and num_samples < len(dataset):
+ dataset = dataset[:num_samples]
+ return dataset
+
+
+def extract_model_response(generated_text: str) -> str:
+ model_response = generated_text
+ if "assistant:" in generated_text:
+ model_response = generated_text.split("assistant:")[-1].strip()
+ return model_response
+
+
+def extract_numbers_and_target(input_text: str, target_text: str) -> Tuple[List[int], int]:
+ numbers, target = None, None
+ if "[" in input_text and "]" in input_text:
+ start_idx = input_text.find("[")
+ end_idx = input_text.find("]")
+ if start_idx != -1 and end_idx != -1:
+ numbers_str = input_text[start_idx+1:end_idx]
+ numbers = [int(n) for n in numbers_str.split() if n.isdigit()]
+ if target_text.isdigit():
+ target = int(target_text)
+ return numbers, target
+
+
+def evaluate_batch_vllm(llm, input_texts: List[str], target_texts: List[str], args, verbose: bool = False) -> List[Dict[str, Any]]:
+ if verbose:
+ print(f"Batch evaluating {len(input_texts)} samples...")
+
+ # Greedy vs sampling setup matches regular script semantics
+ temperature = args.temperature if args.do_sample else 0.0
+ top_p = args.top_p if args.do_sample else 1.0
+ sampling_params = SamplingParams(
+ temperature=temperature,
+ top_p=top_p,
+ max_tokens=args.max_new_tokens,
+ )
+
+ # Run generation
+ outputs = ray.get(llm.generate.remote(input_texts, sampling_params=sampling_params, use_tqdm=False))
+
+ # Reconstruct generated_text as input + completion (to keep extraction behavior identical)
+ all_results = []
+ for i, out in enumerate(outputs):
+ # vLLM returns a RequestOutput with outputs list; take first completion
+ completion_text = out.outputs[0].text if out.outputs else ""
+ generated_text = f"{completion_text}"
+
+ model_response = extract_model_response(generated_text)
+
+ input_text = input_texts[i]
+ target_text = target_texts[i]
+ numbers, target = extract_numbers_and_target(input_text, target_text)
+
+ reward_result = reward_function(model_response, numbers, target)
+ reward = reward_result["reward"]
+ reward_info = reward_result["reward_info"]
+
+ all_results.append({
+ 'input_text': input_text,
+ 'target_text': target_text,
+ 'generated_text': generated_text,
+ 'model_response': model_response,
+ 'numbers': numbers,
+ 'target': target,
+ 'reward': reward,
+ 'reward_info': reward_info
+ })
+ return all_results
+
+
+def evaluate_dataset_vllm(llm, dataset: List[Tuple[str, str]], args, dataset_name: str, batch_size: int = None) -> Dict[str, Any]:
+ print(f"\n=== Evaluating on {dataset_name} dataset ({len(dataset)} samples) ===")
+ if batch_size is None:
+ batch_size = min(1024, len(dataset))
+ print(f"Using batch size: {batch_size}")
+
+ all_results = []
+ total_reward = 0.0
+ total_format_reward = 0.0
+ total_answer_reward = 0.0
+
+ start_time = time.time()
+ for batch_start in range(0, len(dataset), batch_size):
+ batch_end = min(batch_start + batch_size, len(dataset))
+ batch_dataset = dataset[batch_start:batch_end]
+
+ if args.verbose:
+ print(f"Processing batch {batch_start//batch_size + 1}/{(len(dataset)-1)//batch_size + 1} (samples {batch_start+1}-{batch_end})...")
+
+ input_texts = [item[0] for item in batch_dataset]
+ target_texts = [item[1] for item in batch_dataset]
+
+ batch_results = evaluate_batch_vllm(llm, input_texts, target_texts, args, verbose=args.verbose)
+ all_results.extend(batch_results)
+
+ for result in batch_results:
+ total_reward += result['reward']
+ total_format_reward += result['reward_info']['format_reward']
+ total_answer_reward += result['reward_info']['answer_reward']
+
+ if batch_start == 0:
+ for i, result in enumerate(batch_results[:args.show_examples]):
+ print(f"\n--- Example {i+1} ---")
+ print(f"Input: {result['input_text']}")
+ print(f"Target: {result['target_text']}")
+ print(f"Model Response: {result['model_response']}")
+ print(f"Reward: {result['reward']:.4f} (Format: {result['reward_info']['format_reward']:.4f}, Answer: {result['reward_info']['answer_reward']:.4f})")
+
+ eval_time = time.time() - start_time
+
+ # Stats
+ avg_reward = total_reward / len(dataset)
+ avg_format_reward = total_format_reward / len(dataset)
+ avg_answer_reward = total_answer_reward / len(dataset)
+
+ rewards = [r['reward'] for r in all_results]
+ std_reward = np.std(rewards)
+ min_reward = np.min(rewards)
+ max_reward = np.max(rewards)
+
+ high_reward_count = sum(1 for r in rewards if r >= 1.0)
+ high_reward_percentage = high_reward_count / len(dataset) * 100
+
+ answer_rewards = [r['reward_info']['answer_reward'] for r in all_results]
+ correct_count = sum(1 for r in answer_rewards if r > 0)
+ accuracy = correct_count / len(dataset) * 100
+
+ stats = {
+ 'dataset_name': dataset_name,
+ 'num_samples': len(dataset),
+ 'avg_reward': avg_reward,
+ 'avg_format_reward': avg_format_reward,
+ 'avg_answer_reward': avg_answer_reward,
+ 'std_reward': std_reward,
+ 'min_reward': min_reward,
+ 'max_reward': max_reward,
+ 'high_reward_count': high_reward_count,
+ 'high_reward_percentage': high_reward_percentage,
+ 'correct_count': correct_count,
+ 'accuracy': accuracy,
+ 'eval_time': eval_time,
+ 'all_results': all_results
+ }
+
+ print(f"\n=== {dataset_name} Results Summary ===")
+ print(f"Number of samples: {len(dataset)}")
+ print(f"Average reward: {avg_reward:.4f} ± {std_reward:.4f}")
+ print(f" - Format reward: {avg_format_reward:.4f}")
+ print(f" - Answer reward: {avg_answer_reward:.4f}")
+ print(f"Accuracy (answer_reward > 0): {correct_count}/{len(dataset)} ({accuracy:.1f}%)")
+ print(f"High reward samples (≥1.0): {high_reward_count}/{len(dataset)} ({high_reward_percentage:.1f}%)")
+ print(f"Reward range: [{min_reward:.4f}, {max_reward:.4f}]")
+ print(f"Evaluation time: {eval_time:.2f}s ({eval_time/len(dataset):.3f}s per sample)")
+
+ return stats
+
+
+def save_results(results: Dict[str, Any], output_dir: str, args):
+ os.makedirs(output_dir, exist_ok=True)
+
+ summary = {
+ 'model_id': args.model_id,
+ 'eval_stats': {k: v for k, v in results['eval_stats'].items() if k != 'all_results'},
+ 'generation_config': {
+ 'max_new_tokens': args.max_new_tokens,
+ 'do_sample': args.do_sample,
+ 'temperature': args.temperature if args.do_sample else None,
+ 'top_p': args.top_p if args.do_sample else None,
+ 'batch_size': args.batch_size,
+ 'tensor_parallel_size': args.tensor_parallel_size,
+ 'dtype': args.dtype,
+ }
+ }
+
+ summary_path = os.path.join(output_dir, 'summary.json')
+ with open(summary_path, 'w') as f:
+ json.dump(summary, f, indent=2)
+ print(f"Summary saved to: {summary_path}")
+
+ if args.save_responses:
+ eval_details_path = os.path.join(output_dir, 'eval_detailed_results.json')
+ with open(eval_details_path, 'w') as f:
+ json.dump(results['eval_stats']['all_results'], f, indent=2)
+ print(f"Eval detailed results saved to: {eval_details_path}")
+
+
+def main():
+ args = parse_args()
+
+ # no auto connection
+ os.environ.pop("RAY_ADDRESS", None)
+ os.environ.pop("RAY_HEAD_IP", None)
+ os.environ.pop("RAY_GCS_SERVER_ADDRESS", None)
+
+ # This prevents socket/file lock conflicts with other running Ray instances
+ unique_dir = tempfile.mkdtemp(prefix=f"ray_temp_session_{int(time.time())}_")
+
+ ray.init(
+ address="local",
+ include_dashboard=False,
+ ignore_reinit_error=True,
+ _temp_dir=unique_dir,
+ dashboard_port=None
+ )
+
+ global reward_function
+ from src.countdown.countdown_task import reward_function
+
+ if args.output_dir is None:
+ model_name = os.path.basename(args.model_id.rstrip('/'))
+ batch_suffix = f"_batch{args.batch_size}" if args.batch_size else ""
+ args.output_dir = f"./inference_results_vllm_{model_name}{batch_suffix}"
+
+ print("=== vLLM ES Model Inference Script ===")
+ print(f"Model path: {args.model_id}")
+ if args.tokenizer_path:
+ print(f"Tokenizer path: {args.tokenizer_path}")
+ else:
+ print("Warning: --tokenizer_path not set. vLLM will use the tokenizer from model_id. "
+ "If you trained with a custom/saved tokenizer, pass --tokenizer_path so eval matches training.")
+ print(f"Eval data: {args.eval_data_path} (samples: {args.eval_samples}, offset: {args.eval_offset})")
+ print(f"Output directory: {args.output_dir}")
+ print(f"Tensor parallel size: {args.tensor_parallel_size}")
+ print(f"Dtype: {args.dtype}")
+
+ # Initialize vLLM engine
+ llm, _ = launch_engines(
+ num_engines=1,
+ model_name=args.model_id,
+ tokenizer_path=args.tokenizer_path,
+ )
+ llm = llm[0]
+ llm.collective_rpc.remote("load_weights_from_disk", args=(args.trained_model_path,))
+ print(f"vLLM LLM initialized from {args.trained_model_path}")
+
+ # Load dataset
+ eval_dataset = load_data(
+ os.path.join(os.path.dirname(__file__), args.eval_data_path),
+ num_samples=args.eval_samples,
+ offset=args.eval_offset
+ )
+ print(f"Loaded {len(eval_dataset)} evaluation samples")
+
+ # Evaluate
+ eval_stats = evaluate_dataset_vllm(llm, eval_dataset, args, "Eval", batch_size=args.batch_size)
+ results = {'eval_stats': eval_stats}
+ save_results(results, args.output_dir, args)
+
+
+if __name__ == "__main__":
+ main()
\ No newline at end of file
diff --git a/es-fine-tuning-paper/src/scripts/evaluation/eval_gsm8k_es.sh b/es-fine-tuning-paper/src/scripts/evaluation/eval_gsm8k_es.sh
new file mode 100755
index 0000000..4250033
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/evaluation/eval_gsm8k_es.sh
@@ -0,0 +1,72 @@
+#!/bin/bash
+set -x
+
+# Evaluation script for ES Fine-tuned GSM8K model using vLLM
+
+# Usage:
+# ./eval_gsm8k_es.sh
+# Example: ./eval_gsm8k_es.sh es-ft-gsm8k-experiment/gsm8k_nccl_20260227_103714/model_saves/final_model_iteration_1000/pytorch_model.pth
+
+if [ -z "$1" ]; then
+ echo "Error: Please provide the trained model path"
+ echo "Usage: $0 "
+ echo "Example: $0 es-ft-gsm8k-experiment/gsm8k_nccl_20260227_103714/model_saves/final_model_iteration_1000/pytorch_model.pth"
+ exit 1
+fi
+
+TRAINED_MODEL_PATH="$1"
+
+# Check if the trained model exists
+if [ ! -f "$TRAINED_MODEL_PATH" ]; then
+ echo "Error: Trained model not found at $TRAINED_MODEL_PATH"
+ exit 1
+fi
+
+echo "Evaluating trained model from: $TRAINED_MODEL_PATH"
+
+# Default configuration
+MODEL_ID="${MODEL_ID:-qwen/Qwen2.5-3B}"
+EVAL_DATA_PATH="${EVAL_DATA_PATH:-src/data/gsm8k-0.1/test.parquet}"
+EVAL_SAMPLES="${EVAL_SAMPLES:-200}"
+MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-1024}"
+BATCH_SIZE="${BATCH_SIZE:-32}"
+TENSOR_PARALLEL_SIZE="${TENSOR_PARALLEL_SIZE:-1}"
+OUTPUT_DIR="${OUTPUT_DIR:-src/evals/qwen2.5_3b_base_eval_results_gsm8k_es_0.1}"
+TOKENIZER_PATH=${TOKENIZER_PATH:-"./src/tokenizers/qwen2.5-3b-base-chat"}
+
+
+echo "Configuration:"
+echo " Base Model: $MODEL_ID"
+echo " Eval Data: $EVAL_DATA_PATH"
+echo " Eval Samples: $EVAL_SAMPLES"
+echo " Max New Tokens: $MAX_NEW_TOKENS"
+echo " Batch Size: $BATCH_SIZE"
+echo " Tensor Parallel Size: $TENSOR_PARALLEL_SIZE"
+echo " Output Dir: $OUTPUT_DIR"
+echo ""
+
+# Run evaluation
+python3 eval_gsm8k_vllm.py \
+ --model_id "$MODEL_ID" \
+ --trained_model_path "$TRAINED_MODEL_PATH" \
+ --eval_data_path "$EVAL_DATA_PATH" \
+ --eval_samples $EVAL_SAMPLES \
+ --max_new_tokens $MAX_NEW_TOKENS \
+ --batch_size $BATCH_SIZE \
+ --tensor_parallel_size $TENSOR_PARALLEL_SIZE \
+ --output_dir "$OUTPUT_DIR" \
+ --save_responses \
+ --show_examples 10 \
+ --tokenizer_path $TOKENIZER_PATH \
+ --verbose
+
+# Check exit status
+if [ $? -eq 0 ]; then
+ echo ""
+ echo "=== Evaluation completed successfully! ==="
+ echo "Results saved to: $OUTPUT_DIR"
+else
+ echo ""
+ echo "=== Evaluation failed! ==="
+ exit 1
+fi
diff --git a/es-fine-tuning-paper/src/scripts/evaluation/eval_gsm8k_vllm.py b/es-fine-tuning-paper/src/scripts/evaluation/eval_gsm8k_vllm.py
new file mode 100644
index 0000000..1491d8a
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/evaluation/eval_gsm8k_vllm.py
@@ -0,0 +1,383 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
+import json
+import time
+import argparse
+import tempfile
+from typing import List, Tuple, Dict, Any
+
+import numpy as np
+from vllm import LLM, SamplingParams
+
+import ray
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+from ray.util.placement_group import placement_group
+
+def parse_args():
+ parser = argparse.ArgumentParser(description='vLLM evaluation for ES models on GSM8K')
+ parser.add_argument('--model_id', type=str, default="Qwen/Qwen2.5-3B-Instruct",
+ help='HF model name for vLLM')
+ parser.add_argument('--trained_model_path', type=str, required=True,
+ help='Path to the trained model directory')
+ parser.add_argument(
+ "--tokenizer_path",
+ type=str,
+ default=None,
+ help="Path to custom tokenizer (e.g., src/tokenizers/qwen2.5-3b-base-chat)",
+ )
+
+ # Data args
+ parser.add_argument('--eval_data_path', type=str,
+ default='src/data/gsm8k-0.1/test.parquet',
+ help='Path to evaluation data parquet file')
+ parser.add_argument('--eval_samples', type=int, default=200,
+ help='Number of evaluation samples to evaluate')
+ parser.add_argument('--eval_offset', type=int, default=0,
+ help='Offset for evaluation data')
+
+ # Generation args
+ parser.add_argument('--max_new_tokens', type=int, default=1024,
+ help='Maximum number of new tokens to generate')
+ parser.add_argument('--do_sample', action='store_true',
+ help='Whether to use sampling instead of greedy decoding')
+ parser.add_argument('--temperature', type=float, default=0.8,
+ help='Temperature for sampling')
+ parser.add_argument('--top_p', type=float, default=0.9,
+ help='Top-p for nucleus sampling')
+
+ # Batch/engine args
+ parser.add_argument('--batch_size', type=int, default=None,
+ help='Batch size for prompts per vLLM.generate call (default: min(32, dataset_size))')
+ parser.add_argument('--tensor_parallel_size', type=int, default=1,
+ help='Tensor parallelism for vLLM')
+ parser.add_argument('--dtype', type=str, default='float16',
+ choices=['float16', 'bfloat16', 'float32'],
+ help='Model dtype for vLLM')
+
+ # Output/verbosity
+ parser.add_argument('--output_dir', type=str, default=None,
+ help='Directory to save inference results (default: based on model name)')
+ parser.add_argument('--save_responses', action='store_true',
+ help='Save individual responses to file')
+ parser.add_argument('--verbose', action='store_true',
+ help='Print verbose output')
+ parser.add_argument('--show_examples', type=int, default=5,
+ help='Number of examples to show in detail')
+ return parser.parse_args()
+
+
+class ESNcclLLM(LLM):
+ def __init__(self, *args, **kwargs):
+ # Let Ray/PG determine the actual visible device in the actor
+ os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+ os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
+ super().__init__(*args, **kwargs)
+
+def launch_engines(num_engines, model_name, tokenizer_path):
+ # Strict 1-GPU isolation via PGs
+ pgs = [placement_group([{"GPU": 1, "CPU": 0}], lifetime="detached") for _ in range(num_engines)]
+ ray.get([pg.ready() for pg in pgs])
+
+ strategies = [
+ PlacementGroupSchedulingStrategy(
+ placement_group=pg,
+ placement_group_capture_child_tasks=True,
+ placement_group_bundle_index=0,
+ )
+ for pg in pgs
+ ]
+
+ # Prepare vLLM kwargs
+ vllm_kwargs = {
+ "model": model_name,
+ "tensor_parallel_size": 1,
+ "distributed_executor_backend": "ray",
+ "worker_extension_cls": "utils.worker_extn.WorkerExtension",
+ "dtype": "float16",
+ "enable_prefix_caching": False,
+ "enforce_eager": False,
+ }
+
+ # Add tokenizer path if provided
+ if tokenizer_path:
+ vllm_kwargs["tokenizer"] = tokenizer_path
+
+ engines = [
+ ray.remote(num_cpus=0, num_gpus=0, scheduling_strategy=strategy)(ESNcclLLM).remote(
+ **vllm_kwargs
+ )
+ for strategy in strategies
+ ]
+ return engines, pgs
+
+def load_data(data_path: str, num_samples: int = None, offset: int = 0) -> List[Tuple[str, str]]:
+ """Load GSM8K data from parquet file"""
+ if not os.path.exists(data_path):
+ raise FileNotFoundError(f"Data file not found: {data_path}")
+
+ import datasets
+ ds = datasets.load_dataset('parquet', data_files={'test': data_path})
+
+ # Convert to list of (prompt, ground_truth) tuples
+ dataset = []
+ for item in ds['test']:
+ # Extract prompt text from the prompt field
+ prompt_content = item['prompt'][0]['content'] if isinstance(item['prompt'], list) else item['prompt']
+ ground_truth = item['reward_model']['ground_truth']
+ dataset.append((prompt_content, ground_truth))
+
+ # Apply offset and limit
+ if offset < 0:
+ start_idx, end_idx = len(dataset) + offset, len(dataset)
+ else:
+ start_idx, end_idx = offset, len(dataset)
+ dataset = dataset[start_idx:end_idx]
+
+ if num_samples is not None and num_samples < len(dataset):
+ dataset = dataset[:num_samples]
+
+ return dataset
+
+
+def extract_model_response(generated_text: str) -> str:
+ """Extract the model response from generated text"""
+ model_response = generated_text
+ if "assistant:" in generated_text:
+ model_response = generated_text.split("assistant:")[-1].strip()
+ return model_response
+
+
+def evaluate_batch_vllm(llm, input_texts: List[str], target_texts: List[str], args, verbose: bool = False) -> List[Dict[str, Any]]:
+ """Evaluate a batch of prompts using vLLM"""
+ if verbose:
+ print(f"Batch evaluating {len(input_texts)} samples...")
+
+ # Greedy vs sampling setup
+ temperature = args.temperature if args.do_sample else 0.0
+ top_p = args.top_p if args.do_sample else 1.0
+ sampling_params = SamplingParams(
+ temperature=temperature,
+ top_p=top_p,
+ max_tokens=args.max_new_tokens,
+ )
+
+ # Run generation
+ outputs = ray.get(llm.generate.remote(input_texts, sampling_params=sampling_params, use_tqdm=False))
+
+ # Process results
+ all_results = []
+ for i, out in enumerate(outputs):
+ # vLLM returns a RequestOutput with outputs list; take first completion
+ completion_text = out.outputs[0].text if out.outputs else ""
+ generated_text = f"{completion_text}"
+
+ model_response = extract_model_response(generated_text)
+
+ input_text = input_texts[i]
+ ground_truth = target_texts[i]
+
+ # Evaluate using GSM8K reward function
+ reward_result = reward_function(model_response, ground_truth)
+ reward = reward_result["reward"]
+ reward_info = reward_result["reward_info"]
+
+ all_results.append({
+ 'input_text': input_text,
+ 'ground_truth': ground_truth,
+ 'generated_text': generated_text,
+ 'model_response': model_response,
+ 'reward': reward,
+ 'reward_info': reward_info
+ })
+ return all_results
+
+
+def evaluate_dataset_vllm(llm, dataset: List[Tuple[str, str]], args, dataset_name: str, batch_size: int = None) -> Dict[str, Any]:
+ """Evaluate the entire dataset"""
+ print(f"\n=== Evaluating on {dataset_name} dataset ({len(dataset)} samples) ===")
+ if batch_size is None:
+ batch_size = min(1024, len(dataset))
+ print(f"Using batch size: {batch_size}")
+
+ all_results = []
+ total_reward = 0.0
+ total_format_reward = 0.0
+ total_answer_reward = 0.0
+
+ start_time = time.time()
+ for batch_start in range(0, len(dataset), batch_size):
+ batch_end = min(batch_start + batch_size, len(dataset))
+ batch_dataset = dataset[batch_start:batch_end]
+
+ if args.verbose:
+ print(f"Processing batch {batch_start//batch_size + 1}/{(len(dataset)-1)//batch_size + 1} (samples {batch_start+1}-{batch_end})...")
+
+ input_texts = [item[0] for item in batch_dataset]
+ target_texts = [item[1] for item in batch_dataset]
+
+ batch_results = evaluate_batch_vllm(llm, input_texts, target_texts, args, verbose=args.verbose)
+ all_results.extend(batch_results)
+
+ for result in batch_results:
+ total_reward += result['reward']
+ total_format_reward += result['reward_info']['format_reward']
+ total_answer_reward += result['reward_info']['answer_reward']
+
+ # Show examples from first batch only
+ if batch_start == 0:
+ for i, result in enumerate(batch_results[:args.show_examples]):
+ print(f"\n--- Example {i+1} ---")
+ print(f"Input: {result['input_text'][:200]}...")
+ print(f"Ground Truth: {result['ground_truth']}")
+ print(f"Model Response: {result['model_response'][:300]}...")
+ print(f"Reward: {result['reward']:.4f} (Format: {result['reward_info']['format_reward']:.4f}, Answer: {result['reward_info']['answer_reward']:.4f})")
+
+ eval_time = time.time() - start_time
+
+ # Calculate statistics
+ avg_reward = total_reward / len(dataset)
+ avg_format_reward = total_format_reward / len(dataset)
+ avg_answer_reward = total_answer_reward / len(dataset)
+
+ rewards = [r['reward'] for r in all_results]
+ std_reward = np.std(rewards)
+ min_reward = np.min(rewards)
+ max_reward = np.max(rewards)
+
+ high_reward_count = sum(1 for r in rewards if r >= 1.0)
+ high_reward_percentage = high_reward_count / len(dataset) * 100
+
+ answer_rewards = [r['reward_info']['answer_reward'] for r in all_results]
+ correct_count = sum(1 for r in answer_rewards if r > 0)
+ accuracy = correct_count / len(dataset) * 100
+
+ stats = {
+ 'dataset_name': dataset_name,
+ 'num_samples': len(dataset),
+ 'avg_reward': avg_reward,
+ 'avg_format_reward': avg_format_reward,
+ 'avg_answer_reward': avg_answer_reward,
+ 'std_reward': std_reward,
+ 'min_reward': min_reward,
+ 'max_reward': max_reward,
+ 'high_reward_count': high_reward_count,
+ 'high_reward_percentage': high_reward_percentage,
+ 'correct_count': correct_count,
+ 'accuracy': accuracy,
+ 'eval_time': eval_time,
+ 'all_results': all_results
+ }
+
+ print(f"\n=== {dataset_name} Results Summary ===")
+ print(f"Number of samples: {len(dataset)}")
+ print(f"Average reward: {avg_reward:.4f} ± {std_reward:.4f}")
+ print(f" - Format reward: {avg_format_reward:.4f}")
+ print(f" - Answer reward: {avg_answer_reward:.4f}")
+ print(f"Accuracy (answer_reward > 0): {correct_count}/{len(dataset)} ({accuracy:.1f}%)")
+ print(f"High reward samples (≥1.0): {high_reward_count}/{len(dataset)} ({high_reward_percentage:.1f}%)")
+ print(f"Reward range: [{min_reward:.4f}, {max_reward:.4f}]")
+ print(f"Evaluation time: {eval_time:.2f}s ({eval_time/len(dataset):.3f}s per sample)")
+
+ return stats
+
+
+def save_results(results: Dict[str, Any], output_dir: str, args):
+ """Save evaluation results to files"""
+ os.makedirs(output_dir, exist_ok=True)
+
+ summary = {
+ 'model_id': args.model_id,
+ 'eval_stats': {k: v for k, v in results['eval_stats'].items() if k != 'all_results'},
+ 'generation_config': {
+ 'max_new_tokens': args.max_new_tokens,
+ 'do_sample': args.do_sample,
+ 'temperature': args.temperature if args.do_sample else None,
+ 'top_p': args.top_p if args.do_sample else None,
+ 'batch_size': args.batch_size,
+ 'tensor_parallel_size': args.tensor_parallel_size,
+ 'dtype': args.dtype,
+ }
+ }
+
+ summary_path = os.path.join(output_dir, 'summary_gsm8k.json')
+ with open(summary_path, 'w') as f:
+ json.dump(summary, f, indent=2)
+ print(f"\nSummary saved to: {summary_path}")
+
+ if args.save_responses:
+ eval_details_path = os.path.join(output_dir, 'eval_detailed_results_gsm8k.json')
+ with open(eval_details_path, 'w') as f:
+ json.dump(results['eval_stats']['all_results'], f, indent=2)
+ print(f"Eval detailed results saved to: {eval_details_path}")
+
+
+def main():
+ args = parse_args()
+
+ # No auto connection to Ray cluster
+ os.environ.pop("RAY_ADDRESS", None)
+ os.environ.pop("RAY_HEAD_IP", None)
+ os.environ.pop("RAY_GCS_SERVER_ADDRESS", None)
+
+ # This prevents socket/file lock conflicts with other running Ray instances
+ unique_dir = tempfile.mkdtemp(prefix=f"ray_temp_session_{int(time.time())}_")
+
+ ray.init(
+ address="local",
+ include_dashboard=False,
+ ignore_reinit_error=True,
+ _temp_dir=unique_dir,
+ dashboard_port=None
+ )
+
+ # Import GSM8K reward function
+ global reward_function
+ from src.rewards.gsm8k_reward import reward_function
+
+ if args.output_dir is None:
+ model_name = os.path.basename(args.model_id.rstrip('/'))
+ batch_suffix = f"_batch{args.batch_size}" if args.batch_size else ""
+ args.output_dir = f"./inference_results_gsm8k_vllm_{model_name}{batch_suffix}"
+
+ print("=== vLLM ES Model Inference Script for GSM8K ===")
+ print(f"Base model: {args.model_id}")
+ print(f"Trained model path: {args.trained_model_path}")
+ print(f"Eval data: {args.eval_data_path} (samples: {args.eval_samples}, offset: {args.eval_offset})")
+ print(f"Output directory: {args.output_dir}")
+ print(f"Tensor parallel size: {args.tensor_parallel_size}")
+ print(f"Dtype: {args.dtype}")
+
+ # Initialize vLLM engine
+ print("\nInitializing vLLM engine...")
+ llm, _ = launch_engines(
+ num_engines=1,
+ model_name=args.model_id,
+ tokenizer_path=args.tokenizer_path
+ )
+ llm = llm[0]
+
+ # Load trained weights
+ print(f"Loading trained weights from {args.trained_model_path}...")
+ ray.get(llm.collective_rpc.remote("load_weights_from_disk", args=(args.trained_model_path,)))
+ print(f"vLLM LLM initialized with trained weights")
+
+ # Load dataset
+ print("\nLoading evaluation dataset...")
+ eval_dataset = load_data(
+ args.eval_data_path,
+ num_samples=args.eval_samples,
+ offset=args.eval_offset
+ )
+ print(f"Loaded {len(eval_dataset)} evaluation samples")
+
+ # Evaluate
+ eval_stats = evaluate_dataset_vllm(llm, eval_dataset, args, "GSM8K Test", batch_size=args.batch_size)
+ results = {'eval_stats': eval_stats}
+ save_results(results, args.output_dir, args)
+
+ print("\n=== Evaluation Complete ===")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/es-fine-tuning-paper/src/scripts/evaluation/evaluate_countdown.sh b/es-fine-tuning-paper/src/scripts/evaluation/evaluate_countdown.sh
new file mode 100755
index 0000000..246cf6b
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/evaluation/evaluate_countdown.sh
@@ -0,0 +1,230 @@
+#!/bin/bash
+set -e
+
+# Countdown Model Evaluation Script
+# This script merges FSDP checkpoints and evaluates the fine-tuned model on the test set
+
+echo "=========================================="
+echo "Countdown Model Evaluation Pipeline"
+echo "=========================================="
+
+# Configuration
+PROJECT_NAME="verl_grpo_countdown_base_custom"
+EXPERIMENT_NAME="qwen2.5_3b_base_custom_template_lora"
+BASE_MODEL="Qwen/Qwen2.5-3B"
+# BASE_MODEL='meta-llama/Llama-3.2-3b'
+TOKENIZER_PATH="" # Optional custom tokenizer path
+
+TEST_FILE="./src/data/countdown-0.4/test.parquet"
+TASK_TYPE="countdown"
+USE_FSDP="true" # Set to "false" to skip FSDP model merging
+
+# Parse command line arguments
+CHECKPOINT_STEP=""
+CHECKPOINT_PATH=""
+while [[ $# -gt 0 ]]; do
+ case $1 in
+ --step)
+ CHECKPOINT_STEP="$2"
+ shift 2
+ ;;
+ --checkpoint_path)
+ CHECKPOINT_PATH="$2"
+ shift 2
+ ;;
+ --base_model)
+ BASE_MODEL="$2"
+ shift 2
+ ;;
+ --tokenizer_path)
+ TOKENIZER_PATH="$2"
+ shift 2
+ ;;
+ --project_name)
+ PROJECT_NAME="$2"
+ shift 2
+ ;;
+ --experiment_name)
+ EXPERIMENT_NAME="$2"
+ shift 2
+ ;;
+ --use_fsdp)
+ USE_FSDP="$2"
+ shift 2
+ ;;
+ *)
+ echo "Unknown option: $1"
+ echo "Usage: $0 [--step CHECKPOINT_STEP] [--checkpoint_path CKPT_PATH] [--base_model MODEL_PATH] [--tokenizer_path TOKENIZER_PATH] [--project_name NAME] [--experiment_name NAME] [--use_fsdp true|false]"
+ exit 1
+ ;;
+ esac
+done
+
+# Determine checkpoint directory
+if [ -n "$CHECKPOINT_PATH" ]; then
+ # Use directly provided checkpoint path
+ echo "Using provided checkpoint path: $CHECKPOINT_PATH"
+ CHECKPOINT_DIR="$CHECKPOINT_PATH"
+ CHECKPOINT_STEP=$(basename "$CHECKPOINT_PATH" | sed 's/global_step_//' | sed 's/.*step_//' | sed 's/[^0-9]//g')
+ if [ -z "$CHECKPOINT_STEP" ]; then
+ CHECKPOINT_STEP="custom"
+ fi
+else
+ # Auto-detect latest checkpoint if not specified
+ if [ -z "$CHECKPOINT_STEP" ]; then
+ echo "Auto-detecting latest checkpoint..."
+ CHECKPOINT_BASE="checkpoints/${PROJECT_NAME}/${EXPERIMENT_NAME}"
+ if [ ! -d "$CHECKPOINT_BASE" ]; then
+ echo "Error: Checkpoint directory not found: $CHECKPOINT_BASE"
+ exit 1
+ fi
+
+ # Find the latest global_step directory
+ LATEST_STEP=$(ls -d ${CHECKPOINT_BASE}/global_step_* 2>/dev/null | sort -V | tail -1)
+ if [ -z "$LATEST_STEP" ]; then
+ echo "Error: No checkpoint found in $CHECKPOINT_BASE"
+ exit 1
+ fi
+ CHECKPOINT_STEP=$(basename "$LATEST_STEP" | sed 's/global_step_//')
+ echo "Found latest checkpoint: step $CHECKPOINT_STEP"
+ fi
+ CHECKPOINT_DIR="checkpoints/${PROJECT_NAME}/${EXPERIMENT_NAME}/global_step_${CHECKPOINT_STEP}"
+fi
+
+# Set paths
+ACTOR_DIR="${CHECKPOINT_DIR}/actor"
+MERGED_DIR="${CHECKPOINT_DIR}/merged_model"
+OUTPUT_FILE="evals/qwen2.5_3b_base_eval_results_countdown_${CHECKPOINT_STEP}.json"
+
+echo ""
+echo "Configuration:"
+echo " Project: $PROJECT_NAME"
+echo " Experiment: $EXPERIMENT_NAME"
+echo " Checkpoint Step: $CHECKPOINT_STEP"
+echo " Checkpoint Dir: $CHECKPOINT_DIR"
+echo " Base Model: $BASE_MODEL"
+echo " Test File: $TEST_FILE"
+echo " Use FSDP Merge: $USE_FSDP"
+echo " Output File: $OUTPUT_FILE"
+echo ""
+
+# Create evals directory if it doesn't exist
+mkdir -p evals
+
+# Determine model path based on FSDP flag
+if [ "$USE_FSDP" = "true" ]; then
+ # Check if checkpoint exists
+ if [ ! -d "$ACTOR_DIR" ]; then
+ echo "Error: Checkpoint directory not found: $ACTOR_DIR"
+ exit 1
+ fi
+
+ # Step 1: Merge FSDP checkpoint shards
+ echo "=========================================="
+ echo "Step 1: Merging FSDP checkpoint shards"
+ echo "=========================================="
+
+ if [ -d "$MERGED_DIR" ] && [ "$(ls -A $MERGED_DIR)" ]; then
+ echo "Merged model already exists at: $MERGED_DIR"
+ read -p "Do you want to re-merge? (y/N): " -n 1 -r
+ echo
+ if [[ ! $REPLY =~ ^[Yy]$ ]]; then
+ echo "Skipping merge step..."
+ else
+ echo "Removing existing merged model..."
+ rm -rf "$MERGED_DIR"
+ mkdir -p "$MERGED_DIR"
+
+ echo "Merging checkpoint..."
+ python3 verl/scripts/model_merger.py \
+ --backend fsdp \
+ --hf_model_path "$BASE_MODEL" \
+ --local_dir "$ACTOR_DIR" \
+ --target_dir "$MERGED_DIR"
+
+ echo "✓ Merge completed successfully!"
+ fi
+ else
+ mkdir -p "$MERGED_DIR"
+
+ echo "Merging checkpoint..."
+ python3 verl/scripts/model_merger.py \
+ --backend fsdp \
+ --hf_model_path "$BASE_MODEL" \
+ --local_dir "$ACTOR_DIR" \
+ --target_dir "$MERGED_DIR"
+
+ echo "✓ Merge completed successfully!"
+ fi
+
+ MODEL_PATH="$MERGED_DIR"
+ echo ""
+else
+ # Skip FSDP merge - use checkpoint directly
+ echo "=========================================="
+ echo "Skipping FSDP merge (use_fsdp=false)"
+ echo "=========================================="
+
+ # Check if checkpoint exists (file or directory)
+ if [ -f "$CHECKPOINT_DIR" ]; then
+ # If it's a file, use its parent directory
+ MODEL_PATH=$(dirname "$CHECKPOINT_DIR")
+ echo "Checkpoint is a file, using parent directory: $MODEL_PATH"
+ elif [ -d "$CHECKPOINT_DIR" ]; then
+ # If it's a directory, use it directly
+ MODEL_PATH="$CHECKPOINT_DIR"
+ echo "Using checkpoint directory: $MODEL_PATH"
+ else
+ echo "Error: Checkpoint path not found: $CHECKPOINT_DIR"
+ exit 1
+ fi
+ echo ""
+fi
+
+# Step 2: Run evaluation
+echo "=========================================="
+echo "Step 2: Evaluating on test set"
+echo "=========================================="
+
+if [ ! -f "$TEST_FILE" ]; then
+ echo "Error: Test file not found: $TEST_FILE"
+ echo "Please ensure countdown data is prepared."
+ exit 1
+fi
+
+echo "Running evaluation..."
+echo "This may take several minutes depending on GPU speed..."
+echo ""
+
+# Build evaluation command
+EVAL_CMD="python3 evaluate_model.py \
+ --model_path '$MODEL_PATH' \
+ --test_file '$TEST_FILE' \
+ --task_type '$TASK_TYPE' \
+ --output_file '$OUTPUT_FILE'"
+
+# Add tokenizer path if specified
+if [ -n "$TOKENIZER_PATH" ]; then
+ EVAL_CMD="$EVAL_CMD \
+ --tokenizer_path '$TOKENIZER_PATH'"
+ echo "Using custom tokenizer: $TOKENIZER_PATH"
+fi
+
+# Execute evaluation
+eval $EVAL_CMD
+
+echo ""
+echo "=========================================="
+echo "Evaluation Complete!"
+echo "=========================================="
+echo "Results saved to: $OUTPUT_FILE"
+echo ""
+
+# Display results if jq is available
+if command -v jq &> /dev/null; then
+ echo "Quick Summary:"
+ jq -r '"Total: \(.total) | Correct: \(.correct) | Accuracy: \(.accuracy * 100 | floor)%"' "$OUTPUT_FILE"
+else
+ echo "Install 'jq' to see a quick summary here."
+ echo "Otherwise, check the full results in: $OUTPUT_FILE"
+fi
diff --git a/es-fine-tuning-paper/src/scripts/evaluation/evaluate_gsm8k.sh b/es-fine-tuning-paper/src/scripts/evaluation/evaluate_gsm8k.sh
new file mode 100755
index 0000000..9d774be
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/evaluation/evaluate_gsm8k.sh
@@ -0,0 +1,182 @@
+#!/bin/bash
+set -e
+
+# GSM8K Model Evaluation Script
+# This script merges FSDP checkpoints and evaluates the fine-tuned model on the reserved test set
+
+echo "=========================================="
+echo "GSM8K Model Evaluation Pipeline"
+echo "=========================================="
+
+# Configuration
+PROJECT_NAME="verl_grpo_gsm8k_base_custom"
+EXPERIMENT_NAME="qwen2.5_3b_base_custom_template_lora_0.4"
+# BASE_MODEL="meta-llama/Llama-3.2-3B"
+BASE_MODEL="Qwen/Qwen2.5-3B"
+TOKENIZER_PATH="./tokenizers/qwen2.5-3b-base-chat" # Optional custom tokenizer path
+TEST_FILE="./src/data/gsm8k-0.4/test.parquet"
+TASK_TYPE="gsm8k"
+
+# Parse command line arguments
+CHECKPOINT_STEP=""
+while [[ $# -gt 0 ]]; do
+ case $1 in
+ --step)
+ CHECKPOINT_STEP="$2"
+ shift 2
+ ;;
+ --base_model)
+ BASE_MODEL="$2"
+ shift 2
+ ;;
+ --tokenizer_path)
+ TOKENIZER_PATH="$2"
+ shift 2
+ ;;
+ --project_name)
+ PROJECT_NAME="$2"
+ shift 2
+ ;;
+ --experiment_name)
+ EXPERIMENT_NAME="$2"
+ shift 2
+ ;;
+ *)
+ echo "Unknown option: $1"
+ echo "Usage: $0 [--step CHECKPOINT_STEP] [--base_model MODEL_PATH] [--tokenizer_path TOKENIZER_PATH] [--project_name NAME] [--experiment_name NAME]"
+ exit 1
+ ;;
+ esac
+done
+
+# Auto-detect latest checkpoint if not specified
+if [ -z "$CHECKPOINT_STEP" ]; then
+ echo "Auto-detecting latest checkpoint..."
+ CHECKPOINT_BASE="checkpoints/${PROJECT_NAME}/${EXPERIMENT_NAME}"
+ if [ ! -d "$CHECKPOINT_BASE" ]; then
+ echo "Error: Checkpoint directory not found: $CHECKPOINT_BASE"
+ exit 1
+ fi
+
+ # Find the latest global_step directory
+ LATEST_STEP=$(ls -d ${CHECKPOINT_BASE}/global_step_* 2>/dev/null | sort -V | tail -1)
+ if [ -z "$LATEST_STEP" ]; then
+ echo "Error: No checkpoint found in $CHECKPOINT_BASE"
+ exit 1
+ fi
+ CHECKPOINT_STEP=$(basename "$LATEST_STEP" | sed 's/global_step_//')
+ echo "Found latest checkpoint: step $CHECKPOINT_STEP"
+fi
+
+# Set paths
+CHECKPOINT_DIR="checkpoints/${PROJECT_NAME}/${EXPERIMENT_NAME}/global_step_${CHECKPOINT_STEP}"
+ACTOR_DIR="${CHECKPOINT_DIR}/actor"
+MERGED_DIR="${CHECKPOINT_DIR}/merged_model"
+OUTPUT_FILE="evals/eval_results_gsm8k_step${CHECKPOINT_STEP}.json"
+
+echo ""
+echo "Configuration:"
+echo " Project: $PROJECT_NAME"
+echo " Experiment: $EXPERIMENT_NAME"
+echo " Checkpoint Step: $CHECKPOINT_STEP"
+echo " Base Model: $BASE_MODEL"
+echo " Test File: $TEST_FILE"
+echo " Output File: $OUTPUT_FILE"
+echo ""
+
+# Create evals directory if it doesn't exist
+mkdir -p evals
+
+# Check if checkpoint exists
+if [ ! -d "$ACTOR_DIR" ]; then
+ echo "Error: Checkpoint directory not found: $ACTOR_DIR"
+ exit 1
+fi
+
+# Step 1: Merge FSDP checkpoint shards
+echo "=========================================="
+echo "Step 1: Merging FSDP checkpoint shards"
+echo "=========================================="
+
+if [ -d "$MERGED_DIR" ] && [ "$(ls -A $MERGED_DIR)" ]; then
+ echo "Merged model already exists at: $MERGED_DIR"
+ read -p "Do you want to re-merge? (y/N): " -n 1 -r
+ echo
+ if [[ ! $REPLY =~ ^[Yy]$ ]]; then
+ echo "Skipping merge step..."
+ else
+ echo "Removing existing merged model..."
+ rm -rf "$MERGED_DIR"
+ mkdir -p "$MERGED_DIR"
+
+ echo "Merging checkpoint..."
+ python3 verl/scripts/model_merger.py \
+ --backend fsdp \
+ --hf_model_path "$BASE_MODEL" \
+ --local_dir "$ACTOR_DIR" \
+ --target_dir "$MERGED_DIR"
+
+ echo "✓ Merge completed successfully!"
+ fi
+else
+ mkdir -p "$MERGED_DIR"
+
+ echo "Merging checkpoint..."
+ python3 verl/scripts/model_merger.py \
+ --backend fsdp \
+ --hf_model_path "$BASE_MODEL" \
+ --local_dir "$ACTOR_DIR" \
+ --target_dir "$MERGED_DIR"
+
+ echo "✓ Merge completed successfully!"
+fi
+
+echo ""
+
+# Step 2: Run evaluation
+echo "=========================================="
+echo "Step 2: Evaluating on reserved test set"
+echo "=========================================="
+
+if [ ! -f "$TEST_FILE" ]; then
+ echo "Error: Test file not found: $TEST_FILE"
+ echo "Please run: ./prepare_gsm8k_data.sh"
+ exit 1
+fi
+
+echo "Running evaluation..."
+echo "This may take several minutes depending on GPU speed..."
+echo ""
+
+# Build evaluation command
+EVAL_CMD="python3 evaluate_model.py \
+ --model_path '$MERGED_DIR' \
+ --test_file '$TEST_FILE' \
+ --task_type '$TASK_TYPE' \
+ --output_file '$OUTPUT_FILE'"
+
+# Add tokenizer path if specified
+if [ -n "$TOKENIZER_PATH" ]; then
+ EVAL_CMD="$EVAL_CMD \
+ --tokenizer_path '$TOKENIZER_PATH'"
+ echo "Using custom tokenizer: $TOKENIZER_PATH"
+fi
+
+# Execute evaluation
+eval $EVAL_CMD
+
+echo ""
+echo "=========================================="
+echo "Evaluation Complete!"
+echo "=========================================="
+echo "Results saved to: $OUTPUT_FILE"
+echo ""
+
+# Display results if jq is available
+if command -v jq &> /dev/null; then
+ echo "Quick Summary:"
+ jq -r '"Total: \(.total) | Correct: \(.correct) | Accuracy: \(.accuracy * 100 | floor)%"' "$OUTPUT_FILE"
+else
+ echo "Install 'jq' to see a quick summary here."
+ echo "Otherwise, check the full results in: $OUTPUT_FILE"
+fi
\ No newline at end of file
diff --git a/es-fine-tuning-paper/src/scripts/evaluation/evaluate_model.py b/es-fine-tuning-paper/src/scripts/evaluation/evaluate_model.py
new file mode 100644
index 0000000..70ea2f8
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/evaluation/evaluate_model.py
@@ -0,0 +1,268 @@
+#!/usr/bin/env python3
+"""
+Evaluation script for fine-tuned models on GSM8K and Countdown test sets.
+This script loads a saved model checkpoint and evaluates it on the reserved test set.
+"""
+
+import argparse
+import json
+import os
+import re
+from typing import Dict, List, Any
+
+import pandas as pd
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from tqdm import tqdm
+
+
+def extract_solution(solution_str):
+ """Extract numerical solution from GSM8K answer format."""
+ solution = re.search("#### (\\-?[0-9\\.\\,]+)", solution_str)
+ if solution is None:
+ return None
+ final_solution = solution.group(0)
+ final_solution = final_solution.split("#### ")[1].replace(",", "")
+ return final_solution
+
+
+def extract_answer_from_response(response: str, task_type: str = "gsm8k") -> str:
+ """Extract answer from model response based on task type."""
+ if task_type == "gsm8k":
+ # Look for #### format answer - find ALL matches and take the LAST one
+ matches = list(re.finditer(r"#### (\-?[0-9\\.\\,]+)", response))
+ if matches:
+ # Take the last match (the final answer)
+ return matches[-1].group(1).replace(",", "")
+ elif task_type == "countdown":
+ # Look for tags - find ALL matches and take the LAST one
+ matches = re.findall(r"(.*?)", response, re.DOTALL)
+ if matches:
+ return matches[-1].strip()
+ return None
+
+
+def evaluate_gsm8k_answer(predicted: str, ground_truth: str) -> bool:
+ """Check if predicted answer matches ground truth for GSM8K."""
+ if predicted is None or ground_truth is None:
+ return False
+
+ try:
+ pred_float = float(predicted.replace(",", ""))
+ gt_float = float(ground_truth.replace(",", ""))
+ return abs(pred_float - gt_float) < 1e-5
+ except:
+ return False
+
+
+def evaluate_countdown_answer(response: str, numbers: List[int], target: float) -> bool:
+ """Check if countdown answer is correct."""
+ answer_regex = r"(.*?)"
+ all_matches = re.findall(answer_regex, response, re.DOTALL)
+ if not all_matches:
+ return False
+
+ # Take the LAST answer tag (final answer)
+ answer_content = all_matches[-1].strip()
+ if not answer_content:
+ return False
+
+ # Check if the answer uses only allowed characters
+ allowed_chars = r"^[0-9+\-*/() ]+$"
+ if not re.match(allowed_chars, answer_content):
+ return False
+
+ # Check if the answer uses all numbers exactly once
+ used_numbers = [int(n) for n in re.findall(r"\d+", answer_content)]
+ if sorted(used_numbers) != sorted(numbers):
+ return False
+
+ # Check if the answer evaluates to the target
+ try:
+ result = eval(answer_content, {"__builtins__": None}, {})
+ if abs(float(result) - float(target)) < 1e-5:
+ return True
+ except:
+ pass
+
+ return False
+
+
+def load_model(model_path: str, base_model: str = None, tokenizer_path: str = None, device: str = "cuda"):
+ """Load the fine-tuned model (base + LoRA adapter if applicable)."""
+ print(f"Loading model from {model_path}...")
+
+
+ # Load full model
+ model = AutoModelForCausalLM.from_pretrained(
+ model_path,
+ torch_dtype=torch.bfloat16,
+ device_map=device
+ )
+ # Use custom tokenizer path if provided, otherwise use model_path
+ if tokenizer_path is not None:
+ print(f"Loading custom tokenizer from: {tokenizer_path}")
+ tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+ else:
+ tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+ model.eval()
+ return model, tokenizer
+
+
+def generate_response(model, tokenizer, messages: List[Dict], max_new_tokens: int = 512, device: str = "cuda"):
+ """Generate response from model using chat template."""
+ # Apply chat template
+ prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+ inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+ with torch.no_grad():
+ outputs = model.generate(
+ **inputs,
+ max_new_tokens=max_new_tokens,
+ do_sample=False,
+ temperature=None, # Disable temperature when not sampling
+ top_p=None, # Disable top_p when not sampling
+ pad_token_id=tokenizer.eos_token_id,
+ )
+
+ # Decode only the generated part (exclude the input prompt)
+ response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+ return response
+
+
+def evaluate_dataset(
+ model,
+ tokenizer,
+ test_file: str,
+ task_type: str = "gsm8k",
+ max_samples: int = None,
+ device: str = "cuda"
+) -> Dict[str, Any]:
+ """Evaluate model on a test dataset."""
+
+ # Load test data
+ df = pd.read_parquet(test_file)
+
+ if max_samples:
+ df = df.head(max_samples)
+
+ total = len(df)
+ correct = 0
+ results = []
+
+ print(f"\nEvaluating on {total} samples from {test_file}...")
+
+ for idx, row in tqdm(df.iterrows(), total=total):
+ # Get messages from the data (this matches the training format)
+ messages = row['prompt']
+
+ # Generate response using the same format as training
+ response = generate_response(model, tokenizer, messages, device=device)
+
+ # Evaluate based on task type
+ is_correct = False
+ ground_truth_value = None
+ extracted_answer = None
+
+ if task_type == "gsm8k":
+ ground_truth = row['reward_model']['ground_truth']
+ ground_truth_value = ground_truth
+ extracted_answer = extract_answer_from_response(response, "gsm8k")
+ is_correct = evaluate_gsm8k_answer(extracted_answer, ground_truth)
+ elif task_type == "countdown":
+ extra_info = row['extra_info']
+ numbers = extra_info['numbers']
+ target = extra_info['target']
+ ground_truth_value = extra_info["solution"]
+ extracted_answer = extract_answer_from_response(response, "countdown")
+ is_correct = evaluate_countdown_answer(response, numbers, target)
+
+ if is_correct:
+ correct += 1
+
+ # Store the prompt as text for readability in results
+ prompt_text = ""
+ for msg in messages:
+ if msg['role'] == 'system':
+ prompt_text += f"System: {msg['content']}\n\n"
+ elif msg['role'] == 'user':
+ prompt_text += f"User: {msg['content']}\n"
+ elif msg['role'] == 'assistant':
+ prompt_text += f"Assistant: {msg['content']}"
+
+ results.append({
+ 'index': idx,
+ 'prompt': prompt_text,
+ 'response': response,
+ 'extracted_answer': extracted_answer,
+ 'correct': is_correct,
+ 'ground_truth': ground_truth_value
+ })
+
+ accuracy = correct / total if total > 0 else 0
+
+ return {
+ 'total': total,
+ 'correct': correct,
+ 'accuracy': accuracy,
+ 'results': results
+ }
+
+
+def main():
+ parser = argparse.ArgumentParser(description="Evaluate fine-tuned model on test set")
+ parser.add_argument("--model_path", type=str, required=True,
+ help="Path to saved model checkpoint or LoRA adapter")
+ parser.add_argument("--base_model", type=str, default=None,
+ help="Base model path (required if model_path is LoRA adapter)")
+ parser.add_argument("--tokenizer_path", type=str, default=None,
+ help="Path to custom tokenizer (optional, if different from model/base_model)")
+ parser.add_argument("--test_file", type=str, required=True,
+ help="Path to test parquet file")
+ parser.add_argument("--task_type", type=str, choices=["gsm8k", "countdown"],
+ default="gsm8k",
+ help="Type of task to evaluate")
+ parser.add_argument("--max_samples", type=int, default=None,
+ help="Maximum number of samples to evaluate (default: all)")
+ parser.add_argument("--output_file", type=str, default=None,
+ help="Path to save detailed results (JSON)")
+ parser.add_argument("--device", type=str, default="cuda",
+ help="Device to run evaluation on")
+
+ args = parser.parse_args()
+
+ # Load model
+ model, tokenizer = load_model(args.model_path, args.base_model, args.tokenizer_path, args.device)
+
+ # Evaluate
+ eval_results = evaluate_dataset(
+ model,
+ tokenizer,
+ args.test_file,
+ args.task_type,
+ args.max_samples,
+ args.device
+ )
+
+ # Print results
+ print(f"\n{'='*60}")
+ print(f"Evaluation Results - {args.task_type.upper()}")
+ print(f"{'='*60}")
+ print(f"Model: {args.model_path}")
+ print(f"Test file: {args.test_file}")
+ print(f"Total samples: {eval_results['total']}")
+ print(f"Correct: {eval_results['correct']}")
+ print(f"Accuracy: {eval_results['accuracy']:.4f} ({eval_results['accuracy']*100:.2f}%)")
+ print(f"{'='*60}\n")
+
+ # Save detailed results if requested
+ if args.output_file:
+ with open(args.output_file, 'w') as f:
+ json.dump(eval_results, f, indent=2)
+ print(f"Detailed results saved to: {args.output_file}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/es-fine-tuning-paper/src/scripts/grpo/grpo-countdown-custom.sh b/es-fine-tuning-paper/src/scripts/grpo/grpo-countdown-custom.sh
new file mode 100644
index 0000000..157bbce
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/grpo/grpo-countdown-custom.sh
@@ -0,0 +1,75 @@
+set -x
+
+# Check if HF_TOKEN is set
+if [ -z "$HF_TOKEN" ]; then
+ echo "Error: HF_TOKEN environment variable is not set"
+ echo "Please run: export HF_TOKEN=your_huggingface_token"
+ exit 1
+fi
+
+# Login to HuggingFace to access models
+huggingface-cli login --token "$HF_TOKEN"
+
+# Alternative approach using custom_reward_function config
+# This version uses an external reward function file instead of built-in VERL reward scoring
+
+#bsz was changed to 128 from 256 because the verl script dropped off data sample if it is not under the batch size we had set
+
+python3 -m verl.trainer.main_ppo \
+ algorithm.adv_estimator=grpo \
+ trainer.val_before_train=False \
+ data.train_files=./src/data/countdown-0.4/train.parquet \
+ data.val_files=./src/data/countdown-0.4/test.parquet \
+ data.train_batch_size=128 \
+ data.max_prompt_length=256 \
+ data.max_response_length=1024 \
+ data.filter_overlong_prompts=True \
+ data.truncation='error' \
+ data.shuffle=False \
+ actor_rollout_ref.model.path=meta-llama/Llama-3.2-3B-Instruct \
+ +actor_rollout_ref.model.lora_rank=64 \
+ +actor_rollout_ref.model.lora_alpha=32 \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.actor.ppo_mini_batch_size=128 \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.actor.use_kl_loss=True \
+ actor_rollout_ref.actor.kl_loss_coef=0.001 \
+ actor_rollout_ref.actor.kl_loss_type=low_var_kl \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.model.enable_gradient_checkpointing=True \
+ actor_rollout_ref.actor.fsdp_config.param_offload=False \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
+ actor_rollout_ref.rollout.name=vllm \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
+ actor_rollout_ref.rollout.n=8 \
+ actor_rollout_ref.rollout.load_format=safetensors \
+ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=True \
+ algorithm.use_kl_in_reward=False \
+ trainer.critic_warmup=0 \
+ trainer.logger='["console","wandb"]' \
+ trainer.project_name='verl_grpo_countdown' \
+ trainer.experiment_name='qwen2.5_3b_grpo_countdown_lora' \
+ trainer.n_gpus_per_node=8 \
+ trainer.nnodes=1 \
+ trainer.save_freq=100 \
+ trainer.test_freq=100 \
+ trainer.total_epochs=100 \
+ custom_reward_function.path=./countdown_reward.py \
+ custom_reward_function.name=countdown_reward_function
+
+ # actor_rollout_ref.actor.ppo_mini_batch_size=256 \
+ # data.train_batch_size=1024 \
+ # trainer.n_gpus_per_node=8 \
+ # actor_rollout_ref.model.use_shm=True \
+
+# After training completes, evaluate the saved model on the test set:
+# python3 evaluate_model.py \
+# --model_path \
+# --base_model meta-llama/Llama-3.2-3B-Instruct \
+# --test_file ./src/data/countdown-0.1/test.parquet \
+# --task_type countdown \
+# --output_file ./eval_results_countdown.json
diff --git a/es-fine-tuning-paper/src/scripts/grpo/grpo-countdown-llama-base-custom.sh b/es-fine-tuning-paper/src/scripts/grpo/grpo-countdown-llama-base-custom.sh
new file mode 100755
index 0000000..21302a8
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/grpo/grpo-countdown-llama-base-custom.sh
@@ -0,0 +1,77 @@
+#!/bin/bash
+set -x
+
+# GRPO training for Countdown using BASE Llama model with custom chat template tokenizer
+# Model: meta-llama/Llama-3.2-3B (base, not instruct)
+# Tokenizer: Custom tokenizer with "Question: {input} Answer: Let's think step by step." template
+
+# Check if HF_TOKEN is set
+if [ -z "$HF_TOKEN" ]; then
+ echo "Error: HF_TOKEN environment variable is not set"
+ echo "Please run: export HF_TOKEN=your_huggingface_token"
+ exit 1
+fi
+
+# Login to HuggingFace to access models
+huggingface-cli login --token "$HF_TOKEN"
+
+# Prepare custom tokenizer if not already created
+if [ ! -d "./tokenizers/llama-3.2-3b-base-chat" ]; then
+ echo "Creating custom tokenizer for Llama base model..."
+ python3 base_model_tokenizer.py \
+ --model_path meta-llama/Llama-3.2-3B \
+ --save_path ./tokenizers/llama-3.2-3b-base-chat
+fi
+
+# Train base Llama model with custom tokenizer
+python3 -m verl.trainer.main_ppo \
+ algorithm.adv_estimator=grpo \
+ trainer.val_before_train=False \
+ data.train_files=./src/data/countdown-0.1/train.parquet \
+ data.val_files=./src/data/countdown-0.1/test.parquet \
+ data.train_batch_size=128 \
+ data.max_prompt_length=256 \
+ data.max_response_length=1024 \
+ data.filter_overlong_prompts=True \
+ data.truncation='error' \
+ data.shuffle=False \
+ actor_rollout_ref.model.path=meta-llama/Llama-3.2-3B \
+ +actor_rollout_ref.model.tokenizer_path=./tokenizers/llama-3.2-3b-base-chat \
+ +actor_rollout_ref.model.lora_rank=64 \
+ +actor_rollout_ref.model.lora_alpha=32 \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.actor.ppo_mini_batch_size=128 \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.actor.use_kl_loss=True \
+ actor_rollout_ref.actor.kl_loss_coef=0.001 \
+ actor_rollout_ref.actor.kl_loss_type=low_var_kl \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.model.enable_gradient_checkpointing=True \
+ actor_rollout_ref.actor.fsdp_config.param_offload=False \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
+ actor_rollout_ref.rollout.name=vllm \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
+ actor_rollout_ref.rollout.n=8 \
+ actor_rollout_ref.rollout.load_format=safetensors \
+ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=True \
+ algorithm.use_kl_in_reward=False \
+ trainer.critic_warmup=0 \
+ trainer.logger='["console","wandb"]' \
+ trainer.project_name='verl_grpo_countdown_base_custom' \
+ trainer.experiment_name='llama3.2_3b_base_custom_template_lora' \
+ trainer.n_gpus_per_node=2 \
+ trainer.nnodes=1 \
+ trainer.save_freq=100 \
+ trainer.test_freq=100 \
+ trainer.total_epochs=100 \
+ custom_reward_function.path=./countdown_reward.py \
+ custom_reward_function.name=countdown_reward_function
+
+# After training completes, evaluate the saved model on the test set:
+# bash evaluate_countdown.sh --base_model meta-llama/Llama-3.2-3B \
+# --project_name verl_grpo_countdown_base_custom \
+# --experiment_name llama3.2_3b_base_custom_template_lora
diff --git a/es-fine-tuning-paper/src/scripts/grpo/grpo-countdown-qwen-base-custom.sh b/es-fine-tuning-paper/src/scripts/grpo/grpo-countdown-qwen-base-custom.sh
new file mode 100755
index 0000000..48c3420
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/grpo/grpo-countdown-qwen-base-custom.sh
@@ -0,0 +1,78 @@
+#!/bin/bash
+set -x
+
+# GRPO training for Countdown using BASE Qwen model with custom chat template tokenizer
+# Model: Qwen/Qwen2.5-3B (base, not instruct)
+# Tokenizer: Custom tokenizer with "Question: {input} Answer: Let's think step by step." template
+
+# Check if HF_TOKEN is set
+if [ -z "$HF_TOKEN" ]; then
+ echo "Error: HF_TOKEN environment variable is not set"
+ echo "Please run: export HF_TOKEN=your_huggingface_token"
+ exit 1
+fi
+
+# Login to HuggingFace to access models
+huggingface-cli login --token "$HF_TOKEN"
+
+# Prepare custom tokenizer if not already created
+if [ ! -d "./tokenizers/qwen2.5-3b-base-chat" ]; then
+ echo "Creating custom tokenizer for Qwen base model..."
+ python3 base_model_tokenizer.py \
+ --model_path Qwen/Qwen2.5-3B \
+ --save_path ./tokenizers/qwen2.5-3b-base-chat
+fi
+
+# Train base Qwen model with custom tokenizer
+python3 -m verl.trainer.main_ppo \
+ algorithm.adv_estimator=grpo \
+ trainer.val_before_train=False \
+ trainer.resume_mode=disable \
+ data.train_files=./data/countdown-0.4/train.parquet \
+ data.val_files=./data/countdown-0.4/test.parquet \
+ data.train_batch_size=128 \
+ data.max_prompt_length=256 \
+ data.max_response_length=1024 \
+ data.filter_overlong_prompts=True \
+ data.truncation='error' \
+ data.shuffle=False \
+ actor_rollout_ref.model.path=Qwen/Qwen2.5-3B \
+ +actor_rollout_ref.model.tokenizer_path=./tokenizers/qwen2.5-3b-base-chat \
+ +actor_rollout_ref.model.lora_rank=64 \
+ +actor_rollout_ref.model.lora_alpha=32 \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.actor.ppo_mini_batch_size=128 \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.actor.use_kl_loss=True \
+ actor_rollout_ref.actor.kl_loss_coef=0.001 \
+ actor_rollout_ref.actor.kl_loss_type=low_var_kl \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.model.enable_gradient_checkpointing=True \
+ actor_rollout_ref.actor.fsdp_config.param_offload=False \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
+ actor_rollout_ref.rollout.name=vllm \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
+ actor_rollout_ref.rollout.n=8 \
+ actor_rollout_ref.rollout.load_format=safetensors \
+ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=True \
+ algorithm.use_kl_in_reward=False \
+ trainer.critic_warmup=0 \
+ trainer.logger='["console","wandb"]' \
+ trainer.project_name='verl_grpo_countdown_base_custom' \
+ trainer.experiment_name='qwen2.5_3b_base_custom_template_lora' \
+ trainer.n_gpus_per_node=8 \
+ trainer.nnodes=1 \
+ trainer.save_freq=100 \
+ trainer.test_freq=100 \
+ trainer.total_epochs=100 \
+ custom_reward_function.path=./countdown_reward.py \
+ custom_reward_function.name=countdown_reward_function
+
+# After training completes, evaluate the saved model on the test set:
+# bash evaluate_countdown.sh --base_model Qwen/Qwen2.5-3B \
+# --project_name verl_grpo_countdown_base_custom \
+# --experiment_name qwen2.5_3b_base_custom_template_lora
diff --git a/es-fine-tuning-paper/src/scripts/grpo/grpo-countdown.sh b/es-fine-tuning-paper/src/scripts/grpo/grpo-countdown.sh
new file mode 100644
index 0000000..59297cd
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/grpo/grpo-countdown.sh
@@ -0,0 +1,51 @@
+set -x
+
+python3 -m verl.trainer.main_ppo \
+ algorithm.adv_estimator=grpo \
+ trainer.val_before_train=False \
+ data.train_files=../countdown/data/train.parquet \
+ data.val_files=../countdown/data/test.parquet \
+ data.train_batch_size=64 \
+ data.max_prompt_length=256 \
+ data.max_response_length=1024 \
+ data.filter_overlong_prompts=True \
+ data.truncation='error' \
+ data.shuffle=False \
+ actor_rollout_ref.model.path=Qwen/Qwen2.5-3B-Instruct \
+ +actor_rollout_ref.model.lora_rank=64 \
+ +actor_rollout_ref.model.lora_alpha=32 \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.actor.ppo_mini_batch_size=2 \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \
+ actor_rollout_ref.actor.use_kl_loss=True \
+ actor_rollout_ref.actor.kl_loss_coef=0.001 \
+ actor_rollout_ref.actor.kl_loss_type=low_var_kl \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.model.enable_gradient_checkpointing=True \
+ actor_rollout_ref.actor.fsdp_config.param_offload=False \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
+ actor_rollout_ref.rollout.name=vllm \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
+ actor_rollout_ref.rollout.n=8 \
+ actor_rollout_ref.rollout.load_format=safetensors \
+ actor_rollout_ref.rollout.max_num_batched_tokens=65535 \
+ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=True \
+ algorithm.use_kl_in_reward=False \
+ trainer.critic_warmup=0 \
+ trainer.logger='["console","wandb"]' \
+ trainer.project_name='verl_grpo_countdown' \
+ trainer.experiment_name='qwen2.5_3b_grpo_countdown_lora' \
+ trainer.n_gpus_per_node=8 \
+ trainer.nnodes=1 \
+ trainer.save_freq=-1 \
+ trainer.test_freq=50 \
+ trainer.total_epochs=256
+
+ # actor_rollout_ref.actor.ppo_mini_batch_size=256 \
+ # data.train_batch_size=1024 \
+ # trainer.n_gpus_per_node=8 \
+ # actor_rollout_ref.model.use_shm=True \
diff --git a/es-fine-tuning-paper/src/scripts/grpo/grpo-gsm8k-base.sh b/es-fine-tuning-paper/src/scripts/grpo/grpo-gsm8k-base.sh
new file mode 100644
index 0000000..ab7d155
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/grpo/grpo-gsm8k-base.sh
@@ -0,0 +1,57 @@
+set -x
+
+# Check if HF_TOKEN is set
+if [ -z "$HF_TOKEN" ]; then
+ echo "Error: HF_TOKEN environment variable is not set"
+ echo "Please run: export HF_TOKEN=your_huggingface_token"
+ exit 1
+fi
+
+# Login to HuggingFace to access models
+huggingface-cli login --token "$HF_TOKEN"
+
+python3 -m verl.trainer.main_ppo \
+ algorithm.adv_estimator=grpo \
+ trainer.val_before_train=False \
+ data.train_files=./src/data/gsm8k-0.1/train.parquet \
+ data.val_files=./src/data/gsm8k-0.1/validation.parquet \
+ data.train_batch_size=32 \
+ data.max_prompt_length=512 \
+ data.max_response_length=1024 \
+ data.filter_overlong_prompts=True \
+ data.truncation='error' \
+ data.shuffle=False \
+ actor_rollout_ref.model.path=meta-llama/Llama-3.2-3B \
+ +actor_rollout_ref.model.tokenizer_path=meta-llama/Llama-3.2-3B-Instruct \
+ +actor_rollout_ref.model.lora_rank=64 \
+ +actor_rollout_ref.model.lora_alpha=32 \
+ actor_rollout_ref.actor.optim.lr=3e-6 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.actor.ppo_mini_batch_size=32 \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.actor.use_kl_loss=True \
+ actor_rollout_ref.actor.kl_loss_coef=0.001 \
+ actor_rollout_ref.actor.kl_loss_type=low_var_kl \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.model.enable_gradient_checkpointing=True \
+ actor_rollout_ref.actor.fsdp_config.param_offload=False \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
+ actor_rollout_ref.rollout.name=vllm \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
+ actor_rollout_ref.rollout.n=8 \
+ actor_rollout_ref.rollout.load_format=safetensors \
+ actor_rollout_ref.rollout.max_num_batched_tokens=65535 \
+ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=True \
+ algorithm.use_kl_in_reward=False \
+ trainer.critic_warmup=0 \
+ trainer.logger='["console","wandb"]' \
+ trainer.project_name='verl_grpo_gsm8k_base' \
+ trainer.experiment_name='llama3.2_3b_base_grpo_lora' \
+ trainer.n_gpus_per_node=8 \
+ trainer.nnodes=1 \
+ trainer.save_freq=23 \
+ trainer.test_freq=1 \
+ trainer.total_epochs=1
\ No newline at end of file
diff --git a/es-fine-tuning-paper/src/scripts/grpo/grpo-gsm8k-llama-base-custom.sh b/es-fine-tuning-paper/src/scripts/grpo/grpo-gsm8k-llama-base-custom.sh
new file mode 100755
index 0000000..7277fef
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/grpo/grpo-gsm8k-llama-base-custom.sh
@@ -0,0 +1,76 @@
+#!/bin/bash
+set -x
+
+# GRPO training for GSM8K using BASE Llama model with custom chat template tokenizer
+# Model: meta-llama/Llama-3.2-3B (base, not instruct)
+# Tokenizer: Custom tokenizer with "Question: {input} Answer: Let's think step by step." template
+
+# Check if HF_TOKEN is set
+if [ -z "$HF_TOKEN" ]; then
+ echo "Error: HF_TOKEN environment variable is not set"
+ echo "Please run: export HF_TOKEN=your_huggingface_token"
+ exit 1
+fi
+
+# Login to HuggingFace to access models
+huggingface-cli login --token "$HF_TOKEN"
+
+# Prepare custom tokenizer if not already created
+if [ ! -d "./tokenizers/llama-3.2-3b-base-chat" ]; then
+ echo "Creating custom tokenizer for Llama base model..."
+ python3 base_model_tokenizer.py \
+ --model_path meta-llama/Llama-3.2-3B \
+ --save_path ./tokenizers/llama-3.2-3b-base-chat
+fi
+
+# Train base Llama model with custom tokenizer
+python3 -m verl.trainer.main_ppo \
+ algorithm.adv_estimator=grpo \
+ trainer.val_before_train=False \
+ data.train_files=./src/data/gsm8k-0.1/train.parquet \
+ data.val_files=./src/data/gsm8k-0.1/validation.parquet \
+ data.train_batch_size=32 \
+ data.max_prompt_length=512 \
+ data.max_response_length=1024 \
+ data.filter_overlong_prompts=True \
+ data.truncation='error' \
+ data.shuffle=False \
+ actor_rollout_ref.model.path=meta-llama/Llama-3.2-3B \
+ +actor_rollout_ref.model.tokenizer_path=./tokenizers/llama-3.2-3b-base-chat \
+ +actor_rollout_ref.model.lora_rank=64 \
+ +actor_rollout_ref.model.lora_alpha=32 \
+ actor_rollout_ref.actor.optim.lr=3e-6 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.actor.ppo_mini_batch_size=32 \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.actor.use_kl_loss=True \
+ actor_rollout_ref.actor.kl_loss_coef=0.001 \
+ actor_rollout_ref.actor.kl_loss_type=low_var_kl \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.model.enable_gradient_checkpointing=True \
+ actor_rollout_ref.actor.fsdp_config.param_offload=False \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
+ actor_rollout_ref.rollout.name=vllm \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
+ actor_rollout_ref.rollout.n=8 \
+ actor_rollout_ref.rollout.load_format=safetensors \
+ actor_rollout_ref.rollout.max_num_batched_tokens=65535 \
+ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=True \
+ algorithm.use_kl_in_reward=False \
+ trainer.critic_warmup=0 \
+ trainer.logger='["console","wandb"]' \
+ trainer.project_name='verl_grpo_gsm8k_base_custom' \
+ trainer.experiment_name='llama3.2_3b_base_custom_template_lora' \
+ trainer.n_gpus_per_node=2 \
+ trainer.nnodes=1 \
+ trainer.save_freq=23 \
+ trainer.test_freq=23 \
+ trainer.total_epochs=1
+
+# After training completes, evaluate the saved model on the reserved test set:
+# bash evaluate_gsm8k.sh --base_model meta-llama/Llama-3.2-3B \
+# --project_name verl_grpo_gsm8k_base_custom \
+# --experiment_name llama3.2_3b_base_custom_template_lora
diff --git a/es-fine-tuning-paper/src/scripts/grpo/grpo-gsm8k-qwen-base-custom.sh b/es-fine-tuning-paper/src/scripts/grpo/grpo-gsm8k-qwen-base-custom.sh
new file mode 100755
index 0000000..dbe33d2
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/grpo/grpo-gsm8k-qwen-base-custom.sh
@@ -0,0 +1,76 @@
+#!/bin/bash
+set -x
+
+# GRPO training for GSM8K using BASE Qwen model with custom chat template tokenizer
+# Model: Qwen/Qwen2.5-3B (base, not instruct)
+# Tokenizer: Custom tokenizer with "Question: {input} Answer: Let's think step by step." template
+
+# Check if HF_TOKEN is set
+if [ -z "$HF_TOKEN" ]; then
+ echo "Error: HF_TOKEN environment variable is not set"
+ echo "Please run: export HF_TOKEN=your_huggingface_token"
+ exit 1
+fi
+
+# Login to HuggingFace to access models
+huggingface-cli login --token "$HF_TOKEN"
+
+# Prepare custom tokenizer if not already created
+if [ ! -d "./tokenizers/qwen2.5-3b-base-chat" ]; then
+ echo "Creating custom tokenizer for Qwen base model..."
+ python3 base_model_tokenizer.py \
+ --model_path Qwen/Qwen2.5-3B \
+ --save_path ./tokenizers/qwen2.5-3b-base-chat
+fi
+
+# Train base Qwen model with custom tokenizer
+python3 -m verl.trainer.main_ppo \
+ algorithm.adv_estimator=grpo \
+ trainer.val_before_train=False \
+ data.train_files=./src/data/gsm8k-0.4/train.parquet \
+ data.val_files=./src/data/gsm8k-0.4/validation.parquet \
+ data.train_batch_size=32 \
+ data.max_prompt_length=512 \
+ data.max_response_length=1024 \
+ data.filter_overlong_prompts=True \
+ data.truncation='error' \
+ data.shuffle=False \
+ actor_rollout_ref.model.path=Qwen/Qwen2.5-3B \
+ +actor_rollout_ref.model.tokenizer_path=./tokenizers/qwen2.5-3b-base-chat \
+ +actor_rollout_ref.model.lora_rank=64 \
+ +actor_rollout_ref.model.lora_alpha=32 \
+ actor_rollout_ref.actor.optim.lr=3e-6 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.actor.ppo_mini_batch_size=32 \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.actor.use_kl_loss=True \
+ actor_rollout_ref.actor.kl_loss_coef=0.001 \
+ actor_rollout_ref.actor.kl_loss_type=low_var_kl \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.model.enable_gradient_checkpointing=True \
+ actor_rollout_ref.actor.fsdp_config.param_offload=False \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
+ actor_rollout_ref.rollout.name=vllm \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
+ actor_rollout_ref.rollout.n=8 \
+ actor_rollout_ref.rollout.load_format=safetensors \
+ actor_rollout_ref.rollout.max_num_batched_tokens=65535 \
+ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=True \
+ algorithm.use_kl_in_reward=False \
+ trainer.critic_warmup=0 \
+ trainer.logger='["console","wandb"]' \
+ trainer.project_name='verl_grpo_gsm8k_base_custom' \
+ trainer.experiment_name='qwen2.5_3b_base_custom_template_lora' \
+ trainer.n_gpus_per_node=2 \
+ trainer.nnodes=1 \
+ trainer.save_freq=23 \
+ trainer.test_freq=1 \
+ trainer.total_epochs=1
+
+# After training completes, evaluate the saved model on the reserved test set:
+# bash evaluate_gsm8k.sh --base_model Qwen/Qwen2.5-3B \
+# --project_name verl_grpo_gsm8k_base_custom \
+# --experiment_name qwen2.5_3b_base_custom_template_lora
\ No newline at end of file
diff --git a/es-fine-tuning-paper/src/scripts/grpo/grpo-gsm8k.sh b/es-fine-tuning-paper/src/scripts/grpo/grpo-gsm8k.sh
new file mode 100644
index 0000000..4c4507a
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/grpo/grpo-gsm8k.sh
@@ -0,0 +1,72 @@
+set -x
+
+# Check if HF_TOKEN is set
+if [ -z "$HF_TOKEN" ]; then
+ echo "Error: HF_TOKEN environment variable is not set"
+ echo "Please run: export HF_TOKEN=your_huggingface_token"
+ exit 1
+fi
+
+# Login to HuggingFace to access models
+huggingface-cli login --token "$HF_TOKEN"
+
+# First, prepare the data with test set reserved for final evaluation
+# python3 grpo_data_gsm8k.py --local_dir ./src/data/gsm8k-0.1 --train_split 0.1 --test_samples 200
+
+python3 -m verl.trainer.main_ppo \
+ algorithm.adv_estimator=grpo \
+ trainer.val_before_train=False \
+ data.train_files=./src/data/gsm8k-0.1/train.parquet \
+ data.val_files=./src/data/gsm8k-0.1/validation.parquet \
+ data.train_batch_size=32 \
+ data.max_prompt_length=512 \
+ data.max_response_length=1024 \
+ data.filter_overlong_prompts=True \
+ data.truncation='error' \
+ data.shuffle=False \
+ actor_rollout_ref.model.path=Qwen/Qwen2.5-3B-Instruct \
+ +actor_rollout_ref.model.lora_rank=64 \
+ +actor_rollout_ref.model.lora_alpha=32 \
+ actor_rollout_ref.actor.optim.lr=3e-6 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.actor.ppo_mini_batch_size=32 \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.actor.use_kl_loss=True \
+ actor_rollout_ref.actor.kl_loss_coef=0.001 \
+ actor_rollout_ref.actor.kl_loss_type=low_var_kl \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.model.enable_gradient_checkpointing=True \
+ actor_rollout_ref.actor.fsdp_config.param_offload=False \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
+ actor_rollout_ref.rollout.name=vllm \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
+ actor_rollout_ref.rollout.n=8 \
+ actor_rollout_ref.rollout.load_format=safetensors \
+ actor_rollout_ref.rollout.max_num_batched_tokens=65535 \
+ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=True \
+ algorithm.use_kl_in_reward=False \
+ trainer.critic_warmup=0 \
+ trainer.logger='["console","wandb"]' \
+ trainer.project_name='verl_grpo_example_gsm8k' \
+ trainer.experiment_name='qwen2.5_3b_ins_grpo_lora_0.1' \
+ trainer.n_gpus_per_node=8 \
+ trainer.nnodes=1 \
+ trainer.save_freq=23 \
+ trainer.test_freq=23 \
+ trainer.total_epochs=1
+
+ # actor_rollout_ref.actor.ppo_mini_batch_size=256 \
+ # data.train_batch_size=1024 \
+ # trainer.n_gpus_per_node=8 \
+ # actor_rollout_ref.model.use_shm=True \
+
+# After training completes, evaluate the saved model on the reserved test set:
+# python3 evaluate_model.py \
+# --model_path \
+# --base_model qwen/Qwen2.5-3B-Instruct \
+# --test_file ./src/data/gsm8k-0.1/test.parquet \
+# --task_type gsm8k \
+# --output_file ./eval_results_gsm8k.json
diff --git a/es-fine-tuning-paper/src/scripts/grpo/grpo_countdown.sh b/es-fine-tuning-paper/src/scripts/grpo/grpo_countdown.sh
new file mode 100644
index 0000000..4bbf614
--- /dev/null
+++ b/es-fine-tuning-paper/src/scripts/grpo/grpo_countdown.sh
@@ -0,0 +1,51 @@
+set -x
+
+python3 -m verl.trainer.main_ppo \
+ algorithm.adv_estimator=grpo \
+ trainer.val_before_train=False \
+ data.train_files=./src/data/gsm8k-0.9/train.parquet \
+ data.val_files=./src/data/gsm8k-0.9/test.parquet \
+ data.train_batch_size=32 \
+ data.max_prompt_length=512 \
+ data.max_response_length=1024 \
+ data.filter_overlong_prompts=True \
+ data.truncation='error' \
+ data.shuffle=False \
+ actor_rollout_ref.model.path=Qwen/Qwen2.5-3B-Instruct \
+ +actor_rollout_ref.model.lora_rank=64 \
+ +actor_rollout_ref.model.lora_alpha=32 \
+ actor_rollout_ref.actor.optim.lr=3e-6 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.actor.ppo_mini_batch_size=32 \
+ actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.actor.use_kl_loss=True \
+ actor_rollout_ref.actor.kl_loss_coef=0.001 \
+ actor_rollout_ref.actor.kl_loss_type=low_var_kl \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.model.enable_gradient_checkpointing=True \
+ actor_rollout_ref.actor.fsdp_config.param_offload=False \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
+ actor_rollout_ref.rollout.name=vllm \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
+ actor_rollout_ref.rollout.n=8 \
+ actor_rollout_ref.rollout.load_format=safetensors \
+ actor_rollout_ref.rollout.max_num_batched_tokens=65535 \
+ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \
+ actor_rollout_ref.ref.fsdp_config.param_offload=True \
+ algorithm.use_kl_in_reward=False \
+ trainer.critic_warmup=0 \
+ trainer.logger='["console","wandb"]' \
+ trainer.project_name='verl_grpo_example_gsm8k' \
+ trainer.experiment_name='qwen2.5_3b_grpo_lora' \
+ trainer.n_gpus_per_node=8 \
+ trainer.nnodes=1 \
+ trainer.save_freq=-1 \
+ trainer.test_freq=93 \
+ trainer.total_epochs=1
+
+ # actor_rollout_ref.actor.ppo_mini_batch_size=256 \
+ # data.train_batch_size=1024 \
+ # trainer.n_gpus_per_node=8 \
+ # actor_rollout_ref.model.use_shm=True \
diff --git a/es-fine-tuning-paper/src/utils/__pycache__/base_model_tokenizer.cpython-313.pyc b/es-fine-tuning-paper/src/utils/__pycache__/base_model_tokenizer.cpython-313.pyc
new file mode 100644
index 0000000..1c2a673
Binary files /dev/null and b/es-fine-tuning-paper/src/utils/__pycache__/base_model_tokenizer.cpython-313.pyc differ
diff --git a/es-fine-tuning-paper/src/utils/base_model_tokenizer.py b/es-fine-tuning-paper/src/utils/base_model_tokenizer.py
new file mode 100644
index 0000000..c4d721b
--- /dev/null
+++ b/es-fine-tuning-paper/src/utils/base_model_tokenizer.py
@@ -0,0 +1,289 @@
+"""
+Custom tokenizer wrapper for base models (Llama, Qwen) with simple chat template.
+This adds a chat template to base model tokenizers for training without instruction-tuned models.
+
+Template format: "Question: {input} Answer: Let's think step by step."
+"""
+
+from transformers import AutoTokenizer
+from typing import List, Dict, Union, Optional
+import os
+
+
+class BaseModelTokenizer:
+ """Wrapper class that adds a simple chat template to base model tokenizers."""
+
+ # Simple chat template for base models
+ CHAT_TEMPLATE = """{% for message in messages %}{% if message['role'] == 'user' %}Question: {{ message['content'] }}
+Answer: Let's think step by step.{% elif message['role'] == 'assistant' %}{{ message['content'] }}{% endif %}{% endfor %}"""
+
+ def __init__(self, model_path: str, **kwargs):
+ """
+ Initialize tokenizer with custom chat template.
+
+ Args:
+ model_path: Path to the base model (e.g., 'meta-llama/Llama-3.2-3B' or 'Qwen/Qwen2.5-3B')
+ **kwargs: Additional arguments to pass to AutoTokenizer
+ """
+ self.tokenizer = AutoTokenizer.from_pretrained(model_path, **kwargs)
+
+ # Add chat template if not present
+ if self.tokenizer.chat_template is None:
+ self.tokenizer.chat_template = self.CHAT_TEMPLATE
+ print(f"Added custom chat template to tokenizer from {model_path}")
+ else:
+ # Override existing chat template with our simple one
+ print(f"Overriding existing chat template for base model training")
+ self.tokenizer.chat_template = self.CHAT_TEMPLATE
+
+ # Ensure pad token is set
+ if self.tokenizer.pad_token is None:
+ if self.tokenizer.eos_token is not None:
+ self.tokenizer.pad_token = self.tokenizer.eos_token
+ print(f"Set pad_token to eos_token: {self.tokenizer.eos_token}")
+ else:
+ self.tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+ print("Added [PAD] as pad_token")
+
+ def save_pretrained(self, save_directory: str, **kwargs):
+ """Save the tokenizer with the custom chat template."""
+ os.makedirs(save_directory, exist_ok=True)
+ self.tokenizer.save_pretrained(save_directory, **kwargs)
+ print(f"Saved tokenizer with custom chat template to {save_directory}")
+
+ def __getattr__(self, name):
+ """Delegate all other attribute access to the underlying tokenizer."""
+ return getattr(self.tokenizer, name)
+
+ def __call__(self, *args, **kwargs):
+ """Make the wrapper callable like the original tokenizer."""
+ return self.tokenizer(*args, **kwargs)
+
+
+def create_base_tokenizer(model_path: str, save_path: Optional[str] = None, **kwargs) -> BaseModelTokenizer:
+ """
+ Create and optionally save a base model tokenizer with custom chat template.
+
+ Args:
+ model_path: Path to the base model
+ save_path: Optional path to save the tokenizer
+ **kwargs: Additional arguments for AutoTokenizer
+
+ Returns:
+ BaseModelTokenizer instance
+ """
+ tokenizer = BaseModelTokenizer(model_path, **kwargs)
+
+ if save_path:
+ tokenizer.save_pretrained(save_path)
+
+ return tokenizer
+
+
+def test_model_generation(model_path: str, data_path: str, num_samples: int = 30):
+ """
+ Test model generation with countdown dataset samples, similar to VERL's approach.
+
+ Args:
+ model_path: Path to the model
+ data_path: Path to parquet data file (e.g., train.parquet)
+ num_samples: Number of samples to generate
+ """
+ from transformers import AutoModelForCausalLM
+ import torch
+ import pandas as pd
+
+ print(f"\n{'='*60}")
+ print(f"Testing Model Generation (VERL-style)")
+ print(f"{'='*60}")
+ print(f"Model: {model_path}")
+ print(f"Data: {data_path}")
+ print(f"Samples: {num_samples}")
+
+ # Load data
+ print("\n[1/4] Loading data...")
+ df = pd.read_parquet(data_path)
+ print(f"Total samples in dataset: {len(df)}")
+
+ # Take first num_samples
+ samples = df.head(num_samples)
+
+ # Load tokenizer with custom chat template
+ print("\n[2/4] Loading tokenizer...")
+ tokenizer = BaseModelTokenizer(model_path)
+
+ # Load model
+ print("[3/4] Loading model...")
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+ model = AutoModelForCausalLM.from_pretrained(
+ model_path,
+ torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+ device_map="auto" if torch.cuda.is_available() else None
+ )
+ if not torch.cuda.is_available():
+ model = model.to(device)
+ model.eval()
+
+ print(f"Model loaded on: {device}")
+
+ print(f"\n[4/4] Generating answers for {num_samples} samples...")
+ print(f"{'='*60}\n")
+
+ # Generate answers
+ results = []
+ for idx, row in samples.iterrows():
+ # Extract the user question from the prompt
+ prompt_messages = row['prompt']
+ # Find the user message (skip system message if present)
+ user_content = None
+ for msg in prompt_messages:
+ if msg['role'] == 'user':
+ user_content = msg['content']
+ break
+
+ if user_content is None:
+ print(f"Warning: No user message found in sample {idx}, skipping...")
+ continue
+
+ # Format message with chat template
+ messages = [{"role": "user", "content": user_content}]
+
+ # Apply chat template and tokenize (VERL-style)
+ input_ids = tokenizer.apply_chat_template(
+ messages,
+ tokenize=True,
+ add_generation_prompt=False,
+ return_tensors="pt"
+ ).to(device)
+
+ # Generate using model.generate (VERL-style)
+ with torch.no_grad():
+ outputs = model.generate(
+ input_ids,
+ max_new_tokens=256,
+ temperature=0.7,
+ top_p=0.9,
+ do_sample=True,
+ pad_token_id=tokenizer.pad_token_id,
+ eos_token_id=tokenizer.eos_token_id,
+ )
+
+ # Decode the generated text
+ generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+ # Extract just the answer part (after the prompt)
+ prompt_text = tokenizer.decode(input_ids[0], skip_special_tokens=True)
+ answer = generated_text[len(prompt_text):].strip()
+
+ results.append({
+ "question": user_content,
+ "prompt": prompt_text,
+ "answer": answer,
+ "full_output": generated_text
+ })
+
+ # Print progress
+ print(f"Sample {idx + 1}/{num_samples}")
+ print(f"Q: {user_content}")
+ print(f"A: {answer}")
+ print(f"{'-'*60}\n")
+
+ print(f"\n{'='*60}")
+ print(f"✓ Generation test completed successfully!")
+ print(f"Total samples: {len(results)}")
+ print(f"{'='*60}\n")
+
+ return results
+
+
+if __name__ == "__main__":
+ import argparse
+
+ parser = argparse.ArgumentParser(
+ description="Create custom tokenizer for base models with simple chat template"
+ )
+ parser.add_argument(
+ "--model_path",
+ type=str,
+ required=True,
+ help="Path to base model (e.g., meta-llama/Llama-3.2-3B, Qwen/Qwen2.5-3B)"
+ )
+ parser.add_argument(
+ "--save_path",
+ type=str,
+ default=None,
+ help="Directory to save the tokenizer with custom chat template"
+ )
+ parser.add_argument(
+ "--test",
+ action="store_true",
+ help="Test the tokenizer with sample messages"
+ )
+ parser.add_argument(
+ "--test_generation",
+ action="store_true",
+ help="Test model generation with multiple samples (VERL-style)"
+ )
+ parser.add_argument(
+ "--data_path",
+ type=str,
+ default="data/countdown-0.1/train.parquet",
+ help="Path to parquet data file (default: data/countdown-0.1/train.parquet)"
+ )
+ parser.add_argument(
+ "--num_samples",
+ type=int,
+ default=30,
+ help="Number of samples to generate (default: 30)"
+ )
+
+ args = parser.parse_args()
+
+ # Test generation mode
+ if args.test_generation:
+ results = test_model_generation(args.model_path, args.data_path, args.num_samples)
+ else:
+ # Create tokenizer
+ print(f"\nCreating tokenizer from: {args.model_path}")
+ tokenizer = create_base_tokenizer(args.model_path, args.save_path)
+
+ if args.test:
+ # Test the tokenizer
+ print("\n" + "="*50)
+ print("Testing tokenizer with sample messages")
+ print("="*50)
+
+ messages = [
+ {
+ "role": "user",
+ "content": "Using the numbers [3, 5, 7, 10], create an equation that equals 24."
+ }
+ ]
+
+ # Apply chat template
+ formatted_text = tokenizer.apply_chat_template(
+ messages,
+ tokenize=False,
+ add_generation_prompt=False
+ )
+
+ print("\nFormatted text:")
+ print(formatted_text)
+
+ # Tokenize
+ tokens = tokenizer.apply_chat_template(
+ messages,
+ tokenize=True,
+ add_generation_prompt=False,
+ return_tensors="pt"
+ )
+
+ print(f"\nTokens shape: {tokens.shape}")
+ print(f"Number of tokens: {tokens.shape[1]}")
+
+ # Decode back
+ decoded = tokenizer.decode(tokens[0])
+ print(f"\nDecoded text:")
+ print(decoded)
+
+ print("\n✓ Tokenizer test completed successfully!")
diff --git a/es-fine-tuning-paper/utils/worker_extn.py b/es-fine-tuning-paper/utils/worker_extn.py
new file mode 100644
index 0000000..e232b24
--- /dev/null
+++ b/es-fine-tuning-paper/utils/worker_extn.py
@@ -0,0 +1,81 @@
+import gc
+import time
+import torch
+
+def _stateless_init_process_group(master_address, master_port, rank, world_size, device):
+ from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+ from vllm.distributed.utils import StatelessProcessGroup
+ pg = StatelessProcessGroup.create(
+ host=master_address, port=master_port, rank=rank, world_size=world_size
+ )
+ return PyNcclCommunicator(pg, device=device)
+
+class WorkerExtension:
+ """
+ Methods used by the ES trainer:
+ - perturb_self_weights(seed, sigma_or_scale, coeff=1.0, negate=False)
+ - restore_self_weights(seed, SIGMA)
+ - init_inter_engine_group(master_address, master_port, rank, world_size)
+ - broadcast_all_weights(src_rank)
+ - save_self_weights_to_disk(filepath)
+ """
+
+ def perturb_self_weights(self, seed, noise_scale, negate=False):
+ scale = float(noise_scale)
+ sign = -1.0 if negate else 1.0
+ for _, p in self.model_runner.model.named_parameters():
+ gen = torch.Generator(device=p.device)
+ gen.manual_seed(int(seed))
+ noise = torch.randn(p.shape, dtype=p.dtype, device=p.device, generator=gen)
+ p.data.add_(sign * scale * noise)
+ del noise
+ if torch.cuda.is_available():
+ torch.cuda.synchronize()
+ torch.cuda.empty_cache()
+ return True
+
+ def restore_self_weights(self, seed, SIGMA):
+ for _, p in self.model_runner.model.named_parameters():
+ gen = torch.Generator(device=p.device)
+ gen.manual_seed(int(seed))
+ noise = torch.randn(p.shape, dtype=p.dtype, device=p.device, generator=gen)
+ p.data.add_(-float(SIGMA) * noise)
+ del noise
+ if torch.cuda.is_available():
+ torch.cuda.synchronize()
+ torch.cuda.empty_cache()
+ return True
+
+ def init_inter_engine_group(self, master_address: str, master_port: int, rank: int, world_size: int):
+ self.inter_pg = _stateless_init_process_group(
+ master_address, master_port, rank, world_size, self.device
+ )
+ return True
+
+ def broadcast_all_weights(self, src_rank: int):
+ for _, p in self.model_runner.model.named_parameters():
+ self.inter_pg.broadcast(p, src=int(src_rank), stream=torch.cuda.current_stream())
+ if torch.cuda.is_available():
+ torch.cuda.synchronize()
+ return True
+
+ def save_self_weights_to_disk(self, filepath):
+ state_dict_to_save = {}
+ for name, p in self.model_runner.model.named_parameters():
+ state_dict_to_save[name] = p.detach().cpu()
+ torch.save(state_dict_to_save, filepath)
+ gc.collect()
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+ time.sleep(0.1)
+ return True
+
+ def load_weights_from_disk(self, filepath):
+ state_dict = torch.load(filepath, map_location=self.device)
+ for name, p in self.model_runner.model.named_parameters():
+ p.data.copy_(state_dict[name].to(self.device))
+ gc.collect()
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+ time.sleep(0.1)
+ return True
\ No newline at end of file
diff --git a/es-fine-tuning-paper/verl-docker-run.sh b/es-fine-tuning-paper/verl-docker-run.sh
new file mode 100644
index 0000000..b442587
--- /dev/null
+++ b/es-fine-tuning-paper/verl-docker-run.sh
@@ -0,0 +1,177 @@
+#!/bin/bash
+
+# Speedrun script for es-fine-tuning-paper setup and training
+# Combines commands from docs/getting_started/quickstart.md
+
+set -e # Exit on any error
+
+
+# Check if WANDB_API_KEY is set
+if [ -z "$WANDB_API_KEY" ]; then
+ echo "WARNING: WANDB_API_KEY environment variable is not set."
+ echo "Training will proceed without Weights & Biases logging."
+ echo "To enable logging, set: export WANDB_API_KEY='your_key_here'"
+ echo "Get your key from: https://wandb.ai/settings"
+ exit 1
+fi
+
+# Pull Docker image if not already present
+echo "Pulling Docker image..."
+sudo docker pull hiyouga/verl:ngc-th2.6.0-cu126-vllm0.8.3-flashinfer0.2.2-cxx11abi0 || {
+ echo "Failed to pull Docker image. Please check your Docker installation and network."
+
+}
+
+# Start Docker container
+echo "Checking for existing Docker container..."
+
+# Check if container is already running
+if [ "$(sudo docker ps -q -f name=verl-es-fine-tuning-paper)" ]; then
+ echo "Container 'verl-es-fine-tuning-paper' is already running. Reusing existing container."
+# Check if container exists but is stopped
+elif [ "$(sudo docker ps -aq -f name=verl-es-fine-tuning-paper)" ]; then
+ echo "Container 'verl-es-fine-tuning-paper' exists but is stopped. Starting existing container..."
+ sudo docker start verl-es-fine-tuning-paper || {
+ echo "Failed to start existing container. Please check Docker status."
+ exit 1
+ }
+ echo "Container started."
+else
+ echo "Starting new Docker container..."
+ sudo docker run -d --gpus all --name verl-es-fine-tuning-paper \
+ --ipc=host \
+ --ulimit memlock=-1 \
+ --ulimit stack=67108864 \
+ -v $(pwd)/..:/workspace \
+ hiyouga/verl:ngc-th2.6.0-cu126-vllm0.8.3-flashinfer0.2.2-cxx11abi0 \
+ sleep infinity || {
+ echo "Failed to start Docker container. Please check GPU availability and Docker setup."
+ exit 1
+ }
+ echo "Waiting for container to start..."
+ sleep 10
+fi
+
+# # Install es-fine-tuning-paper dependencies
+# echo "Installing es-fine-tuning-paper dependencies..."
+# sudo docker exec verl-es-fine-tuning-paper bash -c "cd /workspace/es-fine-tuning-paper && pip3 install -e ." || {
+# echo "Failed to install es-fine-tuning-paper dependencies."
+# exit 1
+# }
+
+# Clone VERL from official repo
+echo "Cloning VERL from official repo..."
+sudo docker exec verl-es-fine-tuning-paper bash -c "if [ ! -d '/workspace/es-fine-tuning-paper/src/verl' ]; then git config --global --add safe.directory '*' && cd /workspace/es-fine-tuning-paper/src && git clone https://github.com/volcengine/verl.git && cd verl && git checkout a43ead6; else echo 'VERL already exists, skipping clone.'; fi" || {
+ echo "Failed to clone VERL from official repo."
+ exit 1
+}
+
+# Patch main_ppo.py to support separate tokenizer path
+echo "Patching main_ppo.py to support separate tokenizer path..."
+sudo docker exec verl-es-fine-tuning-paper bash -c "
+cd /workspace/es-fine-tuning-paper/src/verl
+# Check if patch is already applied
+if grep -q 'tokenizer_path = config.actor_rollout_ref.model.get' verl/trainer/main_ppo.py; then
+ echo 'Patch already applied to main_ppo.py, skipping.'
+else
+ echo 'Applying tokenizer patch to main_ppo.py...'
+ sed -i '/trust_remote_code = config.data.get(\"trust_remote_code\", False)/,/processor = hf_processor(local_path, use_fast=True)/c\ trust_remote_code = config.data.get(\"trust_remote_code\", False)\n tokenizer_path = config.actor_rollout_ref.model.get(\"tokenizer_path\", None)\n if tokenizer_path is not None:\n tokenizer_local_path = copy_to_local(tokenizer_path)\n print(f\"Using separate tokenizer from: {tokenizer_path}\")\n else:\n tokenizer_local_path = local_path\n tokenizer = hf_tokenizer(tokenizer_local_path, trust_remote_code=trust_remote_code)\n processor = hf_processor(tokenizer_local_path, use_fast=True) # used for multimodal LLM, could be none' verl/trainer/main_ppo.py
+fi
+" || {
+ echo "Failed to patch main_ppo.py."
+ exit 1
+}
+
+# Fix permissions for verl directory
+echo "Fixing file permissions..."
+sudo chown -R ubuntu:ubuntu /home/ubuntu/alphaxiv-sandbox/paper-implementations/es-fine-tuning-paper/src/verl/ || {
+ echo "Warning: Failed to fix permissions, but continuing..."
+}
+
+# Install VERL
+echo "Installing VERL..."
+sudo docker exec verl-es-fine-tuning-paper bash -c "cd /workspace/es-fine-tuning-paper/src/verl && pip3 install -e ." || {
+ echo "Failed to install VERL."
+ exit 1
+}
+
+wait
+
+# # Download and preprocess HotpotQA dataset
+# echo "Downloading and preprocessing HotpotQA dataset..."
+# sudo docker exec verl-es-fine-tuning-paper bash -c "cd /workspace/es-fine-tuning-paper && mkdir -p data/hotpotqa && python src/examples/data_preprocess/hotpotqa.py --local_dir data/hotpotqa" || {
+# echo "Failed to download and preprocess HotpotQA dataset."
+# exit 1
+# }
+
+# # Build HotpotQA search index
+# echo "Building HotpotQA search index..."
+# sudo docker exec verl-es-fine-tuning-paper bash -c "cd /workspace/es-fine-tuning-paper && if [ ! -f 'data/corpus/hotpotqa/hpqa_corpus.jsonl' ]; then mkdir -p data/corpus/hotpotqa && wget -q https://huggingface.co/datasets/BeIR/hotpotqa/resolve/main/corpus.jsonl.gz -O data/corpus/hotpotqa/corpus.jsonl.gz && gunzip -c data/corpus/hotpotqa/corpus.jsonl.gz > data/corpus/hotpotqa/hpqa_corpus.jsonl; else echo 'HotpotQA corpus already exists, skipping download.'; fi" || {
+# echo "Failed to download corpus data."
+# exit 1
+# }
+
+# sudo docker exec verl-es-fine-tuning-paper bash -c "cd /workspace/es-fine-tuning-paper && if [ -f 'data/corpus/hotpotqa/index.bin' ]; then echo 'HotpotQA search index already exists, skipping index build.'; else echo 'Building FAISS search index (this may take some time)...'; cd src/scripts/hotpotqa_search/ && python process_hotpotqa.py; fi" || {
+# echo "Failed to build search index."
+# exit 1
+# }
+
+sudo docker exec verl-es-fine-tuning-paper bash -c "pip install --upgrade wandb"
+# Configure Weights & Biases if API key is set
+if [ ! -z "$WANDB_API_KEY" ]; then
+ echo "Configuring Weights & Biases..."
+ sudo docker exec verl-es-fine-tuning-paper bash -c "wandb login $WANDB_API_KEY" || {
+ echo "Failed to login to Weights & Biases."
+ exit 1
+ }
+fi
+
+# Final comprehensive permission fix for all created files
+echo "Fixing all file permissions..."
+sudo chown -R ubuntu:ubuntu /home/ubuntu/alphaxiv-sandbox/paper-implementations/es-fine-tuning-paper/src/ 2>/dev/null || true
+sudo chown -R ubuntu:ubuntu /home/ubuntu/alphaxiv-sandbox/paper-implementations/es-fine-tuning-paper/data/ 2>/dev/null || true
+echo "Setup complete!"
+
+# # Set up environment variables for Docker exec
+# DOCKER_ENV=""
+# if [ ! -z "$HYDRA_FULL_ERROR" ]; then
+# echo "HYDRA_FULL_ERROR is set, enabling full error traces..."
+# DOCKER_ENV="export HYDRA_FULL_ERROR=1 && "
+# fi
+
+# # Run training based on selected algorithm
+# case "$ALGORITHM" in
+# ppo)
+# echo "=========================================="
+# echo "Starting PPO Training on HotpotQA"
+# echo "This will take approximately 22 hours on 4xH100 80GB GPUs"
+# echo "=========================================="
+
+# sudo docker exec verl-es-fine-tuning-paper bash -c "cd /workspace/es-fine-tuning-paper && ${DOCKER_ENV}cp src/examples/trainer/run_ppo_hotpotqa.sh ./ && bash run_ppo_hotpotqa.sh" || {
+# echo "Training failed."
+# exit 1
+# }
+# ;;
+# grpo)
+# echo "=========================================="
+# echo "Starting GRPO Training on HotpotQA"
+# echo "This will take approximately 20-22 hours on 4xH100 80GB GPUs"
+# echo "=========================================="
+
+# sudo docker exec verl-es-fine-tuning-paper bash -c "cd /workspace/es-fine-tuning-paper && ${DOCKER_ENV}cp src/examples/trainer/run_grpo_hotpotqa.sh ./ && bash run_grpo_hotpotqa.sh" || {
+# echo "Training failed."
+# exit 1
+# }
+# ;;
+# *)
+# echo "Unknown algorithm: $ALGORITHM"
+# exit 1
+# ;;
+# esac
+
+# echo "=========================================="
+# echo "Training Complete!"
+# echo "=========================================="
+# echo "Check the results and logs in the es-fine-tuning-paper directory."
+# echo "You can also check Weights & Biases for training metrics if configured."
+