aakaran / finetuning-with-sampling

# Finetuning with Sampling
### Paper | Project Page
This repo contains the official PyTorch implementation of Finetuning with Sampling.
> **Finetuning with Sampling: SFT Learns Better Than You Think** 
> Aayush Karan, Sitan Chen, Yilun Du
> Harvard 
## Setup
Run the following script to setup environment.
 [code] 
## Sampling
The main directory contains slurm scripts to run projection sampling for chemistry ( [code] ), whose training data .json is included in sci_data, as well as math ( [code] ), whose training data .parquet is in math_data. 
To run projection sampling on chemistry with 15 shard parallelism:
 [code] 
The output is several .jsonl files (based on the shard number) that store the boosted traces, original solution, correctness, etc. Accompanying files are trace_boosted .jsonl files, which track which MCMC candidates actually get accepted throughout the entire MCMC process. These can be used to obtain teacher distillation logits for the full boosted traces if needed. For standard SFT, the trace_boosted .jsonl files are not needed, so the logic can be disabled since the trace_boosted file outputs are memory intensive.
## Training
For SFT training, we refer to the setup in https://github.com/yongliang-wu/DFT. For chemistry and math, use [code] , [code] , and [code] . A sample SFT launch script is provided for reference (in [code] in the DFT codebase, may need to modify the file to convert [code] to a list when the instance is a set for indexing). The output of the sampling process should be one large .jsonl file that contains (among other variables) all prompts and generated rollouts. These can be readily converted into .parquet files that rename these keys to [code] and [code] for SFT.
## Evaluation
To evaluate trained checkpoints on say chemistry, include your hf model checkpoints in the [code] of [code] . Depending on the number of checkpoints, run 
 [code] 
The output .jsonl grading file stores correctness per evaluation task and can be directly parsed to obtain final accuracy. Similar commands for [code] and [code] hold. For MMLU and GPQA, we refer to lm-evaluation-harness (https://github.com/eleutherai/lm-evaluation-harness).