DeepSpeed 数据效率:一个可组合的库,旨在更好地利用数据,提高训练效率并改善模型质量
什么是 DeepSpeed 数据效率: DeepSpeed 数据效率是一个专门构建的库,旨在更好地利用数据、提高训练效率并改善模型质量。
为何使用 DeepSpeed 数据效率: DeepSpeed 数据效率提供了新颖的数据效率技术,以实现更高的训练效率和/或更好的模型质量。DeepSpeed 数据效率考虑了可扩展性、灵活性和可组合性,这使得定制技术、将技术应用于各种训练任务以及组合多种技术变得更加容易。我们强烈建议您也阅读我们的博客,以(从高层面)了解我们为何构建 DeepSpeed 数据效率以及它能为用户带来哪些益处。更多技术细节可在我们的论文中找到,例如描述 random-LTD 技术的“Random-LTD: Random and Layerwise Token Dropping Brings Efficient Training for Large-scale Transformers”,以及描述课程学习技术和整个 DeepSpeed 数据效率框架的“DeepSpeed Data Efficiency: Improving Deep Learning Model Quality and Training Efficiency via Efficient Data Sampling and Routing”。
如何使用 DeepSpeed 数据效率: 在以下教程中,前两节将描述该库支持的数据效率技术。第三节将描述如何组合这两种技术以实现更高的训练效率/模型质量。
1. 课程学习
1.1 什么是课程学习
课程学习(由 Yoshua Bengio 等人提出)旨在通过在训练早期呈现相对更容易或更简单的示例来提高训练收敛速度。构建一个课程学习解决方案通常需要两个组成部分:难度度量(即如何量化每个数据样本的难度)和步调函数(即在采样下一个训练数据批次时如何决定课程难度范围)。
1.2 何时使用课程学习
课程学习已成功应用于各种训练任务(详见例如这篇综述论文),去年我们还发布了一种专门的课程学习技术(序列长度热身),用于 GPT 风格模型的预训练(技术细节请参阅我们发表在 NeurIPS 2022 的论文“The Stability-Efficiency Dilemma: Investigating Sequence Length Warmup for Training GPT Models”以及此旧版课程学习功能的教程)。DeepSpeed 数据效率中这个新的通用课程学习库使用户能够在最大可扩展性下将其模型应用于课程学习:用户可以根据各种可定制的策略轻松分析、索引和采样其训练数据。通过使用此库,我们能够探索 GPT-3 和 BERT 预训练的不同 CL 策略,并确定了最佳解决方案,该方案在保持相似模型质量的同时提供了高达 1.5 倍的数据节省。
1.3 如何使用课程学习
1.3.1 GPT-3 和 BERT 预训练
我们的 Megatron-DeepSpeed 仓库中的 examples_deepspeed/data_efficiency
目录包含了我们如何将课程学习应用于 GPT-3 和 BERT 预训练的示例。总共有 3 个步骤:数据分析、预训练和评估/微调。
数据分析: 课程学习在预训练前需要进行数据分析,以计算每个数据样本的难度(基于用户提供的度量),并构建一个将难度值映射到相应数据样本的索引。(也有例外:例如,基于截断的序列长度度量可以通过数据后处理实现,无需数据分析。)我们提供了一个数据分析器来执行离线、仅限 CPU 的数据分析。
examples_deepspeed/data_efficiency/gpt/ds_analyze_*.sh
和 examples_deepspeed/data_efficiency/bert/ds_analyze_*.sh
是用于 GPT-3 和 BERT 数据分析的示例脚本。我们的数据分析器采用简单的 Map-Reduce 方案。首先,在 Map 阶段使用 ds_analyze_*_data_map.sh
来分割数据集并计算每个数据样本的难度值。用户需要提供一个用于计算度量的函数(我们在 examples_deepspeed/data_efficiency/analyze_data.py
中实现了我们的函数)、原始训练数据集以及其他配置,例如 CPU 节点数量和每个节点的线程数量。然后,数据分析器将根据 worker 数量自动分割数据集,以批处理方式计算难度值,并将结果写入两个索引:一个索引将每个数据样本映射到其难度值,另一个索引将每个不同的难度值映射到相应的样本。其次,在 Reduce 阶段使用 ds_analyze_*_data_reduce.sh
来合并所有 worker 生成的索引文件。需要注意的一点是,为了通过分布式实现加速同时仍能合并所有输出,Map 阶段可能会生成大量输出文件,其数量与 CPU 节点数量、每个节点的线程数量以及可能的度量值数量成正比。因此,为了避免生成过多输出文件,我们建议从较小的节点/线程数量开始(在输出日志中,我们提供了所需时间的估计,供用户判断是否要增加 worker 数量),并且我们建议在设计难度度量时限制可能的难度值数量(我们的经验表明,几千个不同的值足以享受课程学习的好处)。
预训练 examples_deepspeed/data_efficiency/gpt/pretrain
和 examples_deepspeed/data_efficiency/bert/pretrain
包含了启用课程学习功能的预训练示例脚本。在预训练期间启用课程学习需要进行几项更改:(1) 用户需要提供一个包含课程学习配置的 DeepSpeed json 配置文件(详见配置列表)。我们在 examples_deepspeed/data_efficiency/gpt/pretrain/ds_pretrain_gpt_1.3B_dense_run.sh
和 examples_deepspeed/data_efficiency/bert/pretrain/ds_pretrain_bert_336M_run.sh
中提供了经过测试的示例配置。(2) 在通过 deepspeed.initialize
初始化 DeepSpeed 引擎时,用户需要提供训练数据集并使用初始化返回的数据加载器(此数据加载器包含课程学习功能)。我们在 megatron/training.py
的 setup_model_and_optimizer
函数中提供了此更改的示例实现。(3) 如果课程学习度量需要数据后处理(例如基于截断的序列长度),用户需要使用 DeepSpeed 引擎的 set_data_post_process_func
API 来提供后处理函数。我们在 megatron/training.py
、pretrain_bert.py
和 pretrain_gpt.py
中提供了此更改的示例实现。(4) 如果课程学习度量需要自定义调度策略(步调函数),用户需要使用 DeepSpeed 引擎的 set_custom_curriculum_learning_schedule
API 来提供在训练期间更新最大可接受难度的函数。DeepSpeed 引擎将向此回调函数提供一个全局训练步数输入。
评估/微调 examples_deepspeed/data_efficiency/gpt/eval/
和 examples_deepspeed/data_efficiency/bert/finetune
包含了 GPT-3 模型的零/少样本评估和 BERT 模型的微调示例脚本。如果您按照我们的示例脚本执行预训练/评估/微调,我们的论文中包含了参考的评估/微调结果。
1.3.2 GPT-2 微调
我们的 DeepSpeedExamples 仓库中的 data_efficiency/gpt_finetuning
目录包含了我们如何将课程学习应用于 GPT-2 微调的示例。data_efficiency/gpt_finetuning/finetune/ds_finetune_gpt2_run.sh
是微调示例脚本。对于需要数据分析的 CL 度量(例如,词汇稀有度度量),您需要首先使用 data_efficiency/gpt_finetuning/finetune/ds_analyze_gpt_data_*
来分析和索引数据集,类似于上面 1.3.1 中描述的 GPT-3 预训练情况。
2. 随机分层 token 丢弃 (random-LTD)
2.1 什么是 random-LTD
Random-LTD 是一种高效的 token 丢弃方法,以随机分配方式应用于每一层。具体来说,对于每一层,与基线相比,random-LTD 随机选择 token 的一个子集并将其送入 Transformer 层。之后,我们将 Transformer 层的输出与丢弃的 token 结合起来,以恢复完整的序列长度。因此,下一层仍然接收完整的序列并可以重复此过程。更多技术细节请阅读我们的 random-LTD 论文。
2.2 何时使用 random-LTD
当您想要预训练/微调基于 Transformer 的模型时,尝试 random-LTD 始终是一个好主意,因为在相同的计算成本下,它比标准基线训练能获得更好的性能。如果您的资源有限,random-LTD 可以实现与原始基线方法相似的精度,同时理论成本节省高达 33.3%,实际运行时间节省高达 25.6%。特别是,如果您需要训练一个层数 >=24 且序列长度 >=2048 的更大模型,我们的方法将比基线效率高得多。
2.3 如何使用 random-LTD
2.3.1 GPT-3 和 BERT 预训练
我们的 Megatron-DeepSpeed 仓库中的 examples_deepspeed/data_efficiency
目录包含了我们如何将 random-LTD 应用于 GPT-3 和 BERT 预训练的示例。
examples_deepspeed/data_efficiency/gpt/pretrain
和 examples_deepspeed/data_efficiency/bert/pretrain
包含了启用 random-LTD 功能的预训练示例脚本。在预训练期间启用 random-LTD 需要进行几项更改:(1) 用户需要提供一个包含 random-LTD 配置的 DeepSpeed json 配置文件(详见配置列表)。我们在 examples_deepspeed/data_efficiency/gpt/pretrain/ds_pretrain_gpt_1.3B_dense_run.sh
和 examples_deepspeed/data_efficiency/bert/pretrain/ds_pretrain_bert_336M_run.sh
中提供了经过测试的示例配置。(2) 在通过 deepspeed.initialize
初始化 DeepSpeed 引擎后,用户需要使用 convert_to_random_ltd
API 来转换和封装模型层,以启用 random-LTD 功能。我们在 megatron/training.py
的 setup_model_and_optimizer
函数中提供了此更改的示例实现。(3) 为了使 random-LTD 理解前向函数的输入参数映射,用户需要将所有输入参数(除了 hidden_states 输入)更改为关键字/命名参数。例如,在 megatron/model/transformer.py
中,我们将前向函数从 def forward(self, hidden_states, attention_mask, encoder_output=None, enc_dec_attn_mask=None, layer_past=None, get_key_value=False):
更改为 def forward(self, hidden_states, attention_mask=None, encoder_output=None, enc_dec_attn_mask=None, layer_past=None, get_key_value=False):
。(4) 在保存模型检查点时(特别是当状态字典具有非传统结构时),用户需要使用 remove_random_ltd_state_dict
API 将 random-LTD 封装的层转换回原始模型层。我们在 megatron/model/language_model.py
中提供了此更改的示例实现。
对于预训练模型的评估/微调,请参阅上一节关于如何使用我们的示例脚本的说明。
2.3.2 GPT-2 和 ViT 微调
我们的 DeepSpeedExamples 仓库中的 data_efficiency
目录包含了我们如何将 random-LTD 应用于 GPT-2 和 ViT 微调的示例。
与预训练情况类似,启用 random-LTD 进行微调也需要进行类似的更改:(1) DeepSpeed json 配置文件。(2) 使用 convert_to_random_ltd
API 转换和封装模型层。(3) 在保存模型检查点时,使用 remove_random_ltd_state_dict
API 将 random-LTD 封装的层转换回原始模型层。
可以通过以下方式运行我们的 GPT 微调示例:
DeepSpeedExamples/data_efficiency/gpt_finetuning$ pip install -r requirement.txt
DeepSpeedExamples/data_efficiency/gpt_finetuning$ bash ./bash_script/run_base_random_ltd.sh
DeepSpeedExamples/data_efficiency/gpt_finetuning$ bash ./bash_script/run_medium_random_ltd.sh
参考的最终结果是:
For run_base_random_ltd.sh:
End of training epoch 3 step 1344 consumed_token 2148032 best perplexity 22.552324221233757 time 0.17486039188173083 hr
For run_medium_random_ltd.sh:
End of training epoch 3 step 1373 consumed_token 2147024 best perplexity 17.332243199130996 time 0.4661190489927928 hr
可以通过以下方式运行我们的 ViT 微调示例:
DeepSpeedExamples/data_efficiency/vit_finetuning$ pip install -r requirement.txt
DeepSpeedExamples/data_efficiency/vit_finetuning$ bash ./bash_script/run_cifar.sh
DeepSpeedExamples/data_efficiency/vit_finetuning$ bash ./bash_script/run_imagenet.sh
参考的最终结果是:
For run_cifar.sh:
13 epoch at time 480.6546013355255s | reserved_length 197
iter 5474 | LR [0.0001]| val_acc 97.97000122070312 | layer_token 305784192
3. 结合课程学习和 random-LTD 以实现更多
3.1 GPT-3 和 BERT 预训练
我们的 Megatron-DeepSpeed 仓库中的 examples_deepspeed/data_efficiency
目录包含了我们如何组合课程学习 random-LTD,并将两者应用于 GPT-3 和 BERT 预训练的示例。
所需的更改与前两节中描述的相同,因为 DeepSpeed 数据效率已经处理了组合这两种技术时的复杂性。然而,需要注意的是,由于 random-LTD 和一些课程学习度量都会改变序列长度,因此可能需要额外的代码来计算每一步的有效序列长度。我们在 megatron/training.py
的 train
函数中提供了此更改的示例实现,我们在其中计算 actual_seq_length
。
3.2 GPT-2 微调
我们的 DeepSpeedExamples 仓库中的 data_efficiency/gpt_finetuning
目录包含了我们如何组合课程学习 random-LTD 以进行 GPT-2 微调的示例。data_efficiency/gpt_finetuning/finetune/ds_finetune_gpt2_run.sh
是微调示例脚本。