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

  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:启用可逆 LayerNorm 执行(丢弃输入激活)
  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 的微批次大小。如果需要更大的批次大小,我们也可以打开其余的内存优化标志。

更新: