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,在继续训练时就报内存不够了,这个也有待后溪继续了解。

猜你喜欢

转载自blog.csdn.net/gzroy/article/details/132521808