Home

Awesome

<p align="center"> <img src="images/spin_dalle.png" width="30%"> <br> </p> <p align="center"> 🤗 <a href="https://huggingface.co/collections/UCLA-AGI/zephyr-7b-sft-full-spin-65c361dfca65637272a02c40" target="_blank">Models</a> | 🤗 <a href="https://huggingface.co/collections/UCLA-AGI/datasets-spin-65c3624e98d4b589bbc76f3a" target="_blank">Datasets</a> </p>

Self-Play Fine-Tuning (SPIN)

Mistral-7B Open LLM MT-Bench

This repository contains the official code for the paper "Self-Play Fine-Tuning Converts Weak Language Models to Strong Language Models".

Authors: Zixiang Chen*, Yihe Deng*, Huizhuo Yuan*, Kaixuan Ji, Quanquan Gu

[Webpage] [Huggingface]

🔔 News

Note: We notice that the Alignment Handbook has updated their configuration and SFT checkpoint since our experiments. The configuration and SFT model from the Alignment Handbook that we used in our experiments for data generation and fine-tuning are the older version (Config, Model). The model checkpoint on Hugging Face (alignment-handbook/zephyr-7b-sft-full) has been updated once with the new config. If you wish to use the newest SFT model, you need to either load the checkpoint with revision=ac6e600eefcce74f5e8bae1035d4f66019e93190 or generate your own data instead of using the datasets we provided on Hugging Face.

Table of Contents

🌀 About SPIN

SPIN utilizes a self-play mechanism, allowing an LLM to improve itself by playing against its previous iterations, without needing additional human-annotated preference data than the SFT dataset itself. More specifically, the LLM generates its own training data from its previous iterations, refining its policy by discerning these self-generated responses from the original SFT data.

<p align="center"> <img src="images/iter_openllm.png" width="35%"> <br> Average score of <b>SPIN</b> at different iterations on the HuggingFace Open LLM leaderboard. </p> SPIN can significantly enhance the performance of an LLM after SFT across various benchmarks, outperforming the model trained with direct preference optimization (DPO) on labelled preference datasets. The approach is theoretically grounded, ensuring that the LLM aligns with the target data distribution, and empirically validated through extensive evaluations on multiple datasets. <p align="center"> <img src="images/dpo_compare.png" width="80%"> <br> Performance comparison with DPO training across the six benchmark datasets. SPIN at iteration 0 achieves comparable performance to DPO training with 62k new data. At iteration 1, SPIN has already surpassed DPO training on the majority of datasets. </p>

For more details, you can check our paper here.

Setup

The following steps provide the necessary setup to run our codes.

  1. Create a Python virtual environment with Conda:
conda create -n myenv python=3.10
conda activate myenv
  1. Install the following Python dependencies to run the codes.
python -m pip install .
python -m pip install flash-attn --no-build-isolation
  1. Login to your huggingface account for downloading models
huggingface-cli login --token "${your_access_token}"

Data

We provide the data used in our experiments along with the synthetic data we generated in this repo as well as on HuggingFace. These data are converted to .parquet format for fine-tuning.

❗ On 04/04/2024, we've discovered that the datasets previously uploaded are incorrect. We've re-upload the corrected datasets.

DatasetDownload
SPIN_iter0🤗 HuggingFace
SPIN_iter1🤗 HuggingFace
SPIN_iter2🤗 HuggingFace
SPIN_iter3🤗 HuggingFace

The input data for our code is required to be of the same format where each data contains the following attributes, as similar to HuggingFaceH4/ultrafeedback_binarized:

{
    "real": [{"role": "user", "content": <prompt>}, 
               {"role": "assistant", "content": <ground truth>}],
    "generated": [{"role": "user", "content": <prompt>}, 
                 {"role": "assistant", "content": <generation>}]
}

🔍 Note: During data generation, the content for generated response can be empty, as we only uses prompt to generate model responses.

Model

We also provide our model checkpoints at iteration 0,1,2,3 on HuggingFace.

ModelDownload
zephyr-7b-sft-full-SPIN-iter0🤗 HuggingFace
zephyr-7b-sft-full-SPIN-iter1🤗 HuggingFace
zephyr-7b-sft-full-SPIN-iter2🤗 HuggingFace
zephyr-7b-sft-full-SPIN-iter3🤗 HuggingFace

