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 (Bidirectional Encoder Representations from Transformers)， 它广泛用于许多自然语言处理（NLP）任务。BERT的详细信息可以在这里找到：BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding。
我们将介绍如何设置数据管道以及如何运行原始的BERT模型。然后，我们将逐步展示如何修改模型以利用DeepSpeed。最后，我们将展示使用DeepSpeed后的性能评估和内存使用减少情况。
不使用DeepSpeed预训练Bing BERT
我们基于 huggingface/transformers 和 NVIDIA/DeepLearningExamples 的改编进行工作。 我们已经在 DeepSpeedExamples/bing_bert 下分叉了这个仓库， 并在他们的脚本中进行了几处修改：
- 我们采用了NVIDIA的BERT建模代码，位于
bing_bert/nvidia/下。 - 我们扩展了来自Project Turing的数据管道，位于
bing_bert/turing/下。
训练数据设置
注意： 下载和预处理的说明即将发布。
下载Wikipedia和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: 训练参数和检查点保存/加载工具
参数解析
我们首先需要在train.py中添加DeepSpeed的参数解析，使用deepspeed.add_config_arguments()。这一步允许应用程序识别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中，我们在checkpoint_model()函数中使用DeepSpeed的检查点API，如下所示，我们收集客户端模型状态并通过调用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: 使用有效批量大小为4096train_micro_batch_size_per_gpu: 每个GPU有足够的内存来即时适应64的批量大小optimizer: 使用Adam训练优化器fp16: 启用FP16混合精度训练，初始损失缩放因子为2^16。
就是这样！这就是你在使用DeepSpeed时需要做的所有修改。我们已经包含了一个修改后的train.py文件，名为DeepSpeedExamples/bing_bert/deepspeed_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 转换器内核实例化转换器层。
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 Kernel和DeepSpeed Fast-Bert Training。
开始训练
一个在四个节点上启动deepspeed_train.py的示例，每个节点有四个GPU，如下所示：
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训练结果
我们在保持行业竞争力的同时，实现了最快的BERT训练时间，在SQUAD 1.1开发集上达到了90.5或更高的F1分数。请按照BERT微调教程来微调由transformer内核预训练的模型，并重现SQUAD F1分数。
- 我们使用1024个V100 GPU（64个NVIDIA DGX-2节点）在44分钟内完成了BERT预训练。相比之下，NVIDIA之前的SOTA使用1472个V100 GPU需要47分钟。DeepSpeed不仅更快，而且使用的资源减少了30%。使用相同的1024个GPU，NVIDIA BERT比DeepSpeed慢52%，需要67分钟来训练。
- 与谷歌原始的BERT训练时间相比，他们在64个TPU2芯片上大约需要96小时才能达到同等水平，而我们在4个DGX-2节点上的64个V100 GPU上训练不到9小时。
- 在256个GPU上，我们花费了2.4小时，比NVIDIA使用相同数量GPU的超级集群获得的最新结果（3.9小时）更快（link）。
| 节点数量 | 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模型的单GPU性能，如BERT。上图展示了通过DeepSpeed优化的BertBERT-Large的单GPU训练吞吐量，与两个著名的Pytorch实现——NVIDIA BERT和HuggingFace BERT进行了比较。对于序列长度为128和512的情况，DeepSpeed分别达到了64和53 teraflops的吞吐量（对应于每秒272和52个样本），比NVIDIA BERT提高了高达28%的吞吐量，比HuggingFace BERT提高了高达62%。我们还支持高达1.8倍的批量大小，而不会耗尽内存。
有关我们如何实现破纪录的BERT训练时间的更多详细信息，请查看深入探讨DeepSpeed BERT 最快的BERT训练