在之前的博客中，我采用SFT（监督优化训练）的方法训练一个GPT2的模型，使得这个模型可以根据提示语进行回答。具体可见博客召唤神龙打造自己的ChatGPT_gzroy的博客-CSDN博客
Huggingface提供了一个TRL的扩展库，可以对transformer模型进行强化学习，SFT是其中的一个训练步骤，为此我也测试一下如何用Huggingface来进行SFT训练，和Pytorch的训练方式做一个比较。
训练数据
首先是获取训练数据，这里同样是采用Huggingface的chatbot_instruction_prompts的数据集，这个数据集涵盖了不同类型的问答，可以用作我们的模型优化之用。
from datasets import load_dataset
ds = load_dataset("alespalla/chatbot_instruction_prompts")
train_ds = ds['train']
eval_dataset = ds['test']
eval_dataset = eval_dataset.select(range(1024))
训练集总共包括了258042条问答数据，对于验证集我只选取了头1024条记录，因为总的数据集太长，如果在训练过程中全部验证的话耗时太长。
加载GPT2模型
Huggingface提供了很多大模型的训练好的参数，这里我们可以直接加载一个已经训练好的GPT2模型来做优化
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token
然后我们可以定义TRL提供的SFTTrainer来进行训练，首先需要对训练数据处理一下，因为训练数据包括了两列，分别是prompt和response，我们需要把两列的文本合为一起，通过格式化字符来区分，如以下格式化函数：
def formatting_func(example):
text = f"### Prompt: {example['prompt']}\n ### Response: {example['response']}"
return text
定义SFTTrainer的训练参数，具体每个参数的含义可见官网的文档：
args = TrainingArguments(
output_dir='checkpoints_hf_sft',
overwrite_output_dir=True,
per_device_train_batch_size=4,
per_device_eval_batch_size=4,
fp16=True,
torch_compile=True,
evaluation_strategy='steps',
prediction_loss_only=True,
eval_accumulation_steps=1,
learning_rate=0.00006,
weight_decay=0.01,
adam_beta1=0.9,
adam_beta2=0.95,
warmup_steps=1000,
eval_steps=4000,
save_steps=4000,
save_total_limit=4,
dataloader_num_workers=4,
max_steps=12000,
optim='adamw_torch_fused')
最后就可以定义一个trainer来训练了
trainer = SFTTrainer(
model,
args = args,
train_dataset=dataset,
eval_dataset=eval_dataset,
tokenizer=tokenizer,
packing=True,
formatting_func=formatting_func,
max_seq_length=1024
)
trainer.train(resume_from_checkpoint=False)
因为是第一次训练，我设置了resume_from_checkpoint=False，如果是继续训练，把这个参数设为True即可从上次checkpoint目录自动加载最新的checkpoint来训练。
训练结果如下：
[12000/12000 45:05, Epoch 0/1]
| Step |
Training Loss |
Validation Loss |
| 4000 |
2.137900 |
2.262321 |
| 8000 |
2.187800 |
2.232235 |
| 12000 |
2.218500 |
2.210413 |
总共耗时45分钟，比我在pytorch上的训练要快一些（快了10分钟多一些），但是这个训练集的Loss随着Step的增加反而增加了，Validation Loss就减少了，有些奇怪。在Pytorch上我同样训练12000个迭代，最后training loss是去到1.8556的。可能还要再调整一下trainer的参数看看。
测试
最后我们把SFT训练完成的模型，通过huggingface的pipeline就可加载进行测试了。
from transformers import pipeline
model = AutoModelForCausalLM.from_pretrained('checkpoints_hf_sft/checkpoint-12000/')
pipe = pipeline(task='text-generation', model=model, tokenizer=tokenizer, device=0)
pipe('### Prompt: Who is the current president of USA?')
[{'generated_text': '### Prompt: Who is the current president of USA?\n ### Response: Harry K. Busby is currently President of the United States.'}]
回答的语法没问题，不过内容是错的。
再测试另一个问题
### Prompt: How to make a cup of coffee?
[{'generated_text': '### Prompt: How to make a cup of coffee?\n ### Response: 1. Preheat the oven to 350°F (175°C).\n\n2. Boil the coffee beans according to package instructions.\n\n3. In a'}]
这个回答就正确了。
总结
通过用Huggingface可以很方便的对大模型进行强化学习的训练，不过也正因为太方便了，很多训练的细节被包装了，所以训练的结果不太容易优化，不像在Pytorch里面控制的自由度更高一些。当然可能我对huggingface的trainer参数的细节还不太了解，这个有待后续继续了解。
另外我还发现huggingface的一个小的bug，就是模型从头训练的时候没有问题，但是当我从之前的checkpoint继续训练时，会报CUDA OOM的错误，从nvidia-smi命令看到的显存占用率来看，好像trainer定义模型和装载Checkpoint会重复占用了显存，因此同样的batch_size，在继续训练时就报内存不够了，这个也有待后溪继续了解。