🔍 Note: <span style="color:red;">With the provided data, you can directly jump to Step 2: Fine-tuning without doing data generation on your own. You may also start from any iteration to reproduce our results using our open-sourced model checkpoints.</span>

Usage

For SPIN, we generate all synthetic data at once for an iteration, and fine-tune the LLM based on the real and synthetic data pairs.

Step 0 (optional): Reformatting SFT dataset

python spin/reformat.py [options]

Options

🔍 Note: If choosing to use SPIN on the entire dataset of HuggingFaceH4/ultrachat_200k instead of our 50k subset, one can reformat the original data with spin/reformat.py. To use other datasets, simply convert the data into the same format and resume with the following steps.

Step 1: Generation

accelerate launch spin/generate.py [options]

Options

The generated data is in json format where each data contains the following attributes:

{
    "real": [{"role": "user", "content": <prompt>}, 
               {"role": "assistant", "content": <ground truth>}],
    "generated": [{"role": "user", "content": <prompt>}, 
                 {"role": "assistant", "content": <generation>}]
}

Note: the iteration number for data generation solely depends on which model is used for generation (e.g. use original SFT model for iter0 data generation and SPIN iter0 model for iter1 data generation). The generate.py script will only use the prompts/questions for the data model.

Example. The following code generates 8k synthetic data for iteration 0.

bash scripts/generate.sh

🚀 Faster generation with vLLM

Alternatively, you could use the following example script to generate LLM responses with speedup. Larger frac_len can be used with vllm.

bash scripts/generate_vllm.sh

Thanks to @sumo43 for implementing vLLM for generation.

Step 1.5: Gather generations and convert data type

python spin/convert_data.py [options]

Options

The code will generate the two final data files including train_prefs-00000-of-00001.parquet and test_prefs-00000-of-00001.parquet, which will be used for fine-tuning.

Note: make sure to collect the generated data filed into the same directory of --input_dir.

Example.

python spin/convert_data.py --output_dir new_data/iter0 --input_dir generated/iter0 --num_fracs 63

Step 2: Fine-tuning

accelerate launch --config_file configs/multi_gpu.yaml --num_processes=8 --main_process_port 29500 spin/run_spin.py configs/config.yaml
<!-- **[TODO]**: wrap up necessary codes into the folder spin. Add explainations/instructions here. -->

You might need to change the configuration in configs/config.yaml. Here are some key configs you might need to customize:

In our experiments, we do full fine-tuning on a multi-GPU machine with DeepSpeed ZeRO-3 (requires A100 (80GB)).

Example.

bash scripts/finetune.sh

Reproducing Our Results

To help reproducing our results, we have made available the scripts corresponding to all four iterations of our study. These scripts are pre-configured with the exact parameters and model versions used in our paper. For each iteration, the base model is initialized with the version released on 🤗 HuggingFace, which can be found at the following links:

DatasetDownload
SPIN_iter0🤗 HuggingFace
SPIN_iter1🤗 HuggingFace
SPIN_iter2🤗 HuggingFace
SPIN_iter3🤗 HuggingFace

To execute the full pipeline using your locally trained models as the base, modify the model_name_or_path parameter in the configuration files to point to your model's path.

To start the full fine-tuning process, run the corresponding script from your terminal:

bash scripts/finetune.sh
bash scripts/finetune_iter1.sh
bash scripts/finetune_iter2.sh
bash scripts/finetune_iter3.sh

By following these steps, you should be able to reproduce our results.


Evaluation

For our evaluation on the Open LLM Leaderboard, please use this lm-evaluation-harness which is v0.4.0. Also, note that we set the number of few shot examples to be the same as instructed on the Leaderboard. Different evaluation versions results in different scores, but the trend will remain the same.

Star History

Star History Chart

Citation

If you find this repo useful for your research, please consider citing the paper

@misc{chen2024selfplay,
      title={Self-Play Fine-Tuning Converts Weak Language Models to Strong Language Models}, 
      author={Zixiang Chen and Yihe Deng and Huizhuo Yuan and Kaixuan Ji and Quanquan Gu},
      year={2024},
      eprint={2401.01335},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

Acknowledgement

This repo is built upon The Alignment Handbook. We thank the authors for their great work.