DeepSpeed Transformer Kernel

本教程展示了如何启用 DeepSpeed Transformer Kernel 并设置其不同的配置参数。

DeepSpeed Transformer Kernel

Transformer 层在许多近期序列处理模型（例如自然语言处理）中无处不在。因此，训练基于 Transformer 的网络需要高效的性能，以便科学家能够在合理的时间内探索各种应用领域中的不同模型。为此，我们为 Transformer 网络开发了一个新的 Kernel，其中包含针对这些层的多项特定优化，可以提高单 GPU 上的训练吞吐量，并随着 GPU 数量的增加而良好扩展。有关 Transformer Kernel 细节的更多信息，请访问我们关于最快 BERT 训练的最新博客文章。

先决条件

要使用 Transformer Kernel 训练模型，您应该根据开始使用指南将 DeepSpeed 集成到您的训练脚本中。

注意：目前 DeepSpeed Transformer Kernels 不支持稀疏注意力。要使用稀疏注意力，您需要禁用 Transformer Kernels！

集成 Transformer Kernel

首先，您需要将 Transformer Kernel 集成到顶层模型中。这里，我们展示了一个使用 Pre-LN BERT-Large 配置设置实例化 Transformer Kernel 的示例。此配置有 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 Kernel 参数

Transformer Kernel 由多个参数配置，这些参数允许用户探索不同的设置。我们将这些参数分为四类：

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

配置 Transformer Kernel 的通用参数包括：

  1. batch_size: 在每个 GPU 上运行 Kernel 时使用的微批次大小
  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 Kernel 的环境参数包括：

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

高性能优化标志

  1. stochastic_mode: 开启此标志后，训练平均可以快 2%。请注意，此标志具有一定程度的非确定性，并且在不同运行中可能产生不同的结果。然而，我们发现启用它并不会影响 BERT 等预训练任务，并且可以获得较高的准确性。另一方面，对于下游任务（例如微调），我们建议将其关闭，以便通过常规 Kernel 执行来重现相同结果。

内存优化标志包括：

  1. attn_dropout_checkpoint: 启用注意力 dropout 的检查点以节省内存
  2. normalize_invertible: 启用可逆 LayerNorm 执行（丢弃输入激活）
  3. gelu_checkpoint: 启用 Gelu 激活输出的检查点以节省内存

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

内存优化标志

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

通过设置 normalize_invertible 标志，我们强制 Kernel 丢弃 Transformer 归一化层的输入激活。我们可以这样做，因为 Kernel 包含一项优化，仅使用输出激活即可计算此层的参数梯度和输入梯度。

attn_dropout_checkpoint 和 gelu_checkpoint 标志指的是检查点方法，其中我们丢弃 Transformer 层某些部分（注意力 dropout 和 Gelu）的输入，以节省激活内存的重要部分。根据我们的性能分析，重新物化这两者的性能成本可以忽略不计，最终我们从运行更大批次大小中获得的性能收益弥补了这一成本。

下表显示了在具有 32GB 内存的 NVIDIA V100 GPU 上运行 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 Kernel

如前所述，为了使用自定义 DeepSpeed Kernel 运行 Transformer 网络，我们只需在运行训练脚本时传递 deepspeed_transformer_kernel 选项即可。下面，我们展示了除了 BERT 预训练任务的其余参数外，如何将此参数传递给 deepspeed 启动器。

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 Kernel 标志外，我们还可以指定前面讨论的内存优化设置。例如，我们在这里使用 attention_dropout_checkpoint 选项来运行序列长度为 512 的任务，以便在每个 GPU 上运行微批次大小为 16 的任务。如果需要更大的批次大小，我们也可以开启其余的内存优化标志。

更新日期：