🤗 Models | 🤗 Datasets

Self-Play Fine-Tuning (SPIN)

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

• [01/02/2024] Our paper is released on arXiv: https://arxiv.org/abs/2401.01335.

Table of Contents

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

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

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

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.

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.

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

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

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

Dataset	Download
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.

Model	Download
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: 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.

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

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

Step 1: Generation

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

Example

The following code generates 8k synthetic data for iteration 1.

bash scripts/generate.sh

Step 1.5: Gather generations and convert data type

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

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

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

Examples

bash scripts/finetune.sh

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.