BingBertSQuAD 微调

在本教程中,我们将为 SQuAD 微调任务(此后称为“BingBertSquad”)的 BingBert 模型添加 DeepSpeed。我们还将演示性能提升。

概述

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

git clone https://github.com/microsoft/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 的微调版本,我们提供了一个 shell 脚本,它采用与 DeepSpeed 相同的参数,名为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
纪元计数 2

表 1. 微调配置

参数解析

应用 DeepSpeed 的第一步是向 BingBertSquad 添加参数,在主入口点的开头使用deepspeed.add_config_arguments(),如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 有一个初始化函数来包装模型、优化器、LR 调度器和数据加载器。对于 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()来显式处理梯度累积边界处的全简化操作。在 DeepSpeed 中,您只需执行model.backward(loss)即可。

权重更新

在基线脚本中,您需要显式指定优化器为FusedAdam(以及 FP16 中的动态损失缩放处理)和 FP32 中的BertAdam,然后调用optimizer.step()optimizer.zero_grad()。DeepSpeed 在内部处理此操作(通过使用 JSON 配置设置优化器),当调用initialize()时,因此您无需显式编写代码,只需执行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 个纪元的训练。尝试了一组参数(种子和学习率),并选择了最佳参数。所有学习率均为 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 内核

utils.py中已经创建了一个参数--deepspeed_transformer_kernel,我们通过在 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,默认值为DS 用于 DeepSpeed。
  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. HuggingFace 或 TensorFlow 预训练模型无法使用--preln标志,因为它们使用后层归一化。

  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. 使用 PyTorch 和 DeepSpeed Transformer 内核在 NVIDIA V100 (16GB) 上运行 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. 使用 PyTorch 和 DeepSpeed Transformer 内核在 NVIDIA V100 (32GB) 上运行 SQuAD 微调的样本/秒。

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

微批次大小 NVIDIA V100 (32-GB) NVIDIA V100 (16-GB)
> 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 训练 中的配方对预训练的模型进行微调应该会产生 90.5 的 F1 分数,并且如果让预训练时间比教程中建议的时间更长,则预计分数会增加。

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

Dropout 设置

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

对于使用确定性 Transformer 预训练的模型,我们使用预训练中使用的相同 dropout 比率 (0.1)。但是,当对使用随机 Transformer 预训练的模型进行微调时,我们会稍微提高 dropout 比率,以补偿微调期间缺乏随机噪声。

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

更新: