使用渐进式层丢弃加速基于Transformer的语言模型训练

在本教程中，我们将介绍DeepSpeed中的渐进式层丢弃 (PLD) 并提供如何使用PLD的示例。PLD允许在相同样本数量下将BERT等Transformer网络训练速度提高24%，并在下游任务上以2.5倍的速度获得相似的准确性。PLD的详细描述和实验结果可在我们的技术报告中查阅。

为了说明如何在DeepSpeed中使用PLD，我们展示了如何启用PLD来预训练BERT模型，以及如何在GLUE数据集上微调预训练模型。

使用DeepSpeed和PLD运行预训练

要执行预训练，首先需要准备数据集。关于这部分，请参阅我们的BERT预训练文章，其中包含有关数据下载和预处理的详细信息。对于以下实验，我们使用Wikipedia文本和Bookcorpus，类似于Devlin 等人的做法。

预训练的主要部分在deepspeed_train.py中完成，该文件已修改为使用DeepSpeed。ds_train_bert_progressive_layer_drop_bsz4k_seq128.sh是启动DeepSpeed和PLD预训练的shell脚本。

bash ds_train_bert_progressive_layer_drop_bsz4k_seq128.sh

如果您已经学习过BERT预训练教程，那么上述脚本中的大多数标志应该很熟悉。要启用PLD训练，需要在客户端脚本和DeepSpeed引擎中都启用PLD。要在客户端脚本中启用PLD，需要添加以下命令行标志以在Transformer块上启用渐进式层丢弃。

--progressive_layer_drop

要在DeepSpeed中启用PLD，需要使用适当的PLD配置字典更新json配置文件，如下所示：

{
  ...
  "progressive_layer_drop": {
    "enabled": true,
    "theta": 0.5,
    "gamma": 0.001
  }
}

我们推荐PLD的theta值为0.5，gamma值为0.001，因为这些值在我们的实验中效果良好。

进行这些配置更改后，DeepSpeed引擎应打印如下运行时消息：

[INFO] [logging.py:60:log_dist] [Rank 0] Enabled progressive layer dropping (theta = 0.5)

deepspeed_bsz4k_progressive_layer_drop_config_seq128.json文件允许用户指定DeepSpeed选项，包括批大小、微批大小、优化器、学习率、序列长度及其他参数。下面是我们用于运行BERT和PLD的DeepSpeed配置文件。

{
  "train_batch_size": 4096,
  "train_micro_batch_size_per_gpu": 16,
  "steps_per_print": 1000,
  "prescale_gradients": true,
  "gradient_predivide_factor": 8,
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 1e-3,
      "weight_decay": 0.01,
      "bias_correction": false
    }
  },
  "gradient_clipping": 1.0,

  "wall_clock_breakdown": false,

  "fp16": {
    "enabled": true,
    "loss_scale": 0
  },

  "progressive_layer_drop": {
    "enabled": true,
    "theta": 0.5,
    "gamma": 0.001
  }
}

请注意，上述配置假定在64块32GB V100 GPU上进行训练。每块GPU使用16的微批大小，并累积梯度直到有效批大小达到4096。如果您的GPU内存较少，可能需要减小“train_micro_batch_size_per_gpu”。或者，如果您有更多GPU，可以增加“train_batch_size”以提高训练速度。我们为启用PLD的BERT预训练使用以下超参数。

参数 值
有效批大小 4K
每GPU训练微批大小 16
优化器 Adam
峰值学习率 1e-3
序列长度 128
学习率调度器 预热线性衰减指数
预热比例 0.02
衰减率 0.99
衰减步长 1000
权重衰减 0.01
梯度裁剪 1.0

表1. 预训练超参数

注意：DeepSpeed现在支持将PreLayerNorm作为BERT训练的默认方式，因为它能够避免梯度消失、稳定优化并提高性能，如我们关于最快BERT训练的博客文章所述。因此，我们直接在带有PreLayerNorm的BERT上支持可切换的Transformer块。实现代码可在“example\bing_bert\nvidia\modelingpreln_layerdrop.py”中找到。

在GLUE任务上使用DeepSpeed进行微调

我们使用GLUE进行微调任务。GLUE（通用语言理解评估基准）(https://gluebenchmark.com/) 是一系列句子或句子对的自然语言理解任务集合，包括问答、情感分析和文本蕴含。它旨在促进在不同领域内各种语言任务上的样本高效学习和知识迁移。

可以使用提供的辅助脚本下载所有GLUE数据。数据下载完成后，可以设置数据并将其移动到“/data/GlueData”，这是存储GLUE数据的默认位置。然后，我们可以使用PLD预训练的BERT模型检查点来运行微调。

微调的主要部分在run_glue_classifier_bert_base.py中完成，该文件已修改为使用DeepSpeed。在微调之前，需要通过run_glue_classifier_bert_base.py中的以下配置指定BERT模型配置。在这种情况下，它已被修改为与预训练模型的配置相同。

    bert_model_config = {
        "vocab_size_or_config_json_file": 119547,
        "hidden_size": 768,
        "num_hidden_layers": 12,
        "num_attention_heads": 12,
        "intermediate_size": 3072,
        "hidden_act": "gelu",
        "hidden_dropout_prob": 0.1,
        "attention_probs_dropout_prob": 0.1,
        "max_position_embeddings": 512,
        "type_vocab_size": 2,
        "initializer_range": 0.02
    }

接下来，可以使用以下命令加载DeepSpeed风格的检查点，该命令也已添加到脚本中。

model.load_state_dict(checkpoint_state_dict['module'], strict=False)

最后，run_glue_classifier_bert_base.sh脚本调用预训练并设置了几个与微调相关的超参数。

bash run_glue_bert_base_finetune.sh [task] [batch size] [learning rate] [number of epochs] [job name] [checkpoint path]

例如：

bash run_glue_bert_base_finetune.sh MNLI 32 3e-5 5 "fine_tune_MNLI" deepspeed_checkpoint.pt

预期结果

微调结果可以在“logs”目录下找到，下面是PLD在GLUE任务上的预期结果。“学习率”行表示我们为每个任务获得相应准确性结果所使用的学习率。

  RTE MRPC STS-B CoLA SST-2 QNLI QQP MNLI-m/mm GLUE
指标 准确率 F1/准确率 PCC/SCC 准确率 准确率 准确率 F1/准确率 准确率  
Bert_{base} (原始) 66.4 88.9/84.8 87.1/89.2 52.1 93.5 90.5 71.2/89.2 84.6/83.4 80.7
Bert_{base} (我们的实现) 67.8 88.0/86.0 89.5/89.2 52.5 91.2 87.1 89.0/90.6 82.5/83.4 82.1
PLD 69.3 86.6/84.3 90.0/89.6 55.8 91.6 90.7 89.6/91.2 84.1/83.8 82.9
学习率 7e-5 9e-5 7e-5 5e-5 7e-5 9e-5 2e-4 3e-5  

更新于: