Huggingface训练Transformer

在之前的博客中，我采用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，在继续训练时就报内存不够了，这个也有待后溪继续了解。