DeepSpeed 变压器内核

本教程展示了如何启用DeepSpeed转换器内核并设置其不同的配置参数。

DeepSpeed 变压器内核

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

先决条件

要使用transformer内核训练模型，您应该按照入门指南将DeepSpeed集成到您的训练脚本中。

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

集成Transformer内核

首先，你需要将transformer内核集成到顶层模型中。这里，我们展示了一个使用Pre-LN BERT-Large配置设置实例化transformer内核的示例。此配置具有24层，隐藏维度为1024，并使用序列长度为128和批量大小为64。为了添加所有这些层，我们在ModuleList中复制相同的层规范num_hidden_layer次，并使用不同的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 内核参数

变压器内核通过多个参数进行配置，这些参数允许用户探索不同的设置。我们将这些参数分为四类：

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

配置变压器内核的通用参数是：

  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: 变压器层数
  8. pre_layer_norm: 选择Pre-LN或Post-LN变压器架构

变压器内核的环境参数包括：

  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_checkpoint 和 gelu_checkpoint 标志指的是检查点方法，在该方法中，我们丢弃了变压器层的某些部分的输入，即注意力丢弃和Gelu，以节省激活内存的重要部分。根据我们的性能分析，重新生成这两者的性能成本可以忽略不计，最终我们从运行更大的批量大小中获得的性能优势弥补了这一点。

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

微批次大小 128 序列长度 512 序列长度
> 12 - attn_dropout_checkpoint
> 16 - normalize_invertible, gelu_checkpoint
> 80 normalize_invertible 内存溢出
> 112 attn_dropout_checkpoint 内存不足
> 128 gelu_checkpoint 内存不足

启用Transformer内核

如前所述，为了使用自定义的DeepSpeed内核运行变压器网络，我们只需要在运行训练脚本时传递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内核标志外，我们还可以指定之前讨论过的内存优化设置。例如，我们在这里使用attention_dropout_checkpoint选项来运行序列长度为512的情况，以便在每个GPU上运行微批次大小为16。如果需要更大的批次大小，我们也可以打开其余的内存优化标志。

更新: