1.3 KiB
Add Peft support for SFT and Prompts model training
The orginal implementation just adopts the loralib and merges the layers into the final model. The huggingface peft is a better lora model implementation and can be easily training and distributed.
Since reward model is relative small, I just keep it as original one. I suggest train full model to get the proper reward/critic model.
Prelimenary installation
Since the current pypi peft package(0.2) has some bugs, please install the peft package using source.
git clone https://github.com/huggingface/peft
cd peft
pip install .
Usage
For SFT training, just call train_peft_sft.py
Its arguments are almost identical to train_sft.py instead adding a new eval_dataset if you have a eval_dataset file. The data file is just a plain datafile, please check the format in the easy_dataset.py.
For stage-3 rlhf training, call train_peft_prompts.py. Its arguments are almost idential to train_prompts.py. The only difference is that I use text files to indicate the prompt and pretrained data file. The models are included in easy_models.py. Currently only bloom models are tested, but technically gpt2/opt/llama should be supported.
Dataformat
Please refer the formats in test_sft.txt, test_prompts.txt, test_pretrained.txt.