🤗 Models | 🤗 Datasets

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]

• [02/13/2024] We updated our arXiv v2: https://arxiv.org/abs/2401.01335.
• [02/09/2024] Our code is open-sourced!
• [01/02/2024] Our paper is released on arXiv: https://arxiv.org/abs/2401.01335.

❗ 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.

• About SPIN
• Setup
  • Data
  • Model
• Usage
  • Step 1: Generation
    • Faster generation with vLLM
  • Step 1.5: Gather generations and convert data type
  • Step 2: Fine-tuning
• Evaluation
• Citation
• Acknowledgement

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.

Average score of SPIN at different iterations on the HuggingFace Open LLM leaderboard.

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.

For more details, you can check our paper here.

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

2. Install PyTorch v2.1.0 with compatible cuda version, following instructions from PyTorch Installation Page. For example with cuda 11:

pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118

3. Install the following Python dependencies to run the codes.

python -m pip install .
python -m pip install flash-attn --no-build-isolation

4. Login to your huggingface account for downloading models

huggingface-cli login --token "${your_access_token}"

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 is converted to .parquet format for fine-tuning.

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.

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

🔍 Note: 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.

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.

python spin/reformat.py [options]

Options

• --data: directory to the SFT dataset (local or huggingface)
  • default: HuggingFaceH4/ultrachat_200k
• --output_dir: local directory to the reformated data files
  • default: UCLA-AGI/SPIN_iter0

🔍 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.

accelerate launch spin/generate.py [options]

Options

• --model: load model checkpoint for generation.
  • default: alignment-handbook/zephyr-7b-sft-full
• --input_dir: directory to the data files with prompts for generation
  • The code is for generation based on data of the format given below.
  • default: UCLA-AGI/SPIN_iter0
• --output_dir: directory to save the output data.
• --batch_size: per device batch size
  • default: 16
• --data_frac: break data into fractions for generations across server.
  • --frac_len: length of the data fraction. Default is 0 which uses the entire dataset for generation. Set frac_len to a positive number to generate only for a fraction of data. Note: we recommend using a smaller frac_len (e.g. 800) to generate data by small batches to avoid unexpected crashes as data generation can be very time-consuming.
  • Setting data_frac to be 0, 1, 2... to generate for different fractions of length frac_len.
  • Note: maintain the same frac length when doing generation using data_frac. It's recommended to set a smaller frac_len to 800.
• --split: choose the split for data generation
  • default: train

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

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.

python spin/convert_data.py [options]

Options

• --num_fracs: number of files to load in.
• --input_dir: directory to the generated data files.
• --output_dir: directory for the unified data that will be used for fine-tuning.

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

accelerate launch --config_file configs/multi_gpu.yaml --num_processes=8 --main_process_port 29500 spin/run_spin.py configs/config.yaml

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

• model_name_or_path: load model checkpoint for finetuning.
  • default: alignment-handbook/zephyr-7b-sft-full
• dataset_mixer: choose data to mix for fine-tuning.
  • default: UCLA-AGI/SPIN_iter0: 1.0
  • For SPIN at iteration 1 and 2, we included both the current iteration and the previous iteration (e.g. for iteration 1 we included both UCLA-AGI/SPIN_iter0: 1.0 and UCLA-AGI/SPIN_iter1: 1.0, summing to 100k data.)
• output_dir: the output directory of finetuned model and checkpoints.
  • default: outputs
• per_device_train_batch_size: batch size on one GPU.
  • default: 16
• gradient_accumulation_steps: make sure that the product per_device_train_batch_size*num_processes*gradient_accumulation_steps equals to your true batch size.
• num_train_epochs: the training epochs of this iteration.
  • default: 3
• beta: beta in SPIN.
  • default: 0.1

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

For our evaluation on the Open LLM Leaderboard, please use the lm-evaluation-harness repository at 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.

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}
}

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