使用 DeepSpeed 和 Hugging Face Transformer 微调 FLAN-T5 XL/XXL

Scaling Instruction-Finetuned Language Models 论文发布了 FLAN-T5 模型,它是 T5 模型的增强版。FLAN-T5 由很多各种各样的任务微调而得,因此,简单来讲,它就是个方方面面都更优的 T5 模型。相同参数量的条件下,FLAN-T5 的性能相比 T5 而言有两位数的提高。Google 在 Hugging Face 上开源了 5 个 FLAN-T5 的 checkpoints,参数量范围从 8000 万 到 110 亿。
Scaling Instruction-Finetuned Language Models 论文地址:
https://arxiv.org/pdf/2210.11416.pdf关于 FLAN-T5 的模型筛选结果:
https://hf.co/models?other=arxiv:2210.11416
在之前的一篇博文中,我们已经学习了如何 针对聊天对话数据摘要生成任务微调 FLAN-T5,那时我们使用的是 Base (250M 参数) 模型。本文,我们将研究如何将训练从 Base 扩展到 XL (30 亿参数) 或 XXL (110 亿参数)。
针对聊天对话数据摘要生成任务微调 FLAN-T5 指南:
https://www.philschmid.de/fine-tune-flan-t5Base (250M 参数) 模型:
https://hf.co/google/flan-t5-baseXL (30 亿参数) 模型:
https://hf.co/google/flan-t5-xlXXL (110 亿参数) 模型:
https://hf.co/google/flan-t5-xxl
这意味着我们将学习如何利用模型并行、多 GPU 以及 DeepSpeed ZeRO 来微调 FLAN-T5 XL 和 XXL。
DeepSpeed ZeRO 链接:
https://www.deepspeed.ai/tutorials/zero/
除了作为教程的部分之外,我们还跑了一系列实验,这些实验数据可以帮助你选择正确的硬件设置。你可以在 结果和实验 部分找到详细信息。
处理数据集
与 针对聊天对话的摘要生成任务微调 FLAN-T5 一文中类似,我们需要先准备一个用于微调的数据集。本文,我们将在 CNN Dailymail 数据集 上微调 FLAN-T5-XXL。我们不会赘述如何生成数据集,如果你想了解数据集生成的详细步骤,请参阅前文提到的 Fine Tune FLAN-T5。
CNN Dailymail 数据集:
https://hf.co/datasets/cnn_dailymailFLAN-T5-XXL:
https://hf.co/google/flan-t5-xxl
我们定义了一些参数,本文的示例都会基于这些参数,但你可以根据实际需要进行调整。
与 Fine Tune FLAN-T5 不同,这次我们把预处理和训练分开。这样我们就可以在非 GPU 实例上运行预处理。我们先对数据集进行预处理 (即分词) 并将其保存到磁盘,然后训练脚本再从磁盘中加载预处理后的数据集。
我们在配置文件中定义了一个 prompt_template
,其可用于来构建指令提示,以提高我们模型的性能。prompt_template
有“固定”的开始词和结束词,文档放在中间。这意味着我们需要确保 “固定”模板词 + 文档 总长不超过模型支持的最大序列长度。因此我们需要计算模型支持的最大文档长度,稍后我们会根据它来填充或截断模板中的文档。
现在我们知道,模型支持的最大输入文档长度为 500。除了输入之外,我们还需要知道最大“目标”序列长度,我们可以通过遍历数据集中的摘要长度来得到。(代码需要运行几分钟)
使用 deepspeed
微调模型
准备完毕!我们现在可以开始训练模型了!如前所述,我们将使用集成了 DeepSpeed 的 Hugging Face Trainer。因此我们需要创建一个 deespeed_config.json
。DeepSpeed 配置 定义了要使用的 ZeRO 策略以及是否要使用混合精度训练等配置项。Hugging Face Trainer 允许我们从 deepspeed_config.json
中的 TrainingArguments
继承相关配置以避免重复设置,查看 文档了解更多信息。
DeepSpeed 配置:
https://www.deepspeed.ai/docs/config-json/DeepSpeed 文档链接:
https://hf.co/docs/transformers/v4.26.1/en/main_classes/deepspeed
我们创建了 4 组 deepspeed 配置文件用于实验,包括 CPU 卸载
和 混合精度
:
以下文件来自:
https://github.com/philschmid/deep-learning-pytorch-huggingface/blob/main/training/configs/
ds_flan_t5_z3_config.json
ds_flan_t5_z3_config_bf16.json
ds_flan_t5_z3_offload.json
ds_flan_t5_z3_offload_bf16.json
你可以根据你的运行环境选择,例如如果在 NVIDIA V100s 上运行,你就不能使用带 bf16
的配置,因为 V100 不支持 bfloat16
数据类型。
在微调T5
模型时,不能使用fp16
,因为它会导致精度溢出问题,参见问题 #4586,#10830,和拉取请求 #10956仓库地址:
https://github.com/huggingface/transformers/
如开头所述,我们使用的是 p4dn.24xlarge AWS EC2 实例,该实例包含 8 张显存为 40GB 的 NVIDIA A100。这意味着我们可以使用 bf16
,它将减少近一半的模型显存占用,使我们能够在不卸载的情况下高效训练。
我们将使用 ds_flan_t5_z3_config_bf16.json。如果你不想用 auto
值,可以查看 文档。
文档地址:
https://hf.co/docs/transformers/v4.26.1/en/main_classes/deepspeed
现在,该训练脚本上场了。我们根据 Fine Tune FLAN-T5 准备了一个 run_seq2seq_deepspeed.py 训练脚本,它支持我们配置 deepspeed 和其他超参数,包括 google/flan-t5-xxl
的模型 ID。
run_seq2seq_deepspeed.py 链接:
https://github.com/philschmid/deep-learning-pytorch-huggingface/blob/main/training/scripts/run_seq2seq_deepspeed.py
我们使用 deepspeed
启动器触发训练,输入给启动器的参数包括 GPU 数量、deepspeed 配置及其它超参数 (如 google/flan-t5-xxl
的模型 ID)。
DeepSpeed 先将模型加载到 CPU 上,然后将其拆分到 8 张 A100 上然后开始训练。使用 CNN Dailymail 数据集 进行训练大约需要 10 个小时,费用约为 322 美元
。
CNN Dailymail 数据集:
https://hf.co/datasets/cnn_dailymail
结果与实验
为了更好地了解硬件要求,我们对 FLAN-T5 XL 和 XXL 进行了一系列实验,以帮助我们评估和了解硬件需求以及训练这些模型的成本。
下表列出了实验和相关设置的详细信息。
数据集: "cnn_dailymail"
训练样本数:
287113
验证样本数:
13368
超参:
epochs:
3
学习率:
1e-4
运行环境设置:
4x V100 16GB: p3.8xlarge
4x A10G 24GB: g5.24xlarge
8x V100 16GB: p3.16xlarge
8x A100 40GB: p4dn.24xlarge
模型DeepSpeed 卸载硬件GPU每卡batch size精度时长费用

我们可以看到 bf16
与 fp32
相比具有显著优势。FLAN-T5-XXL 能放进 4 张 A10G (24GB),但放不进 8 张 V100 16GB。
我们的实验还表明,如果模型可以无需卸载同时以 batch size 大于 4 的配置跑在 GPU 上,其速度将比卸载模型和减小 batch size 的配置快约 2 倍且更具成本效益。
英文原文: https://www.philschmid.de/fine-tune-flan-t5-deepspeed
原文作者: Philipp Schmid
译者: Matrix Yao (姚伟峰),英特尔深度学习工程师,工作方向为 transformer-family 模型在各模态数据上的应用及大规模模型的训练推理。
审校、排版: zhongdongy (阿东)