DeepSpeed Transformer 内核
本教程演示如何启用 DeepSpeed Transformer 内核并设置其不同的配置参数。
DeepSpeed Transformer 内核
Transformer 层在许多最新的序列处理模型(如自然语言处理)中无处不在。因此,训练基于 Transformer 的网络需要在性能方面具有很高的效率,以便科学家能够在合理的时间内探索各种应用领域的不同模型。为此,我们开发了一种新的 Transformer 网络内核,其中包含针对这些层的若干特定优化,这些优化可以提高单 GPU 上的训练吞吐量,并且随着 GPU 数量的增加而良好扩展。有关 Transformer 内核详细信息的更多信息,请访问我们最近关于最快 BERT 训练的博客文章。
先决条件
要将 Transformer 内核用于模型训练,您应该使用入门指南将 DeepSpeed 集成到您的训练脚本中。
注意:目前 DeepSpeed Transformer 内核不支持稀疏注意力。要使用稀疏注意力,您需要禁用 Transformer 内核!
集成 Transformer 内核
首先,您需要将 Transformer 内核集成到顶级模型中。在这里,我们展示了一个使用 Pre-LN BERT-Large 配置设置实例化 Transformer 内核的示例。此配置具有 24 层,隐藏维度为 1024,并使用 128 的序列长度和 64 的批次大小。要添加所有这些层,我们在 ModuleList 中使用不同的 ID 复制相同的层规范num_hidden_layer
次。
config = DeepSpeedTransformerConfig(batch_size = 64,
max_seq_length = 128,
hidden_size = 1024,
heads = 16,
attn_dropout_ratio = 0.1,
hidden_dropout_ratio = 0.1,
num_hidden_layers = 24,
initializer_range = 0.02,
local_rank = 0,
seed = 1234,
fp16 = True,
pre_layer_norm=True,
attn_dropout_checkpoint=False,
normalize_invertible=False,
gelu_checkpoint=False)
self.layer = nn.ModuleList([
copy.deepcopy(DeepSpeedTransformerLayer(cuda_config))
for _ in range(config.num_hidden_layers)
])
Transformer 内核参数
Transformer 内核由许多参数配置,允许用户探索不同的设置。我们将这些参数分为四类
- 通用配置,由不同类型的 Transformer 层使用
- 环境参数,指定系统的设置
- 高性能标志,使用随机计算优化训练
- 内存优化标志,权衡计算能力以换取内存
配置 Transformer 内核的通用参数为
batch_size
:用于在每个 GPU 上运行内核的微批次大小max_seq_length
:使用 DeepSpeed 训练的模型的序列长度hidden_size
:Transformer 层的隐藏大小heads
:Transformer 层的自注意力中的头部数量attn_dropout_ratio
:注意力的输出的 dropout 比率hidden_dropout_ratio
:Transformer 输出的 dropout 比率num_hidden_layers
:Transformer 层的数量pre_layer_norm
:在 Pre-LN 或 Post-LN Transformer 架构之间进行选择
Transformer 内核的环境参数包括
local_rank
:运行 Transformer 内核的当前 GPU 的排名seed
:dropout 层的随机种子fp16
:启用半精度计算initializer_range
:BERT 的初始化范围
高性能优化标志
stochastic_mode
:通过打开此标志,训练平均可以加快 2%。请注意,此标志具有一定的非确定性,并且在不同的运行中可能会产生不同的结果。但是,我们已经看到,通过启用它,BERT 等预训练任务不受影响,并且可以获得较高的准确率。另一方面,对于下游任务(例如微调),我们建议将其关闭,以便能够通过常规内核执行来重现相同的结果。
内存优化标志包括
attn_dropout_checkpoint
:启用注意力 dropout 的检查点以节省内存normalize_invertible
:启用可逆 LayerNorm 执行(丢弃输入激活)gelu_checkpoint
:启用 Gelu 激活输出的检查点以节省内存
为了说明在模型训练中使用 Transformer 内核所需的模型配置更改,我们使用 BERT 模型并逐步介绍不同的配置,以支持不同的序列长度和批次大小。请参阅BERT 训练教程中的说明。
内存优化标志
我们在 Transformer 内核中提供了多种技术,这些技术可以在层的不同部分节省内存。我们将它们作为可配置设置公开,可以在调用内核时启用它们。通过打开这些优化标志中的每一个,我们可以支持更大的批次大小。即使我们使用其中一些技术以性能换取内存,但通过使用更大的批次大小,端到端训练效率也会提高。
通过设置normalize_invertible
标志,我们强制内核丢弃对 Transformer 的归一化层的输入激活。我们可以这样做,因为内核包含一个优化,可以通过仅使用输出激活来计算参数和此层的输入的梯度。
attn_dropout_checkpoint
和gelu_checkpoint
标志指的是检查点方法,其中我们丢弃对 Transformer 层、注意力 dropout 和 Gelu 的某些部分的输入,以节省重要的激活内存部分。根据我们的性能分析,重新生成这两个部分的性能成本可以忽略不计,并且最终我们从运行更大批次大小中获得的性能优势弥补了这一点。
下表显示了在 NVIDIA V100 GPU(具有 32GB 内存)上运行 BERT-Large 时需要打开哪些内存优化标志,并考虑了不同的微批次大小和序列长度。对于我们在实验中使用的两种序列长度 128 和 512,我们已经看到更大的批次大小可以改善两者的整体训练性能。有关这些配置的性能评估的更多信息,请参阅我们的博客文章。
微批次大小 | 128 序列长度 | 512 序列长度 |
---|---|---|
> 12 | - | attn_dropout_checkpoint |
> 16 | - | normalize_invertible ,gelu_checkpoint |
> 80 | normalize_invertible |
OOM |
> 112 | attn_dropout_checkpoint |
OOM |
> 128 | gelu_checkpoint |
OOM |
启用 Transformer 内核
如前所述,为了使用自定义 DeepSpeed 内核运行 Transformer 网络,我们只需要在运行训练脚本时传递deepspeed_transformer_kernel
选项即可。下面,我们展示了一个将此参数传递给deepspeed
启动器的示例,除了 BERT 预训练任务的其他参数之外。
deepspeed deepspeed_train.py \
--cf bert_large_lamb.json \
--max_seq_length 512 \
--print_steps 100 \
--deepspeed \
--deepspeed_transformer_kernel \
--deepspeed_config deepspeed_bsz32K_lamb_config_seq512.json \
--rewarmup \
--lr_schedule "EE" \
--lr_offset 0.0 \
--attention_dropout_checkpoint \
--load_training_checkpoint ${CHECKPOINT_BASE_PATH} \
--load_checkpoint_id ${CHECKPOINT_EPOCH150_NAME}
除了 Transformer 内核标志之外,我们还可以指定前面讨论的内存优化设置。例如,我们在这里使用attention_dropout_checkpoint
选项运行序列长度 512,以便在每个 GPU 上运行 16 的微批次大小。如果需要更大的批次大小,我们也可以打开其余的内存优化标志。