BERT 预训练
注意: 2022 年 8 月 15 日，我们在 github.com/microsoft/Megatron-DeepSpeed/tree/main/examples_deepspeed/bert_with_pile 添加了另一个 BERT 预训练/微调示例，其中包含一个描述如何使用它的 README.md。与下面描述的示例相比，Megatron-DeepSpeed 中的新示例添加了对 ZeRO 和张量切片模型并行的支持（因此支持更大规模的模型），使用了一个公开且更丰富的 Pile 数据集（用户也可以使用自己的数据），以及 这篇论文 中描述的模型架构和训练超参数的一些更改。因此，新示例训练的 BERT 模型能够提供比原始 BERT 更好的 MNLI 结果，但模型架构略有不同，计算量更大。如果要训练更大规模或更高质量的 BERT 风格模型，建议您遵循 Megatron-DeepSpeed 中的新示例。如果您的目标是严格地复制原始 BERT 模型，建议您按照下面描述的 DeepSpeedExamples/bing_bert 中的示例进行操作。另一方面，下面的教程有助于解释如何将 DeepSpeed 集成到预训练代码库中，无论您使用哪个 BERT 示例。
在本教程中，我们将 DeepSpeed 应用于 BERT（来自 Transformer 的双向编码器表示）的预训练，它广泛用于许多自然语言处理 (NLP) 任务。有关 BERT 的详细信息，请参见此处：BERT：用于语言理解的深度双向 Transformer 的预训练。
我们将介绍如何设置数据管道以及如何运行原始 BERT 模型。然后，我们将逐步展示如何修改模型以利用 DeepSpeed。最后，我们将演示使用 DeepSpeed 的性能评估和内存使用量减少。
在没有 DeepSpeed 的情况下预训练 Bing BERT
我们使用来自 huggingface/transformers 和 NVIDIA/DeepLearningExamples 的改编版本。我们在 DeepSpeedExamples/bing_bert 下分叉了这个仓库，并在他们的脚本中进行了一些修改。
- 我们从 NVIDIA 的 BERT 在
bing_bert/nvidia/下的模型代码中借鉴。 - 我们从 Project Turing 在
bing_bert/turing/下的扩展数据管道。
训练数据设置
注意:下载和预处理说明即将推出。
下载维基百科和 BookCorpus 数据集，并在模型配置文件 DeepSpeedExamples/bing_bert/bert_large_adam_seq128.json 中指定其路径。
{
...
"datasets": {
"wiki_pretrain_dataset": "/data/bert/bnorick_format/128/wiki_pretrain",
"bc_pretrain_dataset": "/data/bert/bnorick_format/128/bookcorpus_pretrain"
},
...
}
运行 Bing BERT 模型
从 DeepSpeedExamples/bing_bert 运行
python train.py \
--cf bert_large_adam_seq128.json \
--train_batch_size 64 \
--max_seq_length 128 \
--gradient_accumulation_steps 1 \
--max_grad_norm 1.0 \
--fp16 \
--loss_scale 0 \
--delay_allreduce \
--max_steps 10 \
--output_dir <path-to-model-output>
启用 DeepSpeed
要使用 DeepSpeed，我们需要编辑两个文件
train.py：训练的主要入口点utils.py：训练参数和检查点保存/加载实用程序
参数解析
首先，我们需要使用 deepspeed.add_config_arguments() 将 DeepSpeed 的参数解析添加到 train.py 中。此步骤允许应用程序识别 DeepSpeed 特定的配置。
def get_arguments():
parser = get_argument_parser()
# Include DeepSpeed configuration arguments
parser = deepspeed.add_config_arguments(parser)
args = parser.parse_args()
return args
初始化和训练
我们修改 train.py 以启用使用 DeepSpeed 训练。
初始化
我们使用 deepspeed.initialize() 来创建模型、优化器和学习率调度器。对于 Bing BERT 模型，我们在下面的 prepare_model_optimizer() 函数中初始化 DeepSpeed，以传递原始模型和优化器（从命令选项指定）。
def prepare_model_optimizer(args):
# Loading Model
model = BertMultiTask(args)
# Optimizer parameters
optimizer_parameters = prepare_optimizer_parameters(args, model)
model.network, optimizer, _, _ = deepspeed.initialize(args=args,
model=model.network,
model_parameters=optimizer_parameters,
dist_init_required=False)
return model, optimizer
请注意，对于 Bing BERT，原始模型保留在 model.network 中，因此我们将 model.network 作为参数传递，而不是只传递模型。
训练
由 deepspeed.initialize 返回的 model 是 DeepSpeed 模型引擎，我们将使用它来使用前向、反向和步骤 API 训练模型。由于模型引擎公开了与 nn.Module 对象相同的正向传递 API，因此正向传递没有变化。因此，我们只修改反向传递和优化器/调度器步骤。
反向传播是通过直接使用模型引擎调用 backward(loss) 来执行的。
# Compute loss
if args.deepspeed:
model.network.backward(loss)
else:
if args.fp16:
optimizer.backward(loss)
else:
loss.backward()
DeepSpeed 引擎中的 step() 函数更新模型参数以及学习率。在每一步更新权重后，DeepSpeed 会自动处理梯度清零。
if args.deepspeed:
model.network.step()
else:
optimizer.step()
optimizer.zero_grad()
检查点保存和加载
DeepSpeed 的模型引擎具有灵活的 API，用于保存和加载检查点，以处理客户端模型状态及其自身的内部状态。
def save_checkpoint(self, save_dir, tag, client_state={})
def load_checkpoint(self, load_dir, tag)
在 train.py 中，我们使用 DeepSpeed 的检查点 API 在下面的 checkpoint_model() 函数中，我们收集客户端模型状态，并通过调用 save_checkpoint() 将它们传递给模型引擎。
def checkpoint_model(PATH, ckpt_id, model, epoch, last_global_step, last_global_data_samples, **kwargs):
"""Utility function for checkpointing model + optimizer dictionaries
The main purpose for this is to be able to resume training from that instant again
"""
checkpoint_state_dict = {'epoch': epoch,
'last_global_step': last_global_step,
'last_global_data_samples': last_global_data_samples}
# Add extra kwargs too
checkpoint_state_dict.update(kwargs)
success = model.network.save_checkpoint(PATH, ckpt_id, checkpoint_state_dict)
return
在 load_training_checkpoint() 函数中，我们使用 DeepSpeed 的加载检查点 API 并返回客户端模型的状态。
def load_training_checkpoint(args, model, PATH, ckpt_id):
"""Utility function for checkpointing model + optimizer dictionaries
The main purpose for this is to be able to resume training from that instant again
"""
_, checkpoint_state_dict = model.network.load_checkpoint(PATH, ckpt_id)
epoch = checkpoint_state_dict['epoch']
last_global_step = checkpoint_state_dict['last_global_step']
last_global_data_samples = checkpoint_state_dict['last_global_data_samples']
del checkpoint_state_dict
return (epoch, last_global_step, last_global_data_samples)
DeepSpeed JSON 配置文件
使用 DeepSpeed 的最后一步是创建一个配置 JSON 文件（例如，deepspeed_bsz4096_adam_config.json）。此文件提供用户定义的 DeepSpeed 特定参数，例如，每个 GPU 的批次大小、优化器及其参数以及是否启用 FP16 训练。
{
"train_batch_size": 4096,
"train_micro_batch_size_per_gpu": 64,
"steps_per_print": 1000,
"optimizer": {
"type": "Adam",
"params": {
"lr": 2e-4,
"max_grad_norm": 1.0,
"weight_decay": 0.01,
"bias_correction": false
}
},
"fp16": {
"enabled": true,
"loss_scale": 0,
"initial_scale_power": 16
}
}
特别是，此示例 json 为 DeepSpeed 指定了以下配置参数。
train_batch_size：使用 4096 的有效批次大小train_micro_batch_size_per_gpu：每个 GPU 都有足够的内存来立即适应 64 的批次大小。optimizer：使用 Adam 训练优化器。fp16：启用 FP16 混合精度训练，初始损失缩放因子为 2^16。
就这样！为了使用 DeepSpeed，您只需要进行这些修改。我们已经包含了一个名为 DeepSpeedExamples/bing_bert/deepspeed_train.py 的修改后的 train.py 文件，其中应用了所有更改。
启用 DeepSpeed 的 Transformer 内核
要启用 Transformer 内核以获得更高的性能，首先在 utils.py 中添加参数 --deepspeed_transformer_kernel，默认情况下，我们可以将其设置为 False，以便于打开/关闭。
parser.add_argument('--deepspeed_transformer_kernel',
default=False,
action='store_true',
help='Use DeepSpeed transformer kernel to accelerate.')
然后，在模型源文件中的 BertEncoder 类中，使用 DeepSpeed Transformer 内核实例化 Transformer 层，如下所示。
if args.deepspeed_transformer_kernel:
from deepspeed import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig, DeepSpeedConfig
if hasattr(args, 'deepspeed_config') and args.deepspeed_config:
ds_config = DeepSpeedConfig(args.deepspeed_config)
else:
raise RuntimeError('deepspeed_config is not found in args.')
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,
local_rank = args.local_rank if hasattr(args, 'local_rank') else -1,
seed = args.seed,
fp16 = ds_config.fp16_enabled,
pre_layer_norm=True,
attn_dropout_checkpoint=args.attention_dropout_checkpoint,
normalize_invertible=args.normalize_invertible,
gelu_checkpoint=args.gelu_checkpoint,
stochastic_mode=True)
layer = DeepSpeedTransformerLayer(cuda_config)
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，而预测使用 12 的批次大小，我们可以将 12 用作 Transformer 内核批次大小，或者使用 “–predict_batch_size” 参数将预测批次大小设置为 8 或更小的数字。local_rank在 DeepSpeedTransformerConfig 中用于将 Transformer 内核分配到正确的设备。由于模型已经在之前运行了 set_device()，因此这里不需要设置。stochastic_mode在启用时具有更高的性能，我们在预训练中启用它，并在微调中禁用它。- Transformer 内核有自己的参数，因此使用 Transformer 内核生成的检查点文件必须由启用了 Transformer 内核的模型加载（例如，在微调中）。
有关 Transformer 内核的更多详细信息，请参见 DeepSpeed Transformer 内核 和 DeepSpeed Fast-Bert 训练。
开始训练
在四个节点上（每个节点有四个 GPU）启动 deepspeed_train.py 的示例将是
deepspeed --num_nodes 4 \
deepspeed_train.py \
--deepspeed \
--deepspeed_config deepspeed_bsz4096_adam_config.json \
--cf /path-to-deepspeed/examples/tests/bing_bert/bert_large_adam_seq128.json \
--train_batch_size 4096 \
--max_seq_length 128 \
--gradient_accumulation_steps 4 \
--max_grad_norm 1.0 \
--fp16 \
--loss_scale 0 \
--delay_allreduce \
--max_steps 32 \
--print_steps 1 \
--deepspeed_transformer_kernel \
--output_dir <output_directory>
有关启动 DeepSpeed 的更多信息，请参见 入门 指南。
使用 DeepSpeed 再现最快的 BERT 训练结果
我们在 SQUAD 1.1 开发集上保持竞争力，同时实现了 90.5 或更好的 F1 分数，从而实现了最快的 BERT 训练时间。请按照 BERT 微调 教程微调您的模型，该模型由 Transformer 内核预训练，并重现 SQUAD F1 分数。
- 我们使用 1024 个 V100 GPU（64 个 NVIDIA DGX-2 节点）在 44 分钟内完成了 BERT 预训练。相比之下，NVIDIA 的先前 SOTA 使用 1472 个 V100 GPU 耗时 47 分钟。DeepSpeed 不仅更快，而且使用的资源也减少了 30%。使用相同的 1024 个 GPUS，NVIDIA BERT 比 DeepSpeed 慢 52%，训练时间为 67 分钟。
- 与 Google 的原始 BERT 训练时间相比，它在 64 个 TPU2 芯片上达到同等水平大约需要 96 小时，我们在 4 个 DGX-2 节点（64 个 V100 GPU）上训练时间不到 9 小时。
- 在 256 个 GPU 上，我们花了 2.4 小时，比 NVIDIA 使用相同数量的 GPU 在其超级节点上获得的最新结果（3.9 小时）快 (链接)。
| 节点数 | V100 GPU 数 | 时间 |
|---|---|---|
| 1 个 DGX-2 | 16 | 33 小时 13 分钟 |
| 4 个 DGX-2 | 64 | 8 小时 41 分钟 |
| 16 个 DGX-2 | 256 | 144 分钟 |
| 64 个 DGX-2 | 1024 | 44 分钟 |
我们上面提到的 BERT 训练结果的配置可以使用我们 DeepSpeedExamples 仓库中的脚本/json 配置进行重现。下表汇总了配置。具体来说，请参阅 DeepSpeedExamples 中的 ds_train_bert_bsz64k_seq128.sh 和 ds_train_bert_bsz32k_seq512.sh 脚本以获取更多详细信息。
| 参数 | 128 个序列 | 512 个序列 |
|---|---|---|
| 总批次大小 | 64K | 32K |
| 每个 GPU 的训练微批次大小 | 64 | 8 |
| 优化器 | LAMB | LAMB |
| 学习率 | 11e-3 | 2e-3 |
初始学习率 (lr_offset) |
10e-4 | 0.0 |
| 最小 Lamb 系数 | 0.01 | 0.01 |
| 最大 Lamb 系数 | 0.3 | 0.3 |
| 学习率调度器 | warmup_exp_decay_exp |
warmup_exp_decay_exp |
| 预热比例 | 0.02 | 0.02 |
| 衰减率 | 0.90 | 0.90 |
| 衰减步长 | 250 | 150 |
| 最大训练步数 | 7500 | 7500 |
| 重新预热学习率 | N/A | True |
| 输出检查点编号 | 150 | 160-162 |
| 样本计数 | 403M | 18-22M |
| 轮次计数 | 150 | 160-162 |
DeepSpeed 单 GPU 吞吐量结果
与 SOTA 相比，DeepSpeed 显著提高了基于 Transformer 的模型（如 BERT）的单 GPU 性能。上图显示了通过 DeepSpeed 优化的 BertBERT-Large 训练的单 GPU 吞吐量，与两个知名的 Pytorch 实现（NVIDIA BERT 和 HuggingFace BERT）进行了比较。DeepSpeed 在序列长度为 128 和 512 时分别达到了高达 64 和 53 teraflops 的吞吐量（对应于每秒 272 和 52 个样本），比 NVIDIA BERT 高出 28% 的吞吐量，比 HuggingFace BERT 高出 62%。我们还支持高达 1.8 倍的更大批次大小，而不会出现内存不足问题。
有关我们如何实现创纪录的 BERT 训练时间的更多详细信息，请查看深入了解 DeepSpeed BERT 最快 BERT 训练