BingBertSQuAD 微调

在本教程中,我们将为 BingBert 模型添加 DeepSpeed,以完成 SQuAD 微调任务,此后称之为“BingBertSquad”。我们还将展示性能提升。

概述

如果您还没有 DeepSpeed 仓库的副本,请立即克隆并检出包含 BingBertSquad 示例(DeepSpeedExamples/training/BingBertSquad)的 DeepSpeedExamples 子模块,我们将在本教程的其余部分介绍该示例。

git clone https://github.com/deepspeedai/DeepSpeed
cd DeepSpeed
git submodule update --init --recursive
cd DeepSpeedExamples/training/BingBertSquad

先决条件

您还需要来自 DeepSpeed、HuggingFaceTensorFlow 的预训练 BERT 模型检查点才能运行微调。关于 DeepSpeed 模型,我们将使用 BERT 预训练教程中的检查点 160。

运行 BingBertSquad

  • 启用 DeepSpeed: 我们提供了一个 shell 脚本,您可以调用它来启动 DeepSpeed 训练,它需要 4 个参数:bash run_squad_deepspeed.sh <NUM_GPUS> <PATH_TO_CHECKPOINT> <PATH_TO_DATA_DIR> <PATH_TO_OUTPUT_DIR>。第一个参数是要训练的 GPU 数量,第二个参数是预训练检查点的路径,第三个参数是训练集和验证集(例如 train-v1.1.json)的路径,第四个参数是结果将保存的输出文件夹的路径。此脚本将调用 nvidia_run_squad_deepspeed.py
  • 未修改的基线 如果您想运行未启用 DeepSpeed 的微调版本,我们提供了一个与 DeepSpeed 版本具有相同参数的 shell 脚本,名为 run_squad_baseline.sh。此脚本将调用 nvidia_run_squad_baseline.py

DeepSpeed 集成

训练的主要部分在 nvidia_run_squad_deepspeed.py 中完成,该文件已修改为使用 DeepSpeed。run_squad_deepspeed.sh 脚本有助于调用训练并设置与训练过程相关的几个不同超参数。在接下来的几节中,我们将介绍我们为启用 DeepSpeed 对基线所做的更改,您无需亲自进行这些更改,因为我们已经为您完成了。

配置

deepspeed_bsz24_config.json 文件使用户能够指定 DeepSpeed 选项,包括批处理大小、微批处理大小、学习率和其他参数。运行 nvidia_run_squad_deepspeed.py 时,除了 --deepspeed 标志以启用 DeepSpeed 外,还必须使用 --deepspeed_config deepspeed_bsz24_config.json 指定相应的 DeepSpeed 配置文件。表 1 显示了我们实验中使用的微调配置。

参数
总批处理大小 24
每 GPU 的训练微批处理大小 3
优化器 Adam
学习率 3e-5
序列长度 384
权重衰减 0.0
epoch 计数 2

表 1. 微调配置

参数解析

应用 DeepSpeed 的第一步是在主入口点开头使用 deepspeed.add_config_arguments() 向 BingBertSquad 添加参数,就像 nvidia_run_squad_deepspeed.py 中的 main() 函数一样。add_config_arguments() 的参数从 utils.py 中的 get_argument_parser() 函数获得。

parser = get_argument_parser()
# Include DeepSpeed configuration arguments
parser = deepspeed.add_config_arguments(parser)
args = parser.parse_args()

与此类似,所有选项及其相应的描述都可在 utils.py 中找到。

训练

初始化

DeepSpeed 有一个初始化函数,用于包装模型、优化器、学习率调度器和数据加载器。对于 BingBertSquad,我们只需使用初始化函数来增强基线脚本,以包装模型并按如下方式创建优化器

model, optimizer, _, _ = deepspeed.initialize(
    args=args,
    model=model,
    model_parameters=optimizer_grouped_parameters
)

前向传播

这在基线和 DeepSpeed 中都是相同的,通过 loss = model(input_ids, segment_ids, input_mask, start_positions, end_positions) 完成。

反向传播

在基线脚本中,您需要通过使用 enable_need_reduction(),后跟 FP16 中的 optimizer.backward(loss) 和 FP32 中的 loss.backward(),显式处理梯度累积边界上的 all-reduce 操作。在 DeepSpeed 中,您只需执行 model.backward(loss)

权重更新

在基线脚本中,您需要显式将优化器指定为 FP16 中的 FusedAdam(以及动态损失缩放的处理)和 FP32 中的 BertAdam,然后调用 optimizer.step()optimizer.zero_grad()。DeepSpeed 在调用 initialize() 时会在内部处理此问题(通过使用 JSON 配置设置优化器),因此您无需显式编写代码,只需执行 model.step()

恭喜!已完成向 DeepSpeed 的移植。

评估

训练完成后,可以通过以下命令获取 EM 和 F1 分数

python evaluate-v1.1.py <PATH_TO_DATA_DIR>/dev-v1.1.json <PATH_TO_DATA_DIR>/predictions.json

微调结果

结果总结表如下所示。在所有情况下(除非另有说明),总批处理大小设置为 24,并在 DGX-2 节点上使用 4 个 GPU 训练 2 个 epoch。我们尝试了一组参数(种子和学习率),并选择了最佳参数。所有学习率均为 3e-5;我们将 HuggingFace 和 TensorFlow 模型的种子分别设置为 9041 和 19068。下表中链接了每个情况使用的检查点。

情况 模型 精度 EM F1
TensorFlow Bert-large-uncased-L-24_H-1024_A-16 FP16 84.13 91.03
HuggingFace Bert-large-uncased-whole-word-masking FP16 87.27 93.33

启用 DeepSpeed 的 Transformer 内核以获得更好的吞吐量

DeepSpeed 优化后的 Transformer 内核可以在微调期间启用,以提高训练吞吐量。除了支持使用 DeepSpeed 预训练的模型外,该内核还可以与 TensorFlow 和 HuggingFace 检查点一起使用。

启用 Transformer 内核

参数 --deepspeed_transformer_kernel 已在 utils.py 中创建,我们通过将其添加到 shell 脚本中来启用 transformer 内核。

parser.add_argument(
    '--deepspeed_transformer_kernel',
    default=False,
    action='store_true',
    help='Use DeepSpeed transformer kernel to accelerate.'
)

在建模源文件的 BertEncoder 类中,当通过使用 --deepspeed_transformer_kernel 参数启用 DeepSpeed transformer 内核时,它会按如下方式创建。

if args.deepspeed_transformer_kernel:
    from deepspeed import DeepSpeedTransformerLayer, \
        DeepSpeedTransformerConfig, DeepSpeedConfig

    ds_config = DeepSpeedConfig(args.deepspeed_config)

    cuda_config = DeepSpeedTransformerConfig(
        batch_size=ds_config.train_micro_batch_size_per_gpu,
        max_seq_length=args.max_seq_length,
        hidden_size=config.hidden_size,
        heads=config.num_attention_heads,
        attn_dropout_ratio=config.attention_probs_dropout_prob,
        hidden_dropout_ratio=config.hidden_dropout_prob,
        num_hidden_layers=config.num_hidden_layers,
        initializer_range=config.initializer_range,
        seed=args.seed,
        fp16=ds_config.fp16_enabled
    )
    self.layer = nn.ModuleList([
        copy.deepcopy(DeepSpeedTransformerLayer(i, cuda_config))
        for i in range(config.num_hidden_layers)
    ])
else:
    layer = BertLayer(config)
    self.layer = nn.ModuleList([
        copy.deepcopy(layer)
        for _ in range(config.num_hidden_layers)
    ])

所有配置设置都来自 DeepSpeed 配置文件和命令行参数,因此我们必须将 args 变量传递到此模型中。

注意:batch_size 是输入数据的最大批处理大小,所有微调训练数据或预测数据都不应超过此阈值,否则将抛出异常。在 DeepSpeed 配置文件中,微批处理大小定义为 train_micro_batch_size_per_gpu,例如,如果设置为 8,则 --predict_batch_size 也应为 8。

有关 Transformer 内核的更多详细信息,请参阅我们的使用教程和关于最快 BERT 训练的技术深入探讨

加载 HuggingFace 和 TensorFlow 预训练模型

BingBertSquad 支持 HuggingFace 和 TensorFlow 预训练模型。这里,我们展示两个模型示例

  1. test/huggingface 包含检查点 Bert-large-uncased-whole-word-maskingbert json 配置
  2. test/tensorflow 来自 Google 的一个检查点 zip 包 Bert-large-uncased-L-24_H-1024_A-16
[test/huggingface]
bert-large-uncased-whole-word-masking-config.json
bert-large-uncased-whole-word-masking-pytorch_model.bin
[test/tensorflow]
bert_config.json
bert_model.ckpt.data-00000-of-00001
bert_model.ckpt.index
bert_model.ckpt.meta

加载这两种检查点需要使用三个参数。

  1. --model_file,指向预训练模型文件。
  2. --ckpt_type,指示检查点类型,TF 代表 Tensorflow,HF 代表 HuggingFace,默认值为 DeepSpeed 的 DS
  3. --origin_bert_config_file,指向 BERT 配置文件,通常保存在与 model_file 相同的文件夹中。

我们可以在 run_squad_deepspeed.sh 中的微调 shell 脚本中添加以下内容,以运行上述 HuggingFace 和 TensorFlow 示例。

[HuggingFace]

--model_file test/huggingface/bert-large-uncased-whole-word-masking-pytorch_model.bin \
--ckpt_type HF \
--origin_bert_config_file test/huggingface/bert-large-uncased-whole-word-masking-config.json \
[TensorFlow]

--model_file /test/tensorflow/bert_model.ckpt \
--ckpt_type TF \
--origin_bert_config_file /test/tensorflow/bert_config.json \

注意

  1. 使用 HuggingFace 或 TensorFlow 预训练模型需要 --deepspeed_transformer_kernel 标志。

  2. --preln 标志不能与 HuggingFace 或 TensorFlow 预训练模型一起使用,因为它们使用后层归一化。

  3. BingBertSquad 将检查预训练模型是否具有相同的词汇量大小,如果存在任何不匹配,将无法运行。我们建议您使用上述风格的模型检查点或 DeepSpeed bing_bert 检查点。

性能调优

为了执行微调,我们将总批处理大小设置为 24,如表 1 所示。但是,我们可以调整每 GPU 的微批处理大小以获得高性能训练。在这方面,我们已经在 NVIDIA V100 上尝试了不同的微批处理大小,使用了 16GB 或 32GB 内存。如表 2 和表 3 所示,我们可以通过增加微批处理来提高性能。与 PyTorch 相比,我们可以在 16GB V100 上实现高达 1.5 倍的加速,同时支持每 GPU 2 倍的更大批处理大小。另一方面,我们可以在 32GB V100 上支持高达 32 的批处理大小(比 PyTorch 高 2.6 倍),同时在端到端微调训练中提供 1.3 倍的加速。请注意,我们使用最佳样本/秒来计算 PyTorch 内存不足 (OOM) 情况下的加速。

微批处理大小 PyTorch DeepSpeed 加速 (x)
4 36.34 50.76 1.4
6 OOM 54.28 1.5
8 OOM 54.16 1.5

表 2. 在 NVIDIA V100 (16GB) 上使用 PyTorch 和 DeepSpeed Transformer 内核运行 SQuAD 微调的样本/秒。

微批处理大小 PyTorch DeepSpeed 加速 (x)
4 37.78 50.82 1.3
6 43.81 55.97 1.3
12 49.32 61.41 1.2
24 OOM 60.70 1.2
32 OOM 63.01 1.3

表 3. 在 NVIDIA V100 (32GB) 上使用 PyTorch 和 DeepSpeed Transformer 内核运行 SQuAD 微调的样本/秒。

如前所述,如果需要更大的批处理大小,我们可以将每 GPU 的微批处理大小从 3 增加到 24 甚至更高。为了支持更大的微批处理大小,我们可能需要为我们的 Transformer 内核启用不同的内存优化标志,如DeepSpeed Transformer 内核教程中所述。表 4 显示了运行不同范围的微批处理大小所需的优化标志。

微批处理大小 NVIDIA V100 (32GB) NVIDIA V100 (16GB)
> 4 - normalize_invertible
> 6 - attn_dropout_checkpoint, gelu_checkpoint
> 12 normalize_invertible, attn_dropout_checkpoint OOM
> 24 gelu_checkpoint OOM

表 4. 16GB 和 32GB V100 上一系列微批处理大小的内存优化标志设置。

使用 DeepSpeed Transformer 内核预训练模型进行微调

使用 DeepSpeed Transformer 和 DeepSpeed Fast-Bert Training 中的方法预训练的模型进行微调应产生 90.5 的 F1 分数,如果您让预训练时间比教程中建议的更长,预计还会增加。

为了获得这些结果,我们需要对 dropout 设置进行一些调整,如下所述

Dropout 设置

对于微调,我们只使用确定性 Transformer 以获得可复现的微调结果。但是,我们根据预训练是使用确定性 Transformer 还是随机 Transformer 来选择不同的 dropout 值(有关选择这两种模式的更多详细信息,请参阅 Transformer 教程)。

对于使用确定性 Transformer 预训练的模型,我们使用与预训练中相同的 dropout 比率 (0.1)。但是,当对使用随机 Transformer 预训练的模型进行微调时,我们会略微增加 dropout 比率,以补偿微调过程中随机噪声的不足。

预训练模式 Dropout 比率
确定性 0.1
随机性 0.12 - 0.14

更新日期: