DeepSpeed 稀疏注意力

在本教程中,我们将介绍如何使用 DeepSpeed 稀疏注意力 (SA) 及其构建块内核。使用 SA 的最简单方法是通过 DeepSpeed 启动器。我们将在 如何使用 DeepSpeed 启动器使用稀疏注意力 部分中通过示例介绍这一点。但在此之前,我们将在 下一部分 中介绍 DeepSpeed SA 提供的模块。

注意:目前,DeepSpeed 稀疏注意力只能在使用 Torch >= 1.6 和 CUDA 10.1、10.2、11.0 或 11.1 的 NVIDIA V100 或 A100 GPU 上使用。

稀疏注意力模块

  • MatMul:此模块处理块稀疏矩阵乘法。目前,它支持 SDD、DSD 和 DDS,如 DeepSpeed 稀疏注意力 部分所述。
  • Softmax:此模块应用块稀疏 softmax。它处理正向和反向传递。
  • SparseSelfAttention:此模块使用 MatMul 和 Softmax 内核,并根据查询、键和值生成上下文层输出。它是任何自注意力层中常见操作的简化版本。它还可以应用
    • 相对位置嵌入
    • 注意力掩码
    • 键填充掩码 在中间注意力分数上。有关自注意力的更多详细信息,请查看 MultiHeadAttention
  • BertSparseSelfAttention:此模块包含一个简化的 BertSelfAttention 层,可以用来代替原始的密集 Bert 自注意力层。我们的实现基于 DeepSpeedExample
  • SparseAttentionUtils:此模块提供一些实用程序函数来处理使用稀疏注意力调整预训练模型
    • replace_model_self_attention_with_sparse_self_attention:如果你当前加载了一个模型,并希望用稀疏自注意力模块替换自注意力模块,你可以简单地使用此函数来处理它。它目前处理 BERT 和 RoBERTa 基于预训练的模型,但如果你的模型类型与这两个模型不同,你可以根据你的模型类型扩展它。你还需要扩展位置嵌入以处理新的序列长度;这可以通过使用 extend_position_embedding 函数来完成。
    • update_tokenizer_model_max_length:此函数只是用新值更新标记器中的最大位置嵌入。
    • extend_position_embedding:此函数根据当前值扩展位置嵌入。例如,如果你有一个最大序列长度为 128 的模型,并将其扩展到 1k 序列长度,它会将当前嵌入复制 8 次以初始化新的嵌入。根据实验,我们发现这种初始化比从头开始初始化效果好得多;导致更快收敛。
    • pad_to_block_size:此函数在序列长度维度上填充输入标记和注意力掩码,使其成为块大小的倍数;这是 SA 的要求。
    • unpad_sequence_output:如果模型的输入已填充,此函数会取消填充序列输出。
  • SparsityConfig:这是一个用于稀疏结构的抽象类。任何稀疏结构都需要扩展此类并编写自己的稀疏模式构造;make_layout 函数。DeepSpeed 目前提供以下结构,将在 如何配置稀疏结构 部分中进行介绍
    • FixedSparsityConfig
    • BSLongformerSparsityConfig
    • BigBirdSparsityConfig
    • VariableSparsityConfig
    • DenseSparsityConfig

注意:目前 DeepSpeed Transformer 内核不支持稀疏注意力。要使用稀疏注意力,你需要禁用 Transformer 内核!

如何使用 DeepSpeed 启动器使用稀疏注意力

在本部分中,我们将介绍如何通过我们的 bing_bert 代码使用 DeepSpeed 稀疏注意力。

  • 更新注意力模块:首先,你需要根据稀疏计算更新你的注意力模块。在这里,我们使用 BertSparseSelfAttention,它是我们 bing_bert 代码中 BertSelfAttention 的稀疏版本。它重写了 BertSelfAttention,其中它用
attention_scores = torch.matmul(query_layer, key_layer)
attention_scores = attention_scores / math.sqrt(
    self.attention_head_size)

# Apply the attention mask is (precomputed for all layers in BertModel forward() function)
attention_scores = attention_scores + attention_mask

pdtype = attention_scores.dtype
# Normalize the attention scores to probabilities.
attention_probs = self.softmax(attention_scores)

# This is actually dropping out entire tokens to attend to, which might
# seem a bit unusual, but is taken from the original Transformer paper.
attention_probs = self.dropout(attention_probs)

context_layer = torch.matmul(attention_probs, value_layer)

替换

context_layer =
  self.sparse_self_attention(
	query_layer,
	key_layer,
	value_layer,
	key_padding_mask=attention_mask)

其中 sparse_self_attentionSparseSelfAttention 的一个实例。此模块通过稀疏注意力计算注意力上下文,用其等效的稀疏版本替换底层矩阵乘法和 softmax。你可以类似地更新任何其他注意力模块。

  • 在模型中设置稀疏注意力配置:你需要设置稀疏注意力配置。在我们的示例中,这是在 BertModel 中完成的。
self.pad_token_id = config.pad_token_id if hasattr(
   config, 'pad_token_id') and config.pad_token_id is not None else 0
# set sparse_attention_config if it has been selected
self.sparse_attention_config = get_sparse_attention_config(
   args, config.num_attention_heads)
self.encoder = BertEncoder(
   config, args, sparse_attention_config=self.sparse_attention_config)
  • 更新编码器模型:此外,你需要更新你的编码器模型,以便在启用 SA 时使用 SA 进行注意力层计算。请查看我们的 bing_bert 示例,其中我们在启用 SA 时使用 BertSparseSelfAttention 而不是 BertSelfAttention
if sparse_attention_config is not None:
    from deepspeed.ops.sparse_attention import BertSparseSelfAttention

    layer.attention.self = BertSparseSelfAttention(
         config, sparsity_config=sparse_attention_config)
  • 填充和取消填充输入数据:你可能还需要将 input_idsattention_mask 的序列维度填充为稀疏块大小的倍数。如上面 模块 部分所述,DeepSpeed 提供了用于填充和取消填充的实用程序函数。请查看我们的 bing_bert 示例,以了解在何处以及如何填充或取消填充模型的输入或输出。
if self.sparse_attention_config is not None:
   pad_len, input_ids, attention_mask, token_type_ids, position_ids, inputs_embeds = SparseAttentionUtils.pad_to_block_size(
      block_size=self.sparse_attention_config.block,
      input_ids=input_ids,
      attention_mask=extended_attention_mask,
      token_type_ids=token_type_ids,
      position_ids=None,
      inputs_embeds=None,
      pad_token_id=self.pad_token_id,
      model_embeddings=self.embeddings)
.
.
.
# If BertEncoder uses sparse attention, and input_ids were padded, sequence output needs to be unpadded to original length
if self.sparse_attention_config is not None and pad_len > 0:
   encoded_layers[-1] = SparseAttentionUtils.unpad_sequence_output(
      pad_len, encoded_layers[-1])
  • *启用稀疏注意力*:要使用 DeepSpeed 稀疏注意力,你需要通过 deepspeed_sparse_attention 参数在启动器脚本中启用它
--deepspeed_sparse_attention

请查看 我们的 bing_bert 运行器脚本 作为如何使用 DeepSpeed 启动器启用 SA 的示例。

"sparse_attention": {
    "mode": "fixed",
    "block": 16,
    "different_layout_per_head": true,
    "num_local_blocks": 4,
    "num_global_blocks": 1,
    "attention": "bidirectional",
    "horizontal_global_attention": false,
    "num_different_global_patterns": 4
}

如何使用单个内核

DeepSpeed 稀疏注意力可以用作 DeepSpeed 的一项功能,如上所述,或者简单地作为独立的自注意力模块与任何 Transformer 模型集成在一起。此外,构建块内核、矩阵乘法和 softmax 可以单独使用。要单独使用稀疏注意力,你可以简单地安装 DeepSpeed 并导入 模块 部分中描述的任何模块;示例

from deepspeed.ops.sparse_attention import SparseSelfAttention

请参阅 Docstrings 了解有关如何单独使用每个模块的详细信息。

如何配置稀疏结构

下面将介绍支持的稀疏结构、它们的设置参数以及在自注意力层上添加任意稀疏模式的灵活性。你可以使用任何支持的稀疏结构更新 DeepSpeed 配置文件,并相应地设置参数。

  • SparsityConfig:此模块是所有稀疏结构的父类,包含所有稀疏结构的共享功能。它接受以下参数
    • num_heads:一个整数,确定层的注意力头的数量。
    • block:一个整数,确定块的大小。当前稀疏自注意力的实现基于块稀疏矩阵。其中,此参数定义此类方形块的大小;块 X 块
    • different_layout_per_head:一个布尔值,确定每个头是否应该分配一个不同的稀疏布局;默认值为 false,并将根据可用性满足此条件。
  • Fixed (FixedSparsityConfig):此结构基于 OpenAI 的 使用稀疏 Transformer 进行生成式建模,其中局部和全局注意力由给定的参数固定
    • num_local_blocks:一个整数,确定局部注意力窗口中块的数量。如在下图(改编自原始论文)中所示,局部窗口中的标记会关注与其局部相关的全部标记。在自回归模型的情况下,如上图所示,标记会关注在局部窗口中出现于其之前的标记。在 BERT 等掩码模型的情况下,注意力是双向的。
    • num_global_blocks:一个整数,确定局部窗口中使用多少个连续块作为窗口的代表用于全局注意力;如在下图中所示。
    • attention:一个字符串,确定注意力类型。注意力可以是 unidirectional(如自回归模型),其中标记只关注在上下文中出现于其之前的标记。考虑到这一点,注意力矩阵的上三角部分为空,如上图所示。或者,它可以是 bidirectional(如 BERT),其中标记可以关注其之前或之后的任何其他标记。然后,注意力矩阵的上三角部分是上图中下三角部分的镜像。
    • horizontal_global_attention:一个布尔值,确定作为局部窗口的全局代表的块是否也关注所有其他块。这仅在注意力类型为 bidirectional 时有效。观察注意力矩阵,这意味着全局注意力不仅包括垂直块,还包括水平块。
    • num_different_global_patterns: 用于确定不同全局注意力布局数量的整数。虽然全局注意力可以通过哪些块代表任何局部窗口来固定,但由于存在多头,每个头可以使用不同的全局代表。例如,如果使用 4 个块构成局部窗口,全局注意力大小为单个块,那么我们可以有 4 个不同的版本,其中每个局部窗口的第一个、第二个、第三个或第四个块可以作为该窗口的全局代表。此参数决定我们希望有多少种这样的模式。当然,根据 num_local_blocksnum_global_blocks 存在限制。此外,如果您将其设置为大于 1,则需要将 different_layout_per_head 设置为 True

Fixed sparsity structure

  • BSLongformer (BSLongformerSparsityConfig): 这种结构是对 Longformer: The Long-Document Transformer 的修改版本,其中我们提供块级令牌稀疏性,而不是单个令牌级稀疏性。定义此模式的参数为
    • num_sliding_window_blocks: 用于确定滑动局部注意力窗口中块数量的整数。
    • global_block_indices: 用于确定哪些块被视为全局注意力的整数列表。给定索引,确定所有其他令牌块都注意到的块,以及它们注意到的所有其他令牌块。请注意,如果设置了 global_block_end_indices 参数,则该参数将用作每个全局窗口的起始索引。
    • global_block_end_indices: 用于确定全局窗口块的结束索引的整数列表。默认情况下,它不使用。但如果设置了,它的大小必须与 global_block_indices 参数相同,并且结合这两个参数,对于每个索引 i,从 global_block_indices[i]global_block_end_indices[i](不包含)的块都被视为全局注意力块。
  • BigBird (BigBirdSparsityConfig): 这种结构基于 Big Bird: Transformers for Longer Sequences。它在某种程度上将 fixedlongformer 模式与随机注意力结合起来。以下参数定义了这种结构
    • num_random_blocks: 用于确定每行块中随机关注的块数量的整数。
    • num_sliding_window_blocks: 用于确定滑动局部注意力窗口中块数量的整数。
    • num_global_blocks: 用于确定从索引 0 开始被视为全局注意力的连续块数量的整数。全局块令牌将被所有其他块令牌关注,并将同样关注所有其他块令牌。
  • Variable (VariableSparsityConfig): 这种结构还结合了局部、全局和随机注意力的概念。此外,它具有定义可变大小局部窗口的灵活性。以下是定义此结构的参数列表
    • num_random_blocks: 用于确定每行块中随机关注的块数量的整数。
    • local_window_blocks: 用于确定每个局部注意力窗口中的块数量的整数列表。它假设第一个数字确定第一个局部窗口中的块数,第二个数字确定第二个窗口中的块数,……,最后一个数字确定剩余局部窗口中的块数。
    • global_block_indices: 用于确定哪些块被视为全局注意力的整数列表。给定索引,确定所有其他令牌块都注意到的块,以及它们注意到的所有其他令牌块。请注意,如果设置了 global_block_end_indices 参数,则该参数将用作每个全局窗口的起始索引。
    • global_block_end_indices: 用于确定全局窗口块的结束索引的整数列表。默认情况下,它不使用。但如果设置了,它的大小必须与 global_block_indices 参数相同,并且结合这两个参数,对于每个索引 i,从 global_block_indices[i]global_block_end_indices[i](不包含)的块都被视为全局注意力块。
    • attention:一个字符串,确定注意力类型。注意力可以是 unidirectional(如自回归模型),其中标记只关注在上下文中出现于其之前的标记。考虑到这一点,注意力矩阵的上三角部分为空,如上图所示。或者,它可以是 bidirectional(如 BERT),其中标记可以关注其之前或之后的任何其他标记。然后,注意力矩阵的上三角部分是上图中下三角部分的镜像。
    • horizontal_global_attention: 用于确定作为局部窗口全局代表的块是否也关注所有其他块的布尔值。这仅在注意力类型为 bidirectional 时有效。查看注意力矩阵,这意味着全局注意力不仅包括垂直块,还包括水平块。下图说明了 variable 稀疏性的一个示例,其中蓝色、橙色和绿色块分别说明了局部、全局和随机注意力块。

Variable sparsity structure

此外,我们提供了一个 dense 模式 (DenseSparsityConfig),它可以用于测试目的,因为它代表了完整的注意力。

如何支持新的用户定义稀疏结构

我们的构建块内核,基于块的 MatMulSoftmax,可以接受任何基于块的稀疏性。这提供了将任何基于块的稀疏性模式应用于注意力得分的灵活性。为了定义和应用新的稀疏性模式,您只需遵循上述任何稀疏性结构即可。您需要添加一个扩展 SparsityConfig 的新类,并根据您的稀疏性结构定义 make_layout 函数。您可以添加任何您可能需要的额外参数,或者只使用父类的默认参数。

更新: