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 的批次大小。要添加所有这些层,我们将相同的层规范复制 num_hidden_layer 次,并在 ModuleList 中使用不同的 ID。

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 内核由许多参数配置,允许用户探索不同的设置。我们将这些参数划分为四个类别

  1. 通用配置,供不同类型的 Transformer 层使用
  2. 环境参数,指定系统的设置
  3. 高性能标志,使用随机计算优化训练
  4. 内存优化标志,权衡计算能力以换取内存

用于配置 Transformer 内核的通用参数是

  1. batch_size:在每个 GPU 上运行内核时使用的微批次大小
  2. max_seq_length:使用 DeepSpeed 训练的模型的序列长度
  3. hidden_size:Transformer 层的隐藏大小
  4. heads:Transformer 层自注意力中的头数
  5. attn_dropout_ratio:注意输出的 dropout 比率
  6. hidden_dropout_ratio:Transformer 输出的 dropout 比率
  7. num_hidden_layers:Transformer 层数
  8. pre_layer_norm:在 Pre-LN 或 Post-LN Transformer 架构之间选择

Transformer 内核的环境参数包括

  1. local_rank:当前运行 Transformer 内核的 GPU 的排名
  2. seed:dropout 层的随机种子
  3. fp16:启用半精度计算
  4. initializer_range:BERT 的初始化范围

高性能优化标志

  1. stochastic_mode:通过启用此标志,训练速度平均可以快 2%。请注意,此标志具有一定程度的非确定性,并且可能会在不同的运行中产生不同的结果。但是,我们已经看到,通过启用它,BERT 等预训练任务不受影响,并且可以获得很高的准确率水平。另一方面,对于下游任务,例如微调,我们建议将其关闭,以便能够通过常规内核执行来重现相同的结果。

内存优化标志包括

  1. attn_dropout_checkpoint:启用注意力 dropout 的检查点以节省内存
  2. normalize_invertible:启用可逆层归一化执行(丢弃输入激活)
  3. gelu_checkpoint:启用 Gelu 激活输出的检查点以节省内存

为了说明在模型训练中使用 Transformer 内核所需的模型配置更改,我们使用 BERT 模型并逐步介绍不同的配置,以支持不同的序列长度和批次大小。请参阅BERT 训练教程中的说明。

内存优化标志

我们在 Transformer 内核中提供了几种技术,可以在层的不同部分节省内存。我们将它们作为可配置设置公开,可以在调用内核时启用。通过启用这些优化标志中的每一个,我们可以支持更大的批次大小。即使我们使用其中一些技术来权衡性能以换取内存,但通过使用更大的批次大小,端到端的训练效率也会提高。

通过设置normalize_invertible标志,我们强制内核丢弃对 Transformer 的规范化层的输入激活。我们可以这样做,因为内核包含一个优化来计算参数的梯度和该层的输入,只需使用输出激活。

attn_dropout_checkpointgelu_checkpoint标志是指检查点方法,在该方法中,我们丢弃对 Transformer 层某些部分的输入,即注意力 dropout 和 Gelu,以节省重要部分的激活内存。根据我们的性能分析,重新实现这两个部分的性能成本可以忽略不计,最终,我们在运行更大批次大小时获得的性能优势弥补了这一点。

下表显示了在 NVIDIA V100 GPU(具有 32GB 内存)上运行 BERT-Large 时,需要启用哪些内存优化标志,并考虑了不同的微批次大小和序列长度。对于我们实验中使用的两种序列长度 128 和 512,我们已经看到更大的批次大小提高了这两种情况下的整体训练性能。有关这些配置的性能评估的更多信息,请参阅我们的博客文章

微批次大小 128 序列长度 512 序列长度
> 12 - attn_dropout_checkpoint
> 16 - normalize_invertiblegelu_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 的微批次大小。如果需要更大的批次大小,我们也可以打开其他内存优化标志。

更新: