As described in the Zephyr technical report, training this model proceeds in two steps:
- Apply SFT to fine-tune Mistral 7B on a filtered version of the UltraChat dataset (link). The result is an SFT model like
zephyr-7b-sft-full
orzephyr-7b-sft-qlora
. - Align the SFT model to AI feedback via DPO on a preprocessed version of the UltraFeedback dataset (link). The result is an DPO model like
zephyr-7b-dpo-full
orzephyr-7b-dpo-qlora
.
Note: after the release of Zephyr, the team at Argilla found that the source UltraFeedback dataset had a few thousand incorrect preference labels from GPT-4. Additionally, TRL's SFTTrainer
had a bug in the learning rate scheduler which terminated training early. Accounting for these changes led us to find a better set of hyperparameters from those described in the technical report. In particular, for DPO training we found that training for 1 epoch with beta=0.01
was suffucient to achieve comparable performance to zephyr-7b-beta
(vs. 3 epochs with beta=0.1
).
See below for commands to train these models using either DeepSpeed ZeRO-3 or LoRA.
You will require 8 GPUs (80GB of VRAM) to train the full model.
# Step 1 - SFT
ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/deepspeed_zero3.yaml scripts/run_sft.py recipes/zephyr-7b-beta/sft/config_full.yaml
# Step 2 - DPO
ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/deepspeed_zero3.yaml scripts/run_dpo.py recipes/zephyr-7b-beta/dpo/config_full.yaml
Train faster with flash-attention 2 (GPU supporting FA2: A100, H100, etc)
# Step 1 - SFT
ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/multi_gpu.yaml --num_processes=1 scripts/run_sft.py recipes/zephyr-7b-beta/sft/config_qlora.yaml --load_in_4bit=true
# Step 2 - DPO
ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/multi_gpu.yaml --num_processes=1 scripts/run_dpo.py recipes/zephyr-7b-beta/dpo/config_qlora.yaml
P.S. Using Flash Attention also allows you to drastically increase the batch size (x2 in my case)
Train without flash-attention:
# Step 1 - SFT
ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/multi_gpu.yaml --num_processes=1 scripts/run_sft.py recipes/zephyr-7b-beta/sft/config_qlora.yaml --load_in_4bit=true --use_flash_attention_2=false
# Step 2 - DPO
ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/multi_gpu.yaml --num_processes=1 scripts/run_dpo.py recipes/zephyr-7b-beta/dpo/config_qlora.yaml --use_flash_attention_2=false